diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b50494272ae..71ea2140e98 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3880,9 +3880,11 @@ dependencies = [ name = "litellm-gateway-auth-core" version = "0.1.0" dependencies = [ + "ipnet", "litellm-secret-value", "rstest", "serde", + "serde_json", "strum", "tokio", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index c52b19b83ea..0e6bdfa9297 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -119,6 +119,7 @@ rustify = "=0.7.0" rustify_derive = "=0.5.5" vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] } serde = { version = "1.0", features = ["derive"] } +ipnet = "2.12.0" serde_json = { version = "1.0", features = ["float_roundtrip"] } serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] } sha2 = "0.10" diff --git a/litellm-rust/crates/gateway-auth-core/AGENTS.md b/litellm-rust/crates/gateway-auth-core/AGENTS.md index db632b422b6..e248707f109 100644 --- a/litellm-rust/crates/gateway-auth-core/AGENTS.md +++ b/litellm-rust/crates/gateway-auth-core/AGENTS.md @@ -4,8 +4,8 @@ The Rust side of `litellm/proxy/auth/core`. It mirrors that package's facets, ho The crate is pure. It depends on no IO, runtime, HTTP, database or Python crate, and a facet reaches outside itself only through its host trait. Each host trait has its own `Error` type, so a pyo3-backed host and a native host can fail in their own terms while the facet code stays the same -Each Python protocol is a trait here. A core hands out its facets through accessors whose types are associated types on `AuthCore` and `TargetAuthCore`, the same composition as the Python accessors. Host methods are `-> impl Future + Send` so facets run under Axum and `tokio::spawn` without `#[async_trait]` boxing +Each Python protocol is a trait here. A core hands out its facets through accessors whose types are associated types on `AuthCore`, the same composition as the Python accessors. Host methods that do IO are `-> impl Future + Send` so facets run under Axum and `tokio::spawn` without `#[async_trait]` boxing. Values a host returns as untyped JSON in Python, such as introspection fields, are `serde_json::Value` here -`NativeCore` implements `TargetAuthCore`. The OAuth2 and custom auth facets are worked examples of how a facet ports. Every other facet is `Unported` and panics if called, the same as `RustCore` in Python. A facet leaves `Unported` only together with the Python change that routes the proxy to it, and both sides then share one test suite of cases +`NativeCore` implements `AuthCore`. Its OAuth2 facet (introspection and proxy-header mode) and custom auth facet are ported. Every other facet is `Unported` and panics if called, the same as `RustCore` in Python. A facet leaves `Unported` only together with the Python change that routes the proxy to it, and both sides then share one test suite of cases Python reaches the crate through `crates/python-bridge/src/auth.rs`, never through pyo3 types here. A value that crosses to Python derives serde, and the bridge converts it. Each facet call is a pyo3 class with one `answer_*` method per host call, so `_native.pyi` types the answer to every step. The bridge's host parks each host call until Python answers it. When a serde shape changes, change its TypedDict in `_native.pyi` with it diff --git a/litellm-rust/crates/gateway-auth-core/Cargo.toml b/litellm-rust/crates/gateway-auth-core/Cargo.toml index 43cf1ebb4ee..e81dbc07945 100644 --- a/litellm-rust/crates/gateway-auth-core/Cargo.toml +++ b/litellm-rust/crates/gateway-auth-core/Cargo.toml @@ -7,8 +7,10 @@ repository.workspace = true description = "The pure request auth core behind litellm/proxy/auth/core: facets, their hosts and the values they exchange" [dependencies] +ipnet.workspace = true litellm-secret-value.workspace = true serde.workspace = true +serde_json.workspace = true strum.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/gateway-auth-core/src/answers.rs b/litellm-rust/crates/gateway-auth-core/src/answers.rs index a69a438a80b..6ad369fc3f4 100644 --- a/litellm-rust/crates/gateway-auth-core/src/answers.rs +++ b/litellm-rust/crates/gateway-auth-core/src/answers.rs @@ -1,6 +1,7 @@ use std::{collections::BTreeMap, time::SystemTime}; use serde::Deserialize; +use serde_json::Value; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum PresentedKey { @@ -76,9 +77,9 @@ pub struct SigningKey { #[derive(Clone, Debug, PartialEq, Eq, Deserialize)] pub struct Introspection { pub active: bool, - pub subject: Option, - pub client_id: Option, - pub scopes: Vec, + pub user_id: Value, + pub role: Value, + pub team_id: Value, } #[derive(Clone, Debug, PartialEq)] diff --git a/litellm-rust/crates/gateway-auth-core/src/facets.rs b/litellm-rust/crates/gateway-auth-core/src/facets.rs index a233ded9627..fe8dfefc921 100644 --- a/litellm-rust/crates/gateway-auth-core/src/facets.rs +++ b/litellm-rust/crates/gateway-auth-core/src/facets.rs @@ -2,15 +2,17 @@ use std::{future::Future, time::SystemTime}; use crate::{ host::{ - AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2Host, - TeamMemberModelBudgetHost, + AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost, + OAuth2ProxyHost, TeamMemberModelBudgetHost, }, inputs::{ - AuthorizeSettings, JwtSettings, KeySettings, OAuth2Settings, PlanSettings, RequestFacts, + AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings, + OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts, }, outcome::{ - AnonymousReason, CustomIdentity, Denial, Identity, ModelBudgetResult, OAuth2Denial, - OAuth2Identity, Outcome, Validator, + Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, Denial, Identity, KeyPlan, + ModelBudgetResult, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial, OAuth2ProxyIdentity, + Outcome, Plan, }, }; @@ -47,11 +49,19 @@ pub trait ModelBudgetCore: Sync { } pub trait AuthPlanner: Sync { - fn plan(&self, facts: &RequestFacts, settings: &PlanSettings) -> Validator; + fn plan(&self, facts: &PlanFacts, settings: &PlanSettings) -> Plan; + fn after_custom_auth( + &self, + facts: &PlanFacts, + settings: &PlanSettings, + issued_key: bool, + ) -> CredentialPlan; + fn after_jwt(&self, facts: &PlanFacts, settings: &PlanSettings, issued_key: bool) -> KeyPlan; + fn oauth2_introspection_applies(&self, facts: &PlanFacts, settings: &PlanSettings) -> bool; } pub trait AnonymousCore: Sync { - fn authenticate(&self, facts: &RequestFacts, settings: &PlanSettings) -> AnonymousReason; + fn authenticate(&self, reason: AnonymousReason) -> Anonymous; } pub trait JwtCore: Sync { @@ -65,12 +75,17 @@ pub trait JwtCore: Sync { } pub trait OAuth2Core: Sync { - fn authenticate( + fn introspect( &self, host: &H, - facts: &RequestFacts, - settings: &OAuth2Settings, + token: &str, + settings: &OAuth2IntrospectionSettings, ) -> impl Future, H::Error>> + Send; + fn proxy( + &self, + host: &H, + settings: &OAuth2ProxySettings, + ) -> Result; } pub trait CustomAuthCore: Sync { @@ -78,7 +93,7 @@ pub trait CustomAuthCore: Sync { &self, host: &H, facts: &RequestFacts, - ) -> impl Future> + Send; + ) -> impl Future> + Send; } pub trait AuthorizeCore: Sync { @@ -94,12 +109,6 @@ pub trait AuthorizeCore: Sync { pub trait AuthCore { type Keys: KeyCore; type ModelBudgets: ModelBudgetCore; - - fn keys(&self) -> &Self::Keys; - fn model_budgets(&self) -> &Self::ModelBudgets; -} - -pub trait TargetAuthCore: AuthCore { type Planner: AuthPlanner; type Anonymous: AnonymousCore; type Jwt: JwtCore; @@ -107,6 +116,8 @@ pub trait TargetAuthCore: AuthCore { type Custom: CustomAuthCore; type Authorize: AuthorizeCore; + fn keys(&self) -> &Self::Keys; + fn model_budgets(&self) -> &Self::ModelBudgets; fn plan(&self) -> &Self::Planner; fn anonymous(&self) -> &Self::Anonymous; fn jwt(&self) -> &Self::Jwt; diff --git a/litellm-rust/crates/gateway-auth-core/src/host.rs b/litellm-rust/crates/gateway-auth-core/src/host.rs index bcda0161900..3605d4474d4 100644 --- a/litellm-rust/crates/gateway-auth-core/src/host.rs +++ b/litellm-rust/crates/gateway-auth-core/src/host.rs @@ -1,4 +1,4 @@ -use std::{collections::BTreeMap, error::Error, future::Future}; +use std::{error::Error, future::Future}; use crate::{ answers::{ @@ -7,7 +7,7 @@ use crate::{ TeamMemberModelBudgetScope, }, inputs::RequestFacts, - outcome::{CustomIdentity, Identity, JwtIdentity, ModelBudgetExceeded}, + outcome::{CustomAuthResult, Identity, JwtIdentity, ModelBudgetExceeded}, }; pub trait ModelBudgetHost: Sync { @@ -133,14 +133,18 @@ pub trait JwtHost: Sync { ) -> impl Future, Self::Error>> + Send; } -pub trait OAuth2Host: Sync { +pub trait OAuth2IntrospectionHost: Sync { type Error: Error + Send + Sync + 'static; fn introspect( &self, token: &str, ) -> impl Future> + Send; - fn proxy_headers(&self) -> &BTreeMap; +} + +pub trait OAuth2ProxyHost { + fn client_ip(&self) -> Option<&str>; + fn header(&self, name: &str) -> Option<&str>; } pub trait CustomAuthHost: Sync { @@ -149,7 +153,7 @@ pub trait CustomAuthHost: Sync { fn custom_auth( &self, facts: &RequestFacts, - ) -> impl Future> + Send; + ) -> impl Future> + Send; } pub trait AuthorizeHost: Sync { diff --git a/litellm-rust/crates/gateway-auth-core/src/inputs.rs b/litellm-rust/crates/gateway-auth-core/src/inputs.rs index 06c9068b289..9c9be370109 100644 --- a/litellm-rust/crates/gateway-auth-core/src/inputs.rs +++ b/litellm-rust/crates/gateway-auth-core/src/inputs.rs @@ -1,7 +1,10 @@ -use std::collections::BTreeSet; +use std::collections::{BTreeMap, BTreeSet}; + +use ipnet::IpNet; use litellm_secret_value::SecretValue; use serde::Deserialize; +use serde_json::{Map, Value}; #[derive(Clone, Debug, PartialEq, Deserialize)] pub struct RequestFacts { @@ -10,13 +13,33 @@ pub struct RequestFacts { pub is_llm_api_route: bool, } +#[derive(Clone, Debug, PartialEq)] +pub struct PlanFacts { + pub request: RequestFacts, + pub is_info_route: bool, + pub unverified_claims: Option>, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct RoutingOverride { + pub iss: Option>, + pub client_id: Option>, + pub scope: Option>, + pub aud: Option>, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct PlanSettings { pub custom_auth: bool, - pub jwt_enabled: bool, + pub premium_user: bool, + pub public_routes: BTreeSet, + pub additional_public_routes: Vec, + pub metrics_require_auth: bool, pub oauth2_introspection: bool, pub oauth2_proxy: bool, - pub public_routes: BTreeSet, + pub jwt_enabled: bool, + pub oauth2_routing_overrides: Vec, + pub unauthenticated_pass_through_routes: BTreeSet, } #[derive(Clone, Debug, PartialEq)] @@ -34,10 +57,16 @@ pub struct JwtSettings { pub leeway_seconds: u64, } -#[derive(Clone, Debug, PartialEq, Eq, Deserialize)] -pub struct OAuth2Settings { - pub user_id_header: Option, - pub team_id_header: Option, +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)] +pub struct OAuth2IntrospectionSettings { + pub endpoint_configured: bool, + pub requires_active: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OAuth2ProxySettings { + pub trusted_networks: Vec, + pub header_mappings: BTreeMap, } #[derive(Clone, Debug, PartialEq)] diff --git a/litellm-rust/crates/gateway-auth-core/src/lib.rs b/litellm-rust/crates/gateway-auth-core/src/lib.rs index 58f89a3d543..3ff9692a8ea 100644 --- a/litellm-rust/crates/gateway-auth-core/src/lib.rs +++ b/litellm-rust/crates/gateway-auth-core/src/lib.rs @@ -12,18 +12,20 @@ pub use answers::{ }; pub use facets::{ AnonymousCore, AuthCore, AuthPlanner, AuthorizeCore, CustomAuthCore, JwtCore, KeyCore, - ModelBudgetCore, OAuth2Core, TargetAuthCore, + ModelBudgetCore, OAuth2Core, }; pub use host::{ - AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2Host, SpendCounter, - TeamMemberModelBudgetHost, + AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost, + OAuth2ProxyHost, SpendCounter, TeamMemberModelBudgetHost, }; pub use inputs::{ - AuthorizeSettings, JwtSettings, KeySettings, OAuth2Settings, PlanSettings, RequestFacts, + AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings, OAuth2ProxySettings, + PlanFacts, PlanSettings, RequestFacts, RoutingOverride, }; pub use native::{NativeCore, NativeCustomAuthCore, NativeOAuth2Core, Unported}; pub use outcome::{ - AnonymousReason, CustomIdentity, Denial, Grant, Identity, JwtIdentity, ModelBudgetExceeded, - ModelBudgetResult, OAuth2Denial, OAuth2Identity, OAuth2Mode, Outcome, Validator, - WithinModelBudgets, + Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, CustomIdentity, Denial, Grant, + Identity, JwtIdentity, KeyPlan, ModelBudgetExceeded, ModelBudgetResult, OAuth2Denial, + OAuth2Identity, OAuth2Mode, OAuth2ProxyDenial, OAuth2ProxyIdentity, Outcome, Plan, + PremiumFeature, WithinModelBudgets, }; diff --git a/litellm-rust/crates/gateway-auth-core/src/native.rs b/litellm-rust/crates/gateway-auth-core/src/native.rs index 1dc7d531026..9912530fbce 100644 --- a/litellm-rust/crates/gateway-auth-core/src/native.rs +++ b/litellm-rust/crates/gateway-auth-core/src/native.rs @@ -1,23 +1,36 @@ -use std::time::SystemTime; +use std::{net::IpAddr, time::SystemTime}; + +use ipnet::IpNet; use crate::{ facets::{ AnonymousCore, AuthCore, AuthPlanner, AuthorizeCore, CustomAuthCore, JwtCore, KeyCore, - ModelBudgetCore, OAuth2Core, TargetAuthCore, + ModelBudgetCore, OAuth2Core, }, host::{ - AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2Host, - TeamMemberModelBudgetHost, + AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost, + OAuth2ProxyHost, TeamMemberModelBudgetHost, }, inputs::{ - AuthorizeSettings, JwtSettings, KeySettings, OAuth2Settings, PlanSettings, RequestFacts, + AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings, + OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts, }, outcome::{ - AnonymousReason, CustomIdentity, Denial, Identity, ModelBudgetResult, OAuth2Denial, - OAuth2Identity, Outcome, Validator, + Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, Denial, Identity, KeyPlan, + ModelBudgetResult, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial, OAuth2ProxyIdentity, + Outcome, Plan, }, }; +const ALLOWED_OAUTH2_PROXY_FIELDS: [&str; 6] = [ + "user_id", + "user_email", + "team_id", + "team_alias", + "org_id", + "models", +]; + #[derive(Clone, Copy, Debug, Default)] pub struct NativeCore { unported: Unported, @@ -28,6 +41,12 @@ pub struct NativeCore { impl AuthCore for NativeCore { type Keys = Unported; type ModelBudgets = Unported; + type Planner = Unported; + type Anonymous = Unported; + type Jwt = Unported; + type OAuth2 = NativeOAuth2Core; + type Custom = NativeCustomAuthCore; + type Authorize = Unported; fn keys(&self) -> &Unported { &self.unported @@ -36,15 +55,6 @@ impl AuthCore for NativeCore { fn model_budgets(&self) -> &Unported { &self.unported } -} - -impl TargetAuthCore for NativeCore { - type Planner = Unported; - type Anonymous = Unported; - type Jwt = Unported; - type OAuth2 = NativeOAuth2Core; - type Custom = NativeCustomAuthCore; - type Authorize = Unported; fn plan(&self) -> &Unported { &self.unported @@ -74,26 +84,96 @@ impl TargetAuthCore for NativeCore { #[derive(Clone, Copy, Debug, Default)] pub struct NativeOAuth2Core; +fn ip_in_networks(client_ip: Option<&str>, networks: &[IpNet]) -> bool { + let Some(Ok(address)) = client_ip.map(|ip| ip.trim().parse::()) else { + return false; + }; + let unmapped = match address { + IpAddr::V6(v6) => v6.to_ipv4_mapped().map_or(address, IpAddr::V4), + IpAddr::V4(_) => address, + }; + networks + .iter() + .any(|network| network.contains(&address) || network.contains(&unmapped)) +} + +fn proxy_identity( + host: &H, + settings: &OAuth2ProxySettings, +) -> OAuth2ProxyIdentity { + let value = |field: &str| { + settings + .header_mappings + .get(field) + .and_then(|header| host.header(header)) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + }; + OAuth2ProxyIdentity { + user_id: value("user_id"), + user_email: value("user_email"), + team_id: value("team_id"), + team_alias: value("team_alias"), + org_id: value("org_id"), + models: value("models").map(|models| { + models + .split(',') + .map(|model| model.trim().to_owned()) + .collect() + }), + } +} + impl OAuth2Core for NativeOAuth2Core { - async fn authenticate( + async fn introspect( &self, host: &H, - facts: &RequestFacts, - _settings: &OAuth2Settings, + token: &str, + settings: &OAuth2IntrospectionSettings, ) -> Result, H::Error> { - let Some(token) = &facts.token else { - return Ok(Err(OAuth2Denial::MissingKey)); - }; - let introspection = host.introspect(token.expose()).await?; - if !introspection.active { + if !settings.endpoint_configured { + return Ok(Err(OAuth2Denial::IntrospectionNotConfigured)); + } + let answer = host.introspect(token).await?; + if settings.requires_active && !answer.active { return Ok(Err(OAuth2Denial::InactiveToken)); } Ok(Ok(OAuth2Identity { - user_id: introspection.subject, - team_id: None, - scopes: introspection.scopes, + token: token.to_owned(), + user_id: answer.user_id, + role: answer.role, + team_id: answer.team_id, })) } + + fn proxy( + &self, + host: &H, + settings: &OAuth2ProxySettings, + ) -> Result { + if settings.trusted_networks.is_empty() { + return Err(OAuth2ProxyDenial::TrustedProxyRangesMissing); + } + let client_ip = host.client_ip(); + if !ip_in_networks(client_ip, &settings.trusted_networks) { + return Err(OAuth2ProxyDenial::UntrustedProxy { + client_ip: client_ip.map(str::to_owned), + }); + } + if settings.header_mappings.is_empty() { + return Err(OAuth2ProxyDenial::ProxyMappingsMissing); + } + let disallowed: Vec = settings + .header_mappings + .keys() + .filter(|field| !ALLOWED_OAUTH2_PROXY_FIELDS.contains(&field.as_str())) + .cloned() + .collect(); + if !disallowed.is_empty() { + return Err(OAuth2ProxyDenial::DisallowedProxyMappings { fields: disallowed }); + } + Ok(proxy_identity(host, settings)) + } } #[derive(Clone, Copy, Debug, Default)] @@ -104,7 +184,7 @@ impl CustomAuthCore for NativeCustomAuthCore { &self, host: &H, facts: &RequestFacts, - ) -> Result { + ) -> Result { host.custom_auth(facts).await } } @@ -158,14 +238,36 @@ impl ModelBudgetCore for Unported { } impl AuthPlanner for Unported { - fn plan(&self, _facts: &RequestFacts, _settings: &PlanSettings) -> Validator { - unimplemented!("the proxy picks the credential path") + fn plan(&self, _facts: &PlanFacts, _settings: &PlanSettings) -> Plan { + unimplemented!("the planner runs in PythonCore") + } + + fn after_custom_auth( + &self, + _facts: &PlanFacts, + _settings: &PlanSettings, + _issued_key: bool, + ) -> CredentialPlan { + unimplemented!("the planner runs in PythonCore") + } + + fn after_jwt( + &self, + _facts: &PlanFacts, + _settings: &PlanSettings, + _issued_key: bool, + ) -> KeyPlan { + unimplemented!("the planner runs in PythonCore") + } + + fn oauth2_introspection_applies(&self, _facts: &PlanFacts, _settings: &PlanSettings) -> bool { + unimplemented!("the planner runs in PythonCore") } } impl AnonymousCore for Unported { - fn authenticate(&self, _facts: &RequestFacts, _settings: &PlanSettings) -> AnonymousReason { - unimplemented!("the proxy decides anonymous requests") + fn authenticate(&self, _reason: AnonymousReason) -> Anonymous { + unimplemented!("the anonymous facet runs in PythonCore") } } @@ -195,137 +297,271 @@ impl AuthorizeCore for Unported { #[cfg(test)] mod tests { - use std::{collections::BTreeMap, convert::Infallible, sync::Mutex}; + use std::{ + collections::BTreeMap, + convert::Infallible, + sync::{Mutex, PoisonError}, + }; use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; + use serde_json::json; use super::*; use crate::answers::Introspection; - struct ScriptedOAuth2Host { - introspection: Introspection, - headers: BTreeMap, + struct ScriptedIntrospection { + answer: Introspection, introspected: Mutex>, } - impl OAuth2Host for ScriptedOAuth2Host { + impl OAuth2IntrospectionHost for ScriptedIntrospection { type Error = Infallible; async fn introspect(&self, token: &str) -> Result { self.introspected .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) + .unwrap_or_else(PoisonError::into_inner) .push(token.to_owned()); - Ok(self.introspection.clone()) - } - - fn proxy_headers(&self) -> &BTreeMap { - &self.headers + Ok(self.answer.clone()) } } - #[fixture] - fn settings() -> OAuth2Settings { - OAuth2Settings { - user_id_header: None, - team_id_header: None, - } - } - - fn host(active: bool) -> ScriptedOAuth2Host { - ScriptedOAuth2Host { - introspection: Introspection { + fn introspection(active: bool) -> ScriptedIntrospection { + ScriptedIntrospection { + answer: Introspection { active, - subject: Some("user-1".to_owned()), - client_id: None, - scopes: vec!["read".to_owned()], + user_id: json!("user-1"), + role: json!("internal_user"), + team_id: json!(7), }, - headers: BTreeMap::new(), introspected: Mutex::new(Vec::new()), } } - fn facts(token: Option<&str>) -> RequestFacts { - RequestFacts { - route: "/chat/completions".to_owned(), - token: token.map(SecretValue::from), - is_llm_api_route: true, + fn identity() -> OAuth2Identity { + OAuth2Identity { + token: "tok".to_owned(), + user_id: json!("user-1"), + role: json!("internal_user"), + team_id: json!(7), } } #[rstest] - #[case::active( - true, - Ok(OAuth2Identity { - user_id: Some("user-1".to_owned()), - team_id: None, - scopes: vec!["read".to_owned()], - }) - )] - #[case::inactive(false, Err(OAuth2Denial::InactiveToken))] + #[case::active(true, true, Ok(identity()))] + #[case::inactive(false, true, Err(OAuth2Denial::InactiveToken))] + #[case::token_info_endpoint(false, false, Ok(identity()))] #[tokio::test] - async fn oauth2_introspects_the_token_and_trusts_only_an_active_answer( - settings: OAuth2Settings, + async fn introspection_trusts_an_inactive_answer_only_from_an_introspection_endpoint( #[case] active: bool, + #[case] requires_active: bool, #[case] expected: Result, ) { - let host = host(active); + let host = introspection(active); + let settings = OAuth2IntrospectionSettings { + endpoint_configured: true, + requires_active, + }; let outcome = NativeCore::default() .oauth2() - .authenticate(&host, &facts(Some("tok-1")), &settings) + .introspect(&host, "tok", &settings) .await; assert_eq!(outcome, Ok(expected)); assert_eq!( *host .introspected .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner), - vec!["tok-1".to_owned()] + .unwrap_or_else(PoisonError::into_inner), + vec!["tok".to_owned()] ); } - #[rstest] #[tokio::test] - async fn oauth2_without_a_token_is_denied_before_any_introspection(settings: OAuth2Settings) { - let host = host(true); + async fn introspection_without_an_endpoint_never_calls_out() { + let host = introspection(true); + let settings = OAuth2IntrospectionSettings { + endpoint_configured: false, + requires_active: true, + }; let outcome = NativeCore::default() .oauth2() - .authenticate(&host, &facts(None), &settings) + .introspect(&host, "tok", &settings) .await; - assert_eq!(outcome, Ok(Err(OAuth2Denial::MissingKey))); + assert_eq!(outcome, Ok(Err(OAuth2Denial::IntrospectionNotConfigured))); assert!( host.introspected .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) + .unwrap_or_else(PoisonError::into_inner) .is_empty() ); } - struct ScriptedCustomAuthHost(CustomIdentity); + struct Headers { + client_ip: Option<&'static str>, + headers: BTreeMap<&'static str, &'static str>, + read: Mutex>, + } - impl CustomAuthHost for ScriptedCustomAuthHost { + impl OAuth2ProxyHost for Headers { + fn client_ip(&self) -> Option<&str> { + self.client_ip + } + + fn header(&self, name: &str) -> Option<&str> { + self.read + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(name.to_owned()); + self.headers.get(name).copied() + } + } + + fn headers( + client_ip: Option<&'static str>, + headers: &[(&'static str, &'static str)], + ) -> Headers { + Headers { + client_ip, + headers: headers.iter().copied().collect(), + read: Mutex::new(Vec::new()), + } + } + + #[fixture] + fn trusted() -> Vec { + vec!["10.0.0.0/8".parse().unwrap()] + } + + fn settings(networks: Vec, mappings: &[(&str, &str)]) -> OAuth2ProxySettings { + OAuth2ProxySettings { + trusted_networks: networks, + header_mappings: mappings + .iter() + .map(|(field, header)| ((*field).to_owned(), (*header).to_owned())) + .collect(), + } + } + + #[rstest] + fn proxy_headers_assert_only_the_mapped_identity_fields_that_are_present(trusted: Vec) { + let host = headers( + Some("10.1.2.3"), + &[ + ("x-user", "user-1"), + ("x-team", ""), + ("x-models", "gpt-a, gpt-b"), + ], + ); + let settings = settings( + trusted, + &[ + ("user_id", "x-user"), + ("team_id", "x-team"), + ("models", "x-models"), + ], + ); + assert_eq!( + NativeCore::default().oauth2().proxy(&host, &settings), + Ok(OAuth2ProxyIdentity { + user_id: Some("user-1".to_owned()), + models: Some(vec!["gpt-a".to_owned(), "gpt-b".to_owned()]), + ..OAuth2ProxyIdentity::default() + }) + ); + } + + #[rstest] + #[case::ipv4("10.1.2.3", true)] + #[case::ipv4_mapped("::ffff:10.1.2.3", true)] + #[case::outside("192.168.1.1", false)] + #[case::not_an_ip("not-an-ip", false)] + fn proxy_headers_trust_only_a_peer_inside_the_ranges( + trusted: Vec, + #[case] client_ip: &'static str, + #[case] is_trusted: bool, + ) { + let host = headers(Some(client_ip), &[("x-user", "user-1")]); + let outcome = NativeCore::default() + .oauth2() + .proxy(&host, &settings(trusted, &[("user_id", "x-user")])); + assert_eq!(outcome.is_ok(), is_trusted); + assert_eq!( + host.read + .lock() + .unwrap_or_else(PoisonError::into_inner) + .is_empty(), + !is_trusted + ); + } + + #[rstest] + fn proxy_headers_refuse_every_non_identity_mapping_in_sorted_order(trusted: Vec) { + let host = headers(Some("10.1.2.3"), &[]); + let settings = settings( + trusted, + &[ + ("user_role", "x-role"), + ("user_id", "x-user"), + ("max_budget", "x-budget"), + ], + ); + assert_eq!( + NativeCore::default().oauth2().proxy(&host, &settings), + Err(OAuth2ProxyDenial::DisallowedProxyMappings { + fields: vec!["max_budget".to_owned(), "user_role".to_owned()], + }) + ); + } + + #[rstest] + #[case::no_ranges(Vec::new(), &[("user_id", "x-user")], OAuth2ProxyDenial::TrustedProxyRangesMissing)] + #[case::no_mappings(trusted(), &[], OAuth2ProxyDenial::ProxyMappingsMissing)] + fn proxy_headers_need_ranges_and_mappings( + #[case] networks: Vec, + #[case] mappings: &[(&str, &str)], + #[case] expected: OAuth2ProxyDenial, + ) { + let host = headers(Some("10.1.2.3"), &[]); + assert_eq!( + NativeCore::default() + .oauth2() + .proxy(&host, &settings(networks, mappings)), + Err(expected) + ); + } + + struct ScriptedCustomAuth(CustomAuthResult); + + impl CustomAuthHost for ScriptedCustomAuth { type Error = Infallible; - async fn custom_auth(&self, _facts: &RequestFacts) -> Result { + async fn custom_auth(&self, _facts: &RequestFacts) -> Result { Ok(self.0.clone()) } } + #[rstest] + #[case::identity(CustomAuthResult::Identity(crate::outcome::CustomIdentity { + user_id: Some("user-1".to_owned()), + team_id: Some("team-1".to_owned()), + hashed_token: None, + }))] + #[case::replaced_key(CustomAuthResult::ReplacedKey { token: "sk-new".to_owned() })] + #[case::no_identity(CustomAuthResult::NoIdentity)] #[tokio::test] - async fn custom_auth_returns_the_identity_the_host_resolved() { - let identity = CustomIdentity { - user_id: Some("user-1".to_owned()), - team_id: Some("team-1".to_owned()), - hashed_token: None, + async fn custom_auth_returns_what_the_custom_function_resolved( + #[case] answer: CustomAuthResult, + ) { + let facts = RequestFacts { + route: "/chat/completions".to_owned(), + token: Some(SecretValue::from("tok-1")), + is_llm_api_route: true, }; let outcome = NativeCore::default() .custom() - .authenticate( - &ScriptedCustomAuthHost(identity.clone()), - &facts(Some("tok-1")), - ) + .authenticate(&ScriptedCustomAuth(answer.clone()), &facts) .await; - assert_eq!(outcome, Ok(identity)); + assert_eq!(outcome, Ok(answer)); } } diff --git a/litellm-rust/crates/gateway-auth-core/src/outcome.rs b/litellm-rust/crates/gateway-auth-core/src/outcome.rs index 1f66de52d03..e3a485461bc 100644 --- a/litellm-rust/crates/gateway-auth-core/src/outcome.rs +++ b/litellm-rust/crates/gateway-auth-core/src/outcome.rs @@ -2,6 +2,7 @@ use std::time::SystemTime; use litellm_secret_value::SecretValue; use serde::{Deserialize, Serialize}; +use serde_json::Value; #[derive(Clone, Debug, PartialEq)] pub enum Grant { @@ -67,7 +68,11 @@ pub type ModelBudgetResult = Result; pub enum AnonymousReason { PublicRoute, PassThroughAuthOff, - NoMasterKey, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Anonymous { + pub reason: AnonymousReason, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -82,9 +87,20 @@ pub struct JwtIdentity { #[derive(Clone, Debug, PartialEq, Eq, Serialize)] pub struct OAuth2Identity { + pub token: String, + pub user_id: Value, + pub role: Value, + pub team_id: Value, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)] +pub struct OAuth2ProxyIdentity { pub user_id: Option, + pub user_email: Option, pub team_id: Option, - pub scopes: Vec, + pub team_alias: Option, + pub org_id: Option, + pub models: Option>, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -94,14 +110,37 @@ pub struct CustomIdentity { pub hashed_token: Option, } +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum CustomAuthResult { + Identity(CustomIdentity), + ReplacedKey { token: String }, + NoIdentity, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub enum Identity { MasterKey, Key { hashed_token: String }, Jwt(JwtIdentity), OAuth2(OAuth2Identity), + OAuth2Proxy(OAuth2ProxyIdentity), Custom(CustomIdentity), - Anonymous(AnonymousReason), + Anonymous(Anonymous), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] +pub enum OAuth2Denial { + InactiveToken, + IntrospectionNotConfigured, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum OAuth2ProxyDenial { + TrustedProxyRangesMissing, + UntrustedProxy { client_ip: Option }, + ProxyMappingsMissing, + DisallowedProxyMappings { fields: Vec }, } #[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] @@ -110,17 +149,28 @@ pub enum OAuth2Mode { Proxy, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Validator { - Key, +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] +pub enum PremiumFeature { + OAuth2, Jwt, - OAuth2(OAuth2Mode), - Custom, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum KeyPlan { + Key { pass_through_key: bool }, Anonymous(AnonymousReason), } -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] -pub enum OAuth2Denial { - MissingKey, - InactiveToken, +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CredentialPlan { + Key(KeyPlan), + Jwt, + OAuth2(OAuth2Mode), + PremiumRequired(PremiumFeature), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Plan { + Credential(CredentialPlan), + Custom, } diff --git a/litellm-rust/crates/python-bridge/src/auth.rs b/litellm-rust/crates/python-bridge/src/auth.rs index 4f02a05eb86..00b856bd035 100644 --- a/litellm-rust/crates/python-bridge/src/auth.rs +++ b/litellm-rust/crates/python-bridge/src/auth.rs @@ -5,7 +5,6 @@ //! facet future stays pending until Python answers it. use std::{ - collections::BTreeMap, future::{Future, poll_fn}, pin::Pin, sync::{Arc, Mutex, MutexGuard, PoisonError}, @@ -14,8 +13,9 @@ use std::{ use litellm_error::Report; use litellm_gateway_auth_core::{ - CustomAuthCore, CustomAuthHost, CustomIdentity, Introspection, NativeCore, OAuth2Core, - OAuth2Denial, OAuth2Host, OAuth2Identity, OAuth2Settings, RequestFacts, TargetAuthCore, + AuthCore, CustomAuthCore, CustomAuthHost, CustomAuthResult, Introspection, NativeCore, + OAuth2Core, OAuth2Denial, OAuth2Identity, OAuth2IntrospectionHost, OAuth2IntrospectionSettings, + RequestFacts, }; use litellm_host_python::{from_py_argument, to_py}; use pyo3::{exceptions::PyValueError, prelude::*}; @@ -33,7 +33,7 @@ enum HostCall { struct Exchange { call: Option, introspection: Option, - custom_identity: Option, + custom_auth: Option, } type Shared = Arc>; @@ -46,14 +46,12 @@ type Slot = fn(&mut Exchange) -> &mut Option; struct StepHost { exchange: Shared, - proxy_headers: BTreeMap, } impl StepHost { - fn new(exchange: &Shared, proxy_headers: BTreeMap) -> Self { + fn new(exchange: &Shared) -> Self { Self { exchange: exchange.clone(), - proxy_headers, } } @@ -73,7 +71,7 @@ impl StepHost { } } -impl OAuth2Host for StepHost { +impl OAuth2IntrospectionHost for StepHost { type Error = std::convert::Infallible; async fn introspect(&self, token: &str) -> Result { @@ -82,20 +80,14 @@ impl OAuth2Host for StepHost { }; Ok(self.ask(call, |exchange| &mut exchange.introspection).await) } - - fn proxy_headers(&self) -> &BTreeMap { - &self.proxy_headers - } } impl CustomAuthHost for StepHost { type Error = std::convert::Infallible; - async fn custom_auth(&self, _facts: &RequestFacts) -> Result { + async fn custom_auth(&self, _facts: &RequestFacts) -> Result { Ok(self - .ask(HostCall::CustomAuth, |exchange| { - &mut exchange.custom_identity - }) + .ask(HostCall::CustomAuth, |exchange| &mut exchange.custom_auth) .await) } } @@ -103,17 +95,17 @@ impl CustomAuthHost for StepHost { #[derive(Debug, PartialEq, Eq, Serialize)] #[serde(tag = "kind", rename_all = "snake_case")] enum OAuth2Verdict { - Identity(OAuth2Identity), - MissingKey, + Identity(Box), InactiveToken, + IntrospectionNotConfigured, } impl From> for OAuth2Verdict { fn from(outcome: Result) -> Self { match outcome { - Ok(identity) => Self::Identity(identity), - Err(OAuth2Denial::MissingKey) => Self::MissingKey, + Ok(identity) => Self::Identity(Box::new(identity)), Err(OAuth2Denial::InactiveToken) => Self::InactiveToken, + Err(OAuth2Denial::IntrospectionNotConfigured) => Self::IntrospectionNotConfigured, } } } @@ -221,24 +213,23 @@ impl Driver { fn oauth2_driver( core: NativeCore, - facts: RequestFacts, - settings: OAuth2Settings, - proxy_headers: BTreeMap, + token: String, + settings: OAuth2IntrospectionSettings, ) -> Driver { let exchange = Shared::default(); - let host = StepHost::new(&exchange, proxy_headers); + let host = StepHost::new(&exchange); Driver::new(exchange, async move { - let Ok(outcome) = core.oauth2().authenticate(&host, &facts, &settings).await; + let Ok(outcome) = core.oauth2().introspect(&host, &token, &settings).await; OAuth2Verdict::from(outcome) }) } -fn custom_driver(core: NativeCore, facts: RequestFacts) -> Driver { +fn custom_driver(core: NativeCore, facts: RequestFacts) -> Driver { let exchange = Shared::default(); - let host = StepHost::new(&exchange, BTreeMap::new()); + let host = StepHost::new(&exchange); Driver::new(exchange, async move { - let Ok(identity) = core.custom().authenticate(&host, &facts).await; - identity + let Ok(result) = core.custom().authenticate(&host, &facts).await; + result }) } @@ -268,17 +259,11 @@ impl AuthCoreHandle { } } - fn oauth2( - &self, - facts: &Bound<'_, PyAny>, - settings: &Bound<'_, PyAny>, - proxy_headers: &Bound<'_, PyAny>, - ) -> PyResult { + fn oauth2(&self, token: String, settings: &Bound<'_, PyAny>) -> PyResult { Ok(OAuth2Call(Mutex::new(oauth2_driver( self.core, - from_py_argument(facts)?, + token, from_py_argument(settings)?, - from_py_argument(proxy_headers)?, )))) } @@ -319,7 +304,7 @@ impl OAuth2Call { /// One custom auth call. #[pyclass(frozen)] -pub(crate) struct CustomAuthCall(Mutex>); +pub(crate) struct CustomAuthCall(Mutex>); #[pymethods] impl CustomAuthCall { @@ -327,18 +312,14 @@ impl CustomAuthCall { step_to_py(py, lock_driver(&self.0).start()) } - fn answer_custom_auth( - &self, - py: Python<'_>, - identity: &Bound<'_, PyAny>, - ) -> PyResult> { - let identity = from_py_argument(identity)?; + fn answer_custom_auth(&self, py: Python<'_>, result: &Bound<'_, PyAny>) -> PyResult> { + let result = from_py_argument(result)?; step_to_py( py, lock_driver(&self.0).answer( "custom_auth", - |exchange| &mut exchange.custom_identity, - identity, + |exchange| &mut exchange.custom_auth, + result, ), ) } @@ -346,39 +327,35 @@ impl CustomAuthCall { #[cfg(test)] mod tests { + use litellm_gateway_auth_core::CustomIdentity; use litellm_secret_value::SecretValue; use rstest::rstest; use serde_json::{Value, json}; use super::*; - fn facts(token: Option<&str>) -> RequestFacts { + fn facts() -> RequestFacts { RequestFacts { route: "/chat/completions".to_owned(), - token: token.map(SecretValue::from), + token: Some(SecretValue::from("tok-1")), is_llm_api_route: true, } } - fn oauth2(token: Option<&str>) -> Driver { - let settings = OAuth2Settings { - user_id_header: None, - team_id_header: None, + fn oauth2(endpoint_configured: bool) -> Driver { + let settings = OAuth2IntrospectionSettings { + endpoint_configured, + requires_active: true, }; - oauth2_driver( - NativeCore::default(), - facts(token), - settings, - BTreeMap::new(), - ) + oauth2_driver(NativeCore::default(), "tok-1".to_owned(), settings) } fn introspection(active: bool) -> Introspection { Introspection { active, - subject: Some("user-1".to_owned()), - client_id: None, - scopes: vec!["read".to_owned()], + user_id: json!("user-1"), + role: json!("internal_user"), + team_id: Value::Null, } } @@ -393,69 +370,78 @@ mod tests { ) } + fn answer_custom_auth( + driver: &mut Driver, + result: CustomAuthResult, + ) -> Result, DriveError> { + driver.answer("custom_auth", |exchange| &mut exchange.custom_auth, result) + } + fn wire(step: &Step) -> Value { serde_json::to_value(step).unwrap() } #[test] fn oauth2_asks_python_to_introspect_the_presented_token() { - let step = oauth2(Some("tok-1")).start().unwrap(); + let step = oauth2(true).start().unwrap(); assert_eq!(wire(&step), json!({"kind": "introspect", "token": "tok-1"})); } #[rstest] #[case::active( true, - json!({"kind": "done", "outcome": {"kind": "identity", "user_id": "user-1", "team_id": null, "scopes": ["read"]}}) + json!({"kind": "done", "outcome": {"kind": "identity", "token": "tok-1", "user_id": "user-1", "role": "internal_user", "team_id": null}}) )] #[case::inactive(false, json!({"kind": "done", "outcome": {"kind": "inactive_token"}}))] fn oauth2_finishes_with_the_verdict_on_python_s_introspection( #[case] active: bool, #[case] expected: Value, ) { - let mut driver = oauth2(Some("tok-1")); + let mut driver = oauth2(true); driver.start().unwrap(); let step = answer_introspection(&mut driver, introspection(active)).unwrap(); assert_eq!(wire(&step), expected); } #[test] - fn oauth2_without_a_token_finishes_without_a_host_call() { - let step = oauth2(None).start().unwrap(); + fn oauth2_without_an_endpoint_finishes_without_a_host_call() { + let step = oauth2(false).start().unwrap(); assert_eq!( wire(&step), - json!({"kind": "done", "outcome": {"kind": "missing_key"}}) + json!({"kind": "done", "outcome": {"kind": "introspection_not_configured"}}) ); } - #[test] - fn custom_auth_finishes_with_the_identity_python_resolved() { - let identity = CustomIdentity { + #[rstest] + #[case::identity( + CustomAuthResult::Identity(CustomIdentity { user_id: Some("user-1".to_owned()), team_id: Some("team-1".to_owned()), hashed_token: None, - }; - let mut driver = custom_driver(NativeCore::default(), facts(Some("tok-1"))); + }), + json!({"kind": "identity", "user_id": "user-1", "team_id": "team-1", "hashed_token": null}) + )] + #[case::replaced_key( + CustomAuthResult::ReplacedKey { token: "sk-new".to_owned() }, + json!({"kind": "replaced_key", "token": "sk-new"}) + )] + #[case::no_identity(CustomAuthResult::NoIdentity, json!({"kind": "no_identity"}))] + fn custom_auth_finishes_with_what_python_resolved( + #[case] result: CustomAuthResult, + #[case] outcome: Value, + ) { + let mut driver = custom_driver(NativeCore::default(), facts()); assert_eq!( wire(&driver.start().unwrap()), json!({"kind": "custom_auth"}) ); - let step = driver - .answer( - "custom_auth", - |exchange| &mut exchange.custom_identity, - identity, - ) - .unwrap(); - assert_eq!( - wire(&step), - json!({"kind": "done", "outcome": {"user_id": "user-1", "team_id": "team-1", "hashed_token": null}}) - ); + let step = answer_custom_auth(&mut driver, result).unwrap(); + assert_eq!(wire(&step), json!({"kind": "done", "outcome": outcome})); } #[test] fn an_answer_before_start_is_refused() { - let mut driver = oauth2(Some("tok-1")); + let mut driver = oauth2(true); assert_eq!( answer_introspection(&mut driver, introspection(true)), Err(DriveError::NotStarted) @@ -464,25 +450,25 @@ mod tests { #[test] fn a_second_start_is_refused() { - let mut driver = oauth2(Some("tok-1")); + let mut driver = oauth2(true); driver.start().unwrap(); assert_eq!(driver.start(), Err(DriveError::Started)); } #[test] fn an_answer_to_a_call_that_was_not_made_is_refused() { - let mut driver = oauth2(Some("tok-1")); - driver.start().unwrap(); - let identity = CustomIdentity { - user_id: None, - team_id: None, - hashed_token: None, - }; + let mut driver = custom_driver(NativeCore::default(), facts()); assert_eq!( - driver.answer( + driver.start().map(|step| wire(&step)), + Ok(json!({"kind": "custom_auth"})) + ); + let mut oauth2 = oauth2(true); + oauth2.start().unwrap(); + assert_eq!( + oauth2.answer( "custom_auth", - |exchange| &mut exchange.custom_identity, - identity + |exchange| &mut exchange.custom_auth, + CustomAuthResult::NoIdentity ), Err(DriveError::NotAsked("custom_auth")) ); @@ -490,7 +476,7 @@ mod tests { #[test] fn an_answer_after_the_call_finished_is_refused() { - let mut driver = oauth2(None); + let mut driver = oauth2(false); driver.start().unwrap(); assert_eq!( answer_introspection(&mut driver, introspection(true)), diff --git a/litellm/proxy/auth/core/AGENTS.md b/litellm/proxy/auth/core/AGENTS.md index 9c3437b0019..7935dae5fea 100644 --- a/litellm/proxy/auth/core/AGENTS.md +++ b/litellm/proxy/auth/core/AGENTS.md @@ -10,19 +10,19 @@ Modules here import only the standard library and each other. `tests/code_covera `core.py` splits the decision into facets, each a protocol with its own host. A core hands each facet out through an accessor, the way `RouterCore` hands out `registry()`. Nothing inherits from another facet. -`AuthCore` lists the facets a core decides today. `keys()` returns the `KeyCore` that authenticates a key credential (master key, virtual key, UI session, cached admin key) and re-checks a key the proxy resolved itself. `model_budgets()` returns the `ModelBudgetCore` that checks per-model budgets for the key and user, the user alone, or a team member. +`AuthCore` has one accessor per facet. `plan()` picks the credential path a request takes, in the proxy's order: custom auth, public routes, OAuth2 introspection, OAuth2 proxy headers, JWT, unauthenticated pass-through routes, then keys. A path that can hand the request on (custom auth issuing a key, a JWT mapped to a virtual key) resumes planning through `after_custom_auth()` or `after_jwt()`, whose narrower return types exclude the paths already taken. `anonymous()`, `oauth2()`, `custom()` and `jwt()` turn a credential into an identity. `keys()` authenticates a key credential (master key, virtual key, UI session, cached admin key) and re-checks a key the proxy resolved itself. `model_budgets()` checks per-model budgets for the key and user, the user alone, or a team member. `authorize()` decides whether an identity may make the request. -`TargetAuthCore` extends `AuthCore` with the facets that still run in the proxy: `plan()` picks the credential path, `anonymous()`, `jwt()`, `oauth2()` and `custom()` turn a credential into an `Identity`, and `authorize()` decides whether that identity may make the request. `PythonCore` implements `AuthCore`. `RustCore` declares `TargetAuthCore` with every facet stubbed. The signatures of the target facets, their hosts, settings and identities are provisional and change when the facet moves into the core. +`litellm/proxy/auth/selected_core.py` picks the core the proxy runs, `PythonCore`. `RustCore` declares every accessor and implements none. `PythonCore` does not implement `jwt()` or `authorize()` yet, so their signatures are provisional and the proxy still runs that logic itself. -Each facet returns only its own results. `KeyCore` returns an `Outcome` (`Allow` with a key `Grant`, or `Deny` with a `Denial`), and `ModelBudgetCore` returns `WithinModelBudgets` or `ModelBudgetExceeded`. The proxy converts each one separately. +Each facet returns only its own results, and the proxy converts each one separately. `KeyCore` returns an `Outcome` (`Allow` with a key `Grant`, or `Deny` with a `Denial`), `ModelBudgetCore` returns `WithinModelBudgets` or `ModelBudgetExceeded`, and the planner returns a validator or `PremiumRequired`. ## Hosts carry the IO -`host.py` lists every read, write and alert a facet may ask for, as typed async methods, one host per facet. A host performs IO and returns data. Anything that compares, chooses or denies belongs in the core. When a decision needs data the core doesn't have, add a host method that returns that data instead of letting the host decide. When one check runs another, its host hands out the other host through an accessor, as `KeyHost.model_budgets()` does. +`host.py` lists every read, write and alert a facet may ask for, as typed methods, one host per kind of IO a facet needs (OAuth2 introspection calls out over HTTP while proxy-header mode only reads the request, so each has its own host). A host performs IO and returns data. Anything that compares, chooses or denies belongs in the core. When a decision needs data the core doesn't have, add a host method that returns that data instead of letting the host decide. When one check runs another, its host hands out the other host through an accessor, as `KeyHost.model_budgets()` does. ## Model the data, not flags -The values that cross the boundary live in one module per direction. `inputs.py` holds what a caller passes in (`RequestFacts` and one settings type per facet), `outcome.py` holds what a facet returns (grants, denials, identities and validators), and `answers.py` holds what host methods return. Each is a frozen dataclass, and each choice is a union of them. Express "absent" as `None` on a typed value (a key without a team membership has `team_member=None`) instead of a boolean next to optional fields. +The values that cross the boundary live in one module per direction. `inputs.py` holds what a caller passes in (`RequestFacts`, `PlanFacts` and the settings each facet reads), `outcome.py` holds what a facet returns (grants, denials, identities and validators), and `answers.py` holds what host methods return. Each is a frozen dataclass, and each choice is a union of them. Express "absent" as `None` on a typed value (a key without a team membership has `team_member=None`) instead of a boolean next to optional fields. ## Order is behavior @@ -34,7 +34,7 @@ Nothing here knows about `UserAPIKeyAuth`, HTTP status codes or the proxy's exce ## Credentials -`credential.py` picks the key from the presented headers, query and WebSocket values and owns the precedence between those sources. `key.py` hashes tokens and compares the master key in constant time. Code that needs either goes through these two modules instead of parsing headers or comparing keys itself. +`credential.py` picks the key from the presented headers, query and WebSocket values, owns the precedence between those sources and tells a JWT from a key. `key.py` hashes tokens and compares the master key in constant time. `network.py` decides whether a peer address falls inside the trusted proxy ranges. Code that needs either goes through these two modules instead of parsing headers or comparing keys itself. ## Tests diff --git a/litellm/proxy/auth/core/answers.py b/litellm/proxy/auth/core/answers.py index ffb6019a5b9..cc203ca5d45 100644 --- a/litellm/proxy/auth/core/answers.py +++ b/litellm/proxy/auth/core/answers.py @@ -85,9 +85,9 @@ class SigningKey: @dataclass(frozen=True, slots=True) class Introspection: active: bool - subject: str | None - client_id: str | None - scopes: tuple[str, ...] + user_id: object + role: object + team_id: object @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/auth/core/core.py b/litellm/proxy/auth/core/core.py index cd6a682366f..b127e80ec3c 100644 --- a/litellm/proxy/auth/core/core.py +++ b/litellm/proxy/auth/core/core.py @@ -7,27 +7,35 @@ from litellm.proxy.auth.core.host import ( JwtHost, KeyHost, ModelBudgetHost, - OAuth2Host, + OAuth2IntrospectionHost, + OAuth2ProxyHost, TeamMemberModelBudgetHost, ) from litellm.proxy.auth.core.inputs import ( AuthorizeSettings, JwtSettings, KeySettings, - OAuth2Settings, + OAuth2IntrospectionSettings, + OAuth2ProxySettings, + PlanFacts, PlanSettings, RequestFacts, ) from litellm.proxy.auth.core.outcome import ( Anonymous, - CustomIdentity, + AnonymousReason, + CredentialPlan, + CustomAuthResult, Denial, Identity, + KeyPlan, ModelBudgetResult, OAuth2Denial, OAuth2Identity, + OAuth2ProxyDenial, + OAuth2ProxyIdentity, Outcome, - Validator, + Plan, ) @@ -48,11 +56,17 @@ class ModelBudgetCore(Protocol): class AuthPlanner(Protocol): - def plan(self, facts: RequestFacts, settings: PlanSettings) -> Validator: ... + def plan(self, facts: PlanFacts, settings: PlanSettings) -> Plan: ... + + def after_custom_auth(self, facts: PlanFacts, settings: PlanSettings, *, issued_key: bool) -> CredentialPlan: ... + + def after_jwt(self, facts: PlanFacts, settings: PlanSettings, *, issued_key: bool) -> KeyPlan: ... + + def oauth2_introspection_applies(self, facts: PlanFacts, settings: PlanSettings) -> bool: ... class AnonymousCore(Protocol): - def authenticate(self, facts: RequestFacts, settings: PlanSettings) -> Anonymous: ... + def authenticate(self, reason: AnonymousReason) -> Anonymous: ... class JwtCore(Protocol): @@ -62,13 +76,17 @@ class JwtCore(Protocol): class OAuth2Core(Protocol): - async def authenticate( - self, host: OAuth2Host, facts: RequestFacts, settings: OAuth2Settings + async def introspect( + self, host: OAuth2IntrospectionHost, token: str, settings: OAuth2IntrospectionSettings ) -> OAuth2Identity | OAuth2Denial: ... + def proxy( + self, host: OAuth2ProxyHost, settings: OAuth2ProxySettings + ) -> OAuth2ProxyIdentity | OAuth2ProxyDenial: ... + class CustomAuthCore(Protocol): - async def authenticate(self, host: CustomAuthHost, facts: RequestFacts) -> CustomIdentity: ... + async def authenticate(self, host: CustomAuthHost, facts: RequestFacts) -> CustomAuthResult: ... class AuthorizeCore(Protocol): @@ -82,8 +100,6 @@ class AuthCore(Protocol): def model_budgets(self) -> ModelBudgetCore: ... - -class TargetAuthCore(AuthCore, Protocol): def plan(self) -> AuthPlanner: ... def anonymous(self) -> AnonymousCore: ... diff --git a/litellm/proxy/auth/core/credential.py b/litellm/proxy/auth/core/credential.py index 105eb4afe73..859af3692af 100644 --- a/litellm/proxy/auth/core/credential.py +++ b/litellm/proxy/auth/core/credential.py @@ -120,6 +120,10 @@ def select_api_key(keys: PresentedKeys) -> SelectedKey | None: return _raw_key(keys) +def is_jwt(token: str | None) -> bool: + return token is not None and len(token.split(".")) == 3 + + def _subprotocol_api_key(subprotocols: str | None) -> str | None: protocols: Final = tuple(protocol.strip() for protocol in (subprotocols or "").split(",")) return next( diff --git a/litellm/proxy/auth/core/host.py b/litellm/proxy/auth/core/host.py index 6336f114d90..da602924615 100644 --- a/litellm/proxy/auth/core/host.py +++ b/litellm/proxy/auth/core/host.py @@ -1,4 +1,3 @@ -from collections.abc import Mapping from dataclasses import dataclass from typing import Protocol, TypeAlias @@ -17,7 +16,7 @@ from litellm.proxy.auth.core.answers import ( TeamMemberModelBudgetScope, ) from litellm.proxy.auth.core.inputs import RequestFacts -from litellm.proxy.auth.core.outcome import CustomIdentity, Identity, JwtIdentity, KeyIdentity, ModelBudgetExceeded +from litellm.proxy.auth.core.outcome import CustomAuthResult, Identity, JwtIdentity, KeyIdentity, ModelBudgetExceeded class ModelBudgetHost(Protocol): @@ -113,14 +112,18 @@ class JwtHost(Protocol): async def mapped_key(self, identity: JwtIdentity) -> KeyIdentity | None: ... -class OAuth2Host(Protocol): +class OAuth2IntrospectionHost(Protocol): async def introspect(self, token: str) -> Introspection: ... - def proxy_headers(self) -> Mapping[str, str]: ... + +class OAuth2ProxyHost(Protocol): + def client_ip(self) -> str | None: ... + + def header(self, name: str) -> str | None: ... class CustomAuthHost(Protocol): - async def custom_auth(self, facts: RequestFacts) -> CustomIdentity: ... + async def custom_auth(self, facts: RequestFacts) -> CustomAuthResult: ... class AuthorizeHost(Protocol): diff --git a/litellm/proxy/auth/core/inputs.py b/litellm/proxy/auth/core/inputs.py index 05cfa43b9f0..77f3c5c68c1 100644 --- a/litellm/proxy/auth/core/inputs.py +++ b/litellm/proxy/auth/core/inputs.py @@ -1,4 +1,7 @@ +from collections.abc import Mapping from dataclasses import dataclass +from ipaddress import IPv4Network, IPv6Network +from typing import TypeAlias @dataclass(frozen=True, slots=True) @@ -8,13 +11,33 @@ class RequestFacts: is_llm_api_route: bool +@dataclass(frozen=True, slots=True) +class PlanFacts: + request: RequestFacts + is_info_route: bool + unverified_claims: Mapping[str, object] | None + + +@dataclass(frozen=True, slots=True) +class RoutingOverride: + iss: tuple[str, ...] | None + client_id: tuple[str, ...] | None + scope: tuple[str, ...] | None + aud: tuple[str, ...] | None + + @dataclass(frozen=True, slots=True) class PlanSettings: custom_auth: bool - jwt_enabled: bool + premium_user: bool + public_routes: frozenset[str] + additional_public_routes: tuple[str, ...] + metrics_require_auth: bool oauth2_introspection: bool oauth2_proxy: bool - public_routes: frozenset[str] + jwt_enabled: bool + oauth2_routing_overrides: tuple[RoutingOverride, ...] + unauthenticated_pass_through_routes: frozenset[str] @dataclass(frozen=True, slots=True) @@ -33,9 +56,18 @@ class JwtSettings: @dataclass(frozen=True, slots=True) -class OAuth2Settings: - user_id_header: str | None - team_id_header: str | None +class OAuth2IntrospectionSettings: + endpoint_configured: bool + requires_active: bool + + +TrustedNetwork: TypeAlias = IPv4Network | IPv6Network + + +@dataclass(frozen=True, slots=True) +class OAuth2ProxySettings: + trusted_networks: tuple[TrustedNetwork, ...] + header_mappings: Mapping[str, str] @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/auth/core/network.py b/litellm/proxy/auth/core/network.py new file mode 100644 index 00000000000..02af95fd356 --- /dev/null +++ b/litellm/proxy/auth/core/network.py @@ -0,0 +1,24 @@ +from collections.abc import Sequence +from ipaddress import IPv4Address, IPv6Address, ip_address +from typing import Final + +from litellm.proxy.auth.core.inputs import TrustedNetwork + + +def _unmapped(address: IPv4Address | IPv6Address) -> IPv4Address | IPv6Address: + if isinstance(address, IPv6Address) and address.ipv4_mapped is not None: + return address.ipv4_mapped + return address + + +def ip_in_networks(client_ip: str | None, networks: Sequence[TrustedNetwork]) -> bool: + if not client_ip or not networks: + return False + try: + address: Final = ip_address(client_ip.strip()) + except ValueError: + return False + candidates: Final = (address, _unmapped(address)) + return any( + candidate in network for network in networks for candidate in candidates + ) # comprehension-ok: two fixed candidates per network diff --git a/litellm/proxy/auth/core/oauth2.py b/litellm/proxy/auth/core/oauth2.py new file mode 100644 index 00000000000..fdc4ed39e17 --- /dev/null +++ b/litellm/proxy/auth/core/oauth2.py @@ -0,0 +1,63 @@ +from typing import Final + +from litellm.proxy.auth.core.host import OAuth2IntrospectionHost, OAuth2ProxyHost +from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings, OAuth2ProxySettings +from litellm.proxy.auth.core.network import ip_in_networks +from litellm.proxy.auth.core.outcome import ( + DisallowedProxyMappings, + InactiveToken, + IntrospectionNotConfigured, + OAuth2Denial, + OAuth2Identity, + OAuth2ProxyDenial, + OAuth2ProxyIdentity, + ProxyMappingsMissing, + TrustedProxyRangesMissing, + UntrustedProxy, +) + +# Header trust asserts identity only. Mapping any other UserAPIKeyAuth field (role, budgets, limits, permissions) +# would let a caller forge its own enforcement by sending the matching header, and an allowlist stays safe as +# fields are added. A trusted upstream that must assert more should use JWT auth, which checks a signature. +ALLOWED_OAUTH2_PROXY_FIELDS: Final = frozenset({"user_id", "user_email", "team_id", "team_alias", "org_id", "models"}) + + +async def introspect( + host: OAuth2IntrospectionHost, token: str, settings: OAuth2IntrospectionSettings +) -> OAuth2Identity | OAuth2Denial: + if not settings.endpoint_configured: + return IntrospectionNotConfigured() + answer: Final = await host.introspect(token) + if settings.requires_active and not answer.active: + return InactiveToken() + return OAuth2Identity(token=token, user_id=answer.user_id, role=answer.role, team_id=answer.team_id) + + +def _proxy_identity(host: OAuth2ProxyHost, settings: OAuth2ProxySettings) -> OAuth2ProxyIdentity: + def value(field: str) -> str | None: + header: Final = settings.header_mappings.get(field) + return (host.header(header) or None) if header is not None else None + + models: Final = value("models") + return OAuth2ProxyIdentity( + user_id=value("user_id"), + user_email=value("user_email"), + team_id=value("team_id"), + team_alias=value("team_alias"), + org_id=value("org_id"), + models=None if models is None else tuple(model.strip() for model in models.split(",")), + ) + + +def proxy(host: OAuth2ProxyHost, settings: OAuth2ProxySettings) -> OAuth2ProxyIdentity | OAuth2ProxyDenial: + if not settings.trusted_networks: + return TrustedProxyRangesMissing() + client_ip: Final = host.client_ip() + if not ip_in_networks(client_ip, settings.trusted_networks): + return UntrustedProxy(client_ip=client_ip) + if not settings.header_mappings: + return ProxyMappingsMissing() + disallowed: Final = tuple(sorted(settings.header_mappings.keys() - ALLOWED_OAUTH2_PROXY_FIELDS)) + if disallowed: + return DisallowedProxyMappings(fields=disallowed) + return _proxy_identity(host, settings) diff --git a/litellm/proxy/auth/core/outcome.py b/litellm/proxy/auth/core/outcome.py index 36a94d58f24..1f81530ff4d 100644 --- a/litellm/proxy/auth/core/outcome.py +++ b/litellm/proxy/auth/core/outcome.py @@ -147,7 +147,6 @@ ModelBudgetResult: TypeAlias = WithinModelBudgets | ModelBudgetExceeded class AnonymousReason(Enum): PUBLIC_ROUTE = "public_route" PASS_THROUGH_AUTH_OFF = "pass_through_auth_off" - NO_MASTER_KEY = "no_master_key" @dataclass(frozen=True, slots=True) @@ -172,9 +171,20 @@ class JwtIdentity: @dataclass(frozen=True, slots=True) class OAuth2Identity: + token: str + user_id: object + role: object + team_id: object + + +@dataclass(frozen=True, slots=True) +class OAuth2ProxyIdentity: user_id: str | None + user_email: str | None team_id: str | None - scopes: tuple[str, ...] + team_alias: str | None + org_id: str | None + models: tuple[str, ...] | None @dataclass(frozen=True, slots=True) @@ -184,6 +194,19 @@ class CustomIdentity: hashed_token: str | None +@dataclass(frozen=True, slots=True) +class ReplacedKey: + token: str + + +@dataclass(frozen=True, slots=True) +class NoCustomIdentity: + pass + + +CustomAuthResult: TypeAlias = CustomIdentity | ReplacedKey | NoCustomIdentity + + @dataclass(frozen=True, slots=True) class Anonymous: reason: AnonymousReason @@ -194,10 +217,42 @@ class InactiveToken: pass -OAuth2Denial: TypeAlias = MissingKey | InactiveToken +@dataclass(frozen=True, slots=True) +class IntrospectionNotConfigured: + pass -Identity: TypeAlias = MasterKeyIdentity | KeyIdentity | JwtIdentity | OAuth2Identity | CustomIdentity | Anonymous +OAuth2Denial: TypeAlias = InactiveToken | IntrospectionNotConfigured + + +@dataclass(frozen=True, slots=True) +class TrustedProxyRangesMissing: + pass + + +@dataclass(frozen=True, slots=True) +class UntrustedProxy: + client_ip: str | None + + +@dataclass(frozen=True, slots=True) +class ProxyMappingsMissing: + pass + + +@dataclass(frozen=True, slots=True) +class DisallowedProxyMappings: + fields: tuple[str, ...] + + +OAuth2ProxyDenial: TypeAlias = ( + TrustedProxyRangesMissing | UntrustedProxy | ProxyMappingsMissing | DisallowedProxyMappings +) + + +Identity: TypeAlias = ( + MasterKeyIdentity | KeyIdentity | JwtIdentity | OAuth2Identity | OAuth2ProxyIdentity | CustomIdentity | Anonymous +) class OAuth2Mode(Enum): @@ -207,7 +262,7 @@ class OAuth2Mode(Enum): @dataclass(frozen=True, slots=True) class KeyValidator: - pass + pass_through_key: bool @dataclass(frozen=True, slots=True) @@ -230,4 +285,18 @@ class AnonymousValidator: reason: AnonymousReason -Validator: TypeAlias = KeyValidator | JwtValidator | OAuth2Validator | CustomValidator | AnonymousValidator +class PremiumFeature(Enum): + OAUTH2 = "oauth2" + JWT = "jwt" + + +@dataclass(frozen=True, slots=True) +class PremiumRequired: + feature: PremiumFeature + + +KeyPlan: TypeAlias = KeyValidator | AnonymousValidator + +CredentialPlan: TypeAlias = KeyPlan | JwtValidator | OAuth2Validator | PremiumRequired + +Plan: TypeAlias = CredentialPlan | CustomValidator diff --git a/litellm/proxy/auth/core/plan.py b/litellm/proxy/auth/core/plan.py new file mode 100644 index 00000000000..92f763afcf5 --- /dev/null +++ b/litellm/proxy/auth/core/plan.py @@ -0,0 +1,132 @@ +from collections.abc import Iterator +from fnmatch import fnmatchcase +from typing import Final + +from litellm.proxy.auth.core.credential import is_jwt +from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RoutingOverride +from litellm.proxy.auth.core.outcome import ( + AnonymousReason, + AnonymousValidator, + CredentialPlan, + CustomValidator, + JwtValidator, + KeyPlan, + KeyValidator, + OAuth2Mode, + OAuth2Validator, + Plan, + PremiumFeature, + PremiumRequired, +) + +_METRICS_ROUTE: Final = "/metrics" +_LENS_ROUTE: Final = "/lens" + + +def _normalized(route: str) -> str: + return route.rstrip("/") if route != "/" and route.endswith("/") else route + + +def _matches_public_pattern(route: str, pattern: str) -> bool: + return route.startswith(pattern[:-1]) if pattern.endswith("*") else route == pattern + + +def is_public_route(route: str, settings: PlanSettings) -> bool: + if _normalized(route) == _METRICS_ROUTE and settings.metrics_require_auth: + return False + if route in settings.public_routes: + return True + return settings.premium_user and any( + _matches_public_pattern(route, pattern) for pattern in settings.additional_public_routes + ) + + +def _claim_values(claim: object, *, split_scope: bool) -> tuple[str, ...]: + if isinstance(claim, list): + return tuple(str(value) for value in claim) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # JWT claims are untyped JSON + if split_scope and isinstance(claim, str) and " " in claim.strip(): + return tuple(value for value in claim.strip().split(" ") if value) + return (str(claim),) + + +def _selector_matches(selector: str, claim: str, *, split_scope: bool) -> bool: + if "*" not in selector and "?" not in selector: + return selector == claim + if not split_scope and any(character.isspace() for character in claim): + return False + return fnmatchcase(claim, selector) + + +def _pairs(selectors: tuple[str, ...], claims: tuple[str, ...]) -> Iterator[tuple[str, str]]: + for selector in selectors: + yield from ((selector, claim) for claim in claims) + + +def _selectors_match(selectors: tuple[str, ...] | None, claim: object, *, split_scope: bool = False) -> bool: + if selectors is None: + return True + if claim is None: + return False + claims: Final = _claim_values(claim, split_scope=split_scope) + return any( + _selector_matches(selector, value, split_scope=split_scope) for selector, value in _pairs(selectors, claims) + ) + + +def _override_matches(override: RoutingOverride, facts: PlanFacts) -> bool: + token_claims: Final = facts.unverified_claims or {} + return ( + _selectors_match(override.iss, token_claims.get("iss")) + and _selectors_match(override.client_id, token_claims.get("client_id")) + and _selectors_match(override.scope, token_claims.get("scope"), split_scope=True) + and _selectors_match(override.aud, token_claims.get("aud")) + ) + + +def _routes_jwt_to_oauth2(facts: PlanFacts, settings: PlanSettings) -> bool: + if facts.unverified_claims is None: + return False + return any(_override_matches(override, facts) for override in settings.oauth2_routing_overrides) + + +def _is_lens_route(route: str) -> bool: + return route == _LENS_ROUTE or route.startswith(f"{_LENS_ROUTE}/") + + +def oauth2_introspection_applies(facts: PlanFacts, settings: PlanSettings) -> bool: + request: Final = facts.request + token_is_jwt: Final = settings.jwt_enabled and is_jwt(request.token) + covered_route: Final = request.is_llm_api_route or facts.is_info_route or _is_lens_route(request.route) + if not covered_route: + return False + if settings.oauth2_introspection and not token_is_jwt: + return True + return token_is_jwt and _routes_jwt_to_oauth2(facts, settings) + + +def _premium(plan: CredentialPlan, feature: PremiumFeature, settings: PlanSettings) -> CredentialPlan: + return plan if settings.premium_user else PremiumRequired(feature) + + +def after_jwt(facts: PlanFacts, settings: PlanSettings, *, issued_key: bool) -> KeyPlan: + if not issued_key and facts.request.route in settings.unauthenticated_pass_through_routes: + return AnonymousValidator(AnonymousReason.PASS_THROUGH_AUTH_OFF) + return KeyValidator(pass_through_key=not issued_key) + + +def after_custom_auth(facts: PlanFacts, settings: PlanSettings, *, issued_key: bool) -> CredentialPlan: + if is_public_route(facts.request.route, settings): + return AnonymousValidator(AnonymousReason.PUBLIC_ROUTE) + if oauth2_introspection_applies(facts, settings): + return _premium(OAuth2Validator(OAuth2Mode.INTROSPECTION), PremiumFeature.OAUTH2, settings) + if settings.oauth2_proxy: + return OAuth2Validator(OAuth2Mode.PROXY) + if settings.jwt_enabled and is_jwt(facts.request.token): + return _premium(JwtValidator(), PremiumFeature.JWT, settings) + return after_jwt(facts, settings, issued_key=issued_key) + + +def plan(facts: PlanFacts, settings: PlanSettings) -> Plan: + if settings.custom_auth: + return CustomValidator() + return after_custom_auth(facts, settings, issued_key=False) diff --git a/litellm/proxy/auth/core/python_core.py b/litellm/proxy/auth/core/python_core.py index efac9954892..ed76e79124c 100644 --- a/litellm/proxy/auth/core/python_core.py +++ b/litellm/proxy/auth/core/python_core.py @@ -3,10 +3,36 @@ from dataclasses import dataclass from datetime import datetime from typing import Final +from litellm.proxy.auth.core import oauth2 as oauth2_rules +from litellm.proxy.auth.core import plan as plan_rules from litellm.proxy.auth.core.answers import AdminKey, KeyFacts, TeamMember -from litellm.proxy.auth.core.core import AuthCore, KeyCore, ModelBudgetCore -from litellm.proxy.auth.core.host import KeyHost, ModelBudgetHost, TeamMemberModelBudgetHost -from litellm.proxy.auth.core.inputs import KeySettings, RequestFacts +from litellm.proxy.auth.core.core import ( + AnonymousCore, + AuthCore, + AuthorizeCore, + AuthPlanner, + CustomAuthCore, + JwtCore, + KeyCore, + ModelBudgetCore, + OAuth2Core, +) +from litellm.proxy.auth.core.host import ( + CustomAuthHost, + KeyHost, + ModelBudgetHost, + OAuth2IntrospectionHost, + OAuth2ProxyHost, + TeamMemberModelBudgetHost, +) +from litellm.proxy.auth.core.inputs import ( + KeySettings, + OAuth2IntrospectionSettings, + OAuth2ProxySettings, + PlanFacts, + PlanSettings, + RequestFacts, +) from litellm.proxy.auth.core.key import hash_token, is_master_key from litellm.proxy.auth.core.model_budgets import ( check_key_and_user_model_budgets, @@ -15,11 +41,16 @@ from litellm.proxy.auth.core.model_budgets import ( ) from litellm.proxy.auth.core.outcome import ( Allow, + Anonymous, + AnonymousReason, CachedAdminKey, + CredentialPlan, + CustomAuthResult, Deny, ExpiredKey, KeyBlocked, KeyBudgetExceeded, + KeyPlan, KeyWithoutToken, MalformedKey, MasterKey, @@ -30,8 +61,13 @@ from litellm.proxy.auth.core.outcome import ( NoDatabase, NoMasterKey, NotAVirtualKey, + OAuth2Denial, + OAuth2Identity, + OAuth2ProxyDenial, + OAuth2ProxyIdentity, Outcome, OwnerDeactivated, + Plan, TeamMemberBudgetExceeded, UserAuthDisabled, UserAuthRoute, @@ -63,6 +99,58 @@ class PythonCore(AuthCore): def model_budgets(self) -> ModelBudgetCore: return PythonModelBudgetCore() + def plan(self) -> AuthPlanner: + return PythonAuthPlanner() + + def anonymous(self) -> AnonymousCore: + return PythonAnonymousCore() + + def jwt(self) -> JwtCore: + raise NotImplementedError + + def oauth2(self) -> OAuth2Core: + return PythonOAuth2Core() + + def custom(self) -> CustomAuthCore: + return PythonCustomAuthCore() + + def authorize(self) -> AuthorizeCore: + raise NotImplementedError + + +class PythonAuthPlanner(AuthPlanner): + def plan(self, facts: PlanFacts, settings: PlanSettings) -> Plan: + return plan_rules.plan(facts, settings) + + def after_custom_auth(self, facts: PlanFacts, settings: PlanSettings, *, issued_key: bool) -> CredentialPlan: + return plan_rules.after_custom_auth(facts, settings, issued_key=issued_key) + + def after_jwt(self, facts: PlanFacts, settings: PlanSettings, *, issued_key: bool) -> KeyPlan: + return plan_rules.after_jwt(facts, settings, issued_key=issued_key) + + def oauth2_introspection_applies(self, facts: PlanFacts, settings: PlanSettings) -> bool: + return plan_rules.oauth2_introspection_applies(facts, settings) + + +class PythonAnonymousCore(AnonymousCore): + def authenticate(self, reason: AnonymousReason) -> Anonymous: + return Anonymous(reason) + + +class PythonOAuth2Core(OAuth2Core): + async def introspect( + self, host: OAuth2IntrospectionHost, token: str, settings: OAuth2IntrospectionSettings + ) -> OAuth2Identity | OAuth2Denial: + return await oauth2_rules.introspect(host, token, settings) + + def proxy(self, host: OAuth2ProxyHost, settings: OAuth2ProxySettings) -> OAuth2ProxyIdentity | OAuth2ProxyDenial: + return oauth2_rules.proxy(host, settings) + + +class PythonCustomAuthCore(CustomAuthCore): + async def authenticate(self, host: CustomAuthHost, facts: RequestFacts) -> CustomAuthResult: + return await host.custom_auth(facts) + class PythonKeyCore(KeyCore): async def authenticate(self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime) -> Outcome: diff --git a/litellm/proxy/auth/core/rust_core.py b/litellm/proxy/auth/core/rust_core.py index 113fc26d9bf..ebef7a94ac7 100644 --- a/litellm/proxy/auth/core/rust_core.py +++ b/litellm/proxy/auth/core/rust_core.py @@ -1,5 +1,6 @@ from litellm.proxy.auth.core.core import ( AnonymousCore, + AuthCore, AuthorizeCore, AuthPlanner, CustomAuthCore, @@ -7,11 +8,10 @@ from litellm.proxy.auth.core.core import ( KeyCore, ModelBudgetCore, OAuth2Core, - TargetAuthCore, ) -class RustCore(TargetAuthCore): +class RustCore(AuthCore): def keys(self) -> KeyCore: raise NotImplementedError diff --git a/litellm/proxy/auth/network.py b/litellm/proxy/auth/network.py index 98f15a20a19..0961e88955a 100644 --- a/litellm/proxy/auth/network.py +++ b/litellm/proxy/auth/network.py @@ -8,6 +8,7 @@ from fastapi import Request from pydantic import Field from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.core.network import ip_in_networks as ip_in_networks # noqa: PLC0414 # public re-export from litellm.types.llms.base import LiteLLMBaseModel TrustedProxyNetwork = ipaddress.IPv4Network | ipaddress.IPv6Network @@ -51,23 +52,6 @@ def parse_trusted_proxy_ranges( return networks -def _unmapped(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> ipaddress.IPv4Address | ipaddress.IPv6Address: - if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None: - return addr.ipv4_mapped - return addr - - -def ip_in_networks(client_ip: str | None, networks: list[TrustedProxyNetwork]) -> bool: - if not client_ip or not networks: - return False - try: - addr: Final = ipaddress.ip_address(client_ip.strip()) - except ValueError: - return False - candidates: Final = (addr, _unmapped(addr)) - return any(candidate in network for candidate in candidates for network in networks) - - def _is_valid_ip(value: str) -> bool: try: ipaddress.ip_address(value) diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index dc724e7afbe..1b7be2e43ac 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -1,6 +1,6 @@ import base64 import os -from typing import Final, cast +from typing import Final import httpx from pydantic import ConfigDict, TypeAdapter @@ -10,7 +10,11 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.core.answers import Introspection +from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings +from litellm.proxy.auth.core.outcome import InactiveToken, IntrospectionNotConfigured, OAuth2Identity +from litellm.proxy.auth.selected_core import AUTH_CORE _TOKEN_INFO_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)) @@ -92,130 +96,94 @@ class Oauth2Handler: """ return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} - @staticmethod - def _extract_user_info( - response_data: dict, - user_id_field_name: str, - user_role_field_name: str, - user_team_id_field_name: str, - ) -> tuple[str | None, str | None, str | None]: - """ - Extract user information from OAuth2 response. - - Args: - response_data: The response data from OAuth2 endpoint - user_id_field_name: Field name for user ID - user_role_field_name: Field name for user role - user_team_id_field_name: Field name for team ID - - Returns: - Tuple of (user_id, user_role, user_team_id) - """ - user_id: Final = response_data.get(user_id_field_name) - user_team_id: Final = response_data.get(user_team_id_field_name) - user_role: Final = response_data.get(user_role_field_name) - - return user_id, user_role, user_team_id - @staticmethod async def check_oauth2_token(token: str) -> UserAPIKeyAuth: - """ - Makes a request to the token introspection endpoint to validate the OAuth2 token. + """Validate ``token`` against the configured token-info or RFC 7662 introspection endpoint. - This function implements OAuth2 token introspection according to RFC 7662. - It supports both generic token info endpoints (GET) and OAuth2 introspection endpoints (POST). + Raises ``ValueError`` when the endpoint is not configured, and ``OAuth2TokenValidationError`` + when the endpoint rejects the token, reports it inactive, or cannot be reached.""" + verbose_proxy_logger.debug("Oauth2 token validation for token=[set=%s]", bool(token)) + endpoint: Final = _TokenInfoEndpoint.from_env() + result: Final = await AUTH_CORE.oauth2().introspect(endpoint, token, endpoint.settings()) + match result: + case OAuth2Identity(user_id=user_id, role=role, team_id=team_id): + return UserAPIKeyAuth.model_validate( + {"api_key": token, "team_id": team_id, "user_id": user_id, "user_role": role} + ) + case IntrospectionNotConfigured(): + raise ValueError("OAUTH_TOKEN_INFO_ENDPOINT environment variable is not set") + case InactiveToken(): + raise OAuth2TokenValidationError("An error occurred during token validation: Token is not active") - Args: - token (str): The OAuth2 token to validate. - Returns: - UserAPIKeyAuth: If the token is valid, containing user information. +class _TokenInfoEndpoint: + def __init__( + self, + url: str | None, + client_id: str | None, + client_secret: str | None, + field_names: tuple[str, str, str], + ) -> None: + self._url: Final = url + self._client_id: Final = client_id + self._client_secret: Final = client_secret + self._field_names: Final = field_names - Raises: - ValueError: If the token is invalid, the request fails, or the token info endpoint is not set. - """ - from litellm.proxy.proxy_server import premium_user - - if premium_user is not True: - raise ValueError( - "Oauth2 token validation is only available for premium users" + CommonProxyErrors.not_premium_user.value - ) - - verbose_proxy_logger.debug("Oauth2 token validation for token=[set=%s]", token is not None) - - # Get the token info endpoint from environment variable - token_info_endpoint: Final = os.getenv("OAUTH_TOKEN_INFO_ENDPOINT") - user_id_field_name: Final = os.environ.get("OAUTH_USER_ID_FIELD_NAME", "sub") - user_role_field_name: Final = os.environ.get("OAUTH_USER_ROLE_FIELD_NAME", "role") - user_team_id_field_name: Final = os.environ.get("OAUTH_USER_TEAM_ID_FIELD_NAME", "team_id") - - # OAuth2 client credentials for introspection endpoint authentication - oauth_client_id: Final = os.environ.get("OAUTH_CLIENT_ID") - oauth_client_secret: Final = os.environ.get("OAUTH_CLIENT_SECRET") - - if not token_info_endpoint: - raise ValueError("OAUTH_TOKEN_INFO_ENDPOINT environment variable is not set") - - client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) - - # Determine if this is an introspection endpoint (requires POST) or token info endpoint (uses GET) - is_introspection_endpoint: Final = Oauth2Handler._is_introspection_endpoint( - token_info_endpoint=token_info_endpoint, - oauth_client_id=oauth_client_id, - oauth_client_secret=oauth_client_secret, + @classmethod + def from_env(cls) -> "_TokenInfoEndpoint": + return cls( + url=os.getenv("OAUTH_TOKEN_INFO_ENDPOINT"), + client_id=os.environ.get("OAUTH_CLIENT_ID"), + client_secret=os.environ.get("OAUTH_CLIENT_SECRET"), + field_names=( + os.environ.get("OAUTH_USER_ID_FIELD_NAME", "sub"), + os.environ.get("OAUTH_USER_ROLE_FIELD_NAME", "role"), + os.environ.get("OAUTH_USER_TEAM_ID_FIELD_NAME", "team_id"), + ), ) + def _is_introspection(self) -> bool: + return self._url is not None and Oauth2Handler._is_introspection_endpoint( # pyright: ignore[reportPrivateUsage] # same module + token_info_endpoint=self._url, + oauth_client_id=self._client_id, + oauth_client_secret=self._client_secret, + ) + + def settings(self) -> OAuth2IntrospectionSettings: + return OAuth2IntrospectionSettings( + endpoint_configured=bool(self._url), requires_active=self._is_introspection() + ) + + async def _response(self, url: str, token: str) -> httpx.Response: + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) + if self._is_introspection(): + verbose_proxy_logger.debug("Using OAuth2 introspection endpoint (POST)") + headers, data = Oauth2Handler._prepare_introspection_request( # pyright: ignore[reportPrivateUsage] # same module + token=token, oauth_client_id=self._client_id, oauth_client_secret=self._client_secret + ) + return await client.post(url, headers=headers, data=data) + verbose_proxy_logger.debug("Using generic token info endpoint (GET)") + return await client.get(url, headers=Oauth2Handler._prepare_token_info_request(token=token)) # pyright: ignore[reportPrivateUsage] # same module + + async def introspect(self, token: str) -> Introspection: + url: Final = self._url or "" try: - if is_introspection_endpoint: - # OAuth2 Token Introspection (RFC 7662) - requires POST with form data - verbose_proxy_logger.debug("Using OAuth2 introspection endpoint (POST)") - - headers, data = Oauth2Handler._prepare_introspection_request( - token=token, - oauth_client_id=oauth_client_id, - oauth_client_secret=oauth_client_secret, - ) - - response = await client.post(token_info_endpoint, headers=headers, data=data) - else: - # Generic token info endpoint - uses GET with Bearer token - verbose_proxy_logger.debug("Using generic token info endpoint (GET)") - headers = Oauth2Handler._prepare_token_info_request(token=token) - response = await client.get(token_info_endpoint, headers=headers) - - # if it's a bad token we expect it to raise an HTTPStatusError + response: Final = await self._response(url, token) response.raise_for_status() - - # If we get here, the request was successful - data = _TOKEN_INFO_RESPONSE.validate_python(response.json()) - - verbose_proxy_logger.debug( - "Oauth2 token validation for token=%s, response from endpoint=%s", - token, - data, - ) - - # For introspection endpoints, check if token is active - if is_introspection_endpoint and not data.get("active", True): - raise ValueError("Token is not active") - - # Extract user information from response - user_id, user_role, user_team_id = Oauth2Handler._extract_user_info( - response_data=data, - user_id_field_name=user_id_field_name, - user_role_field_name=user_role_field_name, - user_team_id_field_name=user_team_id_field_name, - ) - - return UserAPIKeyAuth( - api_key=token, - team_id=user_team_id, - user_id=user_id, - user_role=cast(LitellmUserRoles, user_role), - ) + data: Final = _TOKEN_INFO_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as e: - # This will catch any 4xx or 5xx errors raise OAuth2TokenValidationError(f"Oauth 2.0 Token validation failed: {e}") except Exception as e: - # This will catch any other errors (like network issues) raise OAuth2TokenValidationError(f"An error occurred during token validation: {e}") + verbose_proxy_logger.debug( + "Oauth2 token validation for token=%s, response from endpoint=%s", + token, + data, + ) + user_id_field, role_field, team_id_field = self._field_names + return Introspection( + active=bool(data.get("active", True)), + user_id=data.get(user_id_field), + role=data.get(role_field), + team_id=data.get(team_id_field), + ) diff --git a/litellm/proxy/auth/oauth2_proxy_hook.py b/litellm/proxy/auth/oauth2_proxy_hook.py index 01fd747527d..42fa2092614 100644 --- a/litellm/proxy/auth/oauth2_proxy_hook.py +++ b/litellm/proxy/auth/oauth2_proxy_hook.py @@ -1,102 +1,98 @@ from collections.abc import Mapping +from types import MappingProxyType from typing import Final from fastapi import Request +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.trusted_proxy_utils import require_trusted_proxy_request - -# OAuth2-proxy header trust is for **identity assertion** from a trusted -# upstream auth proxy (oauth2-proxy, Authelia, etc.). The allowlist below -# is the only safe surface — anything else (``user_role``, ``api_key``, -# ``permissions``, ``max_budget``, ``user_max_budget``, -# ``team_tpm_limit``, ``end_user_max_budget``, ``allowed_model_region``, -# and dozens of similar policy fields scattered across the -# ``LiteLLM_VerificationTokenView`` hierarchy) is a privilege grant that -# would let a caller forge their own enforcement parameters by sending -# the matching header. -# -# A denylist of "privileged fields" is unmaintainable in this codebase: -# the auth model has ~50 budget/spend/limit/permission fields and gains -# more with each release. An allowlist scoped to identity assertion is -# default-secure — new fields are blocked automatically. -# -# Operators who need a trusted upstream to assert anything beyond -# identity should switch to JWT authentication, which validates a -# signature on the assertion rather than blindly trusting headers. -ALLOWED_OAUTH2_PROXY_FIELDS: Final[frozenset[str]] = frozenset( - { - "user_id", - "user_email", - "team_id", - "team_alias", - "org_id", - "models", - } +from litellm.proxy.auth.core.inputs import OAuth2ProxySettings +from litellm.proxy.auth.core.oauth2 import ALLOWED_OAUTH2_PROXY_FIELDS +from litellm.proxy.auth.core.outcome import ( + DisallowedProxyMappings, + OAuth2ProxyIdentity, + ProxyMappingsMissing, + TrustedProxyRangesMissing, + UntrustedProxy, ) +from litellm.proxy.auth.network import parse_trusted_proxy_ranges +from litellm.proxy.auth.selected_core import AUTH_CORE +from litellm.proxy.auth.trusted_proxy_utils import ( + TRUSTED_PROXY_RANGES_KEY, + direct_client_ip, + trusted_proxy_ranges_missing, + untrusted_proxy, +) + +_FEATURE_NAME: Final = "OAuth2 proxy auth" +_MAPPINGS: Final = TypeAdapter(dict[str, str]) + + +class _RequestHeaders: + def __init__(self, request: Request) -> None: + self._request: Final = request + + def client_ip(self) -> str | None: + return direct_client_ip(self._request) + + def header(self, name: str) -> str | None: + return self._request.headers.get(name) + + +def _settings(general_settings: Mapping[str, object]) -> OAuth2ProxySettings: + mappings: Final = _MAPPINGS.validate_python(general_settings.get("oauth2_config_mappings") or {}) + return OAuth2ProxySettings( + trusted_networks=tuple( + parse_trusted_proxy_ranges( + general_settings.get(TRUSTED_PROXY_RANGES_KEY), setting_name=TRUSTED_PROXY_RANGES_KEY + ) + ), + header_mappings=MappingProxyType(mappings), + ) + + +def _user_api_key_auth(identity: OAuth2ProxyIdentity) -> UserAPIKeyAuth: + fields: Final[dict[str, object]] = { + "user_id": identity.user_id, + "user_email": identity.user_email, + "team_id": identity.team_id, + "team_alias": identity.team_alias, + "org_id": identity.org_id, + "models": None if identity.models is None else list(identity.models), + } + asserted: Final = {field: value for field, value in fields.items() if value is not None} + verbose_proxy_logger.debug("Auth data before creating UserAPIKeyAuth object: keys=%s", list(asserted)) + return UserAPIKeyAuth.model_validate(asserted) async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: """ - Resolve a ``UserAPIKeyAuth`` from request headers per the admin-set - ``oauth2_config_mappings``. - - The auth model assumes the proxy is deployed behind a trusted OAuth2 - reverse proxy that injects authenticated identity headers (e.g. - oauth2-proxy, Authelia). - - **Identity-only allowlist.** ``oauth2_config_mappings`` maps header - names to ``UserAPIKeyAuth`` fields. Without an allowlist, an admin - who maps the wrong header to ``user_role`` lets any caller send - ``X-User-Role: proxy_admin`` and gain full admin privileges - (Pydantic coerces the string into the enum). Only fields in - ``ALLOWED_OAUTH2_PROXY_FIELDS`` (identity assertion only — see the - constant's comment) may be mapped; any other mapping is rejected at - request time so the misconfiguration surfaces loudly rather than as - a silent privesc. + Resolve a ``UserAPIKeyAuth`` from the identity headers a trusted OAuth2 reverse proxy (oauth2-proxy, + Authelia) injects, per the admin-set ``oauth2_config_mappings``. Only the identity fields in + ``ALLOWED_OAUTH2_PROXY_FIELDS`` may be mapped. """ from litellm.proxy.proxy_server import general_settings verbose_proxy_logger.debug("Handling oauth2 proxy request") - require_trusted_proxy_request( - request=request, - general_settings=general_settings, - feature_name="OAuth2 proxy auth", - ) - - oauth2_config_mappings: Final[dict[str, str]] = general_settings.get("oauth2_config_mappings") or {} - verbose_proxy_logger.debug("Oauth2 config mappings: %s", oauth2_config_mappings) - - if not oauth2_config_mappings: - raise ValueError("Oauth2 config mappings not found in general_settings") - - disallowed: Final = sorted(set(oauth2_config_mappings.keys()) - ALLOWED_OAUTH2_PROXY_FIELDS) - if disallowed: - raise ValueError( - "Oauth2 proxy auth refuses to map non-identity UserAPIKeyAuth " - f"fields from request headers: {disallowed}. Only identity " - f"fields are accepted ({sorted(ALLOWED_OAUTH2_PROXY_FIELDS)}); " - "anything else (privileges, budgets, rate limits, metadata) " - "would let a caller forge enforcement parameters by spoofing " - "the matching header. If you need a trusted upstream to " - "assert anything beyond identity, use JWT auth " - "(signature-validated) instead of header-trust." - ) - - auth_data: Final[Mapping[str, str | list[str]]] = { - key: [model.strip() for model in value.split(",")] if key == "models" else value - for key, header in oauth2_config_mappings.items() - if (value := request.headers.get(header)) - } - - verbose_proxy_logger.debug( - "Auth data before creating UserAPIKeyAuth object: keys=%s", - list(auth_data.keys()), - ) - user_api_key_auth: Final = UserAPIKeyAuth.model_validate(auth_data) - verbose_proxy_logger.debug( - "UserAPIKeyAuth object created with keys: %s", - list(user_api_key_auth.__fields_set__), - ) - return user_api_key_auth + result: Final = AUTH_CORE.oauth2().proxy(_RequestHeaders(request), _settings(general_settings)) + match result: + case OAuth2ProxyIdentity(): + return _user_api_key_auth(result) + case TrustedProxyRangesMissing(): + raise trusted_proxy_ranges_missing(_FEATURE_NAME) + case UntrustedProxy(client_ip=client_ip): + raise untrusted_proxy(_FEATURE_NAME, client_ip) + case ProxyMappingsMissing(): + raise ValueError("Oauth2 config mappings not found in general_settings") + case DisallowedProxyMappings(fields=fields): + raise ValueError( + "Oauth2 proxy auth refuses to map non-identity UserAPIKeyAuth " + f"fields from request headers: {list(fields)}. Only identity " + f"fields are accepted ({sorted(ALLOWED_OAUTH2_PROXY_FIELDS)}); " + "anything else (privileges, budgets, rate limits, metadata) " + "would let a caller forge enforcement parameters by spoofing " + "the matching header. If you need a trusted upstream to " + "assert anything beyond identity, use JWT auth " + "(signature-validated) instead of header-trust." + ) diff --git a/litellm/proxy/auth/selected_core.py b/litellm/proxy/auth/selected_core.py new file mode 100644 index 00000000000..e0367191cef --- /dev/null +++ b/litellm/proxy/auth/selected_core.py @@ -0,0 +1,6 @@ +from typing import Final + +from litellm.proxy.auth.core.core import AuthCore +from litellm.proxy.auth.core.python_core import PythonCore + +AUTH_CORE: Final[AuthCore] = PythonCore() diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py index 04c09abd5dd..8e813063645 100644 --- a/litellm/proxy/auth/trusted_proxy_utils.py +++ b/litellm/proxy/auth/trusted_proxy_utils.py @@ -45,6 +45,30 @@ def _get_direct_client_ip(request: Request) -> str | None: return None +def trusted_proxy_ranges_missing(feature_name: str, setting_name: str = TRUSTED_PROXY_RANGES_KEY) -> ValueError: + return ValueError( + f"{feature_name} requires general_settings.{setting_name} before " + "trusting identity headers from an upstream proxy." + ) + + +def untrusted_proxy(feature_name: str, direct_client_ip: str | None) -> ValueError: + verbose_proxy_logger.warning( + "%s rejected identity headers from untrusted direct client IP %r", + feature_name, + direct_client_ip, + ) + return ValueError( + f"{feature_name} only accepts identity headers from configured " + f"trusted proxy ranges. Direct client IP {direct_client_ip!r} " + "is not trusted." + ) + + +def direct_client_ip(request: Request) -> str | None: + return _get_direct_client_ip(request) + + def require_trusted_proxy_request( *, request: Request, @@ -65,20 +89,8 @@ def require_trusted_proxy_request( trusted_networks: Final = parse_trusted_proxy_ranges(general_settings.get(setting_name), setting_name=setting_name) if not trusted_networks: - raise ValueError( - f"{feature_name} requires general_settings.{setting_name} before " - "trusting identity headers from an upstream proxy." - ) + raise trusted_proxy_ranges_missing(feature_name, setting_name) - direct_client_ip: Final = _get_direct_client_ip(request) - if not ip_in_networks(direct_client_ip, trusted_networks): - verbose_proxy_logger.warning( - "%s rejected identity headers from untrusted direct client IP %r", - feature_name, - direct_client_ip, - ) - raise ValueError( - f"{feature_name} only accepts identity headers from configured " - f"trusted proxy ranges. Direct client IP {direct_client_ip!r} " - "is not trusted." - ) + client_ip: Final = _get_direct_client_ip(request) + if not ip_in_networks(client_ip, trusted_networks): + raise untrusted_proxy(feature_name, client_ip) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 1ec23d51318..ce2156cc147 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -8,8 +8,8 @@ Returns a UserAPIKeyAuth object if the API key is valid """ import asyncio -import fnmatch from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timezone from enum import Enum from functools import partial @@ -137,7 +137,6 @@ from litellm.proxy.auth.core.answers import ( TeamMemberBudget, TeamMemberModelBudgetScope, ) -from litellm.proxy.auth.core.core import AuthCore from litellm.proxy.auth.core.credential import ( PresentedKeys, WebSocketDenial, @@ -147,22 +146,36 @@ from litellm.proxy.auth.core.credential import ( select_websocket_api_key, ) from litellm.proxy.auth.core.host import ModelBudgetHost -from litellm.proxy.auth.core.inputs import KeySettings, RequestFacts +from litellm.proxy.auth.core.inputs import KeySettings, PlanFacts, PlanSettings, RequestFacts, RoutingOverride from litellm.proxy.auth.core.outcome import ( Allow, + Anonymous, + AnonymousReason, + AnonymousValidator, CachedAdminKey, + CredentialPlan, + CustomAuthResult, + CustomIdentity, + CustomValidator, Deny, Grant, + JwtValidator, + KeyValidator, MasterKey, ModelBudgetExceeded, ModelBudgetResult, + NoCustomIdentity, NoMasterKey, + OAuth2Mode, + OAuth2Validator, Outcome, + PremiumFeature, + PremiumRequired, + ReplacedKey, UserAuthRoute, VirtualKey, WithinModelBudgets, ) -from litellm.proxy.auth.core.python_core import PythonCore from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context from litellm.proxy.auth.oauth2_check import Oauth2Handler @@ -178,6 +191,7 @@ from litellm.proxy.auth.resolvers.grants import ( ) from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.selected_core import AUTH_CORE from litellm.proxy.auth.team_grants import team_grants from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator @@ -261,20 +275,6 @@ except ImportError as e: user_api_key_service_logger_obj: Final = ServiceLogging() # used for tracking latency on OTEL -def _normalize_public_auth_route(route: str) -> str: - if route != "/" and route.endswith("/"): - return route.rstrip("/") - return route - - -def _route_requires_auth_despite_public(route: str, general_settings: dict | None) -> bool: - normalized_route: Final = _normalize_public_auth_route(route) - if normalized_route == "/metrics": - return litellm.require_auth_for_metrics_endpoint is not False - - return False - - custom_litellm_key_header: Final = APIKeyHeader( name=SpecialHeaders.custom_litellm_api_key.value, auto_error=False, @@ -1225,57 +1225,72 @@ async def _check_team_member_model_budget_with_fallback( return -def _routing_selector_matches_claim( - selector_value: object, - claim_value: object, - *, - split_space_delimited: bool = False, -) -> bool: - if selector_value is None: - return True +def _selectors(value: str | list[str] | None) -> tuple[str, ...] | None: + if value is None: + return None + return tuple(str(selector) for selector in value) if isinstance(value, list) else (str(value),) - selector_list: Final[list[str]] = ( - [str(v) for v in selector_value] if isinstance(selector_value, list) else [str(selector_value)] + +def _routing_override(override: JWTRoutingOverride) -> RoutingOverride: + return RoutingOverride( + iss=_selectors(override.iss), + client_id=_selectors(override.client_id), + scope=_selectors(override.scope), + aud=_selectors(override.aud), ) - if claim_value is None: - return False - if isinstance(claim_value, list): - claim_list = [str(v) for v in claim_value] - elif split_space_delimited and isinstance(claim_value, str) and " " in claim_value.strip(): - # OAuth/OIDC often sends scope as a single space-delimited string. Only split - # for the scope selector: iss/aud/client_id must stay exact full-string match - # on unverified claims (see routing override security review). The elif guard - # (`" " in claim_value.strip()`) ensures at least two non-empty tokens survive. - claim_list = [v for v in claim_value.strip().split(" ") if v] - else: - claim_list = [str(claim_value)] - - def _selector_matches_claim(selector: str, claim: str) -> bool: - # NOTE: wildcard matching is case-sensitive (fnmatch.fnmatchcase). - if "*" in selector or "?" in selector: - # Without scope splitting, do not let `*` span whitespace: a malformed - # iss like "trusted.example.com evil.com" must not match "trusted.*". - # Scope uses split_space_delimited so each claim token is checked separately. - if not split_space_delimited and any(ch.isspace() for ch in claim): - return False - return fnmatch.fnmatchcase(claim, selector) - return selector == claim - - return any(_selector_matches_claim(selector=s, claim=c) for s in selector_list for c in claim_list) +def _oauth2_routing_overrides(jwt_handler: JWTHandler) -> tuple[RoutingOverride, ...]: + overrides: Final = jwt_handler.litellm_jwtauth.routing_overrides or () + return tuple(_routing_override(override) for override in overrides if override.path == "oauth2") -def _matches_routing_override(token_claims: dict, override: "JWTRoutingOverride") -> bool: - return ( - _routing_selector_matches_claim(override.iss, token_claims.get("iss")) - and _routing_selector_matches_claim(override.client_id, token_claims.get("client_id")) - and _routing_selector_matches_claim( - override.scope, - token_claims.get("scope"), - split_space_delimited=True, - ) - and _routing_selector_matches_claim(override.aud, token_claims.get("aud")) +def _unauthenticated_pass_through_routes(endpoints: object) -> frozenset[str]: + if not isinstance(endpoints, list): + return frozenset() + return frozenset( + str(endpoint.get("path", "")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # raw config dicts + for endpoint in endpoints # pyright: ignore[reportUnknownVariableType] # raw config list + if isinstance(endpoint, dict) and endpoint.get("auth", True) is not True # pyright: ignore[reportUnknownMemberType] # raw config dicts + ) + + +def _additional_public_routes(general_settings: Mapping[str, object]) -> tuple[str, ...]: + routes: Final = general_settings.get("public_routes") or [] + return tuple(route for route in routes if isinstance(route, str)) if isinstance(routes, list) else () # pyright: ignore[reportUnknownVariableType] # raw config list + + +def _plan_settings( + general_settings: Mapping[str, object], jwt_handler: JWTHandler, *, premium_user: bool, custom_auth: bool +) -> PlanSettings: + jwt_enabled: Final = general_settings.get("enable_jwt_auth", False) is True + return PlanSettings( + custom_auth=custom_auth, + premium_user=premium_user, + public_routes=frozenset(LiteLLMRoutes.public_routes.value), + additional_public_routes=_additional_public_routes(general_settings), + metrics_require_auth=litellm.require_auth_for_metrics_endpoint is not False, + oauth2_introspection=general_settings.get("enable_oauth2_auth", False) is True, + oauth2_proxy=general_settings.get("enable_oauth2_proxy_auth", False) is True, + jwt_enabled=jwt_enabled, + oauth2_routing_overrides=_oauth2_routing_overrides(jwt_handler) if jwt_enabled else (), + unauthenticated_pass_through_routes=_unauthenticated_pass_through_routes( + general_settings.get("pass_through_endpoints") + ), + ) + + +def _request_facts(route: str, api_key: str) -> RequestFacts: + return RequestFacts(route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)) + + +def _plan_facts(route: str, api_key: str, settings: PlanSettings, jwt_handler: JWTHandler) -> PlanFacts: + return PlanFacts( + request=_request_facts(route, api_key), + is_info_route=RouteChecks.is_info_route(route=route), + unverified_claims=( + jwt_handler.get_unverified_claims(token=api_key) if settings.oauth2_routing_overrides else None + ), ) @@ -1284,34 +1299,156 @@ def oauth2_introspection_applies( ) -> bool: """Whether LiteLLM validates ``api_key`` on ``route`` by OAuth2 token introspection, on LLM API, info and lens routes only: a non-JWT token when OAuth2 auth is on globally, or a JWT matched by a routing override.""" - is_jwt: Final = enable_jwt_auth and jwt_handler.is_jwt(token=api_key) - # Routing uses unverified JWT claims only to choose auth path. - # Final authentication is enforced by the selected validator. - route_jwt_to_oauth2: Final = is_jwt and _should_route_jwt_to_oauth2_override(token=api_key, jwt_handler=jwt_handler) - covered_route: Final = ( - RouteChecks.is_llm_api_route(route=route) - or RouteChecks.is_info_route(route=route) - or route == "/lens" - or route.startswith("/lens/") + settings: Final = _plan_settings( + {"enable_oauth2_auth": enable_oauth2_auth, "enable_jwt_auth": enable_jwt_auth}, + jwt_handler, + premium_user=False, + custom_auth=False, ) - return covered_route and ((enable_oauth2_auth and not is_jwt) or route_jwt_to_oauth2) + return _AUTH_CORE.plan().oauth2_introspection_applies(_plan_facts(route, api_key, settings, jwt_handler), settings) -def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> bool: - routing_overrides: Final = jwt_handler.litellm_jwtauth.routing_overrides - if not routing_overrides: - return False +@dataclass(frozen=True, slots=True) +class _Planned: + plan: CredentialPlan + facts: PlanFacts + api_key: str + issued_key: bool - token_claims: Final = jwt_handler.get_unverified_claims(token=token) - if token_claims is None: - return False - for override in routing_overrides: - if override.path == "oauth2" and _matches_routing_override(token_claims=token_claims, override=override): - verbose_proxy_logger.debug("JWT routing override matched. Routing token to OAuth2 introspection.") - return True +class _CustomAuthHost: + def __init__(self, request: Request, api_key: str, user_custom_auth: Callable | None) -> None: + self._request: Final = request + self._api_key: Final = api_key + self._user_custom_auth: Final = user_custom_auth + self.resolved: UserAPIKeyAuth | None = None - return False + def _resolve(self, response: object) -> CustomIdentity: + self.resolved = _custom_auth_token_for_this_request(response) + return CustomIdentity( + user_id=self.resolved.user_id, team_id=self.resolved.team_id, hashed_token=self.resolved.token + ) + + async def custom_auth(self, facts: RequestFacts) -> CustomAuthResult: + if enterprise_custom_auth is None: + assert self._user_custom_auth is not None + return self._resolve(await self._user_custom_auth(request=self._request, api_key=self._api_key)) + with tracer.trace("litellm.proxy.auth.enterprise_custom_auth"): + response: Final = await enterprise_custom_auth( + request=self._request, api_key=self._api_key, user_custom_auth=self._user_custom_auth + ) + if isinstance(response, UserAPIKeyAuth): + return self._resolve(response) + if isinstance(response, str): + return ReplacedKey(token=response) + return NoCustomIdentity() + + +async def _custom_auth_user( + validated: UserAPIKeyAuth, + request: Request, + request_data: dict[str, object], + route: str, + parent_otel_span: Span | None, +) -> UserAPIKeyAuth: + checked: Final = ( + await _run_post_custom_auth_checks( + valid_token=validated, + request=request, + request_data=request_data, + route=route, + parent_otel_span=parent_otel_span, + ) + if getattr(litellm, "enable_post_custom_auth_checks", False) + else validated + ) + checked.authenticated_by_custom_auth = True + return checked + + +async def _plan_request( + request: Request, + request_data: dict[str, object], + route: str, + api_key: str, + settings: PlanSettings, + jwt_handler: JWTHandler, + user_custom_auth: Callable | None, + parent_otel_span: Span | None, +) -> UserAPIKeyAuth | _Planned: + facts: Final = _plan_facts(route, api_key, settings, jwt_handler) + first: Final = _AUTH_CORE.plan().plan(facts, settings) + if not isinstance(first, CustomValidator): + return _Planned(plan=first, facts=facts, api_key=api_key, issued_key=False) + host: Final = _CustomAuthHost(request=request, api_key=api_key, user_custom_auth=user_custom_auth) + result: Final = await _AUTH_CORE.custom().authenticate(host, facts.request) + match result: + case CustomIdentity(): + assert host.resolved is not None + return await _custom_auth_user(host.resolved, request, request_data, route, parent_otel_span) + case ReplacedKey(token=token): + issued: Final = _plan_facts(route, token, settings, jwt_handler) + return _Planned( + plan=_AUTH_CORE.plan().after_custom_auth(issued, settings, issued_key=True), + facts=issued, + api_key=token, + issued_key=True, + ) + case NoCustomIdentity(): + return _Planned( + plan=_AUTH_CORE.plan().after_custom_auth(facts, settings, issued_key=False), + facts=facts, + api_key=api_key, + issued_key=False, + ) + + +def _anonymous_auth(anonymous: Anonymous) -> UserAPIKeyAuth: + match anonymous.reason: + case AnonymousReason.PUBLIC_ROUTE: + return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) + case AnonymousReason.PASS_THROUGH_AUTH_OFF: + return UserAPIKeyAuth() + + +def _premium_message(feature: PremiumFeature) -> str: + match feature: + case PremiumFeature.OAUTH2: + return ( + "Oauth2 token validation is only available for premium users. " + + CommonProxyErrors.not_premium_user.value + ) + case PremiumFeature.JWT: + return f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}" + + +def _premium_required(required: PremiumRequired) -> ProxyException: + return ProxyException( + message=_premium_message(required.feature), + type=ProxyErrorTypes.auth_error, + param="premium_user", + code=status.HTTP_403_FORBIDDEN, + ) + + +async def _oauth2_auth(mode: OAuth2Mode, *, api_key: str, request: Request) -> UserAPIKeyAuth: + match mode: + case OAuth2Mode.INTROSPECTION: + return await Oauth2Handler.check_oauth2_token(token=api_key) + case OAuth2Mode.PROXY: + return await handle_oauth2_proxy_request(request=request) + + +async def _authenticate_without_key( + plan: AnonymousValidator | OAuth2Validator | PremiumRequired, *, api_key: str, request: Request +) -> UserAPIKeyAuth: + match plan: + case AnonymousValidator(reason=reason): + return _anonymous_auth(_AUTH_CORE.anonymous().authenticate(reason)) + case PremiumRequired(): + raise _premium_required(plan) + case OAuth2Validator(mode=mode): + return await _oauth2_auth(mode, api_key=api_key, request=request) def get_bearer_token( @@ -1644,7 +1781,7 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( route: str, pass_through_endpoints: list[dict] | None, api_key: str, -) -> UserAPIKeyAuth | str: +) -> str: is_mapped_pass_through_route: bool = False normalized_route: Final = normalize_route_for_root_path(route) if normalized_route is not None: @@ -1658,16 +1795,6 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( if pass_through_endpoints is not None: for endpoint in pass_through_endpoints: if isinstance(endpoint, dict) and endpoint.get("path", "") == route: - ## IF AUTH DISABLED - # Default to True: a config dict with no ``auth`` key - # otherwise produced an unauthenticated forwarder. The - # Pydantic ``PassThroughGenericEndpoint.auth`` default - # is also True, but raw config dicts skip that path — - # so this runtime check has to default to True too. - if endpoint.get("auth", True) is not True: - return UserAPIKeyAuth() - ## IF AUTH ENABLED - ### IF CUSTOM PARSER REQUIRED if endpoint.get("custom_auth_parser") is not None and endpoint.get("custom_auth_parser") == "langfuse": # langfuse returns {'Authorization': 'Basic '} # check the langfuse public key if it contains the litellm api key @@ -2306,7 +2433,7 @@ def _custom_auth_token_for_this_request(response: object) -> UserAPIKeyAuth: return UserAPIKeyAuth.model_validate(response).model_copy() -_AUTH_CORE: Final[AuthCore] = PythonCore() +_AUTH_CORE: Final = AUTH_CORE class _ResolvedEndUser(NamedTuple): @@ -2817,6 +2944,7 @@ async def user_api_key_auth_builder( master_key, model_max_budget_limiter, open_telemetry_logger, + premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -2834,7 +2962,6 @@ async def user_api_key_auth_builder( pass route: Final[str] = get_request_route(request=request) valid_token: UserAPIKeyAuth | None = None - custom_auth_api_key: bool = False try: with tracer.trace("litellm.proxy.auth.pre_db_read_auth_checks"): @@ -2869,340 +2996,266 @@ async def user_api_key_auth_builder( _ensure_parent_otel_span_on_request_state(request) parent_otel_span = getattr(request.state, "parent_otel_span", None) - ### USER-DEFINED AUTH FUNCTION ### - if enterprise_custom_auth is not None: - with tracer.trace("litellm.proxy.auth.enterprise_custom_auth"): - response = await enterprise_custom_auth( - request=request, api_key=api_key, user_custom_auth=user_custom_auth - ) - if response is not None and isinstance(response, UserAPIKeyAuth): - validated = _custom_auth_token_for_this_request(response) - if getattr(litellm, "enable_post_custom_auth_checks", False): - validated = await _run_post_custom_auth_checks( - valid_token=validated, - request=request, - request_data=request_data, - route=route, - parent_otel_span=parent_otel_span, - ) - validated.authenticated_by_custom_auth = True - return validated - elif response is not None and isinstance(response, str): - api_key = response - custom_auth_api_key = True - elif user_custom_auth is not None: - response = await user_custom_auth(request=request, api_key=api_key) - validated = _custom_auth_token_for_this_request(response) - if getattr(litellm, "enable_post_custom_auth_checks", False): - validated = await _run_post_custom_auth_checks( - valid_token=validated, - request=request, - request_data=request_data, - route=route, - parent_otel_span=parent_otel_span, - ) - validated.authenticated_by_custom_auth = True - return validated - - ### LITELLM-DEFINED AUTH FUNCTION ### - #### IF JWT #### - """ - LiteLLM supports using JWTs. - - Enable this in proxy config, by setting - ``` - general_settings: - enable_jwt_auth: true - ``` - """ - - ######## Route Checks Before Reading DB / Cache for "token" ################ - if not _route_requires_auth_despite_public(route=route, general_settings=general_settings) and ( - route in LiteLLMRoutes.public_routes.value or route_in_additonal_public_routes(current_route=route) - ): - # check if public endpoint - return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) - - ########## End of Route Checks Before Reading DB / Cache for "token" ######## - - enable_oauth2_auth: Final = general_settings.get("enable_oauth2_auth", False) is True - enable_jwt_auth: Final = general_settings.get("enable_jwt_auth", False) is True - if oauth2_introspection_applies( - api_key, - route, - enable_oauth2_auth=enable_oauth2_auth, - enable_jwt_auth=enable_jwt_auth, + plan_settings: Final = _plan_settings( + general_settings, + jwt_handler, + premium_user=premium_user is True, + custom_auth=enterprise_custom_auth is not None or user_custom_auth is not None, + ) + planned: Final = await _plan_request( + request=request, + request_data=request_data, + route=route, + api_key=api_key, + settings=plan_settings, jwt_handler=jwt_handler, - ): - from litellm.proxy.proxy_server import premium_user + user_custom_auth=user_custom_auth, + parent_otel_span=parent_otel_span, + ) + if isinstance(planned, UserAPIKeyAuth): + return planned + api_key = planned.api_key + enable_jwt_auth: Final = plan_settings.jwt_enabled + credential_plan: Final = planned.plan + if not isinstance(credential_plan, (KeyValidator, JwtValidator)): + return await _authenticate_without_key(credential_plan, api_key=api_key, request=request) + if isinstance(credential_plan, JwtValidator): + # Try JWT-to-Virtual-Key mapping first to avoid + # unnecessary DB queries in auth_builder + do_standard_jwt_auth = True + pending_auto_register: _PendingAutoRegister | None = None + if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured(): + # Decode JWT to get claims without running full auth_builder + jwt_claims: dict | None = await jwt_handler.auth_jwt(token=api_key) - if premium_user is not True: - raise ProxyException( - message="Oauth2 token validation is only available for premium users. " - + CommonProxyErrors.not_premium_user.value, - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) + from litellm.proxy.agent_endpoints.identity_store import resolve_managed_agent - return await Oauth2Handler.check_oauth2_token(token=api_key) - - if general_settings.get("enable_oauth2_proxy_auth", False) is True: - return await handle_oauth2_proxy_request(request=request) - - if general_settings.get("enable_jwt_auth", False) is True: - is_jwt = jwt_handler.is_jwt(token=api_key) - verbose_proxy_logger.debug("is_jwt: %s", is_jwt) - if is_jwt: - from litellm.proxy.proxy_server import premium_user - - if premium_user is not True: - raise ProxyException( - message=f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}", - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, + if ( + jwt_claims + and await resolve_managed_agent(jwt_claims, prisma_client, cache=user_api_key_cache) is not None + ): + raise HTTPException( + 403, "Managed agents require direct JWT authentication without virtual-key mapping" ) - # Try JWT-to-Virtual-Key mapping first to avoid - # unnecessary DB queries in auth_builder - do_standard_jwt_auth = True - pending_auto_register: _PendingAutoRegister | None = None - if jwt_handler.litellm_jwtauth.is_virtual_key_mapping_configured(): - # Decode JWT to get claims without running full auth_builder - jwt_claims: dict | None - if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not is_jwt: - jwt_claims = await jwt_handler.get_oidc_userinfo(token=api_key) - else: - jwt_claims = await jwt_handler.auth_jwt(token=api_key) - from litellm.proxy.agent_endpoints.identity_store import resolve_managed_agent - - if ( - jwt_claims - and await resolve_managed_agent(jwt_claims, prisma_client, cache=user_api_key_cache) is not None - ): - raise HTTPException( - 403, "Managed agents require direct JWT authentication without virtual-key mapping" + resolve_result: Final = await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if isinstance(resolve_result, UserAPIKeyAuth): + mapped_end_user_id: Final = ( + jwt_handler.get_end_user_id( # pyright: ignore[reportUnknownMemberType] # untyped dict token + token=jwt_claims or {}, default_value=None ) + or None + ) + valid_token = resolve_result.model_copy(update={"end_user_id": mapped_end_user_id}) + api_key = valid_token.token or "" + valid_token.jwt_claims = jwt_claims + do_standard_jwt_auth = False + # Fall through to virtual key checks + if valid_token.user_id is not None and valid_token.user_email is None: + mapped_claims = jwt_claims or {} + mapped_user_email = jwt_handler.get_user_email(token=mapped_claims, default_value=None) + mapped_jwt_user_id: Final = jwt_handler.get_user_id(token=mapped_claims, default_value=None) + if mapped_user_email is not None and mapped_jwt_user_id == valid_token.user_id: + try: + mapped_user_obj: Final = await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + user_email=mapped_user_email, + ) + except Exception as e: + verbose_proxy_logger.debug("JWT mapped-key user_email backfill skipped: %s", e) + else: + if mapped_user_obj is not None: + valid_token.user_email = mapped_user_obj.user_email + elif isinstance(resolve_result, _PendingAutoRegister): + # Run full JWT policy (RBAC, scope, custom_validate, + # email-domain) via auth_builder, then create the key + # from the validated identity below. + pending_auto_register = resolve_result + # else: None → FALLBACK_TEAM_MAPPING, falls through to + # standard JWT auth_builder below - resolve_result: Final = await _resolve_jwt_to_virtual_key( - jwt_claims=jwt_claims, + if do_standard_jwt_auth: + with tracer.trace("litellm.proxy.auth.jwt_auth_builder"): + result: Final = await JWTAuthManager.auth_builder( + request_data=request_data, + general_settings=general_settings, + api_key=api_key, + jwt_handler=jwt_handler, + route=route, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + parent_otel_span=parent_otel_span, + request_headers=safe_get_request_headers(request), + request_method=RouteChecks.get_request_method(request=request), + ) + + is_proxy_admin: Final = result["is_proxy_admin"] + team_id: Final = result["team_id"] + user_id: Final = result["user_id"] + user_email: Final = result["user_email"] + user_object: Final = result["user_object"] + end_user_id = result["end_user_id"] + org_id: Final = result["org_id"] + jwt_claims = result.get("jwt_claims", None) + agent_id: Final[str | None] = result.get("agent_id") + + if ( + user_object is not None + and isinstance(user_object.metadata, dict) + and user_object.metadata.get("scim_active") is False + ): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.", + ) + + if is_proxy_admin: + # Proxy admins authenticate via auth_builder (full + # access), not via a mapped virtual key. If + # AUTO_REGISTER was pending, cache a sentinel so + # future requests from this JWT identity skip the + # DB mapping lookup in _resolve_jwt_to_virtual_key. + # Without this, every proxy-admin request under + # AUTO_REGISTER re-hits get_jwt_key_mapping_object. + if pending_auto_register is not None: + await user_api_key_cache.async_set_cache( + key=pending_auto_register.cache_key, + value=_JWT_PROXY_ADMIN_SENTINEL, + ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, + ) + return JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) + + valid_token = JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) + + # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. + # JWT policy (RBAC, scope, custom_validate, email-domain) + # has now been enforced by auth_builder above. Create the + # mapping + virtual key from the *validated* identity, then + # replace valid_token with the new key so downstream checks + # use the key-scoped path. + auto_registered: Final = ( + await _auto_register_jwt_mapping( + virtual_key_claim_field=pending_auto_register.claim_field, + claim_value=pending_auto_register.claim_value, jwt_handler=jwt_handler, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, + cache_key=pending_auto_register.cache_key, + jwt_issuer=pending_auto_register.jwt_issuer, + team_id=team_id, + user_id=user_id, + org_id=org_id, + end_user_id=end_user_id, + agent_id=agent_id, + behavior=pending_auto_register.behavior, ) - if isinstance(resolve_result, UserAPIKeyAuth): - mapped_end_user_id: Final = ( - jwt_handler.get_end_user_id( # pyright: ignore[reportUnknownMemberType] # untyped dict token - token=jwt_claims or {}, default_value=None - ) - or None - ) - valid_token = resolve_result.model_copy(update={"end_user_id": mapped_end_user_id}) - api_key = valid_token.token or "" - valid_token.jwt_claims = jwt_claims - do_standard_jwt_auth = False - # Fall through to virtual key checks - if valid_token.user_id is not None and valid_token.user_email is None: - mapped_claims = jwt_claims or {} - mapped_user_email = jwt_handler.get_user_email(token=mapped_claims, default_value=None) - mapped_jwt_user_id: Final = jwt_handler.get_user_id(token=mapped_claims, default_value=None) - if mapped_user_email is not None and mapped_jwt_user_id == valid_token.user_id: - try: - mapped_user_obj: Final = await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - user_email=mapped_user_email, - ) - except Exception as e: - verbose_proxy_logger.debug("JWT mapped-key user_email backfill skipped: %s", e) - else: - if mapped_user_obj is not None: - valid_token.user_email = mapped_user_obj.user_email - elif isinstance(resolve_result, _PendingAutoRegister): - # Run full JWT policy (RBAC, scope, custom_validate, - # email-domain) via auth_builder, then create the key - # from the validated identity below. - pending_auto_register = resolve_result - # else: None → FALLBACK_TEAM_MAPPING, falls through to - # standard JWT auth_builder below + if pending_auto_register is not None and prisma_client is not None + else None + ) + if isinstance(auto_registered, _JWTKeyMappingFailure): + raise HTTPException(status_code=403, detail=auto_registered.value) + if auto_registered is not None: + auto_registered.jwt_claims = jwt_claims + auto_registered.user_email = user_email + # The auto-registered token is built from the new key's + # columns, which carry no user budget. Carry over the + # already-loaded user row rather than re-reading it, or + # the budget check below has nothing to enforce. + auto_registered.user_model_max_budget = ( + user_object.model_max_budget if user_object is not None else None + ) + valid_token = auto_registered + api_key = valid_token.token or "" - if do_standard_jwt_auth: - with tracer.trace("litellm.proxy.auth.jwt_auth_builder"): - result: Final = await JWTAuthManager.auth_builder( - request_data=request_data, - general_settings=general_settings, - api_key=api_key, - jwt_handler=jwt_handler, - route=route, + falls_through_to_key_checks: Final = ( + auto_registered is not None + and ( + jwt_handler.litellm_jwtauth.auto_register_map_existing_key + or ( + pending_auto_register is not None + and pending_auto_register.behavior == UnregisteredJWTClientBehavior.LINK_ONLY + ) + ) + and master_key is not None + ) + if not falls_through_to_key_checks: + # Check if model has zero cost - if so, skip all budget checks + model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + skip_budget_checks = request_skips_budget_checks( + route=route, model=model, llm_router=llm_router, request_body=request_data + ) + + # Fetch project object for JWT path if project_id is set + _jwt_project_obj = None + if valid_token.project_id is not None: + _jwt_project_obj = await get_project_object( + project_id=valid_token.project_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, - parent_otel_span=parent_otel_span, - request_headers=safe_get_request_headers(request), - request_method=RouteChecks.get_request_method(request=request), + ) + if _jwt_project_obj is not None: + valid_token.project_metadata = _jwt_project_obj.metadata + valid_token.project_alias = _jwt_project_obj.project_alias + + # JWT auth returns here rather than falling through to the + # virtual-key checks below, so the user's per-model budget + # has to be enforced on this path too. Without it the + # post-call increment still charges the counter and nothing + # ever reads it, which is worse than not tracking at all. + # Guarded by the same flag the virtual-key path uses, or a + # zero-cost model would be refused here and allowed there, + # while the log above claims all budget checks were skipped. + if not skip_budget_checks: + await _enforce_model_budgets( + _AUTH_CORE.model_budgets().check_user, + _ModelBudgetHost( + request=request, + request_data=request_data, + valid_token=cast(UserAPIKeyAuth, valid_token), + route=route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=_proxy_model_list(), + ), ) - is_proxy_admin: Final = result["is_proxy_admin"] - team_id: Final = result["team_id"] - user_id: Final = result["user_id"] - user_email: Final = result["user_email"] - user_object: Final = result["user_object"] - end_user_id = result["end_user_id"] - org_id: Final = result["org_id"] - jwt_claims = result.get("jwt_claims", None) - agent_id: Final[str | None] = result.get("agent_id") - - if ( - user_object is not None - and isinstance(user_object.metadata, dict) - and user_object.metadata.get("scim_active") is False - ): - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=f"User={user_id} has been deactivated via SCIM. Keys owned by this user cannot be used.", - ) - - if is_proxy_admin: - # Proxy admins authenticate via auth_builder (full - # access), not via a mapped virtual key. If - # AUTO_REGISTER was pending, cache a sentinel so - # future requests from this JWT identity skip the - # DB mapping lookup in _resolve_jwt_to_virtual_key. - # Without this, every proxy-admin request under - # AUTO_REGISTER re-hits get_jwt_key_mapping_object. - if pending_auto_register is not None: - await user_api_key_cache.async_set_cache( - key=pending_auto_register.cache_key, - value=_JWT_PROXY_ADMIN_SENTINEL, - ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, - ) - return JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) - - valid_token = JWTAuthManager.user_api_key_auth_from_result(result, parent_otel_span) - - # AUTO_REGISTER deferred from _resolve_jwt_to_virtual_key. - # JWT policy (RBAC, scope, custom_validate, email-domain) - # has now been enforced by auth_builder above. Create the - # mapping + virtual key from the *validated* identity, then - # replace valid_token with the new key so downstream checks - # use the key-scoped path. - auto_registered: Final = ( - await _auto_register_jwt_mapping( - virtual_key_claim_field=pending_auto_register.claim_field, - claim_value=pending_auto_register.claim_value, - jwt_handler=jwt_handler, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - cache_key=pending_auto_register.cache_key, - jwt_issuer=pending_auto_register.jwt_issuer, - team_id=team_id, - user_id=user_id, - org_id=org_id, - end_user_id=end_user_id, - agent_id=agent_id, - behavior=pending_auto_register.behavior, - ) - if pending_auto_register is not None and prisma_client is not None - else None - ) - if isinstance(auto_registered, _JWTKeyMappingFailure): - raise HTTPException(status_code=403, detail=auto_registered.value) - if auto_registered is not None: - auto_registered.jwt_claims = jwt_claims - auto_registered.user_email = user_email - # The auto-registered token is built from the new key's - # columns, which carry no user budget. Carry over the - # already-loaded user row rather than re-reading it, or - # the budget check below has nothing to enforce. - auto_registered.user_model_max_budget = ( - user_object.model_max_budget if user_object is not None else None - ) - valid_token = auto_registered - api_key = valid_token.token or "" - - falls_through_to_key_checks: Final = ( - auto_registered is not None - and ( - jwt_handler.litellm_jwtauth.auto_register_map_existing_key - or ( - pending_auto_register is not None - and pending_auto_register.behavior == UnregisteredJWTClientBehavior.LINK_ONLY - ) - ) - and master_key is not None - ) - if not falls_through_to_key_checks: - # Check if model has zero cost - if so, skip all budget checks - model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - skip_budget_checks = request_skips_budget_checks( - route=route, model=model, llm_router=llm_router, request_body=request_data - ) - - # Fetch project object for JWT path if project_id is set - _jwt_project_obj = None - if valid_token.project_id is not None: - _jwt_project_obj = await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if _jwt_project_obj is not None: - valid_token.project_metadata = _jwt_project_obj.metadata - valid_token.project_alias = _jwt_project_obj.project_alias - - # JWT auth returns here rather than falling through to the - # virtual-key checks below, so the user's per-model budget - # has to be enforced on this path too. Without it the - # post-call increment still charges the counter and nothing - # ever reads it, which is worse than not tracking at all. - # Guarded by the same flag the virtual-key path uses, or a - # zero-cost model would be refused here and allowed there, - # while the log above claims all budget checks were skipped. - if not skip_budget_checks: - await _enforce_model_budgets( - _AUTH_CORE.model_budgets().check_user, - _ModelBudgetHost( - request=request, - request_data=request_data, - valid_token=cast(UserAPIKeyAuth, valid_token), - route=route, - limiter=model_max_budget_limiter, - llm_router=llm_router, - llm_model_list=_proxy_model_list(), - ), - ) - - return cast(UserAPIKeyAuth, valid_token) - - #### ELSE #### - ## CHECK PASS-THROUGH ENDPOINTS ## - if not custom_auth_api_key: - response = await check_api_key_for_custom_headers_or_pass_through_endpoints( + return cast(UserAPIKeyAuth, valid_token) + key_plan: Final = ( + credential_plan + if isinstance(credential_plan, KeyValidator) + else _AUTH_CORE.plan().after_jwt(planned.facts, plan_settings, issued_key=planned.issued_key) + ) + if isinstance(key_plan, AnonymousValidator): + return _anonymous_auth(_AUTH_CORE.anonymous().authenticate(key_plan.reason)) + if key_plan.pass_through_key: + api_key = await check_api_key_for_custom_headers_or_pass_through_endpoints( request=request, route=route, pass_through_endpoints=pass_through_endpoints, api_key=api_key, ) - if isinstance(response, str): - api_key = response - elif isinstance(response, UserAPIKeyAuth): - return response if master_key is not None and not isinstance(master_key, str): raise HTTPException( status_code=500, diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 2f942af99a6..74bbac8c5ad 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1295,20 +1295,15 @@ class _AuthRequestFacts(TypedDict): token: str | None is_llm_api_route: bool -class _AuthOAuth2Settings(TypedDict): - user_id_header: str | None - team_id_header: str | None +class _AuthIntrospectionSettings(TypedDict): + endpoint_configured: bool + requires_active: bool class _AuthIntrospection(TypedDict): active: bool - subject: str | None - client_id: str | None - scopes: Sequence[str] - -class _AuthCustomIdentity(TypedDict): - user_id: str | None - team_id: str | None - hashed_token: str | None + user_id: JsonValue + role: JsonValue + team_id: JsonValue class _AuthIntrospectStep(TypedDict): kind: Literal["introspect"] @@ -1316,23 +1311,39 @@ class _AuthIntrospectStep(TypedDict): class _AuthOAuth2Identity(TypedDict): kind: Literal["identity"] - user_id: str | None - team_id: str | None - scopes: list[str] + token: str + user_id: JsonValue + role: JsonValue + team_id: JsonValue class _AuthOAuth2Denied(TypedDict): - kind: Literal["missing_key", "inactive_token"] + kind: Literal["inactive_token", "introspection_not_configured"] class _AuthOAuth2Done(TypedDict): kind: Literal["done"] outcome: _AuthOAuth2Identity | _AuthOAuth2Denied +class _AuthCustomIdentity(TypedDict): + kind: Literal["identity"] + user_id: str | None + team_id: str | None + hashed_token: str | None + +class _AuthReplacedKey(TypedDict): + kind: Literal["replaced_key"] + token: str + +class _AuthNoCustomIdentity(TypedDict): + kind: Literal["no_identity"] + +_AuthCustomAuthResult = _AuthCustomIdentity | _AuthReplacedKey | _AuthNoCustomIdentity + class _AuthCustomAuthStep(TypedDict): kind: Literal["custom_auth"] class _AuthCustomAuthDone(TypedDict): kind: Literal["done"] - outcome: _AuthCustomIdentity + outcome: _AuthCustomAuthResult _AuthOAuth2Steps = _AuthIntrospectStep | _AuthOAuth2Done _AuthCustomAuthSteps = _AuthCustomAuthStep | _AuthCustomAuthDone @@ -1340,9 +1351,7 @@ _AuthCustomAuthSteps = _AuthCustomAuthStep | _AuthCustomAuthDone @final class AuthCoreHandle: def __new__(cls) -> AuthCoreHandle: ... - def oauth2( - self, facts: _AuthRequestFacts, settings: _AuthOAuth2Settings, proxy_headers: Mapping[str, str] - ) -> OAuth2Call: ... + def oauth2(self, token: str, settings: _AuthIntrospectionSettings) -> OAuth2Call: ... def custom(self, facts: _AuthRequestFacts) -> CustomAuthCall: ... # Each call raises ValueError when started twice, or answered before it starts, after it ends, or @@ -1355,4 +1364,4 @@ class OAuth2Call: @final class CustomAuthCall: def start(self) -> _AuthCustomAuthSteps: ... - def answer_custom_auth(self, identity: _AuthCustomIdentity) -> _AuthCustomAuthSteps: ... + def answer_custom_auth(self, result: _AuthCustomAuthResult) -> _AuthCustomAuthSteps: ... diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index d3a5fad1765..bdb2678982e 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -12,6 +12,7 @@ from typing_extensions import TypeIs # noqa: TID251 # narrows the before-valid from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.team import Member from litellm.models.verification_token import LiteLLM_VerificationToken +from litellm.proxy.auth.core.credential import is_jwt as is_jwt # noqa: PLC0414 # public re-export from litellm.proxy.auth.core.key import hash_token as hash_token # noqa: PLC0414 # public re-export from litellm.types.agents import AgentCaller, AgentResponse from litellm.types.proxy.agent_identity import ManagedAgentContext @@ -32,13 +33,6 @@ else: Span: TypeAlias = Any -def is_jwt(token: str | None) -> bool: - if token is None: - return False - parts: Final = token.split(".") - return len(parts) == 3 - - _UntypedDict: TypeAlias = dict[Any, Any] diff --git a/tests/unit/proxy/auth/core/test_oauth2.py b/tests/unit/proxy/auth/core/test_oauth2.py new file mode 100644 index 00000000000..937b71c45d2 --- /dev/null +++ b/tests/unit/proxy/auth/core/test_oauth2.py @@ -0,0 +1,124 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field +from ipaddress import ip_network +from types import MappingProxyType +from typing import Final + +import pytest + +from litellm.proxy.auth.core.answers import Introspection +from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings, OAuth2ProxySettings +from litellm.proxy.auth.core.outcome import ( + DisallowedProxyMappings, + InactiveToken, + IntrospectionNotConfigured, + OAuth2Identity, + OAuth2ProxyIdentity, + ProxyMappingsMissing, + TrustedProxyRangesMissing, + UntrustedProxy, +) +from litellm.proxy.auth.core.python_core import PythonCore + +CORE: Final = PythonCore() +TRUSTED: Final = (ip_network("10.0.0.0/8"),) + + +@dataclass +class RecordingIntrospection: + answer: Introspection + calls: list[str] = field(default_factory=list) # mutable-ok: records host calls under test + + async def introspect(self, token: str) -> Introspection: + self.calls.append(token) + return self.answer + + +@dataclass +class RecordingProxy: + ip: str | None + headers: Mapping[str, str] + calls: list[str] = field(default_factory=list) # mutable-ok: records host calls under test + + def client_ip(self) -> str | None: + self.calls.append("client_ip") + return self.ip + + def header(self, name: str) -> str | None: + self.calls.append(name) + return self.headers.get(name) + + +def _introspection(active: bool) -> RecordingIntrospection: + return RecordingIntrospection(Introspection(active=active, user_id="user-1", role="internal_user", team_id=7)) + + +@pytest.mark.parametrize( + ("active", "requires_active", "expected"), + [ + (True, True, OAuth2Identity(token="tok", user_id="user-1", role="internal_user", team_id=7)), + (False, True, InactiveToken()), + (False, False, OAuth2Identity(token="tok", user_id="user-1", role="internal_user", team_id=7)), + ], +) +async def test_introspection_trusts_an_inactive_answer_only_from_a_token_info_endpoint( + active: bool, requires_active: bool, expected: OAuth2Identity | InactiveToken +) -> None: + host: Final = _introspection(active) + settings: Final = OAuth2IntrospectionSettings(endpoint_configured=True, requires_active=requires_active) + assert await CORE.oauth2().introspect(host, "tok", settings) == expected + assert host.calls == ["tok"] + + +async def test_introspection_without_an_endpoint_never_calls_out() -> None: + host: Final = _introspection(True) + settings: Final = OAuth2IntrospectionSettings(endpoint_configured=False, requires_active=True) + assert await CORE.oauth2().introspect(host, "tok", settings) == IntrospectionNotConfigured() + assert host.calls == [] + + +def _proxy_settings(mappings: Mapping[str, str], *, trusted: bool = True) -> OAuth2ProxySettings: + return OAuth2ProxySettings( + trusted_networks=TRUSTED if trusted else (), header_mappings=MappingProxyType(dict(mappings)) + ) + + +def test_proxy_headers_assert_only_the_mapped_identity_fields_that_are_present() -> None: + host: Final = RecordingProxy("10.1.2.3", {"x-user": "user-1", "x-team": "", "x-models": "gpt-a, gpt-b"}) + settings: Final = _proxy_settings({"user_id": "x-user", "team_id": "x-team", "models": "x-models"}) + assert CORE.oauth2().proxy(host, settings) == OAuth2ProxyIdentity( + user_id="user-1", user_email=None, team_id=None, team_alias=None, org_id=None, models=("gpt-a", "gpt-b") + ) + assert host.calls[0] == "client_ip" + + +@pytest.mark.parametrize("ip", ["10.1.2.3", "::ffff:10.1.2.3"]) +def test_proxy_headers_trust_a_peer_inside_the_ranges_including_ipv4_mapped_addresses(ip: str) -> None: + host: Final = RecordingProxy(ip, {"x-user": "user-1"}) + identity: Final = CORE.oauth2().proxy(host, _proxy_settings({"user_id": "x-user"})) + assert isinstance(identity, OAuth2ProxyIdentity) + assert identity.user_id == "user-1" + + +@pytest.mark.parametrize("ip", ["192.168.1.1", None, "not-an-ip"]) +def test_proxy_headers_from_an_untrusted_peer_are_never_read(ip: str | None) -> None: + host: Final = RecordingProxy(ip, {"x-user": "admin"}) + assert CORE.oauth2().proxy(host, _proxy_settings({"user_id": "x-user"})) == UntrustedProxy(client_ip=ip) + assert host.calls == ["client_ip"] + + +def test_proxy_headers_need_trusted_ranges_before_anything_else() -> None: + host: Final = RecordingProxy("10.1.2.3", {}) + assert CORE.oauth2().proxy(host, _proxy_settings({"user_id": "x"}, trusted=False)) == TrustedProxyRangesMissing() + assert host.calls == [] + + +def test_proxy_headers_need_mappings() -> None: + assert CORE.oauth2().proxy(RecordingProxy("10.1.2.3", {}), _proxy_settings({})) == ProxyMappingsMissing() + + +def test_proxy_headers_refuse_every_non_identity_mapping_before_reading_headers() -> None: + host: Final = RecordingProxy("10.1.2.3", {"x-role": "proxy_admin"}) + settings: Final = _proxy_settings({"user_role": "x-role", "user_id": "x-user", "max_budget": "x-budget"}) + assert CORE.oauth2().proxy(host, settings) == DisallowedProxyMappings(fields=("max_budget", "user_role")) + assert host.calls == ["client_ip"] diff --git a/tests/unit/proxy/auth/core/test_plan.py b/tests/unit/proxy/auth/core/test_plan.py new file mode 100644 index 00000000000..7a25ec9f741 --- /dev/null +++ b/tests/unit/proxy/auth/core/test_plan.py @@ -0,0 +1,268 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +import pytest + +from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RequestFacts, RoutingOverride +from litellm.proxy.auth.core.outcome import ( + AnonymousReason, + AnonymousValidator, + CustomValidator, + JwtValidator, + KeyValidator, + OAuth2Mode, + OAuth2Validator, + PremiumFeature, + PremiumRequired, +) +from litellm.proxy.auth.core.python_core import PythonCore + +PLANNER: Final = PythonCore().plan() +JWT: Final = "header.payload.signature" +LLM_ROUTE: Final = "/chat/completions" +SETTINGS: Final = PlanSettings( + custom_auth=False, + premium_user=True, + public_routes=frozenset({"/health/liveness", "/metrics"}), + additional_public_routes=(), + metrics_require_auth=False, + oauth2_introspection=False, + oauth2_proxy=False, + jwt_enabled=False, + oauth2_routing_overrides=(), + unauthenticated_pass_through_routes=frozenset(), +) + + +def _facts( + route: str = LLM_ROUTE, + token: str | None = "sk-key", + *, + is_llm_api_route: bool = True, + is_info_route: bool = False, + claims: Mapping[str, object] | None = None, +) -> PlanFacts: + return PlanFacts( + request=RequestFacts(route=route, token=token, is_llm_api_route=is_llm_api_route), + is_info_route=is_info_route, + unverified_claims=claims, + ) + + +def _selectors(value: str | list[str] | None) -> tuple[str, ...] | None: + if value is None: + return None + return tuple(value) if isinstance(value, list) else (value,) + + +def _override( + iss: str | list[str] | None = None, + client_id: str | list[str] | None = None, + scope: str | list[str] | None = None, + aud: str | list[str] | None = None, +) -> RoutingOverride: + return RoutingOverride( + iss=_selectors(iss), client_id=_selectors(client_id), scope=_selectors(scope), aud=_selectors(aud) + ) + + +def _routes_to_oauth2(override: RoutingOverride, claims: Mapping[str, object]) -> bool: + settings: Final = replace(SETTINGS, jwt_enabled=True, oauth2_routing_overrides=(override,)) + return PLANNER.oauth2_introspection_applies(_facts(token=JWT, claims=claims), settings) + + +def test_custom_auth_runs_before_every_other_path() -> None: + settings: Final = replace(SETTINGS, custom_auth=True, oauth2_introspection=True) + assert PLANNER.plan(_facts("/health/liveness"), settings) == CustomValidator() + + +def test_a_public_route_needs_no_credential_even_when_oauth2_is_on() -> None: + settings: Final = replace(SETTINGS, oauth2_introspection=True, jwt_enabled=True) + assert PLANNER.plan(_facts("/health/liveness", JWT), settings) == AnonymousValidator(AnonymousReason.PUBLIC_ROUTE) + + +@pytest.mark.parametrize("route", ["/metrics", "/metrics/"]) +def test_metrics_requires_a_credential_when_configured(route: str) -> None: + settings: Final = replace(SETTINGS, metrics_require_auth=True, public_routes=frozenset({route})) + assert PLANNER.plan(_facts(route, is_llm_api_route=False), settings) == KeyValidator(pass_through_key=True) + + +@pytest.mark.parametrize( + ("route", "premium_user", "expected_public"), + [ + ("/api/users/123", True, True), + ("/api/users/123", False, False), + ("/spend/calculate", True, True), + ("/spend/calculate/more", True, False), + ("/other", True, False), + ], +) +def test_configured_public_routes_match_exactly_or_by_prefix_wildcard_for_premium_only( + route: str, premium_user: bool, expected_public: bool +) -> None: + settings: Final = replace( + SETTINGS, premium_user=premium_user, additional_public_routes=("/spend/calculate", "/api/*") + ) + plan: Final = PLANNER.plan(_facts(route, is_llm_api_route=False), settings) + assert (plan == AnonymousValidator(AnonymousReason.PUBLIC_ROUTE)) is expected_public + + +@pytest.mark.parametrize( + ("route", "is_llm_api_route", "is_info_route", "expected"), + [ + (LLM_ROUTE, True, False, True), + ("/key/info", False, True, True), + ("/lens", False, False, True), + ("/lens/traces", False, False, True), + ("/lensing", False, False, False), + ("/key/generate", False, False, False), + ], +) +def test_introspection_covers_llm_info_and_lens_routes_only( + route: str, is_llm_api_route: bool, is_info_route: bool, expected: bool +) -> None: + facts: Final = _facts(route, is_llm_api_route=is_llm_api_route, is_info_route=is_info_route) + plan: Final = PLANNER.plan(facts, replace(SETTINGS, oauth2_introspection=True)) + assert (plan == OAuth2Validator(OAuth2Mode.INTROSPECTION)) is expected + + +@pytest.mark.parametrize( + ("feature", "settings", "token"), + [ + (PremiumFeature.OAUTH2, replace(SETTINGS, oauth2_introspection=True), "sk-key"), + (PremiumFeature.JWT, replace(SETTINGS, jwt_enabled=True), JWT), + ], +) +def test_oauth2_introspection_and_jwt_need_premium(feature: PremiumFeature, settings: PlanSettings, token: str) -> None: + assert PLANNER.plan(_facts(token=token), replace(settings, premium_user=False)) == PremiumRequired(feature) + + +def test_introspection_wins_over_proxy_headers_and_proxy_headers_over_jwt() -> None: + both: Final = replace(SETTINGS, oauth2_introspection=True, oauth2_proxy=True, jwt_enabled=True) + assert PLANNER.plan(_facts(), both) == OAuth2Validator(OAuth2Mode.INTROSPECTION) + assert PLANNER.plan(_facts(token=JWT), both) == OAuth2Validator(OAuth2Mode.PROXY) + assert PLANNER.plan(_facts(token=JWT), replace(both, oauth2_proxy=False)) == JwtValidator() + + +def test_a_jwt_skips_global_introspection_unless_an_override_routes_it() -> None: + settings: Final = replace(SETTINGS, oauth2_introspection=True, jwt_enabled=True) + assert PLANNER.plan(_facts(token=JWT, claims={"iss": "idp"}), settings) == JwtValidator() + routed: Final = replace(settings, oauth2_routing_overrides=(_override(iss="idp"),)) + assert PLANNER.plan(_facts(token=JWT, claims={"iss": "idp"}), routed) == OAuth2Validator(OAuth2Mode.INTROSPECTION) + + +def test_a_jwt_is_validated_before_an_unauthenticated_pass_through_route() -> None: + settings: Final = replace(SETTINGS, jwt_enabled=True, unauthenticated_pass_through_routes=frozenset({"/webhook"})) + assert PLANNER.plan(_facts("/webhook", JWT, is_llm_api_route=False), settings) == JwtValidator() + assert PLANNER.plan(_facts("/webhook", is_llm_api_route=False), settings) == AnonymousValidator( + AnonymousReason.PASS_THROUGH_AUTH_OFF + ) + + +@pytest.mark.parametrize( + ("issued_key", "expected"), + [ + (False, AnonymousValidator(AnonymousReason.PASS_THROUGH_AUTH_OFF)), + (True, KeyValidator(pass_through_key=False)), + ], +) +def test_a_key_issued_by_custom_auth_skips_pass_through_handling( + issued_key: bool, expected: AnonymousValidator | KeyValidator +) -> None: + settings: Final = replace(SETTINGS, custom_auth=True, unauthenticated_pass_through_routes=frozenset({"/webhook"})) + facts: Final = _facts("/webhook", is_llm_api_route=False) + assert PLANNER.after_custom_auth(facts, settings, issued_key=issued_key) == expected + assert PLANNER.after_jwt(facts, settings, issued_key=issued_key) == expected + + +def test_after_custom_auth_still_plans_public_routes_and_oauth2() -> None: + settings: Final = replace(SETTINGS, custom_auth=True, oauth2_introspection=True) + assert PLANNER.after_custom_auth(_facts("/health/liveness"), settings, issued_key=True) == AnonymousValidator( + AnonymousReason.PUBLIC_ROUTE + ) + assert PLANNER.after_custom_auth(_facts(), settings, issued_key=True) == OAuth2Validator(OAuth2Mode.INTROSPECTION) + + +def test_a_plain_key_request_is_a_key_validation_with_pass_through_handling() -> None: + assert PLANNER.plan(_facts(), SETTINGS) == KeyValidator(pass_through_key=True) + + +@pytest.mark.parametrize( + ("selector", "claim", "expected", "scope"), + [ + (None, "any-value", True, False), + ("issuer.example.com", "issuer.example.com", True, False), + ("issuer.example.com", "other-issuer.example.com", False, False), + ("trusted.example.com", "trusted.example.com attacker.example.com", False, False), + ("trusted.*", "trusted.example.com attacker.example.com", False, False), + ("trusted.*", "trusted.example.com", True, False), + (["issuer-a.example.com", "issuer-b.example.com"], "issuer-b.example.com", True, False), + ("*MID_LITELLM", "STREAM_MID_LITELLM", True, False), + ("*MID_LITELLM", "REDIS_LITELLM", False, False), + ("machine-??", "machine-01", True, False), + ("machine-??", "machine-001", False, False), + ("*litellm", "BATCH_LITELLM", False, False), + ("*LITELLM", "BATCH_LITELLM", True, False), + ("App:LiteLLM", "App:LiteLLM openid", True, True), + ("App:*", "App:LiteLLM openid", True, True), + (["openid", "App:LiteLLM"], "openid profile", True, True), + (["service-*", "batch-*"], "batch-123", True, False), + (["service-*", "batch-*"], "other-123", False, False), + ("App:LiteLLM", ["openid", "App:LiteLLM"], True, False), + ("App:LiteLLM", None, False, False), + ], +) +def test_routing_selectors_match_claims( + selector: str | list[str] | None, claim: object, expected: bool, scope: bool +) -> None: + override: Final = _override(scope=selector) if scope else _override(iss=selector) + claims: Final = {"scope": claim} if scope else {"iss": claim} + assert _routes_to_oauth2(override, claims) is expected + + +@pytest.mark.parametrize( + ("override", "claims", "expected"), + [ + (_override(iss="oauth-issuer.example.com"), {"iss": "oauth-issuer.example.com"}, True), + ( + _override(iss="oauth-issuer.example.com", scope="App:LiteLLM"), + {"iss": "oauth-issuer.example.com", "scope": "App:LiteLLM openid"}, + True, + ), + ( + _override(iss="oauth-issuer.example.com", client_id="*MID_LITELLM"), + {"iss": "oauth-issuer.example.com", "client_id": "BATCH_MID_LITELLM"}, + True, + ), + ( + _override(iss="oauth-issuer.example.com", client_id="*MID_LITELLM"), + {"iss": "oauth-issuer.example.com", "client_id": "BATCH_PORTAL"}, + False, + ), + ( + _override(iss="oauth-issuer.example.com", aud=["api://litellm", "api://fallback"]), + {"iss": "oauth-issuer.example.com", "aud": ["api://other", "api://litellm"]}, + True, + ), + ( + _override(iss="oauth-issuer.example.com", scope="App:LiteLLM", client_id="*MID_LITELLM"), + {"iss": "oauth-issuer.example.com", "scope": "App:LiteLLM openid", "client_id": "BATCH_MID_LITELLM"}, + True, + ), + ( + _override(iss="oauth-issuer.example.com", scope="App:LiteLLM", client_id="*MID_LITELLM"), + {"iss": "oauth-issuer.example.com", "scope": "App:Other openid", "client_id": "BATCH_MID_LITELLM"}, + False, + ), + ], +) +def test_every_selector_of_an_override_must_match( + override: RoutingOverride, claims: Mapping[str, object], expected: bool +) -> None: + assert _routes_to_oauth2(override, claims) is expected + + +def test_an_override_needs_decodable_claims() -> None: + settings: Final = replace(SETTINGS, jwt_enabled=True, oauth2_routing_overrides=(_override(),)) + assert PLANNER.oauth2_introspection_applies(_facts(token=JWT, claims=None), settings) is False diff --git a/tests/unit/proxy/auth/core/test_python_core.py b/tests/unit/proxy/auth/core/test_python_core.py index 39cfc4431ae..1431e3457e4 100644 --- a/tests/unit/proxy/auth/core/test_python_core.py +++ b/tests/unit/proxy/auth/core/test_python_core.py @@ -1,6 +1,6 @@ import math from collections.abc import Awaitable -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from datetime import datetime, timedelta, timezone from typing import Final @@ -18,6 +18,9 @@ from litellm.proxy.auth.core.answers import ( ) from litellm.proxy.auth.core.inputs import RequestFacts, KeySettings from litellm.proxy.auth.core.outcome import ( + CustomIdentity, + NoCustomIdentity, + ReplacedKey, Allow, CachedAdminKey, Deny, @@ -534,3 +537,26 @@ async def test_a_model_budget_denial_stops_the_key_chain_before_its_context_load outcome, calls = await _run(PythonCore().keys().check(handler, _facts("hash"), SETTINGS, NOW), handler) assert outcome == Deny(exceeded) assert calls[-2:] == ("read_model_budget_scope", "check_user_model_budget") + + +@dataclass +class ScriptedCustomAuth: + answer: CustomIdentity | ReplacedKey | NoCustomIdentity + asked: list[RequestFacts] = field(default_factory=list) # mutable-ok: records host calls under test + + async def custom_auth(self, facts: RequestFacts) -> CustomIdentity | ReplacedKey | NoCustomIdentity: + self.asked.append(facts) + return self.answer + + +@pytest.mark.parametrize( + "answer", + [CustomIdentity(user_id="u", team_id="t", hashed_token=None), ReplacedKey(token="sk-new"), NoCustomIdentity()], +) +async def test_custom_auth_returns_what_the_custom_function_resolved_for_the_request( + answer: CustomIdentity | ReplacedKey | NoCustomIdentity, +) -> None: + host: Final = ScriptedCustomAuth(answer) + facts: Final = RequestFacts(route="/chat/completions", token="sk-1", is_llm_api_route=True) + assert await PythonCore().custom().authenticate(host, facts) == answer + assert host.asked == [facts] diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index f1922a9a73a..a38e3906594 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -85,10 +85,8 @@ from litellm.proxy.auth.user_api_key_auth import ( _ensure_parent_otel_span_on_request_state, _load_team_member_default_model_budget, _PendingAutoRegister, - _matches_routing_override, + _plan_settings, _reserve_budget_after_common_checks, - _route_requires_auth_despite_public, - _routing_selector_matches_claim, _authorize_authenticated_request, run_centralized_common_checks, _run_post_custom_auth_checks, @@ -97,6 +95,8 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth, user_api_key_auth_websocket_for_model, ) +from litellm.proxy.auth.core.inputs import PlanSettings +from litellm.proxy.auth.core.plan import is_public_route from litellm.proxy.pass_through_endpoints import pass_through_endpoints as pass_through from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -203,15 +203,18 @@ def test_get_api_key(): ) == (api_key, passed_in_key) -def test_route_requires_auth_despite_public_for_metrics(monkeypatch): +def _builder_plan_settings(general_settings: dict) -> PlanSettings: + return _plan_settings(general_settings, JWTHandler(), premium_user=False, custom_auth=False) + + +def test_planning_reads_whether_metrics_requires_auth(monkeypatch): monkeypatch.setattr(litellm, "require_auth_for_metrics_endpoint", True) - assert _route_requires_auth_despite_public("/metrics", {}) is True - assert _route_requires_auth_despite_public("/metrics/", {}) is True + assert _builder_plan_settings({}).metrics_require_auth is True monkeypatch.setattr(litellm, "require_auth_for_metrics_endpoint", False) - assert _route_requires_auth_despite_public("/metrics", {}) is False + assert _builder_plan_settings({}).metrics_require_auth is False def test_public_ai_hub_routes_remain_public(): @@ -223,7 +226,7 @@ def test_public_ai_hub_routes_remain_public(): "/public/skill_hub", ): assert route in LiteLLMRoutes.public_routes.value - assert _route_requires_auth_despite_public(route, {}) is False + assert is_public_route(route, _builder_plan_settings({})) is True @pytest.mark.asyncio @@ -920,146 +923,6 @@ def _assert_get_api_key_with_custom_litellm_key_header(custom_litellm_key_header ) == (api_key, passed_in_key) -@pytest.mark.parametrize( - "selector_value, claim_value, expected, split_space_delimited", - [ - (None, "any-value", True, False), - ("issuer.example.com", "issuer.example.com", True, False), - ("issuer.example.com", "other-issuer.example.com", False, False), - # iss (and other non-scope claims) must not match via space-split injection - ( - "trusted.example.com", - "trusted.example.com attacker.example.com", - False, - False, - ), - # Wildcard iss must not match space-containing claim strings (fnmatch * spans spaces) - ( - "trusted.*", - "trusted.example.com attacker.example.com", - False, - False, - ), - ("trusted.*", "trusted.example.com", True, False), - ( - ["issuer-a.example.com", "issuer-b.example.com"], - "issuer-b.example.com", - True, - False, - ), - ("*MID_LITELLM", "STREAM_MID_LITELLM", True, False), - ("*MID_LITELLM", "REDIS_LITELLM", False, False), - ("machine-??", "machine-01", True, False), - ("machine-??", "machine-001", False, False), - # Wildcard matching is case-sensitive (fnmatch.fnmatchcase) - ("*litellm", "BATCH_LITELLM", False, False), - ("*LITELLM", "BATCH_LITELLM", True, False), - ("App:LiteLLM", "App:LiteLLM openid", True, True), - ("App:*", "App:LiteLLM openid", True, True), - (["openid", "App:LiteLLM"], "openid profile", True, True), - (["service-*", "batch-*"], "batch-123", True, False), - (["service-*", "batch-*"], "other-123", False, False), - ("App:LiteLLM", ["openid", "App:LiteLLM"], True, False), - ("App:LiteLLM", None, False, False), - ], -) -def test_routing_selector_matches_claim_parametrized(selector_value, claim_value, expected, split_space_delimited): - assert ( - _routing_selector_matches_claim( - selector_value=selector_value, - claim_value=claim_value, - split_space_delimited=split_space_delimited, - ) - is expected - ) - - -@pytest.mark.parametrize( - "override, token_claims, expected", - [ - # Only iss selector is required and should match. - ( - JWTRoutingOverride(iss="oauth-issuer.example.com", path="oauth2"), - {"iss": "oauth-issuer.example.com"}, - True, - ), - # Scope selector narrows the match. - ( - JWTRoutingOverride( - iss="oauth-issuer.example.com", - scope="App:LiteLLM", - path="oauth2", - ), - {"iss": "oauth-issuer.example.com", "scope": "App:LiteLLM openid"}, - True, - ), - # client_id wildcard selector narrows the match. - ( - JWTRoutingOverride( - iss="oauth-issuer.example.com", - client_id="*MID_LITELLM", - path="oauth2", - ), - {"iss": "oauth-issuer.example.com", "client_id": "BATCH_MID_LITELLM"}, - True, - ), - ( - JWTRoutingOverride( - iss="oauth-issuer.example.com", - client_id="*MID_LITELLM", - path="oauth2", - ), - {"iss": "oauth-issuer.example.com", "client_id": "BATCH_PORTAL"}, - False, - ), - # aud selector still works with list claims. - ( - JWTRoutingOverride( - iss="oauth-issuer.example.com", - aud=["api://litellm", "api://fallback"], - path="oauth2", - ), - { - "iss": "oauth-issuer.example.com", - "aud": ["api://other", "api://litellm"], - }, - True, - ), - # All provided selectors are AND-ed. - ( - JWTRoutingOverride( - iss="oauth-issuer.example.com", - scope="App:LiteLLM", - client_id="*MID_LITELLM", - path="oauth2", - ), - { - "iss": "oauth-issuer.example.com", - "scope": "App:LiteLLM openid", - "client_id": "BATCH_MID_LITELLM", - }, - True, - ), - ( - JWTRoutingOverride( - iss="oauth-issuer.example.com", - scope="App:LiteLLM", - client_id="*MID_LITELLM", - path="oauth2", - ), - { - "iss": "oauth-issuer.example.com", - "scope": "App:Other openid", - "client_id": "BATCH_MID_LITELLM", - }, - False, - ), - ], -) -def test_matches_routing_override_parametrized(override, token_claims, expected): - assert _matches_routing_override(token_claims=token_claims, override=override) is expected - - def test_get_api_key_with_custom_litellm_key_header_bearer_prefix(): token = "sk-" + "1" * 8 header = f"Bearer {token}" @@ -2113,7 +1976,7 @@ async def test_standard_jwt_auth_propagates_user_email(): ): result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -2183,7 +2046,7 @@ async def test_jwt_auth_propagates_agent_id_to_user_api_key_auth(is_proxy_admin: ): result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -2748,7 +2611,7 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( ): call = user_api_key_auth( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -2857,7 +2720,7 @@ async def test_auto_register_first_request_propagates_user_email(active: bool) - with pytest.raises(ProxyException, match="deactivated via SCIM") as exc: await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -2869,7 +2732,7 @@ async def test_auto_register_first_request_propagates_user_email(active: bool) - return result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -3089,7 +2952,7 @@ async def test_jwt_auto_register_forwards_bound_agent_id(): ): result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -3336,7 +3199,9 @@ class TestJWTOAuth2Coexistence: async def test_auto_register_passes_validated_org_context_to_generated_key(self): jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" general_settings = {"enable_jwt_auth": True} - user_api_key_cache: Final = _validated_jwt_admission_cache(org_id="validated-org", end_user_id="validated-end-user") + user_api_key_cache: Final = _validated_jwt_admission_cache( + org_id="validated-org", end_user_id="validated-end-user" + ) prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -3405,7 +3270,7 @@ class TestJWTOAuth2Coexistence: ): result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -3486,7 +3351,7 @@ class TestJWTOAuth2Coexistence: ): result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -3562,7 +3427,7 @@ class TestJWTOAuth2Coexistence: ): result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -3634,7 +3499,7 @@ class TestJWTOAuth2Coexistence: ): result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, @@ -12509,7 +12374,7 @@ async def test_auth_flow_enters_virtual_key_mapping_when_only_an_issuer_configur ): result = await user_api_key_auth_builder( request=mock_request, - api_key=jwt_token, + api_key=f"Bearer {jwt_token}", azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_auth_default.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_auth_default.py index 7771ff0274e..006bd3f8f8f 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_auth_default.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_auth_default.py @@ -30,7 +30,9 @@ from starlette.requests import Request from litellm.proxy._types import PassThroughGenericEndpoint, ProxyException +from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import ( + _plan_settings, check_api_key_for_custom_headers_or_pass_through_endpoints, ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( @@ -86,58 +88,23 @@ async def test_register_passthrough_with_auth_true_works_for_oss(monkeypatch: py assert response.json()["error"]["type"] == "auth_error" -@pytest.mark.asyncio -async def test_runtime_check_treats_missing_auth_key_as_authenticated(): - # The runtime dispatch in user_api_key_auth pulls - # pass_through_endpoints from general_settings as raw dicts (the - # Pydantic default never applies). A dict without an ``auth`` key - # must default to "authenticated" — without this, the previous - # behaviour (``endpoint.get("auth") is not True`` -> True -> empty - # auth) ships an unauthenticated forwarder. - request = MagicMock() - request.headers = {} - raw_endpoint_no_auth_key = { - "path": "/forwarder", - "target": "https://example.com", - # ``auth`` deliberately omitted - } - - result = await check_api_key_for_custom_headers_or_pass_through_endpoints( - request=request, - route="/forwarder", - pass_through_endpoints=[raw_endpoint_no_auth_key], - api_key="sk-9876", +@pytest.mark.parametrize( + ("raw_endpoint", "unauthenticated"), + [ + ({"path": "/forwarder", "target": "https://example.com"}, frozenset()), + ({"path": "/public-webhook", "target": "https://example.com", "auth": False}, frozenset({"/public-webhook"})), + ({"path": "/forwarder", "target": "https://example.com", "auth": True}, frozenset()), + ], +) +def test_only_an_explicit_auth_false_makes_a_raw_pass_through_route_unauthenticated( + raw_endpoint: dict[str, object], unauthenticated: frozenset[str] +) -> None: + # The runtime dispatch reads pass_through_endpoints from general_settings as raw dicts, where the + # Pydantic default never applies, so a dict without ``auth`` must still require a key. + settings: Final = _plan_settings( + {"pass_through_endpoints": [raw_endpoint]}, JWTHandler(), premium_user=False, custom_auth=False ) - - # Result is the api_key string (auth is REQUIRED for this endpoint - # — flow continues to normal key validation), NOT an empty - # ``UserAPIKeyAuth()`` (which was the unauthenticated-forwarder - # bug). - assert result == "sk-9876" - - -@pytest.mark.asyncio -async def test_runtime_check_explicit_auth_false_still_skips_validation(): - # Operators who explicitly set ``auth: False`` get the legacy - # behaviour — an empty UserAPIKeyAuth, no key required. - from litellm.proxy._types import UserAPIKeyAuth - - request = MagicMock() - request.headers = {} - raw_endpoint_auth_false = { - "path": "/public-webhook", - "target": "https://example.com", - "auth": False, - } - - result = await check_api_key_for_custom_headers_or_pass_through_endpoints( - request=request, - route="/public-webhook", - pass_through_endpoints=[raw_endpoint_auth_false], - api_key="", - ) - - assert isinstance(result, UserAPIKeyAuth) + assert settings.unauthenticated_pass_through_routes == unauthenticated @pytest.mark.asyncio diff --git a/tests/unit/rust_bridge/test_auth_core.py b/tests/unit/rust_bridge/test_auth_core.py index 4e78b8ac146..f2198fd58f4 100644 --- a/tests/unit/rust_bridge/test_auth_core.py +++ b/tests/unit/rust_bridge/test_auth_core.py @@ -6,7 +6,7 @@ from litellm.rust_bridge._native import AuthCoreHandle if TYPE_CHECKING: from litellm.rust_bridge._native import ( - _AuthCustomIdentity, # pyright: ignore[reportPrivateUsage] # the stub's step shapes + _AuthCustomAuthResult, # pyright: ignore[reportPrivateUsage] # the stub's step shapes _AuthOAuth2Steps, # pyright: ignore[reportPrivateUsage] # see above _AuthRequestFacts, # pyright: ignore[reportPrivateUsage] # see above ) @@ -15,17 +15,15 @@ FACTS: Final["_AuthRequestFacts"] = {"route": "/chat/completions", "token": "tok def _oauth2_after_introspection(active: bool) -> "_AuthOAuth2Steps": - call: Final = AuthCoreHandle().oauth2(FACTS, {"user_id_header": None, "team_id_header": None}, {}) + call: Final = AuthCoreHandle().oauth2("tok-1", {"endpoint_configured": True, "requires_active": True}) assert call.start() == {"kind": "introspect", "token": "tok-1"} - return call.answer_introspection( - {"active": active, "subject": "user-1", "client_id": "client-1", "scopes": ("read", "write")} - ) + return call.answer_introspection({"active": active, "user_id": "user-1", "role": None, "team_id": 7}) @pytest.mark.parametrize( ("active", "outcome"), [ - (True, {"kind": "identity", "user_id": "user-1", "team_id": None, "scopes": ["read", "write"]}), + (True, {"kind": "identity", "token": "tok-1", "user_id": "user-1", "role": None, "team_id": 7}), (False, {"kind": "inactive_token"}), ], ) @@ -33,14 +31,21 @@ def test_oauth2_reads_python_s_introspection_and_returns_the_verdict(active: boo assert _oauth2_after_introspection(active) == {"kind": "done", "outcome": outcome} -def test_custom_auth_returns_the_identity_python_resolved() -> None: +@pytest.mark.parametrize( + "result", + [ + {"kind": "identity", "user_id": "user-1", "team_id": "team-1", "hashed_token": "hashed-1"}, + {"kind": "replaced_key", "token": "sk-new"}, + {"kind": "no_identity"}, + ], +) +def test_custom_auth_returns_what_python_resolved(result: "_AuthCustomAuthResult") -> None: call: Final = AuthCoreHandle().custom(FACTS) assert call.start() == {"kind": "custom_auth"} - identity: Final[_AuthCustomIdentity] = {"user_id": "user-1", "team_id": "team-1", "hashed_token": "hashed-1"} - assert call.answer_custom_auth(identity) == {"kind": "done", "outcome": identity} + assert call.answer_custom_auth(result) == {"kind": "done", "outcome": result} def test_an_answer_before_the_call_starts_is_a_value_error() -> None: call: Final = AuthCoreHandle().custom(FACTS) with pytest.raises(ValueError, match="has not started"): - call.answer_custom_auth({"user_id": None, "team_id": None, "hashed_token": None}) + call.answer_custom_auth({"kind": "no_identity"})