From 158815acbdabfa9cd90316cf81fa4bdfd4c7e1f9 Mon Sep 17 00:00:00 2001 From: whit3rabbit Date: Sat, 23 May 2026 11:49:26 -0500 Subject: [PATCH] Block virtual-key passthrough bypass --- crates/proxy/src/main.rs | 37 ++++-- .../proxy/src/server/generic_passthrough.rs | 19 ++- crates/proxy/src/server/passthrough.rs | 13 ++ crates/proxy/tests/virtual_keys.rs | 112 +++++++++++++++++- 4 files changed, 163 insertions(+), 18 deletions(-) diff --git a/crates/proxy/src/main.rs b/crates/proxy/src/main.rs index 7de2c96..c299852 100644 --- a/crates/proxy/src/main.rs +++ b/crates/proxy/src/main.rs @@ -601,9 +601,9 @@ async fn async_main(args: Vec, data_dir: PathBuf) { // --- Admin setup (enabled only when --webui or --admin flag is passed) --- // Returns Some((SharedState, admin Router, admin TcpListener)) when enabled. // admin_redirect_port: passed to the proxy router to enable GET / redirect. - // admin_startup_info: (admin_url, token_str) printed in the startup banner. + // admin_startup_info: admin_url printed in the startup banner. let mut admin_redirect_port: Option = None; - let mut admin_startup_info: Option<(String, String)> = None; + let mut admin_startup_info: Option = None; let admin_parts = if enable_admin { let admin_port: u16 = match std::env::var("ADMIN_PORT") { Ok(val) => val @@ -1008,7 +1008,7 @@ async fn async_main(args: Vec, data_dir: PathBuf) { &admin_bind }; let admin_ui_url = format!("http://{}:{}/admin/", admin_display_host, admin_port); - admin_startup_info = Some((admin_ui_url, admin_token.as_str().to_owned())); + admin_startup_info = Some(admin_ui_url); // Spawn periodic tasks: log retention and metrics snapshot broadcast. let retention_days: u32 = std::env::var("ADMIN_LOG_RETENTION_DAYS") @@ -1181,16 +1181,9 @@ async fn async_main(args: Vec, data_dir: PathBuf) { .unwrap_or_else(|e| panic!("failed to bind proxy to {proxy_addr}: {e}")); tracing::info!("proxy listening on {proxy_addr}"); - // Print a clear startup banner once both servers are bound, so it appears - // at the bottom of startup output and is easy to find and copy from. - if let Some((admin_url, token)) = &admin_startup_info { - let proxy_display = proxy_addr.replace("0.0.0.0", "localhost"); - let border = "─".repeat(56); - println!("{border}"); - println!(" Proxy API http://{proxy_display}"); - println!(" Admin UI {admin_url}"); - println!(" Token {token}"); - println!("{border}"); + // Print non-secret startup details once both servers are bound. + if let Some(admin_url) = &admin_startup_info { + println!("{}", format_startup_banner(&proxy_addr, admin_url)); } // Warn if API keys are configured and listener is on a non-loopback address. @@ -1268,6 +1261,12 @@ async fn async_main(args: Vec, data_dir: PathBuf) { tracing::info!("server shut down gracefully"); } +fn format_startup_banner(proxy_addr: &str, admin_url: &str) -> String { + let proxy_display = proxy_addr.replace("0.0.0.0", "localhost"); + let border = "─".repeat(56); + format!("{border}\n Proxy API http://{proxy_display}\n Admin UI {admin_url}\n{border}") +} + /// Parse a `.env`-format file and return `(key, value)` pairs to set. /// /// Delegates parsing to `anyllm_proxy::env_parser::parse_env_content` (pure, no side effects). @@ -1653,6 +1652,18 @@ fn run_subcommand(proxy_args: Vec, tool_argv: Vec) -> i32 { mod tests { use super::*; + #[test] + fn startup_banner_does_not_include_admin_token() { + let token = "sentinel-admin-token-0123456789abcdef"; + let banner = format_startup_banner("0.0.0.0:3000", "http://127.0.0.1:3001/admin/"); + + assert!(banner.contains("Proxy API http://localhost:3000")); + assert!(banner.contains("Admin UI http://127.0.0.1:3001/admin/")); + assert!(!banner.contains(token)); + assert!(!banner.contains("Admin token:")); + assert!(!banner.contains("Token ")); + } + #[test] fn parse_env_file_double_quoted_newline_escape() { use std::io::Write; diff --git a/crates/proxy/src/server/generic_passthrough.rs b/crates/proxy/src/server/generic_passthrough.rs index 593734e..7dc88fc 100644 --- a/crates/proxy/src/server/generic_passthrough.rs +++ b/crates/proxy/src/server/generic_passthrough.rs @@ -5,23 +5,38 @@ use crate::backend::BackendClient; use crate::server::state::AppState; +use anyllm_translate::{anthropic, mapping}; use axum::{ body::Bytes, extract::{OriginalUri, Path, State}, http::{HeaderMap, Method, StatusCode}, - response::{IntoResponse, Response}, + response::{IntoResponse, Json, Response}, }; /// Catch-all for `ANY /v1/{*path}` paths without an explicit handler. /// Registered last in the Translate-mode router so explicit routes take priority. pub(crate) async fn v1_generic_passthrough( State(state): State, + vk_ctx: Option>, OriginalUri(uri): OriginalUri, Path(tail): Path, method: Method, headers: HeaderMap, body: Bytes, ) -> Response { + state.metrics.record_request(); + + // Virtual keys must use explicit handlers that enforce per-key policy and + // accounting before forwarding with the proxy's provider credentials. + if vk_ctx.is_some() { + let err = mapping::errors_map::create_anthropic_error( + anthropic::ErrorType::PermissionError, + "This endpoint is not available for virtual API keys.".to_string(), + None, + ); + return (StatusCode::FORBIDDEN, Json(err)).into_response(); + } + // This handler is only registered for OpenAI-compatible (Translate) backends. let client = match &state.backend { BackendClient::OpenAI(c) @@ -39,8 +54,6 @@ pub(crate) async fn v1_generic_passthrough( } }; - state.metrics.record_request(); - // Build backend URL: passthrough_url handles per-backend path rewriting (Azure, Vertex, etc.) let path = format!("/v1/{tail}"); let mut url = client.passthrough_url(&path); diff --git a/crates/proxy/src/server/passthrough.rs b/crates/proxy/src/server/passthrough.rs index aebd72e..66a99d5 100644 --- a/crates/proxy/src/server/passthrough.rs +++ b/crates/proxy/src/server/passthrough.rs @@ -137,6 +137,7 @@ pub(crate) async fn anthropic_passthrough( /// route retains its dedicated streaming/model-peek logic. pub(crate) async fn anthropic_generic_passthrough( State(state): State, + vk_ctx: Option>, OriginalUri(uri): OriginalUri, method: axum::http::Method, headers: axum::http::HeaderMap, @@ -144,6 +145,18 @@ pub(crate) async fn anthropic_generic_passthrough( ) -> Response { state.metrics.record_request(); + // Virtual keys must use policy-aware handlers only. The generic passthrough + // can reach Anthropic-native endpoints that lack per-key authorization, + // resource ownership checks, and usage accounting. + if vk_ctx.is_some() { + let err = mapping::errors_map::create_anthropic_error( + anthropic::ErrorType::PermissionError, + "This endpoint is not available for virtual API keys.".to_string(), + None, + ); + return (StatusCode::FORBIDDEN, Json(err)).into_response(); + } + let client = match &state.backend { BackendClient::Anthropic(c) => c, _ => { diff --git a/crates/proxy/tests/virtual_keys.rs b/crates/proxy/tests/virtual_keys.rs index 6c6b07e..fc31a9b 100644 --- a/crates/proxy/tests/virtual_keys.rs +++ b/crates/proxy/tests/virtual_keys.rs @@ -10,13 +10,16 @@ use anyllm_proxy::server::routes; use axum::body::Body; use axum::extract::connect_info::MockConnectInfo; use axum::http::Request; -use axum::routing::post; +use axum::routing::{any, post}; use axum::Router; use dashmap::DashMap; use reqwest::Client; use serde_json::json; use std::net::SocketAddr; -use std::sync::{Arc, OnceLock}; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, OnceLock, +}; use tokio::net::TcpListener; use tower::ServiceExt; @@ -59,6 +62,29 @@ fn shared_state() -> admin::state::SharedState { state } +fn insert_test_virtual_key(raw_key: &str, key_id: i64, allowed_models: Option>) { + let hash = admin::keys::hmac_hash_key(raw_key, &shared_hmac_secret()); + let hash_bytes = admin::keys::hash_from_hex(&hash).unwrap(); + shared_vk_map().insert( + hash_bytes, + admin::keys::VirtualKeyMeta { + id: key_id, + description: Some("generic-passthrough-test".to_string()), + expires_at: None, + rpm_limit: None, + tpm_limit: None, + rate_state: Arc::new(admin::keys::RateLimitState::new()), + role: admin::keys::KeyRole::Developer, + max_budget_usd: None, + budget_duration: None, + period_start: None, + period_spend_usd: 0.0, + allowed_models, + allowed_routes: None, + }, + ); +} + // --------------------------------------------------------------------------- // Admin API CRUD tests (T038) // --------------------------------------------------------------------------- @@ -426,6 +452,25 @@ fn openai_config_with_base(base_url: &str) -> Config { } } +fn anthropic_config_with_base(base_url: &str) -> Config { + Config { + backend: BackendKind::Anthropic, + openai_api_key: "test-key".to_string(), + openai_base_url: base_url.to_string(), + listen_port: 0, + model_mapping: ModelMapping { + big_model: "claude-sonnet-4-6".into(), + small_model: "claude-haiku-4-5".into(), + }, + tls: anyllm_proxy::config::TlsConfig::default(), + backend_auth: BackendAuth::BearerToken("test-key".into()), + log_bodies: false, + expose_degradation_warnings: false, + openai_api_format: OpenAIApiFormat::Chat, + provider_id: None, + } +} + async fn spawn_mock_backend() -> String { let app = Router::new().route( "/v1/chat/completions", @@ -450,6 +495,69 @@ async fn spawn_mock_backend() -> String { format!("http://{addr}") } +async fn spawn_counting_openai_backend(hits: Arc) -> String { + let app = Router::new().route( + "/v1/responses", + post({ + let hits = hits.clone(); + move || { + let hits = hits.clone(); + async move { + hits.fetch_add(1, Ordering::SeqCst); + axum::Json(json!({"id": "resp_mock", "object": "response"})) + } + } + }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + format!("http://{addr}") +} + +async fn spawn_counting_anthropic_backend(hits: Arc) -> String { + let app = Router::new() + .route( + "/v1/messages", + post({ + let hits = hits.clone(); + move || { + let hits = hits.clone(); + async move { + hits.fetch_add(1, Ordering::SeqCst); + axum::Json(json!({ + "id": "msg_mock", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": {"input_tokens": 1, "output_tokens": 1} + })) + } + } + }), + ) + .route( + "/v1/{*path}", + any({ + let hits = hits.clone(); + move || { + let hits = hits.clone(); + async move { + hits.fetch_add(1, Ordering::SeqCst); + axum::Json(json!({"ok": true})) + } + } + }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + format!("http://{addr}") +} + /// Spawn a proxy backed by the shared VK map so auth middleware can find virtual keys. /// Includes a dummy /admin/api/test route behind auth to test RBAC. async fn spawn_proxy_with_shared_vk(config: Config) -> String {