mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
refactor(auth): fold every facet into AuthCore and plan the credential path in the core
AuthCore now hands out all eight facets. TargetAuthCore is gone, and RustCore stubs every accessor. selected_core.py is the one place the proxy picks PythonCore. PythonCore decides which credential path a request takes (custom auth, public routes, OAuth2 introspection, OAuth2 proxy headers, JWT, unauthenticated pass-through routes, keys), plus OAuth2 introspection, the OAuth2 proxy-header checks and custom auth. user_api_key_auth gathers facts, asks plan() for the path, runs that facet and projects the result. The premium, public-route, routing-override and pass-through auth-off decisions no longer live in the proxy, and is_jwt and ip_in_networks have one implementation in the core. JWT and authorize still run in the proxy. gateway-auth-core mirrors the new signatures, with OAuth2 and custom auth ported, and the python-bridge handle follows them.
This commit is contained in:
Yujong Lee
parent
8eb6fd0f14
commit
988a9e4779
38 changed files with 2210 additions and 1178 deletions
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
pub client_id: Option<String>,
|
||||
pub scopes: Vec<String>,
|
||||
pub user_id: Value,
|
||||
pub role: Value,
|
||||
pub team_id: Value,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
|
|
|
|||
|
|
@ -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<H: OAuth2Host>(
|
||||
fn introspect<H: OAuth2IntrospectionHost>(
|
||||
&self,
|
||||
host: &H,
|
||||
facts: &RequestFacts,
|
||||
settings: &OAuth2Settings,
|
||||
token: &str,
|
||||
settings: &OAuth2IntrospectionSettings,
|
||||
) -> impl Future<Output = Result<Result<OAuth2Identity, OAuth2Denial>, H::Error>> + Send;
|
||||
fn proxy<H: OAuth2ProxyHost>(
|
||||
&self,
|
||||
host: &H,
|
||||
settings: &OAuth2ProxySettings,
|
||||
) -> Result<OAuth2ProxyIdentity, OAuth2ProxyDenial>;
|
||||
}
|
||||
|
||||
pub trait CustomAuthCore: Sync {
|
||||
|
|
@ -78,7 +93,7 @@ pub trait CustomAuthCore: Sync {
|
|||
&self,
|
||||
host: &H,
|
||||
facts: &RequestFacts,
|
||||
) -> impl Future<Output = Result<CustomIdentity, H::Error>> + Send;
|
||||
) -> impl Future<Output = Result<CustomAuthResult, H::Error>> + 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;
|
||||
|
|
|
|||
|
|
@ -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<Output = Result<Option<String>, Self::Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait OAuth2Host: Sync {
|
||||
pub trait OAuth2IntrospectionHost: Sync {
|
||||
type Error: Error + Send + Sync + 'static;
|
||||
|
||||
fn introspect(
|
||||
&self,
|
||||
token: &str,
|
||||
) -> impl Future<Output = Result<Introspection, Self::Error>> + Send;
|
||||
fn proxy_headers(&self) -> &BTreeMap<String, String>;
|
||||
}
|
||||
|
||||
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<Output = Result<CustomIdentity, Self::Error>> + Send;
|
||||
) -> impl Future<Output = Result<CustomAuthResult, Self::Error>> + Send;
|
||||
}
|
||||
|
||||
pub trait AuthorizeHost: Sync {
|
||||
|
|
|
|||
|
|
@ -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<Map<String, Value>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct RoutingOverride {
|
||||
pub iss: Option<Vec<String>>,
|
||||
pub client_id: Option<Vec<String>>,
|
||||
pub scope: Option<Vec<String>>,
|
||||
pub aud: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct PlanSettings {
|
||||
pub custom_auth: bool,
|
||||
pub jwt_enabled: bool,
|
||||
pub premium_user: bool,
|
||||
pub public_routes: BTreeSet<String>,
|
||||
pub additional_public_routes: Vec<String>,
|
||||
pub metrics_require_auth: bool,
|
||||
pub oauth2_introspection: bool,
|
||||
pub oauth2_proxy: bool,
|
||||
pub public_routes: BTreeSet<String>,
|
||||
pub jwt_enabled: bool,
|
||||
pub oauth2_routing_overrides: Vec<RoutingOverride>,
|
||||
pub unauthenticated_pass_through_routes: BTreeSet<String>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
pub team_id_header: Option<String>,
|
||||
#[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<IpNet>,
|
||||
pub header_mappings: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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::<IpAddr>()) 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<H: OAuth2ProxyHost>(
|
||||
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<H: OAuth2Host>(
|
||||
async fn introspect<H: OAuth2IntrospectionHost>(
|
||||
&self,
|
||||
host: &H,
|
||||
facts: &RequestFacts,
|
||||
_settings: &OAuth2Settings,
|
||||
token: &str,
|
||||
settings: &OAuth2IntrospectionSettings,
|
||||
) -> Result<Result<OAuth2Identity, OAuth2Denial>, 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<H: OAuth2ProxyHost>(
|
||||
&self,
|
||||
host: &H,
|
||||
settings: &OAuth2ProxySettings,
|
||||
) -> Result<OAuth2ProxyIdentity, OAuth2ProxyDenial> {
|
||||
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<String> = 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<CustomIdentity, H::Error> {
|
||||
) -> Result<CustomAuthResult, H::Error> {
|
||||
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<String, String>,
|
||||
struct ScriptedIntrospection {
|
||||
answer: Introspection,
|
||||
introspected: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl OAuth2Host for ScriptedOAuth2Host {
|
||||
impl OAuth2IntrospectionHost for ScriptedIntrospection {
|
||||
type Error = Infallible;
|
||||
|
||||
async fn introspect(&self, token: &str) -> Result<Introspection, Infallible> {
|
||||
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<String, String> {
|
||||
&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<OAuth2Identity, OAuth2Denial>,
|
||||
) {
|
||||
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<Vec<String>>,
|
||||
}
|
||||
|
||||
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<IpNet> {
|
||||
vec!["10.0.0.0/8".parse().unwrap()]
|
||||
}
|
||||
|
||||
fn settings(networks: Vec<IpNet>, 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<IpNet>) {
|
||||
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<IpNet>,
|
||||
#[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<IpNet>) {
|
||||
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<IpNet>,
|
||||
#[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<CustomIdentity, Infallible> {
|
||||
async fn custom_auth(&self, _facts: &RequestFacts) -> Result<CustomAuthResult, Infallible> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<WithinModelBudgets, ModelBudgetExceeded>;
|
|||
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<String>,
|
||||
pub user_email: Option<String>,
|
||||
pub team_id: Option<String>,
|
||||
pub scopes: Vec<String>,
|
||||
pub team_alias: Option<String>,
|
||||
pub org_id: Option<String>,
|
||||
pub models: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
|
|
@ -94,14 +110,37 @@ pub struct CustomIdentity {
|
|||
pub hashed_token: Option<String>,
|
||||
}
|
||||
|
||||
#[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<String> },
|
||||
ProxyMappingsMissing,
|
||||
DisallowedProxyMappings { fields: Vec<String> },
|
||||
}
|
||||
|
||||
#[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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<HostCall>,
|
||||
introspection: Option<Introspection>,
|
||||
custom_identity: Option<CustomIdentity>,
|
||||
custom_auth: Option<CustomAuthResult>,
|
||||
}
|
||||
|
||||
type Shared = Arc<Mutex<Exchange>>;
|
||||
|
|
@ -46,14 +46,12 @@ type Slot<T> = fn(&mut Exchange) -> &mut Option<T>;
|
|||
|
||||
struct StepHost {
|
||||
exchange: Shared,
|
||||
proxy_headers: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
impl StepHost {
|
||||
fn new(exchange: &Shared, proxy_headers: BTreeMap<String, String>) -> 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<Introspection, Self::Error> {
|
||||
|
|
@ -82,20 +80,14 @@ impl OAuth2Host for StepHost {
|
|||
};
|
||||
Ok(self.ask(call, |exchange| &mut exchange.introspection).await)
|
||||
}
|
||||
|
||||
fn proxy_headers(&self) -> &BTreeMap<String, String> {
|
||||
&self.proxy_headers
|
||||
}
|
||||
}
|
||||
|
||||
impl CustomAuthHost for StepHost {
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
async fn custom_auth(&self, _facts: &RequestFacts) -> Result<CustomIdentity, Self::Error> {
|
||||
async fn custom_auth(&self, _facts: &RequestFacts) -> Result<CustomAuthResult, Self::Error> {
|
||||
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<OAuth2Identity>),
|
||||
InactiveToken,
|
||||
IntrospectionNotConfigured,
|
||||
}
|
||||
|
||||
impl From<Result<OAuth2Identity, OAuth2Denial>> for OAuth2Verdict {
|
||||
fn from(outcome: Result<OAuth2Identity, OAuth2Denial>) -> 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<T> Driver<T> {
|
|||
|
||||
fn oauth2_driver(
|
||||
core: NativeCore,
|
||||
facts: RequestFacts,
|
||||
settings: OAuth2Settings,
|
||||
proxy_headers: BTreeMap<String, String>,
|
||||
token: String,
|
||||
settings: OAuth2IntrospectionSettings,
|
||||
) -> Driver<OAuth2Verdict> {
|
||||
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<CustomIdentity> {
|
||||
fn custom_driver(core: NativeCore, facts: RequestFacts) -> Driver<CustomAuthResult> {
|
||||
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<OAuth2Call> {
|
||||
fn oauth2(&self, token: String, settings: &Bound<'_, PyAny>) -> PyResult<OAuth2Call> {
|
||||
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<Driver<CustomIdentity>>);
|
||||
pub(crate) struct CustomAuthCall(Mutex<Driver<CustomAuthResult>>);
|
||||
|
||||
#[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<Py<PyAny>> {
|
||||
let identity = from_py_argument(identity)?;
|
||||
fn answer_custom_auth(&self, py: Python<'_>, result: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
|
||||
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<OAuth2Verdict> {
|
||||
let settings = OAuth2Settings {
|
||||
user_id_header: None,
|
||||
team_id_header: None,
|
||||
fn oauth2(endpoint_configured: bool) -> Driver<OAuth2Verdict> {
|
||||
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<CustomAuthResult>,
|
||||
result: CustomAuthResult,
|
||||
) -> Result<Step<CustomAuthResult>, DriveError> {
|
||||
driver.answer("custom_auth", |exchange| &mut exchange.custom_auth, result)
|
||||
}
|
||||
|
||||
fn wire<T: Serialize>(step: &Step<T>) -> 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)),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
24
litellm/proxy/auth/core/network.py
Normal file
24
litellm/proxy/auth/core/network.py
Normal file
|
|
@ -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
|
||||
63
litellm/proxy/auth/core/oauth2.py
Normal file
63
litellm/proxy/auth/core/oauth2.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
132
litellm/proxy/auth/core/plan.py
Normal file
132
litellm/proxy/auth/core/plan.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
)
|
||||
|
|
|
|||
6
litellm/proxy/auth/selected_core.py
Normal file
6
litellm/proxy/auth/selected_core.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 <base64(username:password)>'}
|
||||
# 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,
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
||||
|
|
|
|||
124
tests/unit/proxy/auth/core/test_oauth2.py
Normal file
124
tests/unit/proxy/auth/core/test_oauth2.py
Normal file
|
|
@ -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"]
|
||||
268
tests/unit/proxy/auth/core/test_plan.py
Normal file
268
tests/unit/proxy/auth/core/test_plan.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue