diff --git a/litellm-rust/crates/gateway-auth/AGENTS.md b/litellm-rust/crates/gateway-auth/AGENTS.md index 0a4bc6450ae..7196ee74bf2 100644 --- a/litellm-rust/crates/gateway-auth/AGENTS.md +++ b/litellm-rust/crates/gateway-auth/AGENTS.md @@ -2,7 +2,7 @@ Inbound authentication uses a verifier, identity resolver, and authorizer injected into `Auth`. `Auth::from_config` composes the current master-key implementation. `Auth::new` accepts alternate implementations and an injected clock -`CredentialExtractor` selects a token or transport credential from HTTP request parts. The default `Bearer` extractor is strict. A host can supply another extractor with `Auth::with_extractor` for a separate route profile. A transport credential is only a selection marker, never proof of identity +`CredentialExtractor` selects a token or transport credential from HTTP request parts. The default `Bearer` extractor is strict. `BearerOrXApiKey` reads one `x-api-key` header, the Anthropic SDK's key header, only when the request has no Authorization header, so a malformed Authorization header never falls back to it. A host can supply another extractor with `Auth::with_extractor` for a separate route profile. A transport credential is only a selection marker, never proof of identity `Authenticator` receives that selected credential and HTTP request parts. Transport verifiers must obtain proof from trusted server extensions, not client-supplied identity headers. It returns `VerifiedIdentity` only after verification. Request parts provide headers and trusted transport extensions for integrations. The verifier must validate any client-controlled claims before treating them as identity. There is no fallback chain after verification failure @@ -12,7 +12,7 @@ Inbound authentication uses a verifier, identity resolver, and authorizer inject `Permissions::All`, `None`, and `Only` provide the initial policy vocabulary. `Only` matches exact typed access requests. Model requests contain both the public name and resolved deployment model. MCP requests contain the resolved server ID and upstream operation, including bare tool/prompt names and resource URIs. More expressive policy adapters can be added without changing credential verification -The Axum `authenticate` middleware currently accepts one Authorization header using the existing Bearer format. Missing, duplicate, empty, or malformed credentials fail. Authentication replaces any preexisting caller extension and checks the method and matched route before dispatch. Handlers extract `AuthenticatedRequest` and authorize their parsed operation before calling a provider. Missing authenticated context fails closed +The Axum `authenticate` middleware runs the route profile's extractor. Missing, duplicate, empty, or malformed credentials fail. Authentication replaces any preexisting caller extension and checks the method and matched route before dispatch. Handlers extract `AuthenticatedRequest` and authorize their parsed operation before calling a provider. Missing authenticated context fails closed The gateway also inserts an MCP `Authorization` extension and a `SessionOwner` derived from the verified principal and credential reference. The MCP host checks the injected policy at initialization and before sending resolved upstream operations. Standalone MCP library hosts retain their existing explicitly trusted registry behavior when no policy is installed. Production mounts must install authentication and the policy extension together diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs index 7bb6292b18c..a7f7ac42144 100644 --- a/litellm-rust/crates/gateway-auth/src/error.rs +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -6,7 +6,7 @@ use litellm_error::Source; #[derive(Debug, thiserror::Error)] pub enum Error { - #[error("missing or invalid bearer token")] + #[error("missing or invalid API key")] InvalidToken, #[error("invalid UI credentials or session")] InvalidSession, diff --git a/litellm-rust/crates/gateway-auth/src/http.rs b/litellm-rust/crates/gateway-auth/src/http.rs index ad1718161c0..b0c3419699b 100644 --- a/litellm-rust/crates/gateway-auth/src/http.rs +++ b/litellm-rust/crates/gateway-auth/src/http.rs @@ -1,6 +1,9 @@ use axum::{ extract::{FromRequestParts, MatchedPath, Request, State}, - http::{header::AUTHORIZATION, request::Parts}, + http::{ + header::{AUTHORIZATION, AsHeaderName}, + request::Parts, + }, middleware::Next, response::Response, }; @@ -20,20 +23,40 @@ impl CredentialExtractor for Bearer { } } +pub struct BearerOrXApiKey; + +impl CredentialExtractor for BearerOrXApiKey { + fn extract(&self, parts: &Parts) -> Result { + if parts.headers.contains_key(AUTHORIZATION) { + return Bearer.extract(parts); + } + single(parts, "x-api-key") + .and_then(token) + .map(Credential::Token) + } +} + fn bearer(parts: &Parts) -> Result { - let mut values = parts.headers.get_all(AUTHORIZATION).iter(); + single(parts, AUTHORIZATION)? + .strip_prefix("Bearer ") + .ok_or(Error::InvalidToken) + .and_then(token) +} + +fn single(parts: &Parts, name: impl AsHeaderName) -> Result<&str, Error> { + let mut values = parts.headers.get_all(name).iter(); let value = values.next().ok_or(Error::InvalidToken)?; if values.next().is_some() { return Err(Error::InvalidToken); } - let token = value - .to_str() - .ok() - .and_then(|value| value.strip_prefix("Bearer ")) - .map(str::trim) + value.to_str().map_err(|_| Error::InvalidToken) +} + +fn token(value: &str) -> Result { + Some(value.trim()) .filter(|value| !value.is_empty() && !value.bytes().any(|byte| byte.is_ascii_whitespace())) - .ok_or(Error::InvalidToken)?; - Ok(SecretValue::new(token)) + .map(SecretValue::new) + .ok_or(Error::InvalidToken) } pub async fn authenticate( @@ -88,3 +111,39 @@ impl FromRequestParts for RequireMasterKey { Ok(Self) } } + +#[cfg(test)] +mod tests { + use axum::http::Request; + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::x_api_key(&[("x-api-key", "key")], Some("key"))] + #[case::bearer(&[("authorization", "Bearer key")], Some("key"))] + #[case::bearer_wins(&[("authorization", "Bearer bearer"), ("x-api-key", "x")], Some("bearer"))] + #[case::invalid_bearer_does_not_fall_back(&[("authorization", "Basic key"), ("x-api-key", "x")], None)] + #[case::empty(&[("x-api-key", "")], None)] + #[case::inner_whitespace(&[("x-api-key", "a b")], None)] + #[case::duplicate(&[("x-api-key", "key"), ("x-api-key", "key")], None)] + #[case::missing(&[], None)] + fn bearer_or_x_api_key_reads_x_api_key_only_without_authorization( + #[case] headers: &[(&str, &str)], + #[case] expected: Option<&str>, + ) { + let (parts, ()) = headers + .iter() + .fold(Request::builder(), |request, (name, value)| { + request.header(*name, *value) + }) + .body(()) + .unwrap() + .into_parts(); + let token = match BearerOrXApiKey.extract(&parts) { + Ok(Credential::Token(token)) => Some(token.expose().to_owned()), + _ => None, + }; + assert_eq!(token.as_deref(), expected); + } +} diff --git a/litellm-rust/crates/gateway-auth/src/lib.rs b/litellm-rust/crates/gateway-auth/src/lib.rs index 31ce17b3aa4..a34ccd4e2da 100644 --- a/litellm-rust/crates/gateway-auth/src/lib.rs +++ b/litellm-rust/crates/gateway-auth/src/lib.rs @@ -17,7 +17,7 @@ pub use authorization::{ NoAdditionalPolicy, Permissions, UiAction, }; pub use error::Error; -pub use http::{Bearer, CredentialExtractor, RequireMasterKey, authenticate}; +pub use http::{Bearer, BearerOrXApiKey, CredentialExtractor, RequireMasterKey, authenticate}; pub use identity::{ AuthenticatedCaller, Authentication, AuthenticationMethod, Principal, PrincipalKind, ResolvedIdentity, SharedCaller, VerifiedIdentity, diff --git a/litellm-rust/crates/gateway-mcp/AGENTS.md b/litellm-rust/crates/gateway-mcp/AGENTS.md index c4b00e3acbf..f462e44b9da 100644 --- a/litellm-rust/crates/gateway-mcp/AGENTS.md +++ b/litellm-rust/crates/gateway-mcp/AGENTS.md @@ -2,7 +2,7 @@ Native MCP routing using the official [Rust SDK](https://github.com/modelcontextprotocol/rust-sdk), pinned to `rmcp 3.4.1`. The crate exposes a mountable Axum router and an SDK `ServerHandler` -The HTTP router serves `/mcp`, `/{server}/mcp`, `/mcp/{server}`, `/mcp/{server}/mcp`, legacy `/mcp/sse` and `/mcp/sse/messages`, `/mcp/enabled`, and `/mcp-rest/tools/list` and `/mcp-rest/tools/call`. Streamable HTTP supports stateless requests and legacy initialize/session flows. The caller supplies the shutdown token, allowed hosts, allowed origins, and server identity through `HttpConfig` +The HTTP router serves `/mcp`, `/{server}/mcp`, `/mcp/{server}`, `/mcp/{server}/mcp`, legacy `/mcp/sse` and `/mcp/sse/messages`, `/mcp/enabled`, and `/mcp-rest/tools/list` and `/mcp-rest/tools/call`. `routes` returns the protocol routes and the `/mcp-rest` routes apart so a host can authenticate each group its own way, and `router` merges them. Streamable HTTP supports stateless requests and legacy initialize/session flows. The caller supplies the shutdown token, allowed hosts, allowed origins, and server identity through `HttpConfig` `NativeGateway` aggregates upstream tools, prompts, resources, and resource templates. Tool and prompt calls resolve registered server prefixes, including prefixes containing hyphens. Resource names receive prefixes while their URIs remain unchanged, and reads require exactly one selected server. Tool allowlists apply to discovery and execution. Upstream pagination is drained with cycle detection. Aggregate tool listings retain healthy servers and expose sanitized outcomes under `_meta["litellm.ai/server_outcomes"]`, while scoped REST listings preserve upstream failures diff --git a/litellm-rust/crates/gateway-mcp/src/lib.rs b/litellm-rust/crates/gateway-mcp/src/lib.rs index fcb7fd98383..a9624a7d36c 100644 --- a/litellm-rust/crates/gateway-mcp/src/lib.rs +++ b/litellm-rust/crates/gateway-mcp/src/lib.rs @@ -18,7 +18,7 @@ pub use error::{ConfigFault, ConnectError, Error}; pub use native::{Limits, NativeGateway, Registry, Server, ServerResolver}; pub use relay::RelayClient; pub use rmcp; -pub use routes::{HttpConfig, router}; +pub use routes::{HttpConfig, Routes, router, routes}; pub use server::McpServer; pub use sessions::SessionOwner; diff --git a/litellm-rust/crates/gateway-mcp/src/routes.rs b/litellm-rust/crates/gateway-mcp/src/routes.rs index 56d67a5bc77..30dd7cd9120 100644 --- a/litellm-rust/crates/gateway-mcp/src/routes.rs +++ b/litellm-rust/crates/gateway-mcp/src/routes.rs @@ -43,7 +43,17 @@ impl Default for HttpConfig { } } +pub struct Routes { + pub protocol: Router, + pub rest: Router, +} + pub fn router(operations: Arc, config: HttpConfig) -> Router { + let Routes { protocol, rest } = routes(operations, config); + protocol.merge(rest) +} + +pub fn routes(operations: Arc, config: HttpConfig) -> Routes { let server = McpServer::new(operations.clone()).with_info(config.server_info.clone()); let legacy = crate::legacy::router(operations.clone(), server.clone(), &config); let transport_config = StreamableHttpServerConfig::default() @@ -59,7 +69,7 @@ pub fn router(operations: Arc, config: HttpConfig) -> Router { .route_service("/{server}/mcp", service.clone()) .route_service("/{server}/mcp/", service.clone()) .route_layer(middleware::from_fn(scope_server)); - Router::new() + let protocol = Router::new() .route_service("/mcp", service.clone()) .route_service("/mcp/", service) .merge(scoped) @@ -67,10 +77,12 @@ pub fn router(operations: Arc, config: HttpConfig) -> Router { "/mcp/enabled", get(|| async { Json(Enabled { enabled: true }) }), ) + .merge(legacy); + let rest = Router::new() .route("/mcp-rest/tools/list", get(rest_list)) .route("/mcp-rest/tools/call", post(rest_call)) - .with_state(operations) - .merge(legacy) + .with_state(operations); + Routes { protocol, rest } } async fn scope_server(Path(server): Path, mut request: Request, next: Next) -> Response { diff --git a/litellm-rust/crates/gateway-mcp/src/sessions.rs b/litellm-rust/crates/gateway-mcp/src/sessions.rs index 583a8b86341..67656f354db 100644 --- a/litellm-rust/crates/gateway-mcp/src/sessions.rs +++ b/litellm-rust/crates/gateway-mcp/src/sessions.rs @@ -6,7 +6,7 @@ use axum::{ extract::{Request, State}, http::Method, response::IntoResponse, - routing::any, + routing::{RouterIntoService, any}, }; use moka::future::Cache; use rmcp::{ @@ -31,7 +31,10 @@ struct Transport { owners: Cache, } -pub(crate) fn transport(server: McpServer, config: StreamableHttpServerConfig) -> Router { +pub(crate) fn transport( + server: McpServer, + config: StreamableHttpServerConfig, +) -> RouterIntoService { let manager = Arc::new(LocalSessionManager::default()); let cleanup = manager.clone(); let owners = Cache::builder() @@ -55,11 +58,14 @@ pub(crate) fn transport(server: McpServer, config: StreamableHttpServerConfig) - Arc::new(LocalSessionManager::default()), config.with_legacy_session_mode(false), ); - Router::new().fallback(any(dispatch)).with_state(Transport { - stateful: Router::new().fallback_service(stateful), - stateless: Router::new().fallback_service(stateless), - owners, - }) + Router::new() + .fallback(any(dispatch)) + .with_state(Transport { + stateful: Router::new().fallback_service(stateful), + stateless: Router::new().fallback_service(stateless), + owners, + }) + .into_service() } async fn dispatch( diff --git a/litellm-rust/crates/gateway/AGENTS.md b/litellm-rust/crates/gateway/AGENTS.md index 16714cd9a1f..eb1c8dd8228 100644 --- a/litellm-rust/crates/gateway/AGENTS.md +++ b/litellm-rust/crates/gateway/AGENTS.md @@ -1,5 +1,6 @@ - Keep this crate a thin composition layer: mount endpoint routers and serve the supplied listener - Server lifecycle and shared inbound middleware belong here, including client authentication, rate limiting, and request logging +- Inference and `/mcp-rest` routes authenticate with `BearerOrXApiKey`, since the proxy accepts an Anthropic `x-api-key` there. MCP protocol routes keep `Bearer`, since the proxy's MCP admission never reads `x-api-key` - Endpoint paths, request handling, model resolution, and response encoding belong to the mounted crates. Provider execution belongs to the `inference-` crates and `llms` - Inject shared state and infrastructure. Avoid global runtimes, duplicate client pools, and abstractions for hypothetical endpoint groups - Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 0a0191b0499..7f8ec5261bd 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -17,7 +17,7 @@ use axum::{ use http_body_util::BodyExt; use litellm_config::Config; -use litellm_gateway_auth::Auth; +use litellm_gateway_auth::{Auth, BearerOrXApiKey}; use litellm_gateway_inference::{Gateway, ModelRouter}; use litellm_http::{ ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, @@ -140,17 +140,20 @@ pub fn router( inference: Arc, config: &Config, ui: Option, - mcp: Option, + mcp: Option, ) -> Router { let auth = Auth::from_config(config, inference.secrets.clone()); - let inference = litellm_gateway_inference::router(inference) - .merge(mcp.unwrap_or_default()) - .route_layer(axum::middleware::from_fn(auth::bind_session_owner)) - .route_layer(axum::middleware::from_fn_with_state( - auth, - litellm_gateway_auth::authenticate, - )) - .layer(axum::middleware::from_fn(log_request)); + let (protocol, rest) = mcp.map(|mcp| (mcp.protocol, mcp.rest)).unzip(); + let inference = authenticated( + litellm_gateway_inference::router(inference).merge(rest.unwrap_or_default()), + auth.clone().with_extractor(Arc::new(BearerOrXApiKey)), + ) + .merge( + protocol + .map(|protocol| authenticated(protocol, auth)) + .unwrap_or_default(), + ) + .layer(axum::middleware::from_fn(log_request)); let health = Router::new().route("/health/readiness", axum::routing::get(readiness)); match ui { Some(ui) => inference.merge(ui).merge(health), @@ -158,6 +161,15 @@ pub fn router( } } +fn authenticated(routes: Router, auth: Auth) -> Router { + routes + .route_layer(axum::middleware::from_fn(auth::bind_session_owner)) + .route_layer(axum::middleware::from_fn_with_state( + auth, + litellm_gateway_auth::authenticate, + )) +} + async fn readiness() -> axum::Json { axum::Json(serde_json::json!({"status": "healthy", "db": "Not connected"})) } diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index dc77cfe28eb..ff2e2c7268c 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -90,9 +90,9 @@ async fn main() -> Result<(), Box> { &inference.http, ) .await?; - let mcp_router = mcp.as_ref().map(|gateway| { + let mcp_routes = mcp.as_ref().map(|gateway| { let defaults = litellm_gateway_mcp::HttpConfig::default(); - litellm_gateway_mcp::router( + litellm_gateway_mcp::routes( gateway.operations.clone(), litellm_gateway_mcp::HttpConfig { allowed_hosts: config @@ -113,7 +113,7 @@ async fn main() -> Result<(), Box> { let result = axum::serve( listener, - litellm_gateway::router(inference, &config, ui, mcp_router), + litellm_gateway::router(inference, &config, ui, mcp_routes), ) .with_graceful_shutdown(shutdown.clone().cancelled_owned()) .await; diff --git a/litellm-rust/crates/gateway/tests/mcp.rs b/litellm-rust/crates/gateway/tests/mcp.rs index a33c9d7c60e..e6e6af4a17d 100644 --- a/litellm-rust/crates/gateway/tests/mcp.rs +++ b/litellm-rust/crates/gateway/tests/mcp.rs @@ -158,14 +158,14 @@ mcp_servers: .await .unwrap() .unwrap(); - let mcp_router = litellm_gateway_mcp::router( + let mcp_routes = litellm_gateway_mcp::routes( mcp.operations.clone(), HttpConfig { cancellation_token: shutdown.clone(), ..Default::default() }, ); - let app = litellm_gateway::router(inference, &config, None, Some(mcp_router)); + let app = litellm_gateway::router(inference, &config, None, Some(mcp_routes)); Served { app, mcp, @@ -180,26 +180,50 @@ mcp_servers: #[case::streamable_delete("/mcp", "DELETE")] #[case::legacy_sse("/mcp/sse", "GET")] #[case::legacy_messages("/mcp/sse/messages", "POST")] -#[case::rest_list("/mcp-rest/tools/list", "GET")] #[case::enabled("/mcp/enabled", "GET")] #[tokio::test] -async fn every_mcp_route_requires_admission_credentials( +async fn mcp_protocol_routes_require_bearer_admission( #[future] served: Served, #[case] path: &str, #[case] method: &str, + #[values(&[], &[("x-api-key", "admission-secret")])] headers: &[(&str, &str)], ) { let served = served.await; + let request = headers.iter().fold( + Request::builder() + .method(method) + .uri(path) + .header("host", "localhost"), + |request, (name, value)| request.header(*name, *value), + ); let response = served .app .clone() - .oneshot( - Request::builder() - .method(method) - .uri(path) - .header("host", "localhost") - .body(Body::empty()) - .unwrap(), - ) + .oneshot(request.body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), 401); + served.close().await; +} + +#[rstest] +#[case::missing(&[])] +#[case::wrong_x_api_key(&[("x-api-key", "wrong")])] +#[case::non_bearer_authorization(&[("authorization", "Basic admission-secret"), ("x-api-key", "admission-secret")])] +#[tokio::test] +async fn mcp_rest_routes_reject_a_missing_or_invalid_key( + #[future] served: Served, + #[case] headers: &[(&str, &str)], +) { + let served = served.await; + let request = headers.iter().fold( + Request::get("/mcp-rest/tools/list?server_id=configured-id"), + |request, (name, value)| request.header(*name, *value), + ); + let response = served + .app + .clone() + .oneshot(request.body(Body::empty()).unwrap()) .await .unwrap(); assert_eq!(response.status(), 401); @@ -210,6 +234,8 @@ async fn every_mcp_route_requires_admission_credentials( #[tokio::test] async fn serves_configured_upstream_with_separate_admission_and_outbound_credentials( #[future] served: Served, + #[values(("authorization", "Bearer admission-secret"), ("x-api-key", "admission-secret"))] + admission: (&str, &str), ) { let served = served.await; let app = served.app.clone(); @@ -218,7 +244,7 @@ async fn serves_configured_upstream_with_separate_admission_and_outbound_credent .oneshot( Request::builder() .uri("/mcp-rest/tools/list?server_id=configured-id") - .header("authorization", "Bearer admission-secret") + .header(admission.0, admission.1) .body(Body::empty()) .unwrap(), ) @@ -231,7 +257,7 @@ async fn serves_configured_upstream_with_separate_admission_and_outbound_credent assert_eq!(body["tools"][0]["name"], "echo"); assert_eq!(body["tools"][0]["mcp_info"]["alias"], "knowledge"); - let response = app.clone().oneshot(Request::post("/mcp-rest/tools/call").header("authorization", "Bearer admission-secret").header("content-type", "application/json").body(Body::from(json!({"server_id":"configured-id", "name":"echo", "arguments":{"value":"hello"}}).to_string())).unwrap()).await.unwrap(); + let response = app.clone().oneshot(Request::post("/mcp-rest/tools/call").header(admission.0, admission.1).header("content-type", "application/json").body(Body::from(json!({"server_id":"configured-id", "name":"echo", "arguments":{"value":"hello"}}).to_string())).unwrap()).await.unwrap(); assert_eq!(response.status(), 200); let body: Value = serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap(); @@ -243,7 +269,7 @@ async fn serves_configured_upstream_with_separate_admission_and_outbound_credent let blocked = app .oneshot( Request::post("/mcp-rest/tools/call") - .header("authorization", "Bearer admission-secret") + .header(admission.0, admission.1) .header("content-type", "application/json") .body(Body::from( json!({"server_id":"configured-id", "name":"hidden"}).to_string(), diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 7de7a97c870..bf4d8b8e692 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -34,27 +34,29 @@ fn inference() -> Arc { } #[rstest] -#[case::authorized("/v1/messages", Some("Bearer gateway-key"), Some("gateway-key"), 400)] -#[case::missing_token("/v1/messages", None, Some("gateway-key"), 401)] -#[case::wrong_token("/v1/messages", Some("Bearer wrong"), Some("gateway-key"), 401)] -#[case::ocr("/ocr", None, Some("gateway-key"), 401)] -#[case::chat("/v1/chat/completions", None, Some("gateway-key"), 401)] +#[case::authorized("/v1/messages", &[("authorization", "Bearer gateway-key")], Some("gateway-key"), 400)] +#[case::x_api_key("/v1/messages", &[("x-api-key", "gateway-key")], Some("gateway-key"), 400)] +#[case::wrong_x_api_key("/v1/messages", &[("x-api-key", "wrong")], Some("gateway-key"), 401)] +#[case::missing_token("/v1/messages", &[], Some("gateway-key"), 401)] +#[case::wrong_token("/v1/messages", &[("authorization", "Bearer wrong")], Some("gateway-key"), 401)] +#[case::ocr("/ocr", &[], Some("gateway-key"), 401)] +#[case::chat("/v1/chat/completions", &[], Some("gateway-key"), 401)] #[case::deployment( "/openai/deployments/model/chat/completions", - None, + &[], Some("gateway-key"), 401 )] -#[case::transcription("/audio/transcriptions", None, Some("gateway-key"), 401)] -#[case::unsupported_route("/responses", None, Some("gateway-key"), 401)] -#[case::unknown_path("/unknown", None, Some("gateway-key"), 404)] -#[case::unknown_path_unconfigured("/unknown", None, None, 404)] -#[case::unconfigured("/v1/messages", Some("Bearer gateway-key"), None, 500)] +#[case::transcription("/audio/transcriptions", &[], Some("gateway-key"), 401)] +#[case::unsupported_route("/responses", &[], Some("gateway-key"), 401)] +#[case::unknown_path("/unknown", &[], Some("gateway-key"), 404)] +#[case::unknown_path_unconfigured("/unknown", &[], None, 404)] +#[case::unconfigured("/v1/messages", &[("authorization", "Bearer gateway-key")], None, 500)] #[tokio::test] async fn authenticates_before_serving_mounted_inference_routes( inference: Arc, #[case] path: &str, - #[case] authorization: Option<&str>, + #[case] headers: &[(&str, &str)], #[case] master_key: Option<&str>, #[case] status: u16, ) { @@ -87,10 +89,9 @@ async fn authenticates_before_serving_mounted_inference_routes( .timeout(Duration::from_secs(5)) .header("x-request-id", "gateway-test") .json(&json!({"model": "unconfigured-model"})); - let request = match authorization { - Some(value) => request.header("authorization", value), - None => request, - }; + let request = headers.iter().fold(request, |request, (name, value)| { + request.header(*name, *value) + }); let response = request.send().await.unwrap(); assert_eq!(response.status().as_u16(), status); if status == 400 {