From d85f54e726a4918bce9c522b5fb710dc815f0c59 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 26 Oct 2026 08:32:57 -0700 Subject: [PATCH] refactor(gateway-auth): authenticate into a concrete Caller with Axum extractors Replace the pluggable verifier, identity resolver, authorizer and clock with a concrete Auth that turns a presented secret into a Caller. Credential extractors are plain Axum FromRequestParts types chosen per route profile with authenticate::, and Bearer decodes a single typed Authorization header. The gateway refuses general_settings.custom_auth so the Python proxy serves it, and CustomAuth reserves the seam for running it natively later Co-Authored-By: Claude Opus 5.5 --- litellm-rust/Cargo.lock | 45 ++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/gateway-auth/AGENTS.md | 26 +- litellm-rust/crates/gateway-auth/Cargo.toml | 3 +- .../crates/gateway-auth/src/authentication.rs | 237 +-------- .../crates/gateway-auth/src/authorization.rs | 139 ----- .../crates/gateway-auth/src/caller.rs | 52 ++ .../crates/gateway-auth/src/custom.rs | 21 + litellm-rust/crates/gateway-auth/src/error.rs | 7 +- .../crates/gateway-auth/src/extract.rs | 113 ++++ litellm-rust/crates/gateway-auth/src/http.rs | 87 +--- .../crates/gateway-auth/src/identity.rs | 122 ----- litellm-rust/crates/gateway-auth/src/lib.rs | 23 +- litellm-rust/crates/gateway-auth/src/ui.rs | 94 +--- .../crates/gateway-auth/tests/auth.rs | 484 +++--------------- .../src/audio_transcription.rs | 10 +- .../crates/gateway-inference/src/caching.rs | 17 +- .../gateway-inference/src/chat_completions.rs | 16 +- .../crates/gateway-inference/src/messages.rs | 12 +- .../crates/gateway-inference/src/ocr.rs | 10 +- .../crates/gateway-inference/src/request.rs | 10 +- .../crates/gateway-inference/src/responses.rs | 8 +- .../crates/gateway-inference/tests/caching.rs | 57 --- .../gateway-inference/tests/messages.rs | 19 - .../gateway-inference/tests/responses.rs | 19 - .../gateway-inference/tests/support/mod.rs | 82 +-- litellm-rust/crates/gateway-ui/src/session.rs | 16 +- litellm-rust/crates/gateway/src/auth.rs | 26 +- litellm-rust/crates/gateway/src/error.rs | 2 + litellm-rust/crates/gateway/src/lib.rs | 12 +- litellm-rust/crates/gateway/tests/server.rs | 16 + 31 files changed, 452 insertions(+), 1334 deletions(-) delete mode 100644 litellm-rust/crates/gateway-auth/src/authorization.rs create mode 100644 litellm-rust/crates/gateway-auth/src/caller.rs create mode 100644 litellm-rust/crates/gateway-auth/src/custom.rs create mode 100644 litellm-rust/crates/gateway-auth/src/extract.rs delete mode 100644 litellm-rust/crates/gateway-auth/src/identity.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c868d3e4b8a..b66a08f1548 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -774,6 +774,26 @@ dependencies = [ "tower-service", ] +[[package]] +name = "axum-extra" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be44683b41ccb9ab2d23a5230015c9c3c55be97a25e4428366de8873103f7970" +dependencies = [ + "axum-core", + "bytes", + "futures-core", + "futures-util", + "headers", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "mime", + "pin-project-lite", + "tower-layer", + "tower-service", +] + [[package]] name = "axum-login" version = "0.18.0" @@ -2649,6 +2669,30 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "headers" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc2700e3fe555c3310aa7286cac6167449f2c87e0eb58769c9208a1c58a1d106" +dependencies = [ + "base64 0.22.1", + "bytes", + "headers-core", + "http 1.4.2", + "httpdate", + "mime", + "sha1 0.10.7", +] + +[[package]] +name = "headers-core" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4" +dependencies = [ + "http 1.4.2", +] + [[package]] name = "heck" version = "0.5.0" @@ -3764,6 +3808,7 @@ name = "litellm-gateway-auth" version = "0.1.0" dependencies = [ "axum", + "axum-extra", "axum-login", "futures-util", "litellm-auth-types", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 9e5e0a96b73..148c85207aa 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -73,6 +73,7 @@ litellm-python-compat = { path = "crates/python-compat" } askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] } +axum-extra = { version = "0.12.6", default-features = false, features = ["typed-header"] } axum-login = "0.18.0" tower-sessions = { version = "0.14.0", features = ["memory-store"] } bytes = "1" diff --git a/litellm-rust/crates/gateway-auth/AGENTS.md b/litellm-rust/crates/gateway-auth/AGENTS.md index 58457a3151b..0860eaebb49 100644 --- a/litellm-rust/crates/gateway-auth/AGENTS.md +++ b/litellm-rust/crates/gateway-auth/AGENTS.md @@ -1,27 +1,13 @@ # Gateway authentication -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 +`Auth` turns a presented secret into a `Caller`. Today the only caller is `Caller::Master`, verified by a constant-time comparison against the configured master key, which may be an `os.environ/` reference resolved through the injected `SecretSource`. New credential kinds (virtual keys, JWT) become new `Caller` variants and new branches in `Auth::authenticate`, not pluggable verifier traits. Ports exist only at real infrastructure boundaries such as secrets and, later, key storage -`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 +Credential extractors are plain Axum `FromRequestParts` types that convert into `SecretValue`. `Bearer` takes exactly one RFC 6750 Authorization bearer token through typed `headers` decoders, never hand-scanned strings. A route profile picks its extractor with `from_fn_with_state(state, authenticate::)`, and extractor configuration comes from that state through `FromRef` -`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 +The `authenticate` middleware inserts the `Caller` as a request extension and replaces any preexisting one. Handlers extract `Caller`, which fails closed when the middleware did not run, and call `Caller::authorize` with the typed `AccessRequest` before calling a provider or MCP server. `Caller::owner` scopes MCP sessions and response caches -`IdentityResolver` maps verified identity into an internal principal and permissions. The caller retains the original authentication evidence and credential restrictions independently of that mapping. Principals include an authority, subject, and kind. External subjects must remain authority-scoped unless the resolver explicitly maps them to a shared internal identity +Local UI login uses axum-login and tower-sessions. `UiSession` validates the session and its CSRF bearer and never produces a `Caller`, so a UI session cannot authorize inference -`AuthenticatedRequest::authorize` checks expiry and intersects resolved permissions with credential restrictions before calling the injected authorizer. The authorizer may impose additional policy but cannot expand either permission set. The returned operation grant binds the exact authenticated caller instance and access request. Consuming it for another caller or target fails +`general_settings.custom_auth` is served by the Python proxy, and the gateway refuses to start with it. `CustomAuth` is the reserved seam for running it natively. Its outcome maps Python's modes, where `Caller` and `Reject` are the hook's decision, `Fallthrough` is `auto` mode moving on to built-in auth, and `Rekey` is the enterprise hook returning a replacement key for built-in auth. It runs first in `Auth::authenticate` and receives request parts plus the optional extracted key, so a profile with a hook uses `Option` as its extractor. The runtime behind it (Python extension host, Lua, or other) is undecided -`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 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 - -Local UI login continues using axum-login and tower-sessions. Only after session and CSRF validation does `UiSession` expose an authenticated caller. Its credentials are restricted to session-info and logout operations, so the UI CSRF bearer cannot authorize inference. Future UI management routes must extend that explicit scope - -Authentication evidence and principals contain no raw token, password, request body, or mutable accounting state. Session ownership is scoped by principal authority, subject, verifier, and credential ID, so separate credentials do not silently share an MCP session. Credential rotation may retain ownership when the verifier preserves a stable credential ID. Scope and expiry checks still run for each operation - -Failures distinguish invalid or expired credentials, forbidden operations, unavailable authentication services, and missing server configuration/context. HTTP adapters map those outcomes to status codes. Inference keeps its API-specific error envelopes and MCP translates policy failures to its own protocol - -Virtual-key storage, JWT verification, OAuth2 introspection, trusted-proxy validation, SSO, and custom Python hook adapters are not implemented here yet. They should supply these contracts rather than bypassing the shared authorization boundary. Request-body-dependent custom hooks will need a bounded endpoint adapter after parsing - -Budget reservations, rate limits, usage settlement, retries, and cancellation accounting belong after authorization in the execution lifecycle. This foundation adds no accounting backend or cached authorization decisions. Future reservations must bind the same caller and resolved operation and settle idempotently across streaming completion and cancellation +Errors map to HTTP status codes in `Error::status`. Inference keeps its API-specific error envelopes and MCP translates authorization failures to its own protocol diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml index cf040c6eef3..86795b3c6b2 100644 --- a/litellm-rust/crates/gateway-auth/Cargo.toml +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -6,7 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] -axum = { workspace = true, features = ["matched-path"] } +axum.workspace = true +axum-extra.workspace = true axum-login.workspace = true litellm-auth-types.workspace = true litellm-config.workspace = true diff --git a/litellm-rust/crates/gateway-auth/src/authentication.rs b/litellm-rust/crates/gateway-auth/src/authentication.rs index 4ecfaa375b4..86654993759 100644 --- a/litellm-rust/crates/gateway-auth/src/authentication.rs +++ b/litellm-rust/crates/gateway-auth/src/authentication.rs @@ -1,189 +1,42 @@ -use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; +use std::sync::Arc; -use axum::http::{Request, request::Parts}; use litellm_auth_types::SecretValue; use litellm_config::Config; use litellm_secrets::source::SecretSource; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; -use crate::{ - AuthenticatedCaller, AuthenticatedRequest, Authentication, AuthenticationMethod, Authorizer, - Bearer, CredentialExtractor, Error, NoAdditionalPolicy, Permissions, Principal, PrincipalKind, - ResolvedIdentity, VerifiedIdentity, -}; - -pub type AuthFuture<'a, T> = Pin> + Send + 'a>>; - -pub trait Clock: Send + Sync { - fn now(&self) -> SystemTime; -} - -pub struct SystemClock; - -impl Clock for SystemClock { - fn now(&self) -> SystemTime { - SystemTime::now() - } -} - -pub enum Credential { - Token(SecretValue), - Transport, -} - -pub trait Authenticator: Send + Sync { - fn validate(&self) -> AuthFuture<'_, ()> { - Box::pin(async { Ok(()) }) - } - - fn verify<'a>( - &'a self, - credential: &'a Credential, - request: &'a Parts, - ) -> AuthFuture<'a, VerifiedIdentity>; -} - -pub trait IdentityResolver: Send + Sync { - fn resolve<'a>(&'a self, identity: &'a VerifiedIdentity) -> AuthFuture<'a, ResolvedIdentity>; -} - -pub struct LocalAdministrator; - -impl IdentityResolver for LocalAdministrator { - fn resolve<'a>(&'a self, identity: &'a VerifiedIdentity) -> AuthFuture<'a, ResolvedIdentity> { - Box::pin(async move { - let allowed = match identity.authentication.method { - AuthenticationMethod::MasterKey => identity.principal == master_principal(), - AuthenticationMethod::Session => { - identity.principal.authority() == "litellm:local-ui" - && identity.principal.kind() == PrincipalKind::Human - } - AuthenticationMethod::External(_) => false, - }; - if !allowed { - return Err(Error::Forbidden); - } - Ok(ResolvedIdentity { - principal: identity.principal.clone(), - permissions: Permissions::All, - }) - }) - } -} +use crate::{Caller, Error}; #[derive(Clone)] pub struct Auth { - extractor: Arc, - authenticator: Arc, - identities: Arc, - authorizer: Arc, - clock: Arc, -} - -impl Auth { - pub fn new( - authenticator: Arc, - identities: Arc, - authorizer: Arc, - clock: Arc, - ) -> Self { - Self { - extractor: Arc::new(Bearer), - authenticator, - identities, - authorizer, - clock, - } - } - - pub fn from_config(config: &Config, secrets: Arc) -> Self { - let master = Arc::new(MasterKeyAuthenticator { - key: config.general_settings.master_key.clone(), - secrets, - }); - Self { - extractor: Arc::new(Bearer), - authenticator: master, - identities: Arc::new(LocalAdministrator), - authorizer: Arc::new(NoAdditionalPolicy), - clock: Arc::new(SystemClock), - } - } - - pub fn with_extractor(self, extractor: Arc) -> Self { - Self { extractor, ..self } - } - - pub async fn authenticate_parts(&self, request: &Parts) -> Result { - let credential = self.extractor.extract(request)?; - self.authenticate_request(&credential, request).await - } - - pub async fn validate(&self) -> Result<(), Error> { - self.authenticator.validate().await - } - - pub async fn authenticate( - &self, - credential: &SecretValue, - ) -> Result { - self.authenticate_request( - &Credential::Token(credential.clone()), - &Request::new(()).into_parts().0, - ) - .await - } - - pub async fn authenticate_request( - &self, - credential: &Credential, - request: &Parts, - ) -> Result { - let verified = self.authenticator.verify(credential, request).await?; - resolve( - verified, - self.identities.as_ref(), - self.authorizer.clone(), - self.clock.clone(), - ) - .await - } -} - -pub(crate) async fn resolve( - verified: VerifiedIdentity, - identities: &dyn IdentityResolver, - authorizer: Arc, - clock: Arc, -) -> Result { - if verified - .authentication - .expires_at - .is_some_and(|expiry| expiry <= clock.now()) - { - return Err(Error::Expired); - } - let resolved = identities.resolve(&verified).await?; - Ok(AuthenticatedRequest::new( - Arc::new(AuthenticatedCaller::new(verified, resolved)), - authorizer, - clock, - )) -} - -pub struct MasterKeyAuthenticator { - key: Option, + master_key: Option, secrets: Arc, } -impl MasterKeyAuthenticator { - pub fn new(key: Option, secrets: Arc) -> Self { - Self { key, secrets } +impl Auth { + pub fn from_config(config: &Config, secrets: Arc) -> Self { + Self { + master_key: config.general_settings.master_key.clone(), + secrets, + } } - async fn key(&self) -> Result { - let configured = self.key.as_ref().ok_or(Error::Unconfigured)?; + pub async fn validate(&self) -> Result<(), Error> { + self.master_key().await.map(drop) + } + + pub async fn authenticate(&self, token: &SecretValue) -> Result { + let expected = self.master_key().await?; + if bool::from(Sha256::digest(token.expose()).ct_eq(&Sha256::digest(expected.expose()))) { + Ok(Caller::Master) + } else { + Err(Error::InvalidToken) + } + } + + async fn master_key(&self) -> Result { + let configured = self.master_key.as_ref().ok_or(Error::Unconfigured)?; let resolved = match configured.expose().strip_prefix("os.environ/") { Some(name) if !name.is_empty() => self .secrets @@ -199,45 +52,3 @@ impl MasterKeyAuthenticator { Ok(resolved) } } - -impl Authenticator for MasterKeyAuthenticator { - fn validate(&self) -> AuthFuture<'_, ()> { - Box::pin(async { self.key().await.map(|_| ()) }) - } - - fn verify<'a>( - &'a self, - credential: &'a Credential, - _: &'a Parts, - ) -> AuthFuture<'a, VerifiedIdentity> { - Box::pin(async move { - let expected = self.key().await?; - let Credential::Token(credential) = credential else { - return Err(Error::InvalidToken); - }; - if !bool::from( - Sha256::digest(credential.expose()).ct_eq(&Sha256::digest(expected.expose())), - ) { - return Err(Error::InvalidToken); - } - Ok(VerifiedIdentity { - principal: master_principal(), - authentication: Authentication { - method: AuthenticationMethod::MasterKey, - verifier: "litellm:master-key".into(), - credential_id: "master".into(), - expires_at: None, - }, - restrictions: Permissions::All, - }) - }) - } -} - -fn master_principal() -> Principal { - Principal::new( - "litellm:master-key".into(), - "master".into(), - PrincipalKind::System, - ) -} diff --git a/litellm-rust/crates/gateway-auth/src/authorization.rs b/litellm-rust/crates/gateway-auth/src/authorization.rs deleted file mode 100644 index 522fae64f3e..00000000000 --- a/litellm-rust/crates/gateway-auth/src/authorization.rs +++ /dev/null @@ -1,139 +0,0 @@ -use std::sync::Arc; - -use crate::{AuthFuture, AuthenticatedCaller, Clock, Error, SharedCaller}; - -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum McpAction { - Connect, - ListTools, - CallTool(String), - ListPrompts, - GetPrompt(String), - ListResources, - ListResourceTemplates, - ReadResource(String), -} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum UiAction { - SessionInfo, - Logout, -} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum AccessRequest { - Route { method: String, path: String }, - Model { name: String, deployment: String }, - Mcp { server: String, action: McpAction }, - Ui(UiAction), -} - -#[derive(Clone, Debug, Default)] -pub enum Permissions { - All, - #[default] - None, - Only(Arc<[AccessRequest]>), -} - -impl Permissions { - pub fn allows(&self, request: &AccessRequest) -> bool { - match self { - Self::All => true, - Self::None => false, - Self::Only(allowed) => allowed.contains(request), - } - } -} - -pub trait Authorizer: Send + Sync { - fn authorize<'a>( - &'a self, - caller: &'a AuthenticatedCaller, - request: &'a AccessRequest, - ) -> AuthFuture<'a, ()>; -} - -pub struct NoAdditionalPolicy; - -impl Authorizer for NoAdditionalPolicy { - fn authorize<'a>( - &'a self, - _: &'a AuthenticatedCaller, - _: &'a AccessRequest, - ) -> AuthFuture<'a, ()> { - Box::pin(async { Ok(()) }) - } -} - -#[derive(Clone)] -pub struct AuthenticatedRequest { - caller: SharedCaller, - authorizer: Arc, - clock: Arc, -} - -impl AuthenticatedRequest { - pub(crate) fn new( - caller: SharedCaller, - authorizer: Arc, - clock: Arc, - ) -> Self { - Self { - caller, - authorizer, - clock, - } - } - - pub fn caller(&self) -> &AuthenticatedCaller { - &self.caller - } - - pub async fn authorize(&self, request: AccessRequest) -> Result { - if self - .caller - .authentication() - .expires_at - .is_some_and(|expiry| expiry <= self.clock.now()) - { - return Err(Error::Expired); - } - if !self.caller.restrictions().allows(&request) - || !self.caller.permissions().allows(&request) - { - return Err(Error::Forbidden); - } - self.authorizer.authorize(&self.caller, &request).await?; - Ok(AuthorizedOperation { - caller: self.caller.clone(), - request, - }) - } -} - -pub struct AuthorizedOperation { - caller: SharedCaller, - request: AccessRequest, -} - -impl AuthorizedOperation { - pub fn caller(&self) -> &AuthenticatedCaller { - &self.caller - } - pub fn request(&self) -> &AccessRequest { - &self.request - } - - pub fn consume( - self, - caller: &AuthenticatedCaller, - request: &AccessRequest, - ) -> Result<(), Error> { - if std::ptr::eq(self.caller.as_ref(), caller) && &self.request == request { - Ok(()) - } else { - Err(Error::Forbidden) - } - } -} diff --git a/litellm-rust/crates/gateway-auth/src/caller.rs b/litellm-rust/crates/gateway-auth/src/caller.rs new file mode 100644 index 00000000000..2d321ec5925 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/caller.rs @@ -0,0 +1,52 @@ +use axum::{extract::FromRequestParts, http::request::Parts}; + +use crate::Error; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum McpAction { + Connect, + ListTools, + CallTool(String), + ListPrompts, + GetPrompt(String), + ListResources, + ListResourceTemplates, + ReadResource(String), +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum AccessRequest { + Model { name: String, deployment: String }, + Mcp { server: String, action: McpAction }, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum Caller { + Master, +} + +impl Caller { + pub fn authorize(&self, _: &AccessRequest) -> Result<(), Error> { + match self { + Self::Master => Ok(()), + } + } + + pub fn owner(&self) -> String { + match self { + Self::Master => "master".into(), + } + } +} + +impl FromRequestParts for Caller { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, _: &S) -> Result { + parts + .extensions + .get::() + .cloned() + .ok_or(Error::MissingIdentity) + } +} diff --git a/litellm-rust/crates/gateway-auth/src/custom.rs b/litellm-rust/crates/gateway-auth/src/custom.rs new file mode 100644 index 00000000000..1c54a580a03 --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/custom.rs @@ -0,0 +1,21 @@ +use std::{future::Future, pin::Pin}; + +use axum::http::request::Parts; +use litellm_auth_types::SecretValue; + +use crate::{Caller, Error}; + +pub enum CustomOutcome { + Caller(Caller), + Rekey(SecretValue), + Fallthrough, + Reject(Error), +} + +pub trait CustomAuth: Send + Sync { + fn authenticate<'a>( + &'a self, + request: &'a Parts, + key: Option<&'a SecretValue>, + ) -> Pin + Send + 'a>>; +} diff --git a/litellm-rust/crates/gateway-auth/src/error.rs b/litellm-rust/crates/gateway-auth/src/error.rs index 262a666df80..7b528741a06 100644 --- a/litellm-rust/crates/gateway-auth/src/error.rs +++ b/litellm-rust/crates/gateway-auth/src/error.rs @@ -17,12 +17,8 @@ pub enum KeyError { pub enum Error { #[error("operation is not permitted")] Forbidden, - #[error("credential has expired")] - Expired, #[error("authenticated request context is missing")] MissingIdentity, - #[error("authentication service unavailable")] - Unavailable, #[error("gateway auth not configured")] Unconfigured, #[error("missing or invalid bearer token")] @@ -34,9 +30,8 @@ pub enum Error { impl Error { pub fn status(&self) -> StatusCode { match self { - Self::InvalidToken | Self::Expired => StatusCode::UNAUTHORIZED, + Self::InvalidToken => StatusCode::UNAUTHORIZED, Self::Forbidden => StatusCode::FORBIDDEN, - Self::Unavailable => StatusCode::SERVICE_UNAVAILABLE, Self::MissingIdentity | Self::Unconfigured | Self::Secret(_) => { StatusCode::INTERNAL_SERVER_ERROR } diff --git a/litellm-rust/crates/gateway-auth/src/extract.rs b/litellm-rust/crates/gateway-auth/src/extract.rs new file mode 100644 index 00000000000..74b1886a5fb --- /dev/null +++ b/litellm-rust/crates/gateway-auth/src/extract.rs @@ -0,0 +1,113 @@ +use std::iter; + +use axum::{ + extract::FromRequestParts, + http::{HeaderName, HeaderValue, request::Parts}, +}; +use axum_extra::{ + TypedHeader, + headers::{ + self, Header, + authorization::{Authorization, Credentials}, + }, +}; +use litellm_auth_types::SecretValue; + +use crate::Error; + +pub struct Bearer(pub SecretValue); + +impl From for SecretValue { + fn from(value: Bearer) -> Self { + value.0 + } +} + +impl FromRequestParts for Bearer { + type Rejection = Error; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + let TypedHeader(Single(Authorization(BearerToken(token)))) = + TypedHeader::from_request_parts(parts, state) + .await + .map_err(|_| Error::InvalidToken)?; + Ok(Self(token)) + } +} + +struct BearerToken(SecretValue); + +impl Credentials for BearerToken { + const SCHEME: &'static str = "Bearer"; + + fn decode(value: &HeaderValue) -> Option { + let token = value.to_str().ok()?.get(Self::SCHEME.len()..)?.trim(); + is_token(token).then(|| Self(SecretValue::new(token))) + } + + fn encode(&self) -> HeaderValue { + HeaderValue::from_str(&format!("{} {}", Self::SCHEME, self.0.expose())) + .expect("token was decoded from a valid header value") + } +} + +struct Single(H); + +impl Header for Single { + fn name() -> &'static HeaderName { + H::name() + } + + fn decode<'i, I: Iterator>( + values: &mut I, + ) -> Result { + match (values.next(), values.next()) { + (Some(value), None) => H::decode(&mut iter::once(value)).map(Self), + _ => Err(headers::Error::invalid()), + } + } + + fn encode>(&self, values: &mut E) { + self.0.encode(values); + } +} + +fn is_token(token: &str) -> bool { + !token.is_empty() && !token.contains(char::is_whitespace) +} + +#[cfg(test)] +mod tests { + use axum::http::Request; + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::standard(&["Bearer token"], Some("token"))] + #[case::scheme_is_case_insensitive(&["bearer token"], Some("token"))] + #[case::surrounding_space_is_trimmed(&["Bearer token "], Some("token"))] + #[case::empty(&["Bearer "], None)] + #[case::inner_whitespace(&["Bearer a b"], None)] + #[case::other_scheme(&["Basic token"], None)] + #[case::duplicate(&["Bearer token", "Bearer token"], None)] + #[tokio::test] + async fn bearer_accepts_one_rfc_6750_token( + #[case] headers: &[&str], + #[case] expected: Option<&str>, + ) { + let request = headers + .iter() + .fold(Request::builder(), |request, value| { + request.header("authorization", *value) + }) + .body(()) + .unwrap(); + let (mut parts, ()) = request.into_parts(); + let token = Bearer::from_request_parts(&mut parts, &()) + .await + .ok() + .map(|Bearer(token)| token.expose().to_owned()); + assert_eq!(token.as_deref(), expected); + } +} diff --git a/litellm-rust/crates/gateway-auth/src/http.rs b/litellm-rust/crates/gateway-auth/src/http.rs index ddbf240509a..aab6f9361a9 100644 --- a/litellm-rust/crates/gateway-auth/src/http.rs +++ b/litellm-rust/crates/gateway-auth/src/http.rs @@ -1,90 +1,19 @@ use axum::{ - extract::{FromRequestParts, MatchedPath, Request, State}, - http::{header::AUTHORIZATION, request::Parts}, + extract::{Request, State}, middleware::Next, response::Response, }; use litellm_auth_types::SecretValue; -use crate::{AccessRequest, Auth, AuthenticatedRequest, Credential, Error}; +use crate::{Auth, Error}; -pub trait CredentialExtractor: Send + Sync { - fn extract(&self, parts: &Parts) -> Result; -} - -pub struct Bearer; - -impl CredentialExtractor for Bearer { - fn extract(&self, parts: &Parts) -> Result { - bearer(parts).map(Credential::Token) - } -} - -fn bearer(parts: &Parts) -> Result { - let mut values = parts.headers.get_all(AUTHORIZATION).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) - .filter(|value| !value.is_empty() && !value.bytes().any(|byte| byte.is_ascii_whitespace())) - .ok_or(Error::InvalidToken)?; - Ok(SecretValue::new(token)) -} - -pub async fn authenticate( +pub async fn authenticate>( State(auth): State, - request: Request, + credential: E, + mut request: Request, next: Next, ) -> Result { - let (mut parts, body) = request.into_parts(); - let identity = auth.authenticate_parts(&parts).await?; - let route = AccessRequest::Route { - method: parts.method.to_string(), - path: parts - .extensions - .get::() - .map_or(parts.uri.path(), MatchedPath::as_str) - .into(), - }; - identity - .authorize(route.clone()) - .await? - .consume(identity.caller(), &route)?; - parts.extensions.insert(identity); - Ok(next.run(Request::from_parts(parts, body)).await) -} - -impl FromRequestParts for AuthenticatedRequest { - type Rejection = Error; - - async fn from_request_parts(parts: &mut Parts, _: &S) -> Result { - parts - .extensions - .get::() - .cloned() - .ok_or(Error::MissingIdentity) - } -} - -pub struct RequireMasterKey; - -impl FromRequestParts for RequireMasterKey { - type Rejection = Error; - - async fn from_request_parts(parts: &mut Parts, state: &Auth) -> Result { - state.validate().await?; - let identity = state - .authenticate_request(&Credential::Token(bearer(parts)?), parts) - .await?; - if identity.caller().authentication().method != crate::AuthenticationMethod::MasterKey { - return Err(Error::InvalidToken); - } - parts.extensions.insert(identity); - Ok(Self) - } + let caller = auth.authenticate(&credential.into()).await?; + request.extensions_mut().insert(caller); + Ok(next.run(request).await) } diff --git a/litellm-rust/crates/gateway-auth/src/identity.rs b/litellm-rust/crates/gateway-auth/src/identity.rs deleted file mode 100644 index 286983ae647..00000000000 --- a/litellm-rust/crates/gateway-auth/src/identity.rs +++ /dev/null @@ -1,122 +0,0 @@ -use std::{sync::Arc, time::SystemTime}; - -use sha2::{Digest, Sha256}; - -use crate::authorization::Permissions; - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum PrincipalKind { - Human, - Service, - System, -} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub struct Principal { - authority: String, - subject: String, - kind: PrincipalKind, -} - -impl Principal { - pub fn new(authority: String, subject: String, kind: PrincipalKind) -> Self { - Self { - authority, - subject, - kind, - } - } - - pub fn authority(&self) -> &str { - &self.authority - } - pub fn subject(&self) -> &str { - &self.subject - } - pub fn kind(&self) -> PrincipalKind { - self.kind - } -} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub enum AuthenticationMethod { - MasterKey, - Session, - External(String), -} - -#[derive(Clone, Debug)] -pub struct Authentication { - pub method: AuthenticationMethod, - pub verifier: String, - pub credential_id: String, - pub expires_at: Option, -} - -pub struct VerifiedIdentity { - pub principal: Principal, - pub authentication: Authentication, - pub restrictions: Permissions, -} - -pub struct ResolvedIdentity { - pub principal: Principal, - pub permissions: Permissions, -} - -#[derive(Clone, Debug)] -pub struct AuthenticatedCaller { - verified_principal: Principal, - principal: Principal, - authentication: Authentication, - permissions: Permissions, - restrictions: Permissions, -} - -impl AuthenticatedCaller { - pub(crate) fn new(verified: VerifiedIdentity, resolved: ResolvedIdentity) -> Self { - Self { - verified_principal: verified.principal, - principal: resolved.principal, - authentication: verified.authentication, - permissions: resolved.permissions, - restrictions: verified.restrictions, - } - } - - pub fn verified_principal(&self) -> &Principal { - &self.verified_principal - } - - pub fn principal(&self) -> &Principal { - &self.principal - } - pub fn authentication(&self) -> &Authentication { - &self.authentication - } - pub fn permissions(&self) -> &Permissions { - &self.permissions - } - pub fn restrictions(&self) -> &Permissions { - &self.restrictions - } - - pub fn session_owner(&self) -> String { - let fields = [ - self.verified_principal.authority(), - self.verified_principal.subject(), - self.principal.authority(), - self.principal.subject(), - self.authentication.verifier.as_str(), - self.authentication.credential_id.as_str(), - ]; - let digest = fields.into_iter().fold(Sha256::new(), |mut digest, field| { - digest.update((field.len() as u64).to_be_bytes()); - digest.update(field.as_bytes()); - digest - }); - format!("{:x}", digest.finalize()) - } -} - -pub type SharedCaller = Arc; diff --git a/litellm-rust/crates/gateway-auth/src/lib.rs b/litellm-rust/crates/gateway-auth/src/lib.rs index 1b400a94746..5e8a576cef7 100644 --- a/litellm-rust/crates/gateway-auth/src/lib.rs +++ b/litellm-rust/crates/gateway-auth/src/lib.rs @@ -1,27 +1,20 @@ mod authentication; -mod authorization; +mod caller; +mod custom; mod error; +pub mod extract; mod http; -mod identity; pub mod keys; mod ui; use sha2::{Digest, Sha256}; -pub use authentication::{ - Auth, AuthFuture, Authenticator, Clock, Credential, IdentityResolver, LocalAdministrator, - MasterKeyAuthenticator, SystemClock, -}; -pub use authorization::{ - AccessRequest, AuthenticatedRequest, AuthorizedOperation, Authorizer, McpAction, - NoAdditionalPolicy, Permissions, UiAction, -}; +pub use authentication::Auth; +pub use caller::{AccessRequest, Caller, McpAction}; +pub use custom::{CustomAuth, CustomOutcome}; pub use error::{Error, KeyError, UiAuthError}; -pub use http::{Bearer, CredentialExtractor, RequireMasterKey, authenticate}; -pub use identity::{ - AuthenticatedCaller, Authentication, AuthenticationMethod, Principal, PrincipalKind, - ResolvedIdentity, SharedCaller, VerifiedIdentity, -}; +pub use extract::Bearer; +pub use http::authenticate; pub use ui::{UI_CSRF_KEY, UiAuthSession, UiBackend, UiCredentials, UiSession, UiUser}; pub fn hash_token(token: &str) -> String { diff --git a/litellm-rust/crates/gateway-auth/src/ui.rs b/litellm-rust/crates/gateway-auth/src/ui.rs index 9e6b77a17d0..476ddf65d6b 100644 --- a/litellm-rust/crates/gateway-auth/src/ui.rs +++ b/litellm-rust/crates/gateway-auth/src/ui.rs @@ -1,7 +1,6 @@ -use std::{convert::Infallible, sync::Arc}; +use std::convert::Infallible; -use axum::extract::FromRequestParts; -use axum::http::{header::AUTHORIZATION, request::Parts}; +use axum::{extract::FromRequestParts, http::request::Parts}; use axum_login::{AuthUser, AuthnBackend, UserId}; use litellm_auth_types::SecretValue; use serde::Deserialize; @@ -10,11 +9,7 @@ use subtle::ConstantTimeEq; use tower_sessions::Session; use veil::Redact; -use crate::{ - AccessRequest, AuthenticatedRequest, Authentication, AuthenticationMethod, LocalAdministrator, - NoAdditionalPolicy, Permissions, Principal, PrincipalKind, SystemClock, UiAction, UiAuthError, - VerifiedIdentity, -}; +use crate::{Bearer, UiAuthError}; pub const UI_CSRF_KEY: &str = "litellm.ui.csrf"; pub type UiAuthSession = axum_login::AuthSession; @@ -85,7 +80,6 @@ impl AuthnBackend for UiBackend { pub struct UiSession { pub user: UiUser, - pub identity: AuthenticatedRequest, } impl FromRequestParts for UiSession { @@ -103,84 +97,12 @@ impl FromRequestParts for UiSession { .get::(UI_CSRF_KEY) .await? .ok_or(UiAuthError::Unauthorized)?; - let provided = parts - .headers - .get(AUTHORIZATION) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.strip_prefix("Bearer ")) - .ok_or(UiAuthError::Unauthorized)?; - if !bool::from(Sha256::digest(expected).ct_eq(&Sha256::digest(provided))) { + let Bearer(provided) = Bearer::from_request_parts(parts, state) + .await + .map_err(|_| UiAuthError::Unauthorized)?; + if !bool::from(Sha256::digest(expected).ct_eq(&Sha256::digest(provided.expose()))) { return Err(UiAuthError::Unauthorized); } - let identity = session_identity(&user).await?; - parts.extensions.insert(identity.clone()); - Ok(Self { user, identity }) - } -} - -async fn session_identity(user: &UiUser) -> Result { - crate::authentication::resolve( - VerifiedIdentity { - principal: Principal::new( - "litellm:local-ui".into(), - user.username.clone(), - PrincipalKind::Human, - ), - authentication: Authentication { - method: AuthenticationMethod::Session, - verifier: "litellm:local-ui".into(), - credential_id: user.username.clone(), - expires_at: None, - }, - restrictions: Permissions::Only(Arc::from([ - AccessRequest::Ui(UiAction::SessionInfo), - AccessRequest::Ui(UiAction::Logout), - ])), - }, - &LocalAdministrator, - Arc::new(NoAdditionalPolicy), - Arc::new(SystemClock), - ) - .await - .map_err(|_| UiAuthError::Unauthorized) -} - -#[cfg(test)] -mod tests { - use super::*; - use rstest::rstest; - - #[rstest] - #[tokio::test] - async fn ui_identity_is_scoped_to_session_operations() { - let backend = UiBackend::new("admin".into(), SecretValue::new("password")).unwrap(); - let user = backend.get_user(&"admin".into()).await.unwrap().unwrap(); - let identity = session_identity(&user).await.unwrap(); - assert_eq!(identity.caller().principal().subject(), user.username); - assert_eq!( - identity.caller().authentication().method, - AuthenticationMethod::Session - ); - assert!( - identity - .authorize(AccessRequest::Ui(UiAction::SessionInfo)) - .await - .is_ok() - ); - assert!( - identity - .authorize(AccessRequest::Ui(UiAction::Logout)) - .await - .is_ok() - ); - assert!(matches!( - identity - .authorize(AccessRequest::Model { - name: "model".into(), - deployment: "provider/model".into(), - }) - .await, - Err(crate::Error::Forbidden) - )); + Ok(Self { user }) } } diff --git a/litellm-rust/crates/gateway-auth/tests/auth.rs b/litellm-rust/crates/gateway-auth/tests/auth.rs index 5a2a3ded834..e3638a45d7a 100644 --- a/litellm-rust/crates/gateway-auth/tests/auth.rs +++ b/litellm-rust/crates/gateway-auth/tests/auth.rs @@ -4,13 +4,13 @@ use axum::{ Router, body::{Body, to_bytes}, http::{Request, StatusCode}, - middleware::from_extractor_with_state, + middleware::from_fn_with_state, routing::get, }; use futures_util::future::BoxFuture; use litellm_auth_types::SecretValue; use litellm_config::Config; -use litellm_gateway_auth::{Auth, RequireMasterKey, hash_token}; +use litellm_gateway_auth::{Auth, Bearer, Caller, Error, authenticate, hash_token}; use litellm_secrets::source::SecretSource; use rstest::{fixture, rstest}; use tower::ServiceExt; @@ -53,9 +53,9 @@ fn secrets() -> Arc { #[case::wrong_scheme("literal-key", Some("Basic literal-key"), 401)] #[case::empty_token("literal-key", Some("Bearer "), 401)] #[case::missing_reference("os.environ/MISSING", Some("Bearer os.environ/MISSING"), 500)] -#[case::empty_reference("os.environ/EMPTY", Some("Bearer "), 500)] -#[case::empty_key("", Some("Bearer "), 500)] -#[case::whitespace_key(" ", Some("Bearer "), 500)] +#[case::empty_reference("os.environ/EMPTY", Some("Bearer token"), 500)] +#[case::empty_key("", Some("Bearer token"), 500)] +#[case::whitespace_key(" ", Some("Bearer token"), 500)] #[case::empty_reference_name("os.environ/", Some("Bearer os.environ/"), 500)] #[case::secret_failure("os.environ/ERROR", Some("Bearer private-backend-detail"), 500)] #[tokio::test] @@ -71,8 +71,9 @@ async fn enforces_configured_keys_without_exposing_secrets( .unwrap(); let app = Router::new() .route("/protected", get(|| async { StatusCode::NO_CONTENT })) - .layer(from_extractor_with_state::( + .layer(from_fn_with_state( Auth::from_config(&config, secrets), + authenticate::, )); let request = Request::get("/protected"); let request = match authorization { @@ -99,368 +100,109 @@ fn hash_token_matches_python_sha256_hexdigest() { ); } -use litellm_gateway_auth::{ - AccessRequest, AuthFuture, AuthenticatedCaller, AuthenticatedRequest, Authentication, - AuthenticationMethod, Authenticator, Authorizer, Clock, Error, IdentityResolver, Permissions, - Principal, PrincipalKind, ResolvedIdentity, VerifiedIdentity, authenticate, -}; -use std::{ - sync::atomic::{AtomicU64, AtomicUsize, Ordering}, - time::{Duration, SystemTime}, -}; - -struct TestClock(AtomicU64); - -impl Clock for TestClock { - fn now(&self) -> SystemTime { - SystemTime::UNIX_EPOCH + Duration::from_secs(self.0.load(Ordering::SeqCst)) - } -} - -struct Verifier { - authority: &'static str, - restrictions: Permissions, - failure: Option, -} - -impl Authenticator for Verifier { - fn verify<'a>( - &'a self, - credential: &'a litellm_gateway_auth::Credential, - _: &'a axum::http::request::Parts, - ) -> AuthFuture<'a, VerifiedIdentity> { - Box::pin(async move { - match self.failure { - Some(401) => return Err(Error::InvalidToken), - Some(_) => return Err(Error::Unavailable), - None => (), - } - let litellm_gateway_auth::Credential::Token(credential) = credential else { - return Err(Error::InvalidToken); - }; - if credential.expose() != "test-token" { - return Err(Error::InvalidToken); - } - Ok(VerifiedIdentity { - principal: Principal::new( - self.authority.into(), - "subject".into(), - PrincipalKind::Service, - ), - authentication: Authentication { - method: AuthenticationMethod::External("test".into()), - verifier: "test-verifier".into(), - credential_id: "credential".into(), - expires_at: Some(SystemTime::UNIX_EPOCH + Duration::from_secs(100)), - }, - restrictions: self.restrictions.clone(), - }) - }) - } -} - -struct Identities { - mapped_principal: Option, - calls: AtomicUsize, - permissions: Permissions, -} - -impl IdentityResolver for Identities { - fn resolve<'a>(&'a self, identity: &'a VerifiedIdentity) -> AuthFuture<'a, ResolvedIdentity> { - Box::pin(async move { - self.calls.fetch_add(1, Ordering::SeqCst); - Ok(ResolvedIdentity { - principal: self - .mapped_principal - .clone() - .unwrap_or_else(|| identity.principal.clone()), - permissions: self.permissions.clone(), - }) - }) - } -} - -struct Policy(AtomicUsize); - -impl Authorizer for Policy { - fn authorize<'a>( - &'a self, - _: &'a AuthenticatedCaller, - _: &'a AccessRequest, - ) -> AuthFuture<'a, ()> { - Box::pin(async move { - self.0.fetch_add(1, Ordering::SeqCst); - Ok(()) - }) - } -} - -struct Pipeline { - auth: Auth, - identities: Arc, - policy: Arc, - clock: Arc, -} - -fn pipeline( - authority: &'static str, - permissions: Permissions, - restrictions: Permissions, - failure: Option, -) -> Pipeline { - let identities = Arc::new(Identities { - mapped_principal: None, - calls: AtomicUsize::new(0), - permissions, - }); - let policy = Arc::new(Policy(AtomicUsize::new(0))); - let clock = Arc::new(TestClock(AtomicU64::new(10))); - let auth = Auth::new( - Arc::new(Verifier { - authority, - restrictions, - failure, - }), - identities.clone(), - policy.clone(), - clock.clone(), - ); - Pipeline { - auth, - identities, - policy, - clock, - } -} - -fn model(name: &str) -> AccessRequest { - AccessRequest::Model { - name: name.into(), - deployment: name.into(), - } +fn master_auth(secrets: Arc, key: &str) -> Auth { + let config = Config::from_yaml(&format!( + "model_list: []\ngeneral_settings:\n master_key: {key}\n" + )) + .unwrap(); + Auth::from_config(&config, secrets) } #[rstest] -#[case::invalid_token(401)] -#[case::unavailable_verifier(503)] +#[case::valid(Some("Bearer master-key"), None, 204)] +#[case::duplicate(Some("Bearer master-key"), Some("Bearer master-key"), 401)] +#[case::forged_extension_without_credentials(None, None, 401)] #[tokio::test] -async fn failed_verification_never_resolves_identity_or_runs_the_handler(#[case] status: u16) { - let pipeline = pipeline("issuer-a", Permissions::All, Permissions::All, Some(status)); +async fn middleware_only_trusts_the_presented_credential( + secrets: Arc, + #[case] first: Option<&str>, + #[case] second: Option<&str>, + #[case] status: u16, +) { let app = Router::new() - .route("/protected", get(|| async { StatusCode::NO_CONTENT })) - .layer(axum::middleware::from_fn_with_state( - pipeline.auth, - authenticate, + .route( + "/protected", + get(|caller: Caller| async move { + assert_eq!(caller, Caller::Master); + StatusCode::NO_CONTENT + }), + ) + .layer(from_fn_with_state( + master_auth(secrets, "master-key"), + authenticate::, )); + let request = [first, second] + .into_iter() + .flatten() + .fold(Request::get("/protected"), |request, value| { + request.header("authorization", value) + }) + .extension(Caller::Master); + let response = app + .oneshot(request.body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status().as_u16(), status); +} + +#[rstest] +#[tokio::test] +async fn handlers_fail_closed_without_the_middleware() { + let app = Router::new().route( + "/protected", + get(|_: Caller| async { StatusCode::NO_CONTENT }), + ); let response = app .oneshot( Request::get("/protected") - .header("authorization", "Bearer test-token") + .header("authorization", "Bearer master-key") .body(Body::empty()) .unwrap(), ) .await .unwrap(); - assert_eq!(response.status().as_u16(), status); - assert_eq!(pipeline.identities.calls.load(Ordering::SeqCst), 0); - assert_eq!(pipeline.policy.0.load(Ordering::SeqCst), 0); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); } -#[rstest] -#[case::membership_restricts(Permissions::Only(Arc::from([model("allowed")])), Permissions::All)] -#[case::credential_restricts(Permissions::All, Permissions::Only(Arc::from([model("allowed")])))] -#[tokio::test] -async fn policy_cannot_expand_membership_or_credential_permissions( - #[case] permissions: Permissions, - #[case] restrictions: Permissions, -) { - let pipeline = pipeline("issuer-a", permissions, restrictions, None); - let identity = pipeline - .auth - .authenticate(&SecretValue::new("test-token")) - .await - .unwrap(); - assert!(matches!( - identity.authorize(model("denied")).await, - Err(Error::Forbidden) - )); - assert_eq!(pipeline.policy.0.load(Ordering::SeqCst), 0); - let allowed = model("allowed"); - identity - .authorize(allowed.clone()) - .await - .unwrap() - .consume(identity.caller(), &allowed) - .unwrap(); - assert_eq!(pipeline.policy.0.load(Ordering::SeqCst), 1); -} +struct CustomHeader(SecretValue); -#[rstest] -#[tokio::test] -async fn authorization_grants_bind_the_exact_caller_and_target() { - let first = pipeline("issuer-a", Permissions::All, Permissions::All, None); - let second = pipeline("issuer-b", Permissions::All, Permissions::All, None); - let caller = first - .auth - .authenticate(&SecretValue::new("test-token")) - .await - .unwrap(); - let other = second - .auth - .authenticate(&SecretValue::new("test-token")) - .await - .unwrap(); - assert_ne!(caller.caller().principal(), other.caller().principal()); - assert_ne!( - caller.caller().session_owner(), - other.caller().session_owner() - ); - let access = model("allowed"); - assert!(matches!( - caller - .authorize(access.clone()) - .await - .unwrap() - .consume(other.caller(), &access), - Err(Error::Forbidden) - )); - assert!(matches!( - caller - .authorize(access) - .await - .unwrap() - .consume(caller.caller(), &model("different")), - Err(Error::Forbidden) - )); -} +impl axum::extract::FromRequestParts for CustomHeader { + type Rejection = Error; -#[rstest] -#[tokio::test] -async fn expiry_is_checked_again_for_operations_on_existing_connections() { - let pipeline = pipeline("issuer-a", Permissions::All, Permissions::All, None); - let identity = pipeline - .auth - .authenticate(&SecretValue::new("test-token")) - .await - .unwrap(); - assert!(identity.authorize(model("allowed")).await.is_ok()); - pipeline.clock.0.store(100, Ordering::SeqCst); - assert!(matches!( - identity.authorize(model("allowed")).await, - Err(Error::Expired) - )); - assert!(matches!( - pipeline - .auth - .authenticate(&SecretValue::new("test-token")) - .await, - Err(Error::Expired) - )); - assert_eq!(pipeline.identities.calls.load(Ordering::SeqCst), 1); - assert_eq!(pipeline.policy.0.load(Ordering::SeqCst), 1); -} - -#[rstest] -#[tokio::test] -async fn bearer_extraction_rejects_duplicates_and_does_not_trust_preexisting_identity() { - let pipeline = pipeline("issuer-a", Permissions::All, Permissions::All, None); - let identity = pipeline - .auth - .authenticate(&SecretValue::new("test-token")) - .await - .unwrap(); - let app = Router::new() - .route( - "/protected", - get(|_: AuthenticatedRequest| async { StatusCode::NO_CONTENT }), - ) - .layer(axum::middleware::from_fn_with_state( - pipeline.auth, - authenticate, - )); - let duplicate = Request::get("/protected") - .header("authorization", "Bearer test-token") - .header("authorization", "Bearer test-token") - .body(Body::empty()) - .unwrap(); - assert_eq!(app.clone().oneshot(duplicate).await.unwrap().status(), 401); - let forged = Request::get("/protected") - .extension(identity) - .body(Body::empty()) - .unwrap(); - assert_eq!(app.clone().oneshot(forged).await.unwrap().status(), 401); - let valid = Request::get("/protected") - .header("authorization", "Bearer test-token") - .body(Body::empty()) - .unwrap(); - assert_eq!(app.oneshot(valid).await.unwrap().status(), 204); -} - -#[rstest] -#[tokio::test] -async fn route_permissions_are_checked_before_parsing_or_dispatch() { - let allowed = AccessRequest::Route { - method: "GET".into(), - path: "/allowed/{id}".into(), - }; - let pipeline = pipeline( - "issuer-a", - Permissions::Only(Arc::from([allowed])), - Permissions::All, - None, - ); - let app = Router::new() - .route("/allowed/{id}", get(|| async { StatusCode::NO_CONTENT })) - .route("/denied", get(|| async { StatusCode::NO_CONTENT })) - .layer(axum::middleware::from_fn_with_state( - pipeline.auth, - authenticate, - )); - let allowed = Request::get("/allowed/123") - .header("authorization", "Bearer test-token") - .body(Body::empty()) - .unwrap(); - let denied = Request::get("/denied") - .header("authorization", "Bearer test-token") - .body(Body::empty()) - .unwrap(); - assert_eq!(app.clone().oneshot(allowed).await.unwrap().status(), 204); - assert_eq!(app.oneshot(denied).await.unwrap().status(), 403); -} - -struct CustomHeader; - -impl litellm_gateway_auth::CredentialExtractor for CustomHeader { - fn extract( - &self, - parts: &axum::http::request::Parts, - ) -> Result { + async fn from_request_parts( + parts: &mut axum::http::request::Parts, + _: &Auth, + ) -> Result { let value = parts .headers .get("x-test-key") .and_then(|value| value.to_str().ok()) .ok_or(Error::InvalidToken)?; - Ok(litellm_gateway_auth::Credential::Token(SecretValue::new( - value, - ))) + Ok(Self(SecretValue::new(value))) + } +} + +impl From for SecretValue { + fn from(value: CustomHeader) -> Self { + value.0 } } #[rstest] -#[case::selected_header("test-token", "wrong-token", 204)] -#[case::no_fallback("wrong-token", "test-token", 401)] +#[case::selected_header("master-key", "wrong-key", 204)] +#[case::no_fallback("wrong-key", "master-key", 401)] #[tokio::test] async fn route_profile_uses_only_its_selected_credentials( + secrets: Arc, #[case] custom: &str, #[case] bearer: &str, #[case] status: u16, ) { - let pipeline = pipeline("issuer-a", Permissions::All, Permissions::All, None); let app = Router::new() .route("/protected", get(|| async { StatusCode::NO_CONTENT })) - .layer(axum::middleware::from_fn_with_state( - pipeline.auth.with_extractor(Arc::new(CustomHeader)), - authenticate, + .layer(from_fn_with_state( + master_auth(secrets, "master-key"), + authenticate::, )); let response = app .oneshot( @@ -477,79 +219,17 @@ async fn route_profile_uses_only_its_selected_credentials( #[rstest] #[tokio::test] -async fn master_rotation_preserves_session_owner_but_rejects_the_old_key( - secrets: Arc, -) { - let first = - Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: first-key\n").unwrap(); - let second = - Config::from_yaml("model_list: []\ngeneral_settings:\n master_key: second-key\n").unwrap(); - let original = Auth::from_config(&first, secrets.clone()) - .authenticate(&SecretValue::new("first-key")) - .await - .unwrap(); - let rotated = Auth::from_config(&second, secrets); - let replacement = rotated - .authenticate(&SecretValue::new("second-key")) - .await - .unwrap(); +async fn master_rotation_rejects_the_old_key(secrets: Arc) { + let rotated = master_auth(secrets, "second-key"); assert_eq!( - original.caller().session_owner(), - replacement.caller().session_owner() + rotated + .authenticate(&SecretValue::new("second-key")) + .await + .unwrap(), + Caller::Master ); assert!(matches!( rotated.authenticate(&SecretValue::new("first-key")).await, Err(Error::InvalidToken) )); } - -#[rstest] -#[tokio::test] -async fn mapping_two_issuers_to_one_account_does_not_merge_session_ownership() { - let resolver = Arc::new(Identities { - mapped_principal: Some(Principal::new( - "internal".into(), - "shared-account".into(), - PrincipalKind::Human, - )), - calls: AtomicUsize::new(0), - permissions: Permissions::All, - }); - let first = Auth::new( - Arc::new(Verifier { - authority: "issuer-a", - restrictions: Permissions::All, - failure: None, - }), - resolver.clone(), - Arc::new(Policy(AtomicUsize::new(0))), - Arc::new(TestClock(AtomicU64::new(10))), - ); - let second = Auth::new( - Arc::new(Verifier { - authority: "issuer-b", - restrictions: Permissions::All, - failure: None, - }), - resolver, - Arc::new(Policy(AtomicUsize::new(0))), - Arc::new(TestClock(AtomicU64::new(10))), - ); - let original = first - .authenticate(&SecretValue::new("test-token")) - .await - .unwrap(); - let mapped = second - .authenticate(&SecretValue::new("test-token")) - .await - .unwrap(); - assert_eq!(original.caller().principal(), mapped.caller().principal()); - assert_ne!( - original.caller().verified_principal(), - mapped.caller().verified_principal() - ); - assert_ne!( - original.caller().session_owner(), - mapped.caller().session_owner() - ); -} diff --git a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs index 8d571eccc5b..762abcbfaa5 100644 --- a/litellm-rust/crates/gateway-inference/src/audio_transcription.rs +++ b/litellm-rust/crates/gateway-inference/src/audio_transcription.rs @@ -1,4 +1,4 @@ -use litellm_gateway_auth::AuthenticatedRequest; +use litellm_gateway_auth::Caller; use std::{path::Path, sync::Arc}; use axum::{Json, extract::State, response::IntoResponse}; @@ -13,22 +13,22 @@ use crate::{ pub(crate) async fn create( State(gateway): State>, - identity: AuthenticatedRequest, + caller: Caller, body: InferenceBody, ) -> Result { - handle(&gateway, &identity, body).await.map(Json) + handle(&gateway, &caller, body).await.map(Json) } async fn handle( gateway: &Gateway, - identity: &AuthenticatedRequest, + caller: &Caller, InferenceBody { fields: body, upload, }: InferenceBody, ) -> Result { let deployment = request::resolve_deployment(gateway, &body)?; - request::authorize_model(identity, deployment, &body).await?; + request::authorize_model(caller, deployment, &body)?; let audio = match upload { Some(upload) => { let format = upload diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index 426b8b91182..27df50e1410 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -1,7 +1,7 @@ use std::time::Duration; use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope}; -use litellm_gateway_auth::AuthenticatedRequest; +use litellm_gateway_auth::Caller; use serde::Deserialize; use serde_json::{Map, Value}; @@ -21,10 +21,7 @@ struct Controls { type Prepared = (Map, CacheOptions); -pub(crate) fn prepare( - identity: &AuthenticatedRequest, - body: Map, -) -> Result { +pub(crate) fn prepare(caller: &Caller, body: Map) -> Result { let controls: Controls = match body.get("cache").filter(|value| !value.is_null()) { Some(value) => serde_json::from_value(value.clone()) .map_err(|error| Error::InvalidBody(error.to_string()))?, @@ -36,7 +33,6 @@ pub(crate) fn prepare( .map(|value| serde_json::from_value(value.clone())) .transpose() .map_err(|error| Error::InvalidBody(error.to_string()))?; - let caller = identity.caller(); let options = CacheOptions { policy: CachePolicy { caching, @@ -45,14 +41,7 @@ pub(crate) fn prepare( ttl: controls.ttl.map(duration).transpose()?, max_age: controls.max_age.map(duration).transpose()?, }, - scope: CacheScope::Isolated( - serde_json::json!([ - caller.principal().authority(), - caller.principal().subject(), - caller.authentication().credential_id - ]) - .to_string(), - ), + scope: CacheScope::Isolated(caller.owner()), }; Ok(( body.into_iter() diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 742a373b286..ab981f51f7d 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -1,4 +1,4 @@ -use litellm_gateway_auth::AuthenticatedRequest; +use litellm_gateway_auth::Caller; use std::sync::Arc; use axum::{ @@ -13,15 +13,15 @@ use crate::{Error, Gateway, JsonObject, request}; pub(crate) async fn create( State(gateway): State>, - identity: AuthenticatedRequest, + caller: Caller, JsonObject(body): JsonObject, ) -> Result { - handle(&gateway, &identity, body).await + handle(&gateway, &caller, body).await } pub(crate) async fn create_from_model_path( State(gateway): State>, - identity: AuthenticatedRequest, + caller: Caller, Path(model): Path, JsonObject(body): JsonObject, ) -> Result { @@ -32,17 +32,17 @@ pub(crate) async fn create_from_model_path( .collect(), Some(_) => body, }; - handle(&gateway, &identity, body).await + handle(&gateway, &caller, body).await } async fn handle( gateway: &Gateway, - identity: &AuthenticatedRequest, + caller: &Caller, body: Map, ) -> Result { let deployment = request::resolve_deployment(gateway, &body)?; - request::authorize_model(identity, deployment, &body).await?; - let (body, cache_options) = crate::caching::prepare(identity, body)?; + request::authorize_model(caller, deployment, &body)?; + let (body, cache_options) = crate::caching::prepare(caller, body)?; let route = gateway.chat_completions.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index d612a042e03..821661290c3 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -1,6 +1,6 @@ //! `POST /v1/messages`, as the Python proxy's `anthropic_response` serves it. -use litellm_gateway_auth::AuthenticatedRequest; +use litellm_gateway_auth::Caller; use std::sync::Arc; use axum::{ @@ -28,13 +28,13 @@ const ANTHROPIC_API_HEADER_PROVIDERS: &str = pub async fn create( State(gateway): State>, - identity: AuthenticatedRequest, + caller: Caller, RequestId(request_id): RequestId, headers: HeaderMap, body: Result, ) -> impl IntoResponse { let result = match body { - Ok(JsonObject(body)) => handle(&gateway, &identity, &headers, body).await, + Ok(JsonObject(body)) => handle(&gateway, &caller, &headers, body).await, Err(error) => Err(error), }; result.map_err(|error| (error.status(), Json(error.body(request_id.as_deref())))) @@ -42,13 +42,13 @@ pub async fn create( async fn handle( gateway: &Gateway, - identity: &AuthenticatedRequest, + caller: &Caller, headers: &HeaderMap, body: Map, ) -> Result { let deployment = request::resolve_deployment(gateway, &body)?; - request::authorize_model(identity, deployment, &body).await?; - let (body, cache_options) = crate::caching::prepare(identity, body)?; + request::authorize_model(caller, deployment, &body)?; + let (body, cache_options) = crate::caching::prepare(caller, body)?; let route = gateway.messages.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs index 03d064b6bae..ee9d7959851 100644 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -1,4 +1,4 @@ -use litellm_gateway_auth::AuthenticatedRequest; +use litellm_gateway_auth::Caller; use std::sync::Arc; use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse}; @@ -15,16 +15,16 @@ use crate::{ pub(crate) async fn create( State(gateway): State>, - identity: AuthenticatedRequest, + caller: Caller, headers: HeaderMap, body: InferenceBody, ) -> Result { - handle(&gateway, &identity, &headers, body).await.map(Json) + handle(&gateway, &caller, &headers, body).await.map(Json) } async fn handle( gateway: &Gateway, - identity: &AuthenticatedRequest, + caller: &Caller, headers: &HeaderMap, InferenceBody { fields: body, @@ -36,7 +36,7 @@ async fn handle( .and_then(|value| value.to_str().ok()) .map(str::to_owned); let deployment = request::resolve_deployment(gateway, &body)?; - request::authorize_model(identity, deployment, &body).await?; + request::authorize_model(caller, deployment, &body)?; let document = match upload { Some(upload) => OcrDocumentInput::Bytes { bytes: upload.bytes, diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index b398da2013b..31769655291 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -90,8 +90,8 @@ fn object(body: &[u8]) -> Result, Error> { } } -pub(crate) async fn authorize_model( - identity: &litellm_gateway_auth::AuthenticatedRequest, +pub(crate) fn authorize_model( + caller: &litellm_gateway_auth::Caller, deployment: &Deployment, body: &Map, ) -> Result<(), Error> { @@ -103,11 +103,7 @@ pub(crate) async fn authorize_model( name: name.into(), deployment: deployment.model.clone(), }; - identity - .authorize(access.clone()) - .await? - .consume(identity.caller(), &access)?; - Ok(()) + Ok(caller.authorize(&access)?) } pub(crate) fn resolve_deployment<'a>( diff --git a/litellm-rust/crates/gateway-inference/src/responses.rs b/litellm-rust/crates/gateway-inference/src/responses.rs index 77bebf610e1..ae7a768b8e7 100644 --- a/litellm-rust/crates/gateway-inference/src/responses.rs +++ b/litellm-rust/crates/gateway-inference/src/responses.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use axum::{Json, body::Bytes, extract::State, response::Response}; -use litellm_gateway_auth::AuthenticatedRequest; +use litellm_gateway_auth::Caller; use litellm_host_http::Sse; use litellm_inference_responses::{route::Responses, types::ResponsesCall}; use serde_json::json; @@ -10,12 +10,12 @@ use crate::{Error, Gateway, JsonObject, request}; pub(crate) async fn create( State(gateway): State>, - identity: AuthenticatedRequest, + caller: Caller, JsonObject(body): JsonObject, ) -> Result { let deployment = request::resolve_deployment(&gateway, &body)?; - request::authorize_model(&identity, deployment, &body).await?; - let (body, cache_options) = crate::caching::prepare(&identity, body)?; + request::authorize_model(&caller, deployment, &body)?; + let (body, cache_options) = crate::caching::prepare(&caller, body)?; let route = gateway.responses.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( diff --git a/litellm-rust/crates/gateway-inference/tests/caching.rs b/litellm-rust/crates/gateway-inference/tests/caching.rs index 143bc1ba6db..93cf25e9ff3 100644 --- a/litellm-rust/crates/gateway-inference/tests/caching.rs +++ b/litellm-rust/crates/gateway-inference/tests/caching.rs @@ -126,60 +126,3 @@ async fn all_inference_endpoints_share_native_cache( ); assert_eq!(to_bytes(restored.into_body(), 4096).await.unwrap(), second); } - -#[rstest] -#[case::different_subject("issuer", "tenant-b")] -#[case::different_authority("other-issuer", "tenant-a")] -#[tokio::test] -async fn authenticated_callers_do_not_share_cached_responses( - #[case] authority: &str, - #[case] subject: &str, -) { - use litellm_gateway_auth::{Principal, PrincipalKind}; - - let upstream = MockServer::start().await; - Mock::given(method("POST")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "id":"response-1", "model":"test-model", "status":"completed", "output":[] - }))) - .expect(2) - .mount(&upstream) - .await; - let cache: Arc = Arc::new(ResponseCache::new(Arc::new( - InMemoryCache::new(Some(100), Some(Duration::from_secs(60))), - ))); - let first_caller = support::app_with_cache_for_principal( - "openai/test-model", - &upstream.uri(), - cache.clone(), - Principal::new("issuer".into(), "tenant-a".into(), PrincipalKind::Service), - ); - let second_caller = support::app_with_cache_for_principal( - "openai/test-model", - &upstream.uri(), - cache, - Principal::new(authority.into(), subject.into(), PrincipalKind::Service), - ); - let body = json!({"model":"public/model","input":"same prompt"}); - let first = support::post(first_caller.clone(), "/v1/responses", body.clone()).await; - assert_eq!(first.status(), 200); - assert!(!first.headers().contains_key("x-litellm-cache-key")); - let first_hit = support::post(first_caller.clone(), "/v1/responses", body.clone()).await; - assert_eq!(first_hit.status(), 200); - let first_key = first_hit.headers().get("x-litellm-cache-key").unwrap(); - let second = support::post(second_caller.clone(), "/v1/responses", body.clone()).await; - assert_eq!(second.status(), 200); - assert!(!second.headers().contains_key("x-litellm-cache-key")); - let second_hit = support::post(second_caller, "/v1/responses", body.clone()).await; - assert_eq!(second_hit.status(), 200); - assert_ne!( - second_hit.headers().get("x-litellm-cache-key").unwrap(), - first_key - ); - let first_again = support::post(first_caller, "/v1/responses", body).await; - assert_eq!(first_again.status(), 200); - assert_eq!( - first_again.headers().get("x-litellm-cache-key"), - Some(first_key) - ); -} diff --git a/litellm-rust/crates/gateway-inference/tests/messages.rs b/litellm-rust/crates/gateway-inference/tests/messages.rs index bf11e67e0fc..ffe6dcd0eda 100644 --- a/litellm-rust/crates/gateway-inference/tests/messages.rs +++ b/litellm-rust/crates/gateway-inference/tests/messages.rs @@ -174,25 +174,6 @@ async fn hosted_provider_failure_preserves_status_body_and_request_id() { assert_eq!(body["request_id"], "host-http-request"); } -#[rstest] -#[case::messages("/v1/messages")] -#[case::chat("/v1/chat/completions")] -#[case::model_path("/engines/path-model/chat/completions")] -#[case::ocr("/ocr")] -#[case::transcription("/audio/transcriptions")] -#[tokio::test] -async fn model_permissions_prevent_provider_calls(#[case] path: &str) { - let upstream = MockServer::start().await; - let app = support::app_with_permissions( - "anthropic/test-model", - &upstream.uri(), - litellm_gateway_auth::Permissions::None, - ); - let response = support::post(app, path, json!({"model": "public/model"})).await; - assert_eq!(response.status(), 403); - assert!(upstream.received_requests().await.unwrap().is_empty()); -} - #[rstest] #[case::completion(false)] #[case::streaming(true)] diff --git a/litellm-rust/crates/gateway-inference/tests/responses.rs b/litellm-rust/crates/gateway-inference/tests/responses.rs index ded7d73c5f3..68b40a16472 100644 --- a/litellm-rust/crates/gateway-inference/tests/responses.rs +++ b/litellm-rust/crates/gateway-inference/tests/responses.rs @@ -97,22 +97,3 @@ async fn responses_preserve_upstream_errors_without_retry() { .contains("slow down") ); } - -#[rstest] -#[tokio::test] -async fn responses_authorize_the_public_model_before_provider_execution() { - let upstream = MockServer::start().await; - let app = support::app_with_permissions( - "openai/test-model", - &upstream.uri(), - litellm_gateway_auth::Permissions::None, - ); - let response = support::post( - app, - "/v1/responses", - json!({"model": "public/model", "input": "hello"}), - ) - .await; - assert_eq!(response.status(), 403); - assert!(upstream.received_requests().await.unwrap().is_empty()); -} diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 0e37bb0aff9..f1859608f5d 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -29,15 +29,7 @@ impl SecretSource for NoSecrets { } pub fn app(model: &str, api_base: &str) -> Router { - app_with_permissions(model, api_base, litellm_gateway_auth::Permissions::All) -} - -pub fn app_with_permissions( - model: &str, - api_base: &str, - permissions: litellm_gateway_auth::Permissions, -) -> Router { - configured_app(model, api_base, permissions, None, None) + configured_app(model, api_base, None) } pub fn app_with_cache( @@ -45,36 +37,13 @@ pub fn app_with_cache( api_base: &str, cache: Arc, ) -> Router { - configured_app( - model, - api_base, - litellm_gateway_auth::Permissions::All, - Some(cache), - None, - ) -} - -pub fn app_with_cache_for_principal( - model: &str, - api_base: &str, - cache: Arc, - principal: litellm_gateway_auth::Principal, -) -> Router { - configured_app( - model, - api_base, - litellm_gateway_auth::Permissions::All, - Some(cache), - Some(principal), - ) + configured_app(model, api_base, Some(cache)) } fn configured_app( model: &str, api_base: &str, - permissions: litellm_gateway_auth::Permissions, cache: Option>, - principal: Option, ) -> Router { let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); let http = Resolution::from(&HttpSettings::default()).config; @@ -102,34 +71,16 @@ fn configured_app( Some(cache) => gateway.with_cache(cache), None => gateway, }; - router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state( - (permissions, principal), - test_identity, - )) + router(Arc::new(gateway)).layer(axum::middleware::from_fn(master_caller)) } -async fn test_identity( - axum::extract::State((permissions, principal)): axum::extract::State<( - litellm_gateway_auth::Permissions, - Option, - )>, +async fn master_caller( mut request: axum::extract::Request, next: axum::middleware::Next, ) -> Response { - let auth = litellm_gateway_auth::Auth::new( - Arc::new(litellm_gateway_auth::MasterKeyAuthenticator::new( - Some(SecretValue::new("test-inbound-key")), - Arc::new(NoSecrets), - )), - Arc::new(TestPermissions(permissions, principal)), - Arc::new(litellm_gateway_auth::NoAdditionalPolicy), - Arc::new(litellm_gateway_auth::SystemClock), - ); - let identity = auth - .authenticate(&SecretValue::new("test-inbound-key")) - .await - .unwrap(); - request.extensions_mut().insert(identity); + request + .extensions_mut() + .insert(litellm_gateway_auth::Caller::Master); next.run(request).await } @@ -147,22 +98,3 @@ pub async fn post(app: Router, path: &str, body: Value) -> Response { pub async fn json(response: Response) -> Value { serde_json::from_slice(&to_bytes(response.into_body(), 1024 * 1024).await.unwrap()).unwrap() } - -struct TestPermissions( - litellm_gateway_auth::Permissions, - Option, -); - -impl litellm_gateway_auth::IdentityResolver for TestPermissions { - fn resolve<'a>( - &'a self, - identity: &'a litellm_gateway_auth::VerifiedIdentity, - ) -> litellm_gateway_auth::AuthFuture<'a, litellm_gateway_auth::ResolvedIdentity> { - Box::pin(async move { - Ok(litellm_gateway_auth::ResolvedIdentity { - principal: self.1.clone().unwrap_or_else(|| identity.principal.clone()), - permissions: self.0.clone(), - }) - }) - } -} diff --git a/litellm-rust/crates/gateway-ui/src/session.rs b/litellm-rust/crates/gateway-ui/src/session.rs index b80f40bb376..618b30869c6 100644 --- a/litellm-rust/crates/gateway-ui/src/session.rs +++ b/litellm-rust/crates/gateway-ui/src/session.rs @@ -10,9 +10,7 @@ use tower_cookies::{Cookie, Cookies}; use tower_sessions::{Expiry, Session, cookie::SameSite}; use crate::Error; -use litellm_gateway_auth::{ - AccessRequest, UI_CSRF_KEY, UiAction, UiAuthError, UiAuthSession, UiCredentials, UiSession, -}; +use litellm_gateway_auth::{UI_CSRF_KEY, UiAuthError, UiAuthSession, UiCredentials, UiSession}; pub struct State { pub signing_key: EncodingKey, @@ -109,11 +107,6 @@ pub async fn login( pub async fn info(session: Result) -> Result { let session = session?; - session - .identity - .authorize(AccessRequest::Ui(UiAction::SessionInfo)) - .await - .map_err(|_| UiAuthError::Unauthorized)?; Ok(Json(SessionInfo { user_id: session.user.username, user_role: "proxy_admin", @@ -125,12 +118,7 @@ pub async fn logout( mut auth: UiAuthSession, cookies: Cookies, ) -> Result { - let validated = validated?; - validated - .identity - .authorize(AccessRequest::Ui(UiAction::Logout)) - .await - .map_err(|_| UiAuthError::Unauthorized)?; + validated?; auth.logout().await?; cookies.add( Cookie::build(("token", "")) diff --git a/litellm-rust/crates/gateway/src/auth.rs b/litellm-rust/crates/gateway/src/auth.rs index 72892be783c..7d40426d4a9 100644 --- a/litellm-rust/crates/gateway/src/auth.rs +++ b/litellm-rust/crates/gateway/src/auth.rs @@ -1,27 +1,27 @@ use std::sync::Arc; use axum::{extract::Request, middleware::Next, response::Response}; -use litellm_gateway_auth::{AccessRequest, AuthenticatedRequest, Error, McpAction}; +use litellm_gateway_auth::{AccessRequest, Caller, Error, McpAction}; use litellm_gateway_mcp::{ Authorization, GatewayFuture, OperationAuthorizer, ServerInfo, SessionOwner, rmcp::model::ClientRequest, }; pub(crate) async fn bind_session_owner( - identity: AuthenticatedRequest, + caller: Caller, mut request: Request, next: Next, ) -> Response { request .extensions_mut() - .insert(SessionOwner(identity.caller().session_owner())); + .insert(SessionOwner(caller.owner())); request .extensions_mut() - .insert(Authorization(Arc::new(McpPolicy(identity)))); + .insert(Authorization(Arc::new(McpPolicy(caller)))); next.run(request).await } -struct McpPolicy(AuthenticatedRequest); +struct McpPolicy(Caller); impl OperationAuthorizer for McpPolicy { fn authorize<'a>( @@ -49,15 +49,11 @@ impl OperationAuthorizer for McpPolicy { } Some(_) => return Err(litellm_gateway_mcp::Error::Forbidden), }; - let access = AccessRequest::Mcp { - server: server.server_id.clone(), - action, - }; self.0 - .authorize(access.clone()) - .await - .map_err(mcp_error)? - .consume(self.0.caller(), &access) + .authorize(&AccessRequest::Mcp { + server: server.server_id.clone(), + action, + }) .map_err(mcp_error) }) } @@ -65,9 +61,7 @@ impl OperationAuthorizer for McpPolicy { fn mcp_error(error: Error) -> litellm_gateway_mcp::Error { match error { - Error::InvalidToken | Error::Expired | Error::Forbidden => { - litellm_gateway_mcp::Error::Forbidden - } + Error::InvalidToken | Error::Forbidden => litellm_gateway_mcp::Error::Forbidden, _ => litellm_gateway_mcp::Error::Configuration("MCP authorization unavailable".into()), } } diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index 75dc04d179d..c08b9f0f014 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -4,6 +4,8 @@ pub enum Error { Http(#[from] litellm_http::Error), #[error("environment_variables values must be strings, numbers or booleans")] Environment, + #[error("general_settings.custom_auth is only supported by the Python proxy")] + CustomAuth, #[error("MCP host does not yet support configured setting {0}")] McpSetting(String), #[error(transparent)] diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 0b3aba81181..5a3af29992d 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -15,7 +15,7 @@ use axum::{ use http_body_util::BodyExt; use litellm_config::Config; -use litellm_gateway_auth::Auth; +use litellm_gateway_auth::{Auth, Bearer}; use litellm_gateway_inference::{Gateway, ModelRouter}; use litellm_http::{ ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, @@ -27,6 +27,14 @@ use litellm_tracing::ByteChunk; use uuid::Uuid; pub fn build_inference(config: &Config) -> Result, Error> { + if config + .general_settings + .additional_fields + .get("custom_auth") + .is_some_and(|value| !value.is_null()) + { + return Err(Error::CustomAuth); + } let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); let environment = secrets::environment_values(&config.environment_variables)?; let lookup = |name: &str| { @@ -125,7 +133,7 @@ pub fn router( .route_layer(axum::middleware::from_fn(auth::bind_session_owner)) .route_layer(axum::middleware::from_fn_with_state( auth, - litellm_gateway_auth::authenticate, + litellm_gateway_auth::authenticate::, )) .layer(axum::middleware::from_fn(log_request)); match ui { diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 71edd3c3650..fd408ed14ee 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -33,6 +33,22 @@ fn inference() -> Arc { litellm_gateway::build_inference(&Config::from_yaml("model_list: []").unwrap()).unwrap() } +#[rstest] +#[case::configured("custom_auth: auth.user_api_key_auth", false)] +#[case::null("custom_auth: null", true)] +#[case::absent("master_key: key", true)] +fn custom_auth_is_left_to_the_python_proxy(#[case] setting: &str, #[case] builds: bool) { + let config = + Config::from_yaml(&format!("model_list: []\ngeneral_settings:\n {setting}\n")).unwrap(); + match litellm_gateway::build_inference(&config) { + Ok(_) => assert!(builds), + Err(error) => { + assert!(!builds); + assert!(matches!(error, litellm_gateway::Error::CustomAuth)); + } + } +} + #[rstest] #[case::authorized("/v1/messages", Some("Bearer gateway-key"), Some("gateway-key"), 400)] #[case::missing_token("/v1/messages", None, Some("gateway-key"), 401)]