// Integration tests for virtual key admin API (T038), rate limiting (T051), // budget enforcement (US5), and RBAC (US6). // // Admin routes require CSRF double-submit cookie protection on POST/PUT/DELETE. // Tests inject a fixed test token via X-CSRF-Token header + Cookie to satisfy the middleware. use anyllm_proxy::admin; use anyllm_proxy::config::{ BackendAuth, BackendConfig, BackendKind, Config, ModelMapping, MultiConfig, OpenAIApiFormat, TlsConfig, }; use anyllm_proxy::server::routes; use axum::body::Body; use axum::extract::connect_info::MockConnectInfo; use axum::http::Request; use axum::response::{IntoResponse, Response}; use axum::routing::{any, post}; use axum::Router; use dashmap::DashMap; use indexmap::IndexMap; use reqwest::Client; use serde_json::json; use std::net::SocketAddr; use std::sync::{ atomic::{AtomicUsize, Ordering}, Arc, Mutex, OnceLock, }; use tokio::net::TcpListener; use tower::ServiceExt; // --------------------------------------------------------------------------- // Shared DashMap for tests that need the proxy auth middleware. // `set_virtual_keys` uses a global OnceLock — whichever test runs first wins. // All proxy-auth tests share this one Arc so the middleware always // looks at the same map that the tests populate. // --------------------------------------------------------------------------- static TEST_VK_MAP: OnceLock>> = OnceLock::new(); static TEST_HMAC_SECRET: OnceLock>> = OnceLock::new(); static GEMINI_NATIVE_ENV_LOCK: OnceLock> = OnceLock::new(); fn shared_vk_map() -> Arc> { TEST_VK_MAP .get_or_init(|| { let map = Arc::new(DashMap::new()); anyllm_proxy::server::middleware::set_virtual_keys(map.clone()); map }) .clone() } fn shared_hmac_secret() -> Arc> { TEST_HMAC_SECRET .get_or_init(|| { // Use a fixed test secret so all tests agree on hash values. let secret = Arc::new(b"test-hmac-secret-for-integration".to_vec()); anyllm_proxy::server::middleware::set_hmac_secret(secret.clone()); secret }) .clone() } /// Build a SharedState whose `virtual_keys` is the shared test map. fn shared_state() -> admin::state::SharedState { let mut state = admin::state::SharedState::new_for_test(); state.virtual_keys = shared_vk_map(); state.hmac_secret = shared_hmac_secret(); state } fn insert_test_virtual_key(raw_key: &str, key_id: i64, allowed_models: Option>) { insert_test_virtual_key_with_tpm(raw_key, key_id, None, allowed_models); } fn insert_test_virtual_key_with_tpm( raw_key: &str, key_id: i64, tpm_limit: Option, 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, 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, }, ); } async fn recv_request_completed( rx: &mut tokio::sync::broadcast::Receiver, ) -> admin::state::RequestLogEntry { tokio::time::timeout(std::time::Duration::from_secs(2), async { loop { match rx.recv().await { Ok(admin::state::AdminEvent::RequestCompleted(entry)) => return entry, Ok(_) => {} Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} Err(e) => panic!("request log channel closed: {e}"), } } }) .await .expect("request log event") } // --------------------------------------------------------------------------- // Admin API CRUD tests (T038) // --------------------------------------------------------------------------- /// CSRF token used by unit-level tests (oneshot). Must match the cookie value below. const TEST_CSRF_TOKEN: &str = "0000000000000000000000000000000000000000000000000000000000000001"; /// Cookie header value that satisfies the CSRF double-submit check for the token above. const TEST_CSRF_COOKIE: &str = "csrf_token=0000000000000000000000000000000000000000000000000000000000000001"; fn test_admin_router() -> (Router, admin::state::SharedState) { // Raise admin rate limit so parallel tests from 127.0.0.1 don't starve each other. admin::routes::set_admin_rpm(10_000); let state = shared_state(); // Pre-register the test CSRF token so the first mutating request passes. // Tests that make multiple mutations must call reinsert_csrf(&state) before // each additional mutation (one-time tokens are consumed on use). state .issued_csrf_tokens .insert(TEST_CSRF_TOKEN.to_string(), ()); let token = Arc::new(zeroize::Zeroizing::new("test-admin-token".to_string())); let router = admin::routes::admin_router(state.clone(), token) // ConnectInfo extractor requires the service to be wrapped with // into_make_service_with_connect_info in production. In tests we use // MockConnectInfo so handlers can extract a fake peer address. .layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 0)))); (router, state) } fn test_admin_router_with_state(state: admin::state::SharedState) -> Router { admin::routes::set_admin_rpm(10_000); state .issued_csrf_tokens .insert(TEST_CSRF_TOKEN.to_string(), ()); let token = Arc::new(zeroize::Zeroizing::new("test-admin-token".to_string())); admin::routes::admin_router(state, token) .layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 0)))) } /// Re-register the test CSRF token before a subsequent mutating request. /// The CSRF middleware consumes tokens on first use; call this between mutations. fn reinsert_csrf(state: &admin::state::SharedState) { state .issued_csrf_tokens .insert(TEST_CSRF_TOKEN.to_string(), ()); } #[tokio::test] async fn create_key_returns_201_with_raw_key() { let (app, _state) = test_admin_router(); let req = Request::post("/admin/api/keys") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({"description": "test key", "rpm_limit": 60})).unwrap(), )) .unwrap(); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), 201); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); assert!(body["key"].as_str().unwrap().starts_with("sk-vk")); assert!(body["id"].as_i64().is_some()); assert_eq!(body["description"], "test key"); assert_eq!(body["rpm_limit"], 60); } #[tokio::test] async fn list_keys_returns_created_keys() { let (app, state) = test_admin_router(); let create_req = Request::post("/admin/api/keys") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({"description": "list-test"})).unwrap(), )) .unwrap(); let _ = app.clone().oneshot(create_req).await.unwrap(); // GET does not consume a CSRF token; no reinsert needed here. let _ = &state; // satisfy compiler: state used for reinsert if needed let list_req = Request::get("/admin/api/keys") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .body(Body::empty()) .unwrap(); let resp = app.oneshot(list_req).await.unwrap(); assert_eq!(resp.status(), 200); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); let keys = body["keys"].as_array().unwrap(); assert!(!keys.is_empty()); } #[tokio::test] async fn revoke_key_removes_from_dashmap() { let (app, state) = test_admin_router(); let create_req = Request::post("/admin/api/keys") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({"description": "revoke-test"})).unwrap(), )) .unwrap(); let resp = app.clone().oneshot(create_req).await.unwrap(); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); let id = body["id"].as_i64().unwrap(); let raw_key = body["key"].as_str().unwrap().to_string(); let hash = admin::keys::hmac_hash_key(&raw_key, &state.hmac_secret); let hash_bytes = admin::keys::hash_from_hex(&hash).unwrap(); assert!(state.virtual_keys.contains_key(&hash_bytes)); reinsert_csrf(&state); let revoke_req = Request::delete(format!("/admin/api/keys/{id}")) .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::empty()) .unwrap(); let resp = app.clone().oneshot(revoke_req).await.unwrap(); assert_eq!(resp.status(), 200); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); assert_eq!(body["status"], "revoked"); assert!(!state.virtual_keys.contains_key(&hash_bytes)); } #[tokio::test] async fn revoke_nonexistent_key_returns_404() { let (app, _state) = test_admin_router(); let req = Request::delete("/admin/api/keys/9999") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::empty()) .unwrap(); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), 404); } // --------------------------------------------------------------------------- // Update key (PUT /admin/api/keys/{id}) tests // --------------------------------------------------------------------------- #[tokio::test] async fn update_key_returns_200_with_updated_fields() { let (app, state) = test_admin_router(); // Create a key first. let create_req = Request::post("/admin/api/keys") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({"description": "update-test", "rpm_limit": 10})).unwrap(), )) .unwrap(); let resp = app.clone().oneshot(create_req).await.unwrap(); assert_eq!(resp.status(), 201); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); let id = body["id"].as_i64().unwrap(); reinsert_csrf(&state); // Update description and rpm_limit. let update_req = Request::put(format!("/admin/api/keys/{id}")) .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({ "description": "updated-desc", "rpm_limit": 200, "allowed_models": ["gpt-4o", "claude-*"] })) .unwrap(), )) .unwrap(); let resp = app.clone().oneshot(update_req).await.unwrap(); assert_eq!(resp.status(), 200); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); assert_eq!(body["description"], "updated-desc"); assert_eq!(body["rpm_limit"], 200); let models = body["allowed_models"].as_array().unwrap(); assert_eq!(models.len(), 2); assert_eq!(models[0], "gpt-4o"); assert_eq!(models[1], "claude-*"); } #[tokio::test] async fn update_nonexistent_key_returns_404() { let (app, _state) = test_admin_router(); let req = Request::put("/admin/api/keys/99999") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({"description": "no-such-key"})).unwrap(), )) .unwrap(); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), 404); } #[tokio::test] async fn update_revoked_key_returns_404() { let (app, state) = test_admin_router(); // Create then revoke. let create_req = Request::post("/admin/api/keys") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({"description": "revoke-then-update"})).unwrap(), )) .unwrap(); let resp = app.clone().oneshot(create_req).await.unwrap(); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); let id = body["id"].as_i64().unwrap(); reinsert_csrf(&state); let revoke_req = Request::delete(format!("/admin/api/keys/{id}")) .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::empty()) .unwrap(); let resp = app.clone().oneshot(revoke_req).await.unwrap(); assert_eq!(resp.status(), 200); reinsert_csrf(&state); // Update should fail with 404. let update_req = Request::put(format!("/admin/api/keys/{id}")) .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({"description": "should-fail"})).unwrap(), )) .unwrap(); let resp = app.oneshot(update_req).await.unwrap(); assert_eq!(resp.status(), 404); } #[tokio::test] async fn update_key_refreshes_dashmap() { let (app, state) = test_admin_router(); // Create key. let create_req = Request::post("/admin/api/keys") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({"description": "dashmap-update-test", "rpm_limit": 10})) .unwrap(), )) .unwrap(); let resp = app.clone().oneshot(create_req).await.unwrap(); assert_eq!(resp.status(), 201); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); let id = body["id"].as_i64().unwrap(); let raw_key = body["key"].as_str().unwrap().to_string(); reinsert_csrf(&state); // Update rpm_limit and allowed_routes via PUT. let update_req = Request::put(format!("/admin/api/keys/{id}")) .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({ "rpm_limit": 500, "allowed_routes": ["route-allowed"] })) .unwrap(), )) .unwrap(); let resp = app.oneshot(update_req).await.unwrap(); assert_eq!(resp.status(), 200); // Verify DashMap entry was updated. let hash = admin::keys::hmac_hash_key(&raw_key, &state.hmac_secret); let hash_bytes = admin::keys::hash_from_hex(&hash).unwrap(); let meta = state .virtual_keys .get(&hash_bytes) .expect("key should exist in DashMap"); assert_eq!(meta.rpm_limit, Some(500)); assert_eq!( meta.allowed_routes.as_ref().unwrap(), &vec!["route-allowed".to_string()] ); } // --------------------------------------------------------------------------- // Helpers for proxy-level tests // --------------------------------------------------------------------------- fn openai_config_with_base(base_url: &str) -> Config { Config { backend: BackendKind::OpenAI, openai_api_key: "test-key".to_string(), openai_base_url: base_url.to_string(), listen_port: 0, model_mapping: ModelMapping { big_model: "gpt-4o".into(), small_model: "gpt-4o-mini".into(), }, tls: anyllm_proxy::config::TlsConfig::default(), backend_auth: BackendAuth::BearerToken("test-key".into()), log_bodies: false, redact_secrets: false, anthropic_thinking_repair: false, pxpipe_compress: false, expose_degradation_warnings: false, openai_api_format: OpenAIApiFormat::Chat, provider_id: None, } } 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, redact_secrets: false, anthropic_thinking_repair: false, pxpipe_compress: false, expose_degradation_warnings: false, openai_api_format: OpenAIApiFormat::Chat, provider_id: None, } } fn gemini_config_with_base(base_url: &str) -> Config { Config { backend: BackendKind::Gemini, openai_api_key: "test-key".to_string(), openai_base_url: base_url.to_string(), listen_port: 0, model_mapping: ModelMapping { big_model: "gemini-2.5-pro".into(), small_model: "gemini-2.5-flash".into(), }, tls: anyllm_proxy::config::TlsConfig::default(), backend_auth: BackendAuth::GoogleApiKey("test-key".into()), log_bodies: false, redact_secrets: false, anthropic_thinking_repair: false, pxpipe_compress: 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", post(|| async { axum::Json(json!({ "id": "chatcmpl-mock", "object": "chat.completion", "created": 1700000000, "model": "gpt-4o", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop" }], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} })) }), ); 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_openai_streaming_backend() -> String { let app = Router::new().route( "/v1/chat/completions", post(|| async { Response::builder() .status(200) .header("content-type", "text/event-stream") .body(Body::from(concat!( "data: {\"id\":\"chatcmpl-mock\",\"object\":\"chat.completion.chunk\",\"created\":1700000000,\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"}}]}\n\n", "data: {\"id\":\"chatcmpl-mock\",\"object\":\"chat.completion.chunk\",\"created\":1700000000,\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hello\"}}]}\n\n", "data: {\"id\":\"chatcmpl-mock\",\"object\":\"chat.completion.chunk\",\"created\":1700000000,\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\n", "data: [DONE]\n\n" ))) .unwrap() }), ); 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_hanging_openai_streaming_backend() -> (String, Arc) { let release = Arc::new(tokio::sync::Notify::new()); let app = Router::new().route( "/v1/chat/completions", post({ let release = release.clone(); move || { let release = release.clone(); async move { let stream = futures::stream::once(async move { release.notified().await; Ok::<_, std::convert::Infallible>(bytes::Bytes::from_static( b"data: [DONE]\n\n", )) }); Response::builder() .status(200) .header("content-type", "text/event-stream") .body(Body::from_stream(stream)) .unwrap() } } }), ); 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}"), release) } async fn spawn_gemini_native_backend() -> String { let app = Router::new().route( "/v1beta/models/{model_action}", post( |axum::extract::Path(model_action): axum::extract::Path| async move { let payload = r#"{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3,"totalTokenCount":10}}"#; if model_action.ends_with(":streamGenerateContent") { Response::builder() .status(200) .header("content-type", "text/event-stream") .body(Body::from(format!("data: {payload}\n\n"))) .unwrap() } else { ([(axum::http::header::CONTENT_TYPE, "application/json")], payload) .into_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_anthropic_streaming_backend() -> String { let app = Router::new().route( "/v1/messages", post(|| async { Response::builder() .status(200) .header("content-type", "text/event-stream") .body(Body::from(concat!( "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_mock\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-haiku-4-5\",\"usage\":{\"input_tokens\":2,\"output_tokens\":0}}}\n\n", "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"ok\"}}\n\n", "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":2}}\n\n", "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" ))) .unwrap() }), ); 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_chat_backend() -> (String, Arc) { let hits = Arc::new(AtomicUsize::new(0)); let app = Router::new().route( "/v1/chat/completions", post({ let hits = hits.clone(); move || { let hits = hits.clone(); async move { hits.fetch_add(1, Ordering::SeqCst); axum::Json(json!({ "id": "chatcmpl-mock", "object": "chat.completion", "created": 1700000000, "model": "gpt-4o", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop" }], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} })) } } }), ); 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}"), hits) } fn openai_backend_config_with_base(base_url: &str) -> BackendConfig { BackendConfig { kind: BackendKind::OpenAI, provider_id: None, api_key: "test-key".to_string(), base_url: base_url.to_string(), api_format: OpenAIApiFormat::Chat, model_mapping: ModelMapping { big_model: "gpt-4o".into(), small_model: "gpt-4o-mini".into(), }, tls: anyllm_proxy::config::TlsConfig::default(), backend_auth: BackendAuth::BearerToken("test-key".into()), log_bodies: false, omit_stream_options: false, stream_timeout_secs: 900, bedrock_credentials: None, allow_local_ssrf: false, } } fn multi_config_with_backend_bases(allowed_base: &str, denied_base: &str) -> MultiConfig { let mut backends = IndexMap::new(); backends.insert( "allowed".to_string(), openai_backend_config_with_base(allowed_base), ); backends.insert( "denied".to_string(), openai_backend_config_with_base(denied_base), ); MultiConfig { listen_port: 0, log_bodies: false, redact_secrets: false, anthropic_thinking_repair: false, pxpipe_compress: false, forward_client_auth: false, default_backend: "allowed".to_string(), backends, expose_degradation_warnings: false, } } fn seed_route_scope( state: &admin::state::SharedState, route_id: &str, allowed_base: &str, denied_base: &str, ) { let conn = state.db.lock().unwrap_or_else(|e| e.into_inner()); let now = admin::db::now_iso8601(); for (id, name, base) in [ ("backend-allowed", "allowed", allowed_base), ("backend-denied", "denied", denied_base), ] { admin::db::insert_managed_backend( &conn, &admin::db::ManagedBackendRow { id: id.to_string(), name: name.to_string(), provider_id: "openai".to_string(), api_key: Some("test-key".to_string()), api_base: Some(base.to_string()), deployment: None, api_version: None, project: None, region: None, aws_access_key_id: None, aws_secret_access_key: None, aws_session_token: None, rpm: None, tpm: None, enabled: true, created_at: now.clone(), updated_at: now.clone(), }, ) .unwrap(); } admin::db::insert_route( &conn, &admin::db::RouteRow { id: route_id.to_string(), name: "allowed-route".to_string(), description: None, strategy: "failover".to_string(), rpm: None, tpm: None, budget_usd: None, enabled: true, guardrail_mode: None, pxpipe_compress: None, pxpipe_models: None, redact_secrets: None, position: 0, created_at: now.clone(), updated_at: now, }, ) .unwrap(); admin::db::add_route_provider( &conn, route_id, "backend-allowed", &["*".to_string()], 0, true, ) .unwrap(); } async fn create_route_scoped_key(app: Router, route_id: &str) -> String { let req = Request::post("/admin/api/keys") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({ "description": "route-scope-test", "allowed_routes": [route_id] })) .unwrap(), )) .unwrap(); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), 201); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); body["key"].as_str().unwrap().to_string() } async fn spawn_proxy_with_multi( config: MultiConfig, state: admin::state::SharedState, model_router: Option>>, ) -> String { let app = routes::app_multi_with_shared(config, Some(state), model_router, None, None, None); 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_gemini_native_proxy_with_state( base_url: &str, state: admin::state::SharedState, ) -> String { let app = { let _guard = GEMINI_NATIVE_ENV_LOCK .get_or_init(|| Mutex::new(())) .lock() .unwrap_or_else(|e| e.into_inner()); let previous = std::env::var_os("GEMINI_API_FORMAT"); std::env::set_var("GEMINI_API_FORMAT", "native"); let native_base_url = format!("{}/v1beta", base_url.trim_end_matches('/')); let app = routes::app_multi_with_shared( MultiConfig::from_single_config(&gemini_config_with_base(&native_base_url)), Some(state), None, None, None, None, ); match previous { Some(value) => std::env::set_var("GEMINI_API_FORMAT", value), None => std::env::remove_var("GEMINI_API_FORMAT"), } app }; 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_openai_backend(hits: Arc) -> String { let count = { 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 app = Router::new() .route("/v1/responses", post(count.clone())) .route( "/v1/{*path}", any(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 { let state = shared_state(); // must call before building app to ensure set_virtual_keys fires let multi = anyllm_proxy::config::MultiConfig::from_single_config(&config); let base_app = routes::app_multi_with_shared(multi, Some(state), None, None, None, None); // Add a test /admin/ route behind the same auth middleware so RBAC can be tested. let admin_test = Router::new() .route( "/admin/api/test", axum::routing::get(|| async { axum::Json(json!({"ok": true})) }), ) .layer(axum::middleware::from_fn( anyllm_proxy::server::middleware::validate_auth, )); let app = base_app.merge(admin_test); 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}") } #[tokio::test] async fn route_allowlist_denies_unassigned_backend_prefix() { let (allowed_base, allowed_hits) = spawn_counting_chat_backend().await; let (denied_base, denied_hits) = spawn_counting_chat_backend().await; let state = shared_state(); let route_id = "route-prefix-allowed"; seed_route_scope(&state, route_id, &allowed_base, &denied_base); let admin_app = test_admin_router_with_state(state.clone()); let raw_key = create_route_scoped_key(admin_app, route_id).await; let proxy_url = spawn_proxy_with_multi( multi_config_with_backend_bases(&allowed_base, &denied_base), state, None, ) .await; let client = Client::new(); let msg = json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] }); let resp = client .post(format!("{proxy_url}/allowed/v1/messages")) .header("x-api-key", &raw_key) .json(&msg) .send() .await .unwrap(); assert_eq!(resp.status(), 200); assert_eq!(allowed_hits.load(Ordering::SeqCst), 1); let resp = client .post(format!("{proxy_url}/denied/v1/messages")) .header("x-api-key", &raw_key) .json(&msg) .send() .await .unwrap(); assert_eq!(resp.status(), 403); assert_eq!(denied_hits.load(Ordering::SeqCst), 0); } #[tokio::test] async fn route_allowlist_denies_model_router_cross_backend_dispatch() { use anyllm_proxy::config::model_router::{Deployment, ModelRouter}; use std::collections::HashMap; use std::sync::RwLock; let (allowed_base, allowed_hits) = spawn_counting_chat_backend().await; let (denied_base, denied_hits) = spawn_counting_chat_backend().await; let state = shared_state(); let route_id = "route-model-router-allowed"; seed_route_scope(&state, route_id, &allowed_base, &denied_base); let admin_app = test_admin_router_with_state(state.clone()); let raw_key = create_route_scoped_key(admin_app, route_id).await; let mut model_routes = HashMap::new(); model_routes.insert( "claude-sonnet-4-20250514".to_string(), vec![Arc::new(Deployment::new( "denied".to_string(), "gpt-4o".to_string(), None, None, ))], ); let model_router = Some(Arc::new(RwLock::new(ModelRouter::new(model_routes)))); let proxy_url = spawn_proxy_with_multi( multi_config_with_backend_bases(&allowed_base, &denied_base), state, model_router, ) .await; let resp = Client::new() .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!(resp.status(), 403); assert_eq!(allowed_hits.load(Ordering::SeqCst), 0); assert_eq!(denied_hits.load(Ordering::SeqCst), 0); } #[tokio::test] async fn gemini_input_generate_content_records_virtual_key_usage_and_tpm() { let mock = spawn_mock_backend().await; let state = shared_state(); let mut events = state.events_tx.subscribe(); let proxy_url = spawn_proxy_with_multi( MultiConfig::from_single_config(&openai_config_with_base(&mock)), state, None, ) .await; let raw_key = "sk-vkgeminiinputusage"; insert_test_virtual_key_with_tpm(raw_key, 91_001, Some(5), Some(vec!["claude-*".to_string()])); let client = Client::new(); let body = json!({"contents": [{"role": "user", "parts": [{"text": "Hi"}]}]}); let resp = client .post(format!( "{proxy_url}/v1beta/models/claude-sonnet-4-20250514:generateContent" )) .header("x-goog-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 200); let entry = recv_request_completed(&mut events).await; assert_eq!(entry.key_id, Some(91_001)); assert_eq!(entry.input_tokens, Some(10)); assert_eq!(entry.output_tokens, Some(5)); assert!(entry.cost_usd.unwrap_or(0.0) > 0.0); let resp = client .post(format!( "{proxy_url}/v1beta/models/claude-sonnet-4-20250514:generateContent" )) .header("x-goog-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 429); } #[tokio::test] async fn gemini_input_stream_holds_concurrency_permits_until_body_drop() { let (mock, release) = spawn_hanging_openai_streaming_backend().await; let state = shared_state(); let proxy_url = spawn_proxy_with_multi( MultiConfig::from_single_config(&openai_config_with_base(&mock)), state, None, ) .await; let raw_key = "sk-vkgeminiinputstreampermit"; insert_test_virtual_key(raw_key, 91_002, Some(vec!["claude-*".to_string()])); let client = Client::new(); let body = json!({"contents": [{"role": "user", "parts": [{"text": "Hi"}]}]}); let mut held = Vec::new(); for _ in 0..100 { let resp = client .post(format!( "{proxy_url}/v1beta/models/claude-sonnet-4-20250514:streamGenerateContent" )) .header("x-goog-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 200); held.push(resp); } let resp = client .post(format!( "{proxy_url}/v1beta/models/claude-sonnet-4-20250514:streamGenerateContent" )) .header("x-goog-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 429); drop(held); release.notify_waiters(); } #[tokio::test] async fn translated_messages_stream_records_virtual_key_tpm() { let mock = spawn_openai_streaming_backend().await; let state = shared_state(); let mut events = state.events_tx.subscribe(); let proxy_url = spawn_proxy_with_multi( MultiConfig::from_single_config(&openai_config_with_base(&mock)), state, None, ) .await; let raw_key = "sk-vkmessagesstreamusage"; insert_test_virtual_key_with_tpm(raw_key, 91_003, Some(5), Some(vec!["claude-*".to_string()])); let client = Client::new(); let body = json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "stream": true, "messages": [{"role": "user", "content": "Hi"}] }); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 200); let text = resp.text().await.unwrap(); assert!(text.contains("message_delta")); let entry = recv_request_completed(&mut events).await; assert_eq!(entry.key_id, Some(91_003)); assert_eq!(entry.input_tokens, Some(10)); assert_eq!(entry.output_tokens, Some(5)); assert!(entry.is_streaming); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 429); } #[tokio::test] async fn openai_chat_completions_stream_records_virtual_key_tpm() { let mock = spawn_openai_streaming_backend().await; let state = shared_state(); let mut events = state.events_tx.subscribe(); let proxy_url = spawn_proxy_with_multi( MultiConfig::from_single_config(&openai_config_with_base(&mock)), state, None, ) .await; let raw_key = "sk-vkchatstreamusage"; insert_test_virtual_key_with_tpm(raw_key, 91_004, Some(5), Some(vec!["gpt-4o".to_string()])); let client = Client::new(); let body = json!({ "model": "gpt-4o", "max_tokens": 100, "stream": true, "messages": [{"role": "user", "content": "Hi"}] }); let resp = client .post(format!("{proxy_url}/v1/chat/completions")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 200); let text = resp.text().await.unwrap(); assert!(text.contains("[DONE]")); let entry = recv_request_completed(&mut events).await; assert_eq!(entry.key_id, Some(91_004)); assert_eq!(entry.input_tokens, Some(10)); assert_eq!(entry.output_tokens, Some(5)); assert!(entry.is_streaming); let resp = client .post(format!("{proxy_url}/v1/chat/completions")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 429); } #[tokio::test] async fn gemini_native_messages_records_virtual_key_usage_and_tpm() { let mock = spawn_gemini_native_backend().await; let state = shared_state(); let mut events = state.events_tx.subscribe(); let proxy_url = spawn_gemini_native_proxy_with_state(&mock, state).await; let raw_key = "sk-vkgemininativeusage"; insert_test_virtual_key_with_tpm(raw_key, 91_005, Some(3), Some(vec!["claude-*".to_string()])); let client = Client::new(); let body = json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] }); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 200); let entry = recv_request_completed(&mut events).await; assert_eq!(entry.key_id, Some(91_005)); assert_eq!(entry.input_tokens, Some(7)); assert_eq!(entry.output_tokens, Some(3)); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 429); } #[tokio::test] async fn gemini_native_messages_stream_records_virtual_key_usage_and_tpm() { let mock = spawn_gemini_native_backend().await; let state = shared_state(); let mut events = state.events_tx.subscribe(); let proxy_url = spawn_gemini_native_proxy_with_state(&mock, state).await; let raw_key = "sk-vkgemininativestreamusage"; insert_test_virtual_key_with_tpm(raw_key, 91_006, Some(3), Some(vec!["claude-*".to_string()])); let client = Client::new(); let body = json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "stream": true, "messages": [{"role": "user", "content": "Hi"}] }); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 200); let text = resp.text().await.unwrap(); assert!(text.contains("message_delta")); let entry = recv_request_completed(&mut events).await; assert_eq!(entry.key_id, Some(91_006)); assert_eq!(entry.input_tokens, Some(7)); assert_eq!(entry.output_tokens, Some(3)); assert!(entry.is_streaming); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 429); } #[tokio::test] async fn anthropic_passthrough_messages_records_virtual_key_usage_and_tpm() { let hits = Arc::new(AtomicUsize::new(0)); let mock = spawn_counting_anthropic_backend(hits).await; let state = shared_state(); let mut events = state.events_tx.subscribe(); let proxy_url = spawn_proxy_with_multi( MultiConfig::from_single_config(&anthropic_config_with_base(&mock)), state, None, ) .await; let raw_key = "sk-vkanthropicusage"; insert_test_virtual_key_with_tpm( raw_key, 91_007, Some(1), Some(vec!["claude-haiku-4-5".to_string()]), ); let client = Client::new(); let body = json!({ "model": "claude-haiku-4-5", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] }); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 200); let entry = recv_request_completed(&mut events).await; assert_eq!(entry.key_id, Some(91_007)); assert_eq!(entry.input_tokens, Some(1)); assert_eq!(entry.output_tokens, Some(1)); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 429); } #[tokio::test] async fn anthropic_passthrough_messages_stream_records_virtual_key_usage_and_tpm() { let mock = spawn_anthropic_streaming_backend().await; let state = shared_state(); let mut events = state.events_tx.subscribe(); let proxy_url = spawn_proxy_with_multi( MultiConfig::from_single_config(&anthropic_config_with_base(&mock)), state, None, ) .await; let raw_key = "sk-vkanthropicstreamusage"; insert_test_virtual_key_with_tpm( raw_key, 91_008, Some(2), Some(vec!["claude-haiku-4-5".to_string()]), ); let client = Client::new(); let body = json!({ "model": "claude-haiku-4-5", "max_tokens": 100, "stream": true, "messages": [{"role": "user", "content": "Hi"}] }); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 200); let text = resp.text().await.unwrap(); assert!(text.contains("message_delta")); let entry = recv_request_completed(&mut events).await; assert_eq!(entry.key_id, Some(91_008)); assert_eq!(entry.input_tokens, Some(2)); assert_eq!(entry.output_tokens, Some(2)); assert!(entry.is_streaming); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 429); } // --------------------------------------------------------------------------- // Virtual key auth lifecycle (T038): create → use → revoke → rejected // --------------------------------------------------------------------------- #[tokio::test] async fn virtual_key_auth_and_revocation_lifecycle() { let mock = spawn_mock_backend().await; let proxy_url = spawn_proxy_with_shared_vk(openai_config_with_base(&mock)).await; // Admin server uses shared VK map so create/revoke affect the same DashMap // the middleware checks. let state = shared_state(); let admin_app = admin::routes::admin_router( state, Arc::new(zeroize::Zeroizing::new("admin-token".to_string())), ); let admin_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let admin_port = admin_listener.local_addr().unwrap().port(); let admin_url = format!("http://127.0.0.1:{admin_port}"); tokio::spawn(async move { axum::serve( admin_listener, admin_app.into_make_service_with_connect_info::(), ) .await .unwrap() }); let client = Client::new(); // 1. Create a virtual key (fetch fresh CSRF token for each mutation) let csrf = fetch_csrf(&client, &admin_url, admin_port, "admin-token").await; let resp = client .post(format!("{admin_url}/admin/api/keys")) .header("host", format!("localhost:{admin_port}")) .header("authorization", "Bearer admin-token") .header("x-csrf-token", &csrf) .header("cookie", format!("csrf_token={csrf}")) .json(&json!({"description": "lifecycle-test"})) .send() .await .unwrap(); assert_eq!(resp.status(), 201); let body: serde_json::Value = resp.json().await.unwrap(); let raw_key = body["key"].as_str().unwrap().to_string(); let key_id = body["id"].as_i64().unwrap(); // 2. Use the virtual key to authenticate let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!(resp.status(), 200, "virtual key should authenticate"); // 3. Revoke the key let csrf = fetch_csrf(&client, &admin_url, admin_port, "admin-token").await; let resp = client .delete(format!("{admin_url}/admin/api/keys/{key_id}")) .header("host", format!("localhost:{admin_port}")) .header("authorization", "Bearer admin-token") .header("x-csrf-token", &csrf) .header("cookie", format!("csrf_token={csrf}")) .send() .await .unwrap(); assert_eq!(resp.status(), 200); // 4. Revoked key must be rejected let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!(resp.status(), 401, "revoked key should be rejected"); } #[tokio::test] async fn anthropic_generic_passthrough_rejects_virtual_key_without_upstream_call() { let upstream_hits = Arc::new(AtomicUsize::new(0)); let mock = spawn_counting_anthropic_backend(upstream_hits.clone()).await; let proxy_url = spawn_proxy_with_shared_vk(anthropic_config_with_base(&mock)).await; let raw_key = "sk-vkanthropicgenericdenied"; insert_test_virtual_key(raw_key, 90_001, Some(vec!["claude-haiku-4-5".to_string()])); let client = Client::new(); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", raw_key) .json(&json!({ "model": "claude-haiku-4-5", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!( resp.status(), 200, "virtual key should still work on policy-aware /v1/messages" ); assert_eq!(upstream_hits.load(Ordering::SeqCst), 1); let resp = client .post(format!("{proxy_url}/v1/messages/batches")) .header("x-api-key", raw_key) .json(&json!({ "requests": [{ "custom_id": "req-1", "params": { "model": "claude-opus-4-5", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] } }] })) .send() .await .unwrap(); assert_eq!(resp.status(), 403); let body: serde_json::Value = resp.json().await.unwrap(); assert_eq!(body["error"]["type"], "permission_error"); assert_eq!( body["error"]["message"], "This endpoint is not available for virtual API keys." ); assert_eq!( upstream_hits.load(Ordering::SeqCst), 1, "denied generic passthrough must not reach upstream" ); } #[tokio::test] async fn translate_generic_passthrough_rejects_virtual_key_without_upstream_call() { let upstream_hits = Arc::new(AtomicUsize::new(0)); let mock = spawn_counting_openai_backend(upstream_hits.clone()).await; let proxy_url = spawn_proxy_with_shared_vk(openai_config_with_base(&mock)).await; let raw_key = "sk-vktranslategenericdenied"; insert_test_virtual_key(raw_key, 90_002, Some(vec!["claude-*".to_string()])); let client = Client::new(); let resp = client .post(format!("{proxy_url}/v1/responses")) .header("x-api-key", raw_key) .json(&json!({ "model": "claude-sonnet-4-6", "input": "Hi" })) .send() .await .unwrap(); assert_eq!(resp.status(), 403); let body: serde_json::Value = resp.json().await.unwrap(); assert_eq!(body["error"]["type"], "permission_error"); assert_eq!( body["error"]["message"], "This endpoint is not available for virtual API keys." ); assert_eq!( upstream_hits.load(Ordering::SeqCst), 0, "denied generic passthrough must not reach upstream" ); } #[tokio::test] async fn translate_model_passthrough_routes_enforce_virtual_key_model_allowlist() { let upstream_hits = Arc::new(AtomicUsize::new(0)); let mock = spawn_counting_openai_backend(upstream_hits.clone()).await; let proxy_url = spawn_proxy_with_shared_vk(openai_config_with_base(&mock)).await; let raw_key = "sk-vkmodelpassthroughdenied"; insert_test_virtual_key(raw_key, 90_003, Some(vec!["allowed-model".to_string()])); let client = Client::new(); for (path, body) in [ ( "/v1/embeddings", json!({"model": "denied-model", "input": "hello"}), ), ("/v1/images/generations", json!({"model": "denied-model"})), ( "/v1/audio/speech", json!({"model": "denied-model", "input": "hello", "voice": "alloy"}), ), ] { let resp = client .post(format!("{proxy_url}{path}")) .header("x-api-key", raw_key) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 403, "path {path} should be denied"); let body: serde_json::Value = resp.json().await.unwrap(); assert_eq!(body["error"]["type"], "permission_error"); } let resp = client .post(format!("{proxy_url}/v1/audio/transcriptions")) .header("x-api-key", raw_key) .header("content-type", "multipart/form-data; boundary=test") .body("--test\r\n--test--\r\n") .send() .await .unwrap(); assert_eq!(resp.status(), 403); assert_eq!( upstream_hits.load(Ordering::SeqCst), 0, "denied model passthrough routes must not reach upstream" ); } // --------------------------------------------------------------------------- // Bedrock native routes (POST /model/{modelId}/...) must enforce the virtual // key's model allowlist, same as bedrock_passthrough. The modelId comes from // the URL path here, so the check runs before any AWS call — a denied model is // rejected 403 without dummy credentials ever being signed/sent. // --------------------------------------------------------------------------- fn bedrock_backend_config() -> BackendConfig { BackendConfig { kind: BackendKind::Bedrock, provider_id: Some("bedrock".to_string()), api_key: String::new(), base_url: "us-east-1".to_string(), // region is stored in base_url for Bedrock api_format: OpenAIApiFormat::Chat, model_mapping: ModelMapping { big_model: "anthropic.claude-3-5-sonnet-20241022-v2:0".into(), small_model: "anthropic.claude-3-5-haiku-20241022-v1:0".into(), }, tls: TlsConfig::default(), backend_auth: BackendAuth::BearerToken(String::new()), log_bodies: false, omit_stream_options: false, stream_timeout_secs: 900, // Dummy static credentials — never actually used because the denied // model is rejected before signing. bedrock_credentials: Some(aws_credential_types::Credentials::new( "AKIDTEST", "secret", None, None, "test", )), allow_local_ssrf: true, } } async fn spawn_bedrock_proxy_with_shared_vk() -> String { let state = shared_state(); // fires set_virtual_keys before app build let mut backends = IndexMap::new(); backends.insert("bedrock".to_string(), bedrock_backend_config()); let multi = MultiConfig { listen_port: 0, log_bodies: false, redact_secrets: false, anthropic_thinking_repair: false, pxpipe_compress: false, forward_client_auth: false, default_backend: "bedrock".to_string(), backends, expose_degradation_warnings: false, }; let app = routes::app_multi_with_shared(multi, Some(state), None, None, None, None); 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}") } #[tokio::test] async fn bedrock_native_routes_enforce_virtual_key_model_allowlist() { let proxy_url = spawn_bedrock_proxy_with_shared_vk().await; let raw_key = "sk-vkbedrocknativedenied"; insert_test_virtual_key( raw_key, 90_010, Some(vec!["anthropic.claude-3-5-sonnet-20241022-v2:0".to_string()]), ); let client = Client::new(); // Denied model in the URL path across all four native endpoints. let denied = "anthropic.claude-3-opus-20240229-v1:0"; for suffix in [ "converse", "converse-stream", "invoke", "invoke-with-response-stream", ] { let resp = client .post(format!("{proxy_url}/model/{denied}/{suffix}")) .header("x-api-key", raw_key) .json(&json!({"messages": [{"role": "user", "content": "hi"}]})) .send() .await .unwrap(); assert_eq!(resp.status(), 403, "native {suffix} must deny scoped model"); let body: serde_json::Value = resp.json().await.unwrap(); assert_eq!(body["error"]["type"], "permission_error"); } } // --------------------------------------------------------------------------- // RPM rate limiting (T051): create key with rpm_limit:2, 3rd request → 429 // --------------------------------------------------------------------------- #[tokio::test] async fn rpm_limit_returns_429_after_exceeded() { let mock = spawn_mock_backend().await; let proxy_url = spawn_proxy_with_shared_vk(openai_config_with_base(&mock)).await; let state = shared_state(); let admin_app = admin::routes::admin_router( state, Arc::new(zeroize::Zeroizing::new("admin-token2".to_string())), ); let admin_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let admin_port = admin_listener.local_addr().unwrap().port(); let admin_url = format!("http://127.0.0.1:{admin_port}"); tokio::spawn(async move { axum::serve( admin_listener, admin_app.into_make_service_with_connect_info::(), ) .await .unwrap() }); let client = Client::new(); // Create a key with rpm_limit: 2 let csrf = fetch_csrf(&client, &admin_url, admin_port, "admin-token2").await; let resp = client .post(format!("{admin_url}/admin/api/keys")) .header("host", format!("localhost:{admin_port}")) .header("authorization", "Bearer admin-token2") .header("x-csrf-token", &csrf) .header("cookie", format!("csrf_token={csrf}")) .json(&json!({"description": "rate-limit-test", "rpm_limit": 2})) .send() .await .unwrap(); assert_eq!(resp.status(), 201); let body: serde_json::Value = resp.json().await.unwrap(); let raw_key = body["key"].as_str().unwrap().to_string(); let msg = json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] }); // First 2 requests should succeed for _ in 0..2 { let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&msg) .send() .await .unwrap(); assert_eq!(resp.status(), 200); } // 3rd request must be rate-limited let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&msg) .send() .await .unwrap(); assert_eq!(resp.status(), 429); assert!( resp.headers().get("retry-after").is_some(), "429 must include retry-after header" ); } // --------------------------------------------------------------------------- // Budget enforcement tests (US5: T046) // --------------------------------------------------------------------------- /// Fetch a fresh server-issued CSRF token from the real admin server. async fn fetch_csrf( client: &Client, admin_url: &str, admin_port: u16, admin_token: &str, ) -> String { let resp = client .get(format!("{admin_url}/admin/csrf-token")) .header("host", format!("localhost:{admin_port}")) .header("authorization", format!("Bearer {admin_token}")) .send() .await .unwrap(); let body: serde_json::Value = resp.json().await.unwrap(); body["csrf_token"].as_str().unwrap().to_string() } /// Helper: create a key via admin API and return (raw_key, key_id). async fn create_key_via_admin( admin_url: &str, admin_port: u16, admin_token: &str, body: serde_json::Value, ) -> (String, i64) { let client = Client::new(); // Fetch a fresh server-issued CSRF token (one-time use). let csrf_resp = client .get(format!("{admin_url}/admin/csrf-token")) .header("host", format!("localhost:{admin_port}")) .header("authorization", format!("Bearer {admin_token}")) .send() .await .unwrap(); let csrf_body: serde_json::Value = csrf_resp.json().await.unwrap(); let csrf_token = csrf_body["csrf_token"].as_str().unwrap().to_string(); let csrf_cookie = format!("csrf_token={csrf_token}"); let resp = client .post(format!("{admin_url}/admin/api/keys")) .header("host", format!("localhost:{admin_port}")) .header("authorization", format!("Bearer {admin_token}")) .header("x-csrf-token", &csrf_token) .header("cookie", &csrf_cookie) .json(&body) .send() .await .unwrap(); assert_eq!(resp.status(), 201, "create key failed"); let body: serde_json::Value = resp.json().await.unwrap(); let raw_key = body["key"].as_str().unwrap().to_string(); let key_id = body["id"].as_i64().unwrap(); (raw_key, key_id) } /// Helper: spawn admin + proxy servers, returns (proxy_url, admin_url, admin_port). async fn spawn_test_servers(admin_token: &str) -> (String, String, u16) { // Raise admin rate limit so parallel tests from 127.0.0.1 don't starve each other. admin::routes::set_admin_rpm(10_000); let mock = spawn_mock_backend().await; let proxy_url = spawn_proxy_with_shared_vk(openai_config_with_base(&mock)).await; let state = shared_state(); let admin_app = admin::routes::admin_router( state, Arc::new(zeroize::Zeroizing::new(admin_token.to_string())), ); let admin_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let admin_port = admin_listener.local_addr().unwrap().port(); let admin_url = format!("http://127.0.0.1:{admin_port}"); tokio::spawn(async move { axum::serve( admin_listener, admin_app.into_make_service_with_connect_info::(), ) .await .unwrap() }); (proxy_url, admin_url, admin_port) } #[tokio::test] async fn budget_exceeded_returns_429_with_budget_exceeded_type() { let (proxy_url, admin_url, admin_port) = spawn_test_servers("budget-token1").await; // Create key with a tiny budget ($0.0001) and no duration (lifetime budget) let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "budget-token1", json!({"description": "budget-test", "max_budget_usd": 0.0001}), ) .await; // Manually set the key's period_spend above the limit in the DashMap let vk_map = shared_vk_map(); let hash = admin::keys::hmac_hash_key(&raw_key, &shared_hmac_secret()); let hash_bytes = admin::keys::hash_from_hex(&hash).unwrap(); if let Some(mut meta) = vk_map.get_mut(&hash_bytes) { meta.period_spend_usd = 1.0; // Way over the $0.0001 limit } let client = Client::new(); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!(resp.status(), 429, "budget exceeded should return 429"); let body: serde_json::Value = resp.json().await.unwrap(); assert_eq!( body["error"]["type"], "budget_exceeded", "error type must be budget_exceeded, not rate_limit_exceeded" ); assert!( body["error"]["budget_limit_usd"].as_f64().is_some(), "response must include budget_limit_usd" ); assert!( body["error"]["period_spend_usd"].as_f64().is_some(), "response must include period_spend_usd" ); } #[tokio::test] async fn budget_not_exceeded_allows_request() { let (proxy_url, admin_url, admin_port) = spawn_test_servers("budget-token2").await; // Create key with generous budget let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "budget-token2", json!({"description": "budget-ok-test", "spend_limit": 100.0, "budget_duration": "monthly"}), ) .await; let client = Client::new(); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!(resp.status(), 200, "under-budget key should succeed"); } #[tokio::test] async fn budget_resets_on_new_period() { let (proxy_url, admin_url, admin_port) = spawn_test_servers("budget-token3").await; // Create key with daily budget let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "budget-token3", json!({"description": "budget-reset-test", "max_budget_usd": 0.0001, "budget_duration": "daily"}), ) .await; // Manually set the key's spend above limit AND set period_start to yesterday let vk_map = shared_vk_map(); let hash = admin::keys::hmac_hash_key(&raw_key, &shared_hmac_secret()); let hash_bytes = admin::keys::hash_from_hex(&hash).unwrap(); if let Some(mut meta) = vk_map.get_mut(&hash_bytes) { meta.period_spend_usd = 1.0; // Over budget meta.period_start = Some("2020-01-01T00:00:00Z".to_string()); // Long past } // The lazy reset should kick in and allow the request let client = Client::new(); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!(resp.status(), 200, "budget should reset for new period"); } #[tokio::test] async fn no_duration_budget_stays_blocked() { let (proxy_url, admin_url, admin_port) = spawn_test_servers("budget-token4").await; // Create key with lifetime budget (no duration) let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "budget-token4", json!({"description": "lifetime-budget-test", "max_budget_usd": 0.0001}), ) .await; // Set spend above limit; no duration means no reset let vk_map = shared_vk_map(); let hash = admin::keys::hmac_hash_key(&raw_key, &shared_hmac_secret()); let hash_bytes = admin::keys::hash_from_hex(&hash).unwrap(); if let Some(mut meta) = vk_map.get_mut(&hash_bytes) { meta.period_spend_usd = 1.0; } let client = Client::new(); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!(resp.status(), 429, "lifetime budget should stay blocked"); let body: serde_json::Value = resp.json().await.unwrap(); assert_eq!(body["error"]["type"], "budget_exceeded"); // No budget_duration = lifetime (null in response) assert!(body["error"]["budget_duration"].is_null()); } // --------------------------------------------------------------------------- // RBAC tests (US6: T050) // --------------------------------------------------------------------------- #[tokio::test] async fn developer_key_succeeds_on_v1_messages() { let (proxy_url, admin_url, admin_port) = spawn_test_servers("rbac-token1").await; // Create developer key (default role) let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "rbac-token1", json!({"description": "dev-key-test"}), ) .await; let client = Client::new(); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!( resp.status(), 200, "developer key should succeed on /v1/messages" ); } #[tokio::test] async fn developer_key_gets_403_on_admin() { let (proxy_url, admin_url, admin_port) = spawn_test_servers("rbac-token2").await; // Create developer key explicitly let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "rbac-token2", json!({"description": "dev-admin-test", "role": "developer"}), ) .await; // Access the test /admin/ route on the proxy (added by spawn_proxy_with_shared_vk). // The auth middleware checks RBAC before the route handler runs. let client = Client::new(); let resp = client .get(format!("{proxy_url}/admin/api/test")) .header("x-api-key", &raw_key) .send() .await .unwrap(); assert_eq!( resp.status(), 403, "developer key should get 403 on /admin/ path" ); let body: serde_json::Value = resp.json().await.unwrap(); assert_eq!(body["error"]["type"], "permission_denied"); } #[tokio::test] async fn admin_key_succeeds_on_v1_messages() { let (proxy_url, admin_url, admin_port) = spawn_test_servers("rbac-token3").await; // Create admin key let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "rbac-token3", json!({"description": "admin-key-test", "role": "admin"}), ) .await; let client = Client::new(); let resp = client .post(format!("{proxy_url}/v1/messages")) .header("x-api-key", &raw_key) .json(&json!({ "model": "claude-sonnet-4-20250514", "max_tokens": 100, "messages": [{"role": "user", "content": "Hi"}] })) .send() .await .unwrap(); assert_eq!( resp.status(), 200, "admin key should succeed on /v1/messages" ); } #[tokio::test] async fn admin_key_not_blocked_on_admin_path() { let (proxy_url, admin_url, admin_port) = spawn_test_servers("rbac-token4").await; // Create admin key let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "rbac-token4", json!({"description": "admin-path-test", "role": "admin"}), ) .await; let client = Client::new(); let resp = client .get(format!("{proxy_url}/admin/api/test")) .header("x-api-key", &raw_key) .send() .await .unwrap(); // Admin key should NOT get 403. The test route returns 200. assert_eq!( resp.status(), 200, "admin key should not get 403 on /admin/ path" ); } #[tokio::test] async fn new_key_defaults_to_developer_role() { let (_proxy_url, admin_url, admin_port) = spawn_test_servers("rbac-token5").await; // Create key with no explicit role let (raw_key, _key_id) = create_key_via_admin( &admin_url, admin_port, "rbac-token5", json!({"description": "default-role-test"}), ) .await; // Check that the in-memory meta has developer role let vk_map = shared_vk_map(); let hash = admin::keys::hmac_hash_key(&raw_key, &shared_hmac_secret()); let hash_bytes = admin::keys::hash_from_hex(&hash).unwrap(); let meta = vk_map.get(&hash_bytes).expect("key should exist in map"); assert_eq!( meta.role, admin::keys::KeyRole::Developer, "new keys should default to developer role" ); // Also check via list endpoint let client = Client::new(); let resp = client .get(format!("{admin_url}/admin/api/keys")) .header("host", format!("localhost:{admin_port}")) .header("authorization", "Bearer rbac-token5") .send() .await .unwrap(); assert_eq!(resp.status(), 200); let body: serde_json::Value = resp.json().await.unwrap(); let keys = body["keys"].as_array().unwrap(); // Find our key by description let our_key = keys .iter() .find(|k| k["description"] == "default-role-test") .expect("our key should appear in list"); assert_eq!(our_key["role"], "developer"); } // --------------------------------------------------------------------------- // Model management validation tests // --------------------------------------------------------------------------- #[tokio::test] async fn add_model_rejects_unknown_backend() { use anyllm_proxy::config::model_router::ModelRouter; use std::collections::HashMap; use std::sync::RwLock; admin::routes::set_admin_rpm(10_000); let mut state = admin::state::SharedState::new_for_test(); // Give it a model_router so the handler passes the "no model router active" guard. state.model_router = Some(Arc::new(RwLock::new(ModelRouter::new(HashMap::new())))); // backend_metrics is empty by default — any backend_name should be rejected. state .issued_csrf_tokens .insert(TEST_CSRF_TOKEN.to_string(), ()); let token = Arc::new(zeroize::Zeroizing::new("test-admin-token".to_string())); let app = admin::routes::admin_router(state, token) .layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 0)))); let req = Request::post("/admin/api/models") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .header("content-type", "application/json") .header("x-csrf-token", TEST_CSRF_TOKEN) .header("cookie", TEST_CSRF_COOKIE) .body(Body::from( serde_json::to_string(&json!({ "model_name": "my-model", "backend_name": "nonexistent", "actual_model": "gpt-4o" })) .unwrap(), )) .unwrap(); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), 400); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); assert!( body["error"] .as_str() .unwrap_or("") .contains("unknown backend"), "expected 'unknown backend' in error, got: {:?}", body ); } // --------------------------------------------------------------------------- // Pagination has_more field tests // --------------------------------------------------------------------------- #[tokio::test] async fn get_requests_response_has_has_more_field() { let (app, _state) = test_admin_router(); let req = Request::get("/admin/api/requests?limit=10&offset=0") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .body(Body::empty()) .unwrap(); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), 200); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); assert!( body.get("has_more").is_some(), "response should include has_more field" ); assert_eq!(body["has_more"], serde_json::Value::Bool(false)); } #[tokio::test] async fn get_audit_response_has_has_more_field() { let (app, _state) = test_admin_router(); let req = Request::get("/admin/api/audit?limit=10&offset=0") .header("host", "localhost:9090") .header("authorization", "Bearer test-admin-token") .body(Body::empty()) .unwrap(); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), 200); let body: serde_json::Value = serde_json::from_slice( &axum::body::to_bytes(resp.into_body(), 1 << 20) .await .unwrap(), ) .unwrap(); assert!( body.get("has_more").is_some(), "response should include has_more field" ); assert_eq!(body["has_more"], serde_json::Value::Bool(false)); }