From 9a79b2b2a9d4c5135423dbaf4d46939bfa29604e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 10 Nov 2026 22:38:27 +0900 Subject: [PATCH] refactor(auth): split key authentication from authorization keys() now answers only who the credential belongs to, returning an identity or the reason the key is unusable, and the new authorize() facet decides whether that key may make the request: master-key-only routes, model access, team member, key and per-model budgets, and the proxy budget alert. outcome.py is split into identity.py (identities and authentication denials), outcome.py (the authorization verdict) and validators.py (the planner's choices), and gateway-auth-core mirrors the same split. Because every credential check now runs before any authorization check, three answers change. An expired key whose team member is over budget gets the 401 expired-key error instead of the budget error, and no team member budget alert fires. A blocked or unknown virtual key on a master-key-only route gets the blocked or invalid-key 401 instead of the master-key-only 401. A virtual key on a master-key-only route against a proxy without a database gets the 400 no-db error instead of the master-key-only 401. --- .../crates/gateway-auth-core/AGENTS.md | 4 +- .../crates/gateway-auth-core/src/answers.rs | 15 - .../crates/gateway-auth-core/src/facets.rs | 23 +- .../crates/gateway-auth-core/src/host.rs | 52 +-- .../crates/gateway-auth-core/src/identity.rs | 126 +++++++ .../crates/gateway-auth-core/src/inputs.rs | 7 +- .../crates/gateway-auth-core/src/lib.rs | 22 +- .../crates/gateway-auth-core/src/native.rs | 23 +- .../crates/gateway-auth-core/src/outcome.rs | 152 +------- .../gateway-auth-core/src/validators.rs | 33 ++ litellm/proxy/auth/auth_denial.py | 15 +- litellm/proxy/auth/core/AGENTS.md | 12 +- litellm/proxy/auth/core/answers.py | 15 - litellm/proxy/auth/core/core.py | 41 +- litellm/proxy/auth/core/host.py | 53 +-- litellm/proxy/auth/core/identity.py | 191 ++++++++++ litellm/proxy/auth/core/inputs.py | 7 +- litellm/proxy/auth/core/model_budgets.py | 4 +- litellm/proxy/auth/core/oauth2.py | 6 +- litellm/proxy/auth/core/outcome.py | 260 +------------ litellm/proxy/auth/core/plan.py | 6 +- litellm/proxy/auth/core/python_core.py | 184 +++++---- litellm/proxy/auth/core/validators.py | 52 +++ litellm/proxy/auth/oauth2_check.py | 6 +- litellm/proxy/auth/oauth2_proxy_hook.py | 6 +- litellm/proxy/auth/user_api_key_auth.py | 132 ++++--- .../proxy/auth/core/test_model_budgets.py | 6 +- tests/unit/proxy/auth/core/test_oauth2.py | 4 +- tests/unit/proxy/auth/core/test_plan.py | 8 +- .../unit/proxy/auth/core/test_python_core.py | 355 ++++++++++-------- 30 files changed, 949 insertions(+), 871 deletions(-) create mode 100644 litellm-rust/crates/gateway-auth-core/src/identity.rs create mode 100644 litellm-rust/crates/gateway-auth-core/src/validators.rs create mode 100644 litellm/proxy/auth/core/identity.py create mode 100644 litellm/proxy/auth/core/validators.py diff --git a/litellm-rust/crates/gateway-auth-core/AGENTS.md b/litellm-rust/crates/gateway-auth-core/AGENTS.md index e5a6b238c56..3d2c89d2bb6 100644 --- a/litellm-rust/crates/gateway-auth-core/AGENTS.md +++ b/litellm-rust/crates/gateway-auth-core/AGENTS.md @@ -1,9 +1,9 @@ # gateway-auth-core -The Rust side of `litellm/proxy/auth/core`. It mirrors that package's facets, hosts, inputs, outcomes and answers type for type, so a facet can move from `PythonCore` into Rust without changing what the proxy hands in or gets back. Read `litellm/proxy/auth/core/AGENTS.md` first, its rules apply here unchanged +The Rust side of `litellm/proxy/auth/core`. It mirrors that package's facets, hosts, inputs, identities, outcomes, validators and answers type for type, so a facet can move from `PythonCore` into Rust without changing what the proxy hands in or gets back. Read `litellm/proxy/auth/core/AGENTS.md` first, its rules apply here unchanged 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`, 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 +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. A Python union of an identity and a denial is a `Result` here, as in `KeyAuthentication` `NativeCore` implements `AuthCore` with every facet `Unported`, which panics if called, the same as `RustCore` in Python. The crate exists so the types stay in step with 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 diff --git a/litellm-rust/crates/gateway-auth-core/src/answers.rs b/litellm-rust/crates/gateway-auth-core/src/answers.rs index 6ad369fc3f4..efa927f6e38 100644 --- a/litellm-rust/crates/gateway-auth-core/src/answers.rs +++ b/litellm-rust/crates/gateway-auth-core/src/answers.rs @@ -81,18 +81,3 @@ pub struct Introspection { pub role: Value, pub team_id: Value, } - -#[derive(Clone, Debug, PartialEq)] -pub struct TeamFacts { - pub blocked: bool, - pub models: Vec, - pub spend: f64, - pub max_budget: Option, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct OrganizationFacts { - pub models: Vec, - pub spend: f64, - pub max_budget: Option, -} diff --git a/litellm-rust/crates/gateway-auth-core/src/facets.rs b/litellm-rust/crates/gateway-auth-core/src/facets.rs index fe8dfefc921..c5295485004 100644 --- a/litellm-rust/crates/gateway-auth-core/src/facets.rs +++ b/litellm-rust/crates/gateway-auth-core/src/facets.rs @@ -5,15 +5,17 @@ use crate::{ AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost, OAuth2ProxyHost, TeamMemberModelBudgetHost, }, + identity::{ + Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, Identity, KeyAuthentication, + KeyDenial, KeyIdentity, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial, + OAuth2ProxyIdentity, VirtualKey, + }, inputs::{ AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings, OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts, }, - outcome::{ - Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, Denial, Identity, KeyPlan, - ModelBudgetResult, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial, OAuth2ProxyIdentity, - Outcome, Plan, - }, + outcome::{ModelBudgetResult, Outcome}, + validators::{CredentialPlan, KeyPlan, Plan}, }; pub trait KeyCore: Sync { @@ -23,14 +25,13 @@ pub trait KeyCore: Sync { facts: &RequestFacts, settings: &KeySettings, now: SystemTime, - ) -> impl Future> + Send; - fn check( + ) -> impl Future> + Send; + fn verify( &self, host: &H, facts: &RequestFacts, - settings: &KeySettings, now: SystemTime, - ) -> impl Future> + Send; + ) -> impl Future, H::Error>> + Send; } pub trait ModelBudgetCore: Sync { @@ -71,7 +72,7 @@ pub trait JwtCore: Sync { facts: &RequestFacts, settings: &JwtSettings, now: SystemTime, - ) -> impl Future, H::Error>> + Send; + ) -> impl Future, H::Error>> + Send; } pub trait OAuth2Core: Sync { @@ -100,7 +101,7 @@ pub trait AuthorizeCore: Sync { fn authorize( &self, host: &H, - identity: &Identity, + identity: &KeyIdentity, facts: &RequestFacts, settings: &AuthorizeSettings, ) -> impl Future> + Send; diff --git a/litellm-rust/crates/gateway-auth-core/src/host.rs b/litellm-rust/crates/gateway-auth-core/src/host.rs index 3605d4474d4..45357bd7e58 100644 --- a/litellm-rust/crates/gateway-auth-core/src/host.rs +++ b/litellm-rust/crates/gateway-auth-core/src/host.rs @@ -2,12 +2,12 @@ use std::{error::Error, future::Future}; use crate::{ answers::{ - FallbackTarget, Introspection, KeyFacts, ModelBudgetScope, OrganizationFacts, OwnerFacts, - PresentedKey, SigningKey, TeamFacts, TeamMember, TeamMemberBudget, - TeamMemberModelBudgetScope, + FallbackTarget, Introspection, KeyFacts, ModelBudgetScope, OwnerFacts, PresentedKey, + SigningKey, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope, }, + identity::{CustomAuthResult, JwtIdentity}, inputs::RequestFacts, - outcome::{CustomAuthResult, Identity, JwtIdentity, ModelBudgetExceeded}, + outcome::ModelBudgetExceeded, }; pub trait ModelBudgetHost: Sync { @@ -60,9 +60,6 @@ pub trait TeamMemberModelBudgetHost: Sync { pub trait KeyHost: Sync { type Error: Error + Send + Sync + 'static; - type ModelBudgets: ModelBudgetHost; - - fn model_budgets(&self) -> Self::ModelBudgets; fn resolve_end_user( &self, @@ -82,11 +79,19 @@ pub trait KeyHost: Sync { loaded_from_db: bool, ) -> impl Future> + Send; fn inspect_key(&self) -> impl Future> + Send; - fn enforce_model_access(&self) -> impl Future> + Send; fn load_key_owner( &self, user_id: &str, ) -> impl Future, Self::Error>> + Send; +} + +pub trait AuthorizeHost: Sync { + type Error: Error + Send + Sync + 'static; + type ModelBudgets: ModelBudgetHost; + + fn model_budgets(&self) -> Self::ModelBudgets; + + fn enforce_model_access(&self) -> impl Future> + Send; fn check_budget_exemption(&self) -> impl Future> + Send; fn read_team_member_budget( &self, @@ -112,14 +117,6 @@ pub trait KeyHost: Sync { ) -> impl Future> + Send; } -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum SpendCounter { - Key { hashed_token: String }, - Team { team_id: String }, - Organization { org_id: String }, - TeamMember(TeamMember), -} - pub trait JwtHost: Sync { type Error: Error + Send + Sync + 'static; @@ -127,7 +124,7 @@ pub trait JwtHost: Sync { &self, issuer: &str, ) -> impl Future, Self::Error>> + Send; - fn mapped_key( + fn mapped_key_hash( &self, identity: &JwtIdentity, ) -> impl Future, Self::Error>> + Send; @@ -155,24 +152,3 @@ pub trait CustomAuthHost: Sync { facts: &RequestFacts, ) -> impl Future> + Send; } - -pub trait AuthorizeHost: Sync { - type Error: Error + Send + Sync + 'static; - - fn team( - &self, - team_id: &str, - ) -> impl Future, Self::Error>> + Send; - fn organization( - &self, - org_id: &str, - ) -> impl Future, Self::Error>> + Send; - fn allowed_models( - &self, - identity: &Identity, - ) -> impl Future, Self::Error>> + Send; - fn spend( - &self, - counter: &SpendCounter, - ) -> impl Future> + Send; -} diff --git a/litellm-rust/crates/gateway-auth-core/src/identity.rs b/litellm-rust/crates/gateway-auth-core/src/identity.rs new file mode 100644 index 00000000000..01839bd1709 --- /dev/null +++ b/litellm-rust/crates/gateway-auth-core/src/identity.rs @@ -0,0 +1,126 @@ +use std::time::SystemTime; + +use litellm_secret_value::SecretValue; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::answers::KeyFacts; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] +pub enum AnonymousReason { + PublicRoute, + PassThroughAuthOff, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Anonymous { + pub reason: AnonymousReason, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct JwtIdentity { + pub user_id: Option, + pub team_id: Option, + pub org_id: Option, + pub end_user_id: Option, + pub role: Option, + pub scopes: Vec, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct VirtualKey { + pub key: KeyFacts, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum KeyIdentity { + NoMasterKey { token: Option }, + UserAuthRoute, + MasterKey, + CachedAdminKey, + VirtualKey(VirtualKey), +} + +#[derive(Clone, Debug, PartialEq)] +pub enum KeyDenial { + MissingKey, + MalformedKey, + NoDatabase, + NotAVirtualKey, + KeyBlocked, + OwnerDeactivated { + user_id: String, + }, + ExpiredKey { + expires: SystemTime, + now: SystemTime, + key_label: SecretValue, + }, + KeyWithoutToken, +} + +pub type KeyAuthentication = Result; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct OAuth2Identity { + pub token: String, + pub user_id: Value, + pub role: Value, + pub team_id: Value, +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)] +pub struct OAuth2ProxyIdentity { + pub user_id: Option, + pub user_email: Option, + pub team_id: Option, + pub team_alias: Option, + pub org_id: Option, + pub models: Option>, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct CustomIdentity { + pub user_id: Option, + pub team_id: Option, + pub hashed_token: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum CustomAuthResult { + Identity(CustomIdentity), + ReplacedKey { token: String }, + NoIdentity, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum Identity { + Key(KeyIdentity), + Jwt(JwtIdentity), + OAuth2(OAuth2Identity), + OAuth2Proxy(OAuth2ProxyIdentity), + Custom(CustomIdentity), + Anonymous(Anonymous), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] +pub enum OAuth2Denial { + InactiveToken, + IntrospectionNotConfigured, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum OAuth2ProxyDenial { + TrustedProxyRangesMissing, + UntrustedProxy { client_ip: Option }, + ProxyMappingsMissing, + DisallowedProxyMappings { fields: Vec }, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum AuthDenial { + Key(KeyDenial), + OAuth2(OAuth2Denial), + OAuth2Proxy(OAuth2ProxyDenial), +} diff --git a/litellm-rust/crates/gateway-auth-core/src/inputs.rs b/litellm-rust/crates/gateway-auth-core/src/inputs.rs index 9c9be370109..a1bfc2aa0aa 100644 --- a/litellm-rust/crates/gateway-auth-core/src/inputs.rs +++ b/litellm-rust/crates/gateway-auth-core/src/inputs.rs @@ -45,10 +45,7 @@ pub struct PlanSettings { #[derive(Clone, Debug, PartialEq)] pub struct KeySettings { pub master_key: Option, - pub allow_user_auth: bool, pub has_db: bool, - pub master_key_only_routes: BTreeSet, - pub proxy_max_budget: f64, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -71,6 +68,8 @@ pub struct OAuth2ProxySettings { #[derive(Clone, Debug, PartialEq)] pub struct AuthorizeSettings { - pub admin_only_routes: BTreeSet, + pub allow_user_auth: bool, + pub has_db: bool, + pub master_key_only_routes: BTreeSet, pub proxy_max_budget: f64, } diff --git a/litellm-rust/crates/gateway-auth-core/src/lib.rs b/litellm-rust/crates/gateway-auth-core/src/lib.rs index 13f216ba5c6..21e3b53c9d2 100644 --- a/litellm-rust/crates/gateway-auth-core/src/lib.rs +++ b/litellm-rust/crates/gateway-auth-core/src/lib.rs @@ -1,14 +1,15 @@ mod answers; mod facets; mod host; +mod identity; mod inputs; mod native; mod outcome; +mod validators; pub use answers::{ - FallbackTarget, Introspection, KeyFacts, KeyModelBudgets, ModelBudgetScope, OrganizationFacts, - OwnerFacts, PresentedKey, SigningKey, TeamFacts, TeamMember, TeamMemberBudget, - TeamMemberModelBudgetScope, + FallbackTarget, Introspection, KeyFacts, KeyModelBudgets, ModelBudgetScope, OwnerFacts, + PresentedKey, SigningKey, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope, }; pub use facets::{ AnonymousCore, AuthCore, AuthPlanner, AuthorizeCore, CustomAuthCore, JwtCore, KeyCore, @@ -16,16 +17,17 @@ pub use facets::{ }; pub use host::{ AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost, - OAuth2ProxyHost, SpendCounter, TeamMemberModelBudgetHost, + OAuth2ProxyHost, TeamMemberModelBudgetHost, +}; +pub use identity::{ + Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, CustomIdentity, Identity, + JwtIdentity, KeyAuthentication, KeyDenial, KeyIdentity, OAuth2Denial, OAuth2Identity, + OAuth2ProxyDenial, OAuth2ProxyIdentity, VirtualKey, }; pub use inputs::{ AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings, OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts, RoutingOverride, }; pub use native::{NativeCore, Unported}; -pub use outcome::{ - Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, CustomIdentity, Denial, Grant, - Identity, JwtIdentity, KeyPlan, ModelBudgetExceeded, ModelBudgetResult, OAuth2Denial, - OAuth2Identity, OAuth2Mode, OAuth2ProxyDenial, OAuth2ProxyIdentity, Outcome, Plan, - PremiumFeature, WithinModelBudgets, -}; +pub use outcome::{Denial, ModelBudgetExceeded, ModelBudgetResult, Outcome, WithinModelBudgets}; +pub use validators::{CredentialPlan, KeyPlan, OAuth2Mode, Plan, PremiumFeature}; diff --git a/litellm-rust/crates/gateway-auth-core/src/native.rs b/litellm-rust/crates/gateway-auth-core/src/native.rs index 02c593e544f..fc1fc531cf1 100644 --- a/litellm-rust/crates/gateway-auth-core/src/native.rs +++ b/litellm-rust/crates/gateway-auth-core/src/native.rs @@ -9,15 +9,17 @@ use crate::{ AuthorizeHost, CustomAuthHost, JwtHost, KeyHost, ModelBudgetHost, OAuth2IntrospectionHost, OAuth2ProxyHost, TeamMemberModelBudgetHost, }, + identity::{ + Anonymous, AnonymousReason, AuthDenial, CustomAuthResult, Identity, KeyAuthentication, + KeyDenial, KeyIdentity, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial, + OAuth2ProxyIdentity, VirtualKey, + }, inputs::{ AuthorizeSettings, JwtSettings, KeySettings, OAuth2IntrospectionSettings, OAuth2ProxySettings, PlanFacts, PlanSettings, RequestFacts, }, - outcome::{ - Anonymous, AnonymousReason, CredentialPlan, CustomAuthResult, Denial, Identity, KeyPlan, - ModelBudgetResult, OAuth2Denial, OAuth2Identity, OAuth2ProxyDenial, OAuth2ProxyIdentity, - Outcome, Plan, - }, + outcome::{ModelBudgetResult, Outcome}, + validators::{CredentialPlan, KeyPlan, Plan}, }; #[derive(Clone, Copy, Debug, Default)] @@ -78,17 +80,16 @@ impl KeyCore for Unported { _facts: &RequestFacts, _settings: &KeySettings, _now: SystemTime, - ) -> Result { + ) -> Result { unimplemented!("{RUNS_IN_PYTHON}") } - async fn check( + async fn verify( &self, _host: &H, _facts: &RequestFacts, - _settings: &KeySettings, _now: SystemTime, - ) -> Result { + ) -> Result, H::Error> { unimplemented!("{RUNS_IN_PYTHON}") } } @@ -157,7 +158,7 @@ impl JwtCore for Unported { _facts: &RequestFacts, _settings: &JwtSettings, _now: SystemTime, - ) -> Result, H::Error> { + ) -> Result, H::Error> { unimplemented!("{RUNS_IN_PYTHON}") } } @@ -195,7 +196,7 @@ impl AuthorizeCore for Unported { async fn authorize( &self, _host: &H, - _identity: &Identity, + _identity: &KeyIdentity, _facts: &RequestFacts, _settings: &AuthorizeSettings, ) -> Result { diff --git a/litellm-rust/crates/gateway-auth-core/src/outcome.rs b/litellm-rust/crates/gateway-auth-core/src/outcome.rs index e3a485461bc..9ad26ce3d32 100644 --- a/litellm-rust/crates/gateway-auth-core/src/outcome.rs +++ b/litellm-rust/crates/gateway-auth-core/src/outcome.rs @@ -1,18 +1,3 @@ -use std::time::SystemTime; - -use litellm_secret_value::SecretValue; -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -#[derive(Clone, Debug, PartialEq)] -pub enum Grant { - NoMasterKey { token: Option }, - UserAuthRoute, - MasterKey, - CachedAdminKey, - VirtualKey { budget_throttle_pct: Option }, -} - #[derive(Clone, Debug, PartialEq)] pub struct ModelBudgetExceeded { pub message: String, @@ -22,23 +7,15 @@ pub struct ModelBudgetExceeded { pub entity_id: Option, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct WithinModelBudgets; + +pub type ModelBudgetResult = Result; + #[derive(Clone, Debug, PartialEq)] pub enum Denial { - MissingKey, - MalformedKey, UserAuthDisabled, - ExpiredKey { - expires: SystemTime, - now: SystemTime, - key_label: SecretValue, - }, MasterKeyOnlyRoute, - NoDatabase, - NotAVirtualKey, - KeyBlocked, - OwnerDeactivated { - user_id: String, - }, TeamMemberBudgetExceeded { user_id: String, team_id: String, @@ -49,128 +26,11 @@ pub enum Denial { spend: f64, max_budget: f64, }, - KeyWithoutToken, ModelBudgetExceeded(ModelBudgetExceeded), } #[derive(Clone, Debug, PartialEq)] pub enum Outcome { - Allow(Grant), + Allow { budget_throttle_pct: Option }, Deny(Denial), } - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct WithinModelBudgets; - -pub type ModelBudgetResult = Result; - -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] -pub enum AnonymousReason { - PublicRoute, - PassThroughAuthOff, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Anonymous { - pub reason: AnonymousReason, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct JwtIdentity { - pub user_id: Option, - pub team_id: Option, - pub org_id: Option, - pub end_user_id: Option, - pub role: Option, - pub scopes: Vec, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize)] -pub struct OAuth2Identity { - pub token: String, - pub user_id: Value, - pub role: Value, - pub team_id: Value, -} - -#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)] -pub struct OAuth2ProxyIdentity { - pub user_id: Option, - pub user_email: Option, - pub team_id: Option, - pub team_alias: Option, - pub org_id: Option, - pub models: Option>, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct CustomIdentity { - pub user_id: Option, - pub team_id: Option, - pub hashed_token: Option, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum CustomAuthResult { - Identity(CustomIdentity), - ReplacedKey { token: String }, - NoIdentity, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum Identity { - MasterKey, - Key { hashed_token: String }, - Jwt(JwtIdentity), - OAuth2(OAuth2Identity), - OAuth2Proxy(OAuth2ProxyIdentity), - Custom(CustomIdentity), - Anonymous(Anonymous), -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] -pub enum OAuth2Denial { - InactiveToken, - IntrospectionNotConfigured, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum OAuth2ProxyDenial { - TrustedProxyRangesMissing, - UntrustedProxy { client_ip: Option }, - ProxyMappingsMissing, - DisallowedProxyMappings { fields: Vec }, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] -pub enum OAuth2Mode { - Introspection, - Proxy, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] -pub enum PremiumFeature { - OAuth2, - Jwt, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum KeyPlan { - Key { pass_through_key: bool }, - Anonymous(AnonymousReason), -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum CredentialPlan { - Key(KeyPlan), - Jwt, - OAuth2(OAuth2Mode), - PremiumRequired(PremiumFeature), -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Plan { - Credential(CredentialPlan), - Custom, -} diff --git a/litellm-rust/crates/gateway-auth-core/src/validators.rs b/litellm-rust/crates/gateway-auth-core/src/validators.rs new file mode 100644 index 00000000000..1f2f2870aa1 --- /dev/null +++ b/litellm-rust/crates/gateway-auth-core/src/validators.rs @@ -0,0 +1,33 @@ +use crate::identity::AnonymousReason; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] +pub enum OAuth2Mode { + Introspection, + Proxy, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::VariantArray)] +pub enum PremiumFeature { + OAuth2, + Jwt, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum KeyPlan { + Key { pass_through_key: bool }, + Anonymous(AnonymousReason), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CredentialPlan { + Key(KeyPlan), + Jwt, + OAuth2(OAuth2Mode), + PremiumRequired(PremiumFeature), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Plan { + Credential(CredentialPlan), + Custom, +} diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py index 313365b4302..5d7dbb41321 100644 --- a/litellm/proxy/auth/auth_denial.py +++ b/litellm/proxy/auth/auth_denial.py @@ -7,19 +7,22 @@ from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER, INVALID_VIRTUAL_ from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import key_budget_exceeded_error from litellm.proxy.auth.auth_utils import abbreviate_api_key -from litellm.proxy.auth.core.outcome import ( - Denial, +from litellm.proxy.auth.core.identity import ( ExpiredKey, KeyBlocked, - KeyBudgetExceeded, + KeyDenial, KeyWithoutToken, MalformedKey, - MasterKeyOnlyRoute, MissingKey, - ModelBudgetExceeded, NoDatabase, NotAVirtualKey, OwnerDeactivated, +) +from litellm.proxy.auth.core.outcome import ( + Denial, + KeyBudgetExceeded, + MasterKeyOnlyRoute, + ModelBudgetExceeded, TeamMemberBudgetExceeded, UserAuthDisabled, ) @@ -45,7 +48,7 @@ def _not_a_virtual_key(token: str, jwt_auth_enabled: bool) -> HTTPException: def denial_exception( - denial: Denial, + denial: KeyDenial | Denial, token: str | None, route: str, jwt_auth_enabled: bool, diff --git a/litellm/proxy/auth/core/AGENTS.md b/litellm/proxy/auth/core/AGENTS.md index 95d149b346b..02028a4c330 100644 --- a/litellm/proxy/auth/core/AGENTS.md +++ b/litellm/proxy/auth/core/AGENTS.md @@ -10,19 +10,19 @@ Modules here import only the standard library and each other. A core does no IO `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` 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. +`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 verifies a key the proxy resolved itself. `authorize()` decides whether a key identity may make the request: master-key-only routes, model access and the key, team member and per-model budgets. `model_budgets()` checks per-model budgets for the key and user, the user alone, or a team member. -`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. +`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()` yet, so its signature is provisional and the proxy still validates JWTs itself. `authorize()` covers key identities only, and every identity still passes through the proxy's common checks afterwards. -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`. +Authentication and authorization are separate steps with separate results. An authenticating facet returns an identity or the reason the credential is unusable (`keys()` returns a `KeyAuthentication`). `authorize()` takes that identity and returns an `Outcome`, `Allow` or `Deny` with a `Denial`. A check about the credential itself (missing, malformed, blocked, expired, owner deactivated) belongs in authentication, and a check about what the identity may do (routes, models, budgets) belongs in `authorize()`, so a credential that is unusable is always refused as such before any budget is read or alert sent. `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 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. +`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 `AuthorizeHost.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`, `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. +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), `identity.py` holds what authentication returns (identities and why a credential was refused), `outcome.py` holds the authorization verdict, `validators.py` holds the planner's choices, 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 @@ -30,7 +30,7 @@ The order in which a core calls host methods is part of its contract. It decides ## Projection stays outside -Nothing here knows about `UserAPIKeyAuth`, HTTP status codes or the proxy's exception types. The proxy turns each grant into its token object and maps every `Denial` to the public error contract in one place, `litellm/proxy/auth/auth_denial.py`. A new denial needs a case there. +Nothing here knows about `UserAPIKeyAuth`, HTTP status codes or the proxy's exception types. The proxy turns each identity into its token object and maps every `KeyDenial` and `Denial` to the public error contract in one place, `litellm/proxy/auth/auth_denial.py`. A new denial needs a case there. ## Credentials diff --git a/litellm/proxy/auth/core/answers.py b/litellm/proxy/auth/core/answers.py index cc203ca5d45..c2eeb9fec71 100644 --- a/litellm/proxy/auth/core/answers.py +++ b/litellm/proxy/auth/core/answers.py @@ -88,18 +88,3 @@ class Introspection: user_id: object role: object team_id: object - - -@dataclass(frozen=True, slots=True) -class TeamFacts: - blocked: bool - models: tuple[str, ...] - spend: float - max_budget: float | None - - -@dataclass(frozen=True, slots=True) -class OrganizationFacts: - models: tuple[str, ...] - spend: float - max_budget: float | None diff --git a/litellm/proxy/auth/core/core.py b/litellm/proxy/auth/core/core.py index b127e80ec3c..3c49376733f 100644 --- a/litellm/proxy/auth/core/core.py +++ b/litellm/proxy/auth/core/core.py @@ -11,6 +11,21 @@ from litellm.proxy.auth.core.host import ( OAuth2ProxyHost, TeamMemberModelBudgetHost, ) +from litellm.proxy.auth.core.identity import ( + Anonymous, + AnonymousReason, + AuthDenial, + CustomAuthResult, + Identity, + KeyAuthentication, + KeyDenial, + KeyIdentity, + OAuth2Denial, + OAuth2Identity, + OAuth2ProxyDenial, + OAuth2ProxyIdentity, + VirtualKey, +) from litellm.proxy.auth.core.inputs import ( AuthorizeSettings, JwtSettings, @@ -21,30 +36,16 @@ from litellm.proxy.auth.core.inputs import ( PlanSettings, RequestFacts, ) -from litellm.proxy.auth.core.outcome import ( - Anonymous, - AnonymousReason, - CredentialPlan, - CustomAuthResult, - Denial, - Identity, - KeyPlan, - ModelBudgetResult, - OAuth2Denial, - OAuth2Identity, - OAuth2ProxyDenial, - OAuth2ProxyIdentity, - Outcome, - Plan, -) +from litellm.proxy.auth.core.outcome import ModelBudgetResult, Outcome +from litellm.proxy.auth.core.validators import CredentialPlan, KeyPlan, Plan class KeyCore(Protocol): async def authenticate( self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime - ) -> Outcome: ... + ) -> KeyAuthentication: ... - async def check(self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime) -> Outcome: ... + async def verify(self, host: KeyHost, facts: RequestFacts, now: datetime) -> VirtualKey | KeyDenial: ... class ModelBudgetCore(Protocol): @@ -72,7 +73,7 @@ class AnonymousCore(Protocol): class JwtCore(Protocol): async def authenticate( self, host: JwtHost, facts: RequestFacts, settings: JwtSettings, now: datetime - ) -> Identity | Denial: ... + ) -> Identity | AuthDenial: ... class OAuth2Core(Protocol): @@ -91,7 +92,7 @@ class CustomAuthCore(Protocol): class AuthorizeCore(Protocol): async def authorize( - self, host: AuthorizeHost, identity: Identity, facts: RequestFacts, settings: AuthorizeSettings + self, host: AuthorizeHost, identity: KeyIdentity, facts: RequestFacts, settings: AuthorizeSettings ) -> Outcome: ... diff --git a/litellm/proxy/auth/core/host.py b/litellm/proxy/auth/core/host.py index da602924615..65d223bfc3d 100644 --- a/litellm/proxy/auth/core/host.py +++ b/litellm/proxy/auth/core/host.py @@ -1,22 +1,20 @@ -from dataclasses import dataclass -from typing import Protocol, TypeAlias +from typing import Protocol from litellm.proxy.auth.core.answers import ( FallbackTarget, Introspection, KeyFacts, ModelBudgetScope, - OrganizationFacts, OwnerFacts, PresentedKey, SigningKey, - TeamFacts, TeamMember, TeamMemberBudget, TeamMemberModelBudgetScope, ) +from litellm.proxy.auth.core.identity import CustomAuthResult, JwtIdentity from litellm.proxy.auth.core.inputs import RequestFacts -from litellm.proxy.auth.core.outcome import CustomAuthResult, Identity, JwtIdentity, KeyIdentity, ModelBudgetExceeded +from litellm.proxy.auth.core.outcome import ModelBudgetExceeded class ModelBudgetHost(Protocol): @@ -44,8 +42,6 @@ class TeamMemberModelBudgetHost(Protocol): class KeyHost(Protocol): - def model_budgets(self) -> ModelBudgetHost: ... - async def resolve_end_user(self, prefetch_key: bool) -> None: ... async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: ... @@ -58,10 +54,14 @@ class KeyHost(Protocol): async def inspect_key(self) -> KeyFacts: ... - async def enforce_model_access(self) -> None: ... - async def load_key_owner(self, user_id: str) -> OwnerFacts | None: ... + +class AuthorizeHost(Protocol): + def model_budgets(self) -> ModelBudgetHost: ... + + async def enforce_model_access(self) -> None: ... + async def check_budget_exemption(self) -> bool: ... async def read_team_member_budget(self, member: TeamMember) -> TeamMemberBudget | None: ... @@ -83,33 +83,10 @@ class KeyHost(Protocol): async def alert_proxy_budget(self, spend: float) -> None: ... -@dataclass(frozen=True, slots=True) -class KeySpend: - hashed_token: str - - -@dataclass(frozen=True, slots=True) -class TeamSpend: - team_id: str - - -@dataclass(frozen=True, slots=True) -class OrganizationSpend: - org_id: str - - -@dataclass(frozen=True, slots=True) -class TeamMemberSpend: - member: TeamMember - - -SpendCounter: TypeAlias = KeySpend | TeamSpend | OrganizationSpend | TeamMemberSpend - - class JwtHost(Protocol): async def signing_keys(self, issuer: str) -> tuple[SigningKey, ...]: ... - async def mapped_key(self, identity: JwtIdentity) -> KeyIdentity | None: ... + async def mapped_key_hash(self, identity: JwtIdentity) -> str | None: ... class OAuth2IntrospectionHost(Protocol): @@ -124,13 +101,3 @@ class OAuth2ProxyHost(Protocol): class CustomAuthHost(Protocol): async def custom_auth(self, facts: RequestFacts) -> CustomAuthResult: ... - - -class AuthorizeHost(Protocol): - async def team(self, team_id: str) -> TeamFacts | None: ... - - async def organization(self, org_id: str) -> OrganizationFacts | None: ... - - async def allowed_models(self, identity: Identity) -> frozenset[str]: ... - - async def spend(self, counter: SpendCounter) -> float: ... diff --git a/litellm/proxy/auth/core/identity.py b/litellm/proxy/auth/core/identity.py new file mode 100644 index 00000000000..c915c970540 --- /dev/null +++ b/litellm/proxy/auth/core/identity.py @@ -0,0 +1,191 @@ +from dataclasses import dataclass +from datetime import datetime +from enum import Enum +from typing import TypeAlias + +from litellm.proxy.auth.core.answers import KeyFacts + + +class AnonymousReason(Enum): + PUBLIC_ROUTE = "public_route" + PASS_THROUGH_AUTH_OFF = "pass_through_auth_off" + + +@dataclass(frozen=True, slots=True) +class Anonymous: + reason: AnonymousReason + + +@dataclass(frozen=True, slots=True) +class NoMasterKey: + token: str | None + + +@dataclass(frozen=True, slots=True) +class UserAuthRoute: + pass + + +@dataclass(frozen=True, slots=True) +class MasterKey: + pass + + +@dataclass(frozen=True, slots=True) +class CachedAdminKey: + pass + + +@dataclass(frozen=True, slots=True) +class VirtualKey: + key: KeyFacts + + +KeyIdentity: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey + + +@dataclass(frozen=True, slots=True) +class JwtIdentity: + user_id: str | None + team_id: str | None + org_id: str | None + end_user_id: str | None + role: str | None + scopes: tuple[str, ...] + + +@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 + team_alias: str | None + org_id: str | None + models: tuple[str, ...] | None + + +@dataclass(frozen=True, slots=True) +class CustomIdentity: + user_id: str | None + team_id: str | None + hashed_token: str | None + + +Identity: TypeAlias = KeyIdentity | JwtIdentity | OAuth2Identity | OAuth2ProxyIdentity | CustomIdentity | Anonymous + + +@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 MissingKey: + pass + + +@dataclass(frozen=True, slots=True) +class MalformedKey: + pass + + +@dataclass(frozen=True, slots=True) +class NoDatabase: + pass + + +@dataclass(frozen=True, slots=True) +class NotAVirtualKey: + pass + + +@dataclass(frozen=True, slots=True) +class KeyBlocked: + pass + + +@dataclass(frozen=True, slots=True) +class OwnerDeactivated: + user_id: str + + +@dataclass(frozen=True, slots=True) +class ExpiredKey: + expires: datetime + now: datetime + key_label: str + + +@dataclass(frozen=True, slots=True) +class KeyWithoutToken: + pass + + +KeyDenial: TypeAlias = ( + MissingKey + | MalformedKey + | NoDatabase + | NotAVirtualKey + | KeyBlocked + | OwnerDeactivated + | ExpiredKey + | KeyWithoutToken +) + +KeyAuthentication: TypeAlias = KeyIdentity | KeyDenial + + +@dataclass(frozen=True, slots=True) +class InactiveToken: + pass + + +@dataclass(frozen=True, slots=True) +class IntrospectionNotConfigured: + pass + + +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 +) + +AuthDenial: TypeAlias = KeyDenial | OAuth2Denial | OAuth2ProxyDenial diff --git a/litellm/proxy/auth/core/inputs.py b/litellm/proxy/auth/core/inputs.py index 77f3c5c68c1..2a9511023ba 100644 --- a/litellm/proxy/auth/core/inputs.py +++ b/litellm/proxy/auth/core/inputs.py @@ -43,10 +43,7 @@ class PlanSettings: @dataclass(frozen=True, slots=True) class KeySettings: master_key: str | None - allow_user_auth: bool has_db: bool - master_key_only_routes: frozenset[str] - proxy_max_budget: float @dataclass(frozen=True, slots=True) @@ -72,5 +69,7 @@ class OAuth2ProxySettings: @dataclass(frozen=True, slots=True) class AuthorizeSettings: - admin_only_routes: frozenset[str] + allow_user_auth: bool + has_db: bool + master_key_only_routes: frozenset[str] proxy_max_budget: float diff --git a/litellm/proxy/auth/core/model_budgets.py b/litellm/proxy/auth/core/model_budgets.py index b3a70b645c6..1c9fb32f989 100644 --- a/litellm/proxy/auth/core/model_budgets.py +++ b/litellm/proxy/auth/core/model_budgets.py @@ -3,7 +3,9 @@ from typing import Final from litellm.proxy.auth.core.answers import FallbackTarget, ModelBudgetScope, TeamMemberModelBudgetScope from litellm.proxy.auth.core.host import ModelBudgetHost, TeamMemberModelBudgetHost -from litellm.proxy.auth.core.outcome import ModelBudgetExceeded +from litellm.proxy.auth.core.outcome import ( + ModelBudgetExceeded, +) async def check_key_and_user_model_budgets(host: ModelBudgetHost) -> ModelBudgetExceeded | None: diff --git a/litellm/proxy/auth/core/oauth2.py b/litellm/proxy/auth/core/oauth2.py index fdc4ed39e17..0585d595b5c 100644 --- a/litellm/proxy/auth/core/oauth2.py +++ b/litellm/proxy/auth/core/oauth2.py @@ -1,9 +1,7 @@ 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 ( +from litellm.proxy.auth.core.identity import ( DisallowedProxyMappings, InactiveToken, IntrospectionNotConfigured, @@ -15,6 +13,8 @@ from litellm.proxy.auth.core.outcome import ( TrustedProxyRangesMissing, UntrustedProxy, ) +from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings, OAuth2ProxySettings +from litellm.proxy.auth.core.network import ip_in_networks # 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 diff --git a/litellm/proxy/auth/core/outcome.py b/litellm/proxy/auth/core/outcome.py index 1f81530ff4d..7457f7e53c6 100644 --- a/litellm/proxy/auth/core/outcome.py +++ b/litellm/proxy/auth/core/outcome.py @@ -1,89 +1,17 @@ from dataclasses import dataclass -from datetime import datetime -from enum import Enum from typing import TypeAlias -@dataclass(frozen=True, slots=True) -class NoMasterKey: - token: str | None - - -@dataclass(frozen=True, slots=True) -class UserAuthRoute: - pass - - -@dataclass(frozen=True, slots=True) -class MasterKey: - pass - - -@dataclass(frozen=True, slots=True) -class CachedAdminKey: - pass - - -@dataclass(frozen=True, slots=True) -class VirtualKey: - budget_throttle_pct: float | None - - -@dataclass(frozen=True, slots=True) -class WithinModelBudgets: - pass - - -Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey - - -@dataclass(frozen=True, slots=True) -class MissingKey: - pass - - -@dataclass(frozen=True, slots=True) -class MalformedKey: - pass - - @dataclass(frozen=True, slots=True) class UserAuthDisabled: pass -@dataclass(frozen=True, slots=True) -class ExpiredKey: - expires: datetime - now: datetime - key_label: str - - @dataclass(frozen=True, slots=True) class MasterKeyOnlyRoute: pass -@dataclass(frozen=True, slots=True) -class NoDatabase: - pass - - -@dataclass(frozen=True, slots=True) -class NotAVirtualKey: - pass - - -@dataclass(frozen=True, slots=True) -class KeyBlocked: - pass - - -@dataclass(frozen=True, slots=True) -class OwnerDeactivated: - user_id: str - - @dataclass(frozen=True, slots=True) class TeamMemberBudgetExceeded: user_id: str @@ -98,11 +26,6 @@ class KeyBudgetExceeded: max_budget: float -@dataclass(frozen=True, slots=True) -class KeyWithoutToken: - pass - - @dataclass(frozen=True, slots=True) class ModelBudgetExceeded: message: str @@ -112,26 +35,21 @@ class ModelBudgetExceeded: entity_id: str | None +@dataclass(frozen=True, slots=True) +class WithinModelBudgets: + pass + + +ModelBudgetResult: TypeAlias = WithinModelBudgets | ModelBudgetExceeded + Denial: TypeAlias = ( - MissingKey - | MalformedKey - | UserAuthDisabled - | ExpiredKey - | MasterKeyOnlyRoute - | NoDatabase - | NotAVirtualKey - | KeyBlocked - | OwnerDeactivated - | TeamMemberBudgetExceeded - | KeyBudgetExceeded - | KeyWithoutToken - | ModelBudgetExceeded + UserAuthDisabled | MasterKeyOnlyRoute | TeamMemberBudgetExceeded | KeyBudgetExceeded | ModelBudgetExceeded ) @dataclass(frozen=True, slots=True) class Allow: - grant: Grant + budget_throttle_pct: float | None @dataclass(frozen=True, slots=True) @@ -140,163 +58,3 @@ class Deny: Outcome: TypeAlias = Allow | Deny - -ModelBudgetResult: TypeAlias = WithinModelBudgets | ModelBudgetExceeded - - -class AnonymousReason(Enum): - PUBLIC_ROUTE = "public_route" - PASS_THROUGH_AUTH_OFF = "pass_through_auth_off" - - -@dataclass(frozen=True, slots=True) -class MasterKeyIdentity: - pass - - -@dataclass(frozen=True, slots=True) -class KeyIdentity: - hashed_token: str - - -@dataclass(frozen=True, slots=True) -class JwtIdentity: - user_id: str | None - team_id: str | None - org_id: str | None - end_user_id: str | None - role: str | None - scopes: tuple[str, ...] - - -@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 - team_alias: str | None - org_id: str | None - models: tuple[str, ...] | None - - -@dataclass(frozen=True, slots=True) -class CustomIdentity: - user_id: str | None - team_id: str | None - 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 - - -@dataclass(frozen=True, slots=True) -class InactiveToken: - pass - - -@dataclass(frozen=True, slots=True) -class IntrospectionNotConfigured: - pass - - -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): - INTROSPECTION = "introspection" - PROXY = "proxy" - - -@dataclass(frozen=True, slots=True) -class KeyValidator: - pass_through_key: bool - - -@dataclass(frozen=True, slots=True) -class JwtValidator: - pass - - -@dataclass(frozen=True, slots=True) -class OAuth2Validator: - mode: OAuth2Mode - - -@dataclass(frozen=True, slots=True) -class CustomValidator: - pass - - -@dataclass(frozen=True, slots=True) -class AnonymousValidator: - reason: AnonymousReason - - -class PremiumFeature(Enum): - OAUTH2 = "oauth2" - JWT = "jwt" - - -@dataclass(frozen=True, slots=True) -class PremiumRequired: - feature: PremiumFeature - - -KeyPlan: TypeAlias = KeyValidator | AnonymousValidator - -CredentialPlan: TypeAlias = KeyPlan | JwtValidator | OAuth2Validator | PremiumRequired - -Plan: TypeAlias = CredentialPlan | CustomValidator diff --git a/litellm/proxy/auth/core/plan.py b/litellm/proxy/auth/core/plan.py index 92f763afcf5..147f7356dd2 100644 --- a/litellm/proxy/auth/core/plan.py +++ b/litellm/proxy/auth/core/plan.py @@ -3,9 +3,11 @@ 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 ( +from litellm.proxy.auth.core.identity import ( AnonymousReason, +) +from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RoutingOverride +from litellm.proxy.auth.core.validators import ( AnonymousValidator, CredentialPlan, CustomValidator, diff --git a/litellm/proxy/auth/core/python_core.py b/litellm/proxy/auth/core/python_core.py index ed76e79124c..cad7a6946ff 100644 --- a/litellm/proxy/auth/core/python_core.py +++ b/litellm/proxy/auth/core/python_core.py @@ -18,6 +18,7 @@ from litellm.proxy.auth.core.core import ( OAuth2Core, ) from litellm.proxy.auth.core.host import ( + AuthorizeHost, CustomAuthHost, KeyHost, ModelBudgetHost, @@ -25,7 +26,33 @@ from litellm.proxy.auth.core.host import ( OAuth2ProxyHost, TeamMemberModelBudgetHost, ) +from litellm.proxy.auth.core.identity import ( + Anonymous, + AnonymousReason, + CachedAdminKey, + CustomAuthResult, + ExpiredKey, + KeyAuthentication, + KeyBlocked, + KeyDenial, + KeyIdentity, + KeyWithoutToken, + MalformedKey, + MasterKey, + MissingKey, + NoDatabase, + NoMasterKey, + NotAVirtualKey, + OAuth2Denial, + OAuth2Identity, + OAuth2ProxyDenial, + OAuth2ProxyIdentity, + OwnerDeactivated, + UserAuthRoute, + VirtualKey, +) from litellm.proxy.auth.core.inputs import ( + AuthorizeSettings, KeySettings, OAuth2IntrospectionSettings, OAuth2ProxySettings, @@ -41,52 +68,22 @@ 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, MasterKeyOnlyRoute, - MissingKey, ModelBudgetExceeded, ModelBudgetResult, - NoDatabase, - NoMasterKey, - NotAVirtualKey, - OAuth2Denial, - OAuth2Identity, - OAuth2ProxyDenial, - OAuth2ProxyIdentity, Outcome, - OwnerDeactivated, - Plan, TeamMemberBudgetExceeded, UserAuthDisabled, - UserAuthRoute, - VirtualKey, WithinModelBudgets, ) +from litellm.proxy.auth.core.validators import CredentialPlan, KeyPlan, Plan _USER_AUTH_ROUTE: Final = "/user/auth" _VIRTUAL_KEY_PREFIX: Final = "sk-" -@dataclass(frozen=True, slots=True) -class _KeyCheck: - settings: KeySettings - now: datetime - is_llm_api_route: bool - key_label: str - - @dataclass(frozen=True, slots=True) class _WithinKeyBudget: budget_throttle_pct: float | None @@ -115,7 +112,7 @@ class PythonCore(AuthCore): return PythonCustomAuthCore() def authorize(self) -> AuthorizeCore: - raise NotImplementedError + return PythonAuthorizeCore() class PythonAuthPlanner(AuthPlanner): @@ -153,24 +150,39 @@ class PythonCustomAuthCore(CustomAuthCore): class PythonKeyCore(KeyCore): - async def authenticate(self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime) -> Outcome: + async def authenticate( + self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime + ) -> KeyAuthentication: token: Final = facts.token if settings.master_key is None: - return Allow(NoMasterKey(token=token)) + return NoMasterKey(token=token) if token is None: - return Deny(MissingKey()) + return MissingKey() if token == "": - return Deny(MalformedKey()) + return MalformedKey() if facts.route == _USER_AUTH_ROUTE: - return Allow(UserAuthRoute()) if settings.allow_user_auth else Deny(UserAuthDisabled()) + return UserAuthRoute() if is_master_key(token, settings.master_key): await host.resolve_end_user(prefetch_key=False) - return Allow(MasterKey()) + return MasterKey() await host.resolve_end_user(prefetch_key=True) - return await _presented_key(host, token, facts, settings, now) + return await _presented_key(host, token, settings, now) - async def check(self, host: KeyHost, facts: RequestFacts, settings: KeySettings, now: datetime) -> Outcome: - return await _check_key(host, _KeyCheck(settings, now, facts.is_llm_api_route, key_label=facts.token or "")) + async def verify(self, host: KeyHost, facts: RequestFacts, now: datetime) -> VirtualKey | KeyDenial: + return await _verify_key(host, now, key_label=facts.token or "") + + +class PythonAuthorizeCore(AuthorizeCore): + async def authorize( + self, host: AuthorizeHost, identity: KeyIdentity, facts: RequestFacts, settings: AuthorizeSettings + ) -> Outcome: + match identity: + case UserAuthRoute(): + return Allow(budget_throttle_pct=None) if settings.allow_user_auth else Deny(UserAuthDisabled()) + case NoMasterKey() | MasterKey() | CachedAdminKey(): + return Allow(budget_throttle_pct=None) + case VirtualKey(key=key): + return await _authorize_virtual_key(host, key, facts, settings) class PythonModelBudgetCore(ModelBudgetCore): @@ -188,56 +200,37 @@ def _within_model_budgets(exceeded: ModelBudgetExceeded | None) -> ModelBudgetRe return WithinModelBudgets() if exceeded is None else exceeded -async def _presented_key( - host: KeyHost, token: str, facts: RequestFacts, settings: KeySettings, now: datetime -) -> Outcome: +async def _presented_key(host: KeyHost, token: str, settings: KeySettings, now: datetime) -> KeyAuthentication: hashed_token: Final = hash_token(token) presented: Final = await host.resolve_presented_key(hashed_token) if isinstance(presented, AdminKey): if presented.expires is None or presented.expires >= now: - return Allow(CachedAdminKey()) + return CachedAdminKey() await host.evict_presented_key(hashed_token) - return Deny(ExpiredKey(expires=presented.expires, now=now, key_label=token)) - if facts.route in settings.master_key_only_routes: - return Deny(MasterKeyOnlyRoute()) + return ExpiredKey(expires=presented.expires, now=now, key_label=token) if not settings.has_db: - return Deny(NoDatabase()) + return NoDatabase() if presented is not None: await host.prepare_key(loaded_from_db=False) - return await _check_key(host, _KeyCheck(settings, now, facts.is_llm_api_route, key_label=token)) + return await _verify_key(host, now, key_label=token) if not token.startswith(_VIRTUAL_KEY_PREFIX): - return Deny(NotAVirtualKey()) + return NotAVirtualKey() await host.load_key(hashed_token) await host.prepare_key(loaded_from_db=True) - return await _check_key(host, _KeyCheck(settings, now, facts.is_llm_api_route, key_label=hashed_token)) + return await _verify_key(host, now, key_label=hashed_token) -async def _check_key(host: KeyHost, check: _KeyCheck) -> Outcome: +async def _verify_key(host: KeyHost, now: datetime, key_label: str) -> VirtualKey | KeyDenial: key: Final = await host.inspect_key() if key.blocked: - return Deny(KeyBlocked()) - await host.enforce_model_access() + return KeyBlocked() if key.user_id is not None and await _owner_deactivated(host, key.user_id): - return Deny(OwnerDeactivated(user_id=key.user_id)) - exempt: Final = await host.check_budget_exemption() - if not exempt and check.settings.has_db and key.team_member is not None: - team_member_exceeded: Final = await _check_team_member_budget(host, key.team_member) - if team_member_exceeded is not None: - return Deny(team_member_exceeded) - if key.expires is not None and key.expires < check.now: - return Deny(ExpiredKey(expires=key.expires, now=check.now, key_label=check.key_label)) - if exempt: - return await _grant(host, check, key, budget_throttle_pct=None) - await host.send_key_budget_alerts() - key_budget: Final = await _check_key_budget(host, check, key) - if isinstance(key_budget, KeyBudgetExceeded): - return Deny(key_budget) - if key.soft_budget and key.spend >= key.soft_budget: - await host.alert_key_soft_budget() - model_budget_exceeded: Final = await check_key_and_user_model_budgets(host.model_budgets()) - if model_budget_exceeded is not None: - return Deny(model_budget_exceeded) - return await _grant(host, check, key, budget_throttle_pct=key_budget.budget_throttle_pct) + return OwnerDeactivated(user_id=key.user_id) + if key.expires is not None and key.expires < now: + return ExpiredKey(expires=key.expires, now=now, key_label=key_label) + if not key.has_token: + return KeyWithoutToken() + return VirtualKey(key=key) async def _owner_deactivated(host: KeyHost, user_id: str) -> bool: @@ -245,7 +238,32 @@ async def _owner_deactivated(host: KeyHost, user_id: str) -> bool: return owner is not None and owner.deactivated -async def _check_team_member_budget(host: KeyHost, member: TeamMember) -> TeamMemberBudgetExceeded | None: +async def _authorize_virtual_key( + host: AuthorizeHost, key: KeyFacts, facts: RequestFacts, settings: AuthorizeSettings +) -> Outcome: + if facts.route in settings.master_key_only_routes: + return Deny(MasterKeyOnlyRoute()) + await host.enforce_model_access() + exempt: Final = await host.check_budget_exemption() + if not exempt and settings.has_db and key.team_member is not None: + team_member_exceeded: Final = await _check_team_member_budget(host, key.team_member) + if team_member_exceeded is not None: + return Deny(team_member_exceeded) + if exempt: + return await _grant(host, settings, budget_throttle_pct=None) + await host.send_key_budget_alerts() + key_budget: Final = await _check_key_budget(host, key, is_llm_api_route=facts.is_llm_api_route) + if isinstance(key_budget, KeyBudgetExceeded): + return Deny(key_budget) + if key.soft_budget and key.spend >= key.soft_budget: + await host.alert_key_soft_budget() + model_budget_exceeded: Final = await check_key_and_user_model_budgets(host.model_budgets()) + if model_budget_exceeded is not None: + return Deny(model_budget_exceeded) + return await _grant(host, settings, budget_throttle_pct=key_budget.budget_throttle_pct) + + +async def _check_team_member_budget(host: AuthorizeHost, member: TeamMember) -> TeamMemberBudgetExceeded | None: budget: Final = await host.read_team_member_budget(member) if budget is None or budget.spend < budget.max_budget: return None @@ -255,9 +273,11 @@ async def _check_team_member_budget(host: KeyHost, member: TeamMember) -> TeamMe ) -async def _check_key_budget(host: KeyHost, check: _KeyCheck, key: KeyFacts) -> _WithinKeyBudget | KeyBudgetExceeded: +async def _check_key_budget( + host: AuthorizeHost, key: KeyFacts, *, is_llm_api_route: bool +) -> _WithinKeyBudget | KeyBudgetExceeded: max_budget: Final = key.max_budget - if not check.is_llm_api_route or max_budget is None: + if not is_llm_api_route or max_budget is None: return _WithinKeyBudget(budget_throttle_pct=None) spend: Final = await host.read_key_spend(max_budget) await host.alert_key_budget(spend) @@ -269,12 +289,10 @@ async def _check_key_budget(host: KeyHost, check: _KeyCheck, key: KeyFacts) -> _ return KeyBudgetExceeded(spend=spend, max_budget=max_budget) -async def _grant(host: KeyHost, check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Outcome: +async def _grant(host: AuthorizeHost, settings: AuthorizeSettings, budget_throttle_pct: float | None) -> Outcome: await host.load_key_context() - if check.settings.proxy_max_budget > 0 and check.settings.has_db: + if settings.proxy_max_budget > 0 and settings.has_db: proxy_spend: Final = await host.read_proxy_spend() if proxy_spend is not None: await host.alert_proxy_budget(proxy_spend) - if not key.has_token: - return Deny(KeyWithoutToken()) - return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct)) + return Allow(budget_throttle_pct=budget_throttle_pct) diff --git a/litellm/proxy/auth/core/validators.py b/litellm/proxy/auth/core/validators.py new file mode 100644 index 00000000000..a9dafa1cec6 --- /dev/null +++ b/litellm/proxy/auth/core/validators.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass +from enum import Enum +from typing import TypeAlias + +from litellm.proxy.auth.core.identity import AnonymousReason + + +class OAuth2Mode(Enum): + INTROSPECTION = "introspection" + PROXY = "proxy" + + +@dataclass(frozen=True, slots=True) +class KeyValidator: + pass_through_key: bool + + +@dataclass(frozen=True, slots=True) +class JwtValidator: + pass + + +@dataclass(frozen=True, slots=True) +class OAuth2Validator: + mode: OAuth2Mode + + +@dataclass(frozen=True, slots=True) +class CustomValidator: + pass + + +@dataclass(frozen=True, slots=True) +class AnonymousValidator: + reason: AnonymousReason + + +class PremiumFeature(Enum): + OAUTH2 = "oauth2" + JWT = "jwt" + + +@dataclass(frozen=True, slots=True) +class PremiumRequired: + feature: PremiumFeature + + +KeyPlan: TypeAlias = KeyValidator | AnonymousValidator + +CredentialPlan: TypeAlias = KeyPlan | JwtValidator | OAuth2Validator | PremiumRequired + +Plan: TypeAlias = CredentialPlan | CustomValidator diff --git a/litellm/proxy/auth/oauth2_check.py b/litellm/proxy/auth/oauth2_check.py index 1b7be2e43ac..ae193e17171 100644 --- a/litellm/proxy/auth/oauth2_check.py +++ b/litellm/proxy/auth/oauth2_check.py @@ -12,8 +12,12 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.core.answers import Introspection +from litellm.proxy.auth.core.identity import ( + InactiveToken, + IntrospectionNotConfigured, + OAuth2Identity, +) 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)) diff --git a/litellm/proxy/auth/oauth2_proxy_hook.py b/litellm/proxy/auth/oauth2_proxy_hook.py index 42fa2092614..64bf68a7c97 100644 --- a/litellm/proxy/auth/oauth2_proxy_hook.py +++ b/litellm/proxy/auth/oauth2_proxy_hook.py @@ -7,15 +7,15 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth -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 ( +from litellm.proxy.auth.core.identity import ( DisallowedProxyMappings, OAuth2ProxyIdentity, ProxyMappingsMissing, TrustedProxyRangesMissing, UntrustedProxy, ) +from litellm.proxy.auth.core.inputs import OAuth2ProxySettings +from litellm.proxy.auth.core.oauth2 import ALLOWED_OAUTH2_PROXY_FIELDS 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 ( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ce2156cc147..c619edd2c50 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -146,36 +146,49 @@ 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, PlanFacts, PlanSettings, RequestFacts, RoutingOverride -from litellm.proxy.auth.core.outcome import ( - Allow, +from litellm.proxy.auth.core.identity import ( Anonymous, AnonymousReason, - AnonymousValidator, CachedAdminKey, - CredentialPlan, CustomAuthResult, CustomIdentity, - CustomValidator, - Deny, - Grant, - JwtValidator, - KeyValidator, + KeyAuthentication, + KeyDenial, + KeyIdentity, MasterKey, - ModelBudgetExceeded, - ModelBudgetResult, NoCustomIdentity, NoMasterKey, - OAuth2Mode, - OAuth2Validator, - Outcome, - PremiumFeature, - PremiumRequired, ReplacedKey, UserAuthRoute, VirtualKey, +) +from litellm.proxy.auth.core.inputs import ( + AuthorizeSettings, + KeySettings, + PlanFacts, + PlanSettings, + RequestFacts, + RoutingOverride, +) +from litellm.proxy.auth.core.outcome import ( + Allow, + Deny, + ModelBudgetExceeded, + ModelBudgetResult, + Outcome, WithinModelBudgets, ) +from litellm.proxy.auth.core.validators import ( + AnonymousValidator, + CredentialPlan, + CustomValidator, + JwtValidator, + KeyValidator, + OAuth2Mode, + OAuth2Validator, + PremiumFeature, + PremiumRequired, +) 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 @@ -2525,9 +2538,12 @@ def _proxy_model_list() -> list[object] | None: _MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value) -def _core_settings(master_key: str | None, allow_user_auth: bool, prisma_client: PrismaClient | None) -> KeySettings: - return KeySettings( - master_key=master_key, +def _key_settings(master_key: str | None, prisma_client: PrismaClient | None) -> KeySettings: + return KeySettings(master_key=master_key, has_db=prisma_client is not None) + + +def _authorize_settings(allow_user_auth: bool, prisma_client: PrismaClient | None) -> AuthorizeSettings: + return AuthorizeSettings( allow_user_auth=allow_user_auth, has_db=prisma_client is not None, master_key_only_routes=_MASTER_KEY_ONLY_ROUTES, @@ -3272,21 +3288,24 @@ async def user_api_key_auth_builder( proxy_logging_obj=proxy_logging_obj, route=route, ) + key_facts: Final = RequestFacts( + route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route) + ) try: - outcome: Final = await _AUTH_CORE.keys().authenticate( - host, - RequestFacts(route=route, token=api_key, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)), - _core_settings( - master_key=master_key, - allow_user_auth=general_settings.get("allow_user_auth", False) is True, - prisma_client=prisma_client, + authorized: Final = await _authorize_key( + await _AUTH_CORE.keys().authenticate( + host, key_facts, _key_settings(master_key, prisma_client), datetime.now(timezone.utc) + ), + host, + key_facts, + _authorize_settings( + allow_user_auth=general_settings.get("allow_user_auth", False) is True, prisma_client=prisma_client ), - datetime.now(timezone.utc), ) finally: valid_token = host.resolved_key # rebind-ok: the exception handler below reports the resolved identity return await _authenticated( - outcome, + authorized, host, admin_user_id=litellm_proxy_admin_name, route=route, @@ -3422,27 +3441,54 @@ async def _master_key_auth(host: _KeyHost, admin_user_id: str, route: str, start return with_end_user +@dataclass(frozen=True, slots=True) +class _AuthorizedKey: + identity: KeyIdentity + outcome: Outcome + + +async def _authorize_key( + authentication: KeyAuthentication, host: _KeyHost, facts: RequestFacts, settings: AuthorizeSettings +) -> _AuthorizedKey | KeyDenial: + match authentication: + case NoMasterKey() | UserAuthRoute() | MasterKey() | CachedAdminKey() | VirtualKey(): + return _AuthorizedKey( + authentication, await _AUTH_CORE.authorize().authorize(host, authentication, facts, settings) + ) + case _: + return authentication + + async def _authenticated( - outcome: Outcome, + authorized: _AuthorizedKey | KeyDenial, host: _KeyHost, admin_user_id: str, route: str, start_time: datetime, jwt_auth_enabled: bool, ) -> UserAPIKeyAuth: - match outcome: + if not isinstance(authorized, _AuthorizedKey): + raise denial_exception( + authorized, token=host.token, route=route, jwt_auth_enabled=jwt_auth_enabled, key=host.resolved_key + ) + match authorized.outcome: case Deny(denial=denial): raise denial_exception( denial, token=host.token, route=route, jwt_auth_enabled=jwt_auth_enabled, key=host.resolved_key ) - case Allow(grant=grant): - return await _granted(grant, host, admin_user_id, route, start_time) + case Allow(budget_throttle_pct=budget_throttle_pct): + return await _granted(authorized.identity, budget_throttle_pct, host, admin_user_id, route, start_time) async def _granted( - grant: Grant, host: _KeyHost, admin_user_id: str, route: str, start_time: datetime + identity: KeyIdentity, + budget_throttle_pct: float | None, + host: _KeyHost, + admin_user_id: str, + route: str, + start_time: datetime, ) -> UserAPIKeyAuth: - match grant: + match identity: case NoMasterKey(token=token): return UserAPIKeyAuth( api_key=token, user_role=LitellmUserRoles.INTERNAL_USER, parent_otel_span=host.parent_otel_span @@ -3453,10 +3499,10 @@ async def _granted( return await _master_key_auth(host, admin_user_id, route, start_time) case CachedAdminKey(): return _with_end_user(host.key, host.end_user, host.parent_otel_span) - case VirtualKey(budget_throttle_pct=budget_throttle_pct): + case VirtualKey(): return await _virtual_key_auth(host, budget_throttle_pct, route, start_time) case _: - assert_never(grant) + assert_never(identity) async def _virtual_key_auth( @@ -5001,7 +5047,6 @@ async def _authorize_resolved_key( ) -> UserAPIKeyAuth: from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, - master_key, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -5019,14 +5064,15 @@ async def _authorize_resolved_key( proxy_logging_obj=proxy_logging_obj, route=route, ) - outcome: Final = await _AUTH_CORE.keys().check( + facts: Final = RequestFacts(route=route, token=key_hash, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)) + authorized: Final = await _authorize_key( + await _AUTH_CORE.keys().verify(host, facts, start_time), host, - RequestFacts(route=route, token=key_hash, is_llm_api_route=RouteChecks.is_llm_api_route(route=route)), - _core_settings(master_key=master_key, allow_user_auth=False, prisma_client=prisma_client), - start_time, + facts, + _authorize_settings(allow_user_auth=False, prisma_client=prisma_client), ) return await _authenticated( - outcome, + authorized, host, admin_user_id=litellm_proxy_admin_name, route=route, diff --git a/tests/unit/proxy/auth/core/test_model_budgets.py b/tests/unit/proxy/auth/core/test_model_budgets.py index 025e833cbb4..d951cfc4885 100644 --- a/tests/unit/proxy/auth/core/test_model_budgets.py +++ b/tests/unit/proxy/auth/core/test_model_budgets.py @@ -10,7 +10,11 @@ from litellm.proxy.auth.core.answers import ( ModelBudgetScope, TeamMemberModelBudgetScope, ) -from litellm.proxy.auth.core.outcome import ModelBudgetExceeded, ModelBudgetResult, WithinModelBudgets +from litellm.proxy.auth.core.outcome import ( + ModelBudgetExceeded, + ModelBudgetResult, + WithinModelBudgets, +) from litellm.proxy.auth.core.python_core import PythonCore Call = tuple[str, str | None] diff --git a/tests/unit/proxy/auth/core/test_oauth2.py b/tests/unit/proxy/auth/core/test_oauth2.py index 937b71c45d2..8bb13058667 100644 --- a/tests/unit/proxy/auth/core/test_oauth2.py +++ b/tests/unit/proxy/auth/core/test_oauth2.py @@ -7,8 +7,7 @@ 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 ( +from litellm.proxy.auth.core.identity import ( DisallowedProxyMappings, InactiveToken, IntrospectionNotConfigured, @@ -18,6 +17,7 @@ from litellm.proxy.auth.core.outcome import ( TrustedProxyRangesMissing, UntrustedProxy, ) +from litellm.proxy.auth.core.inputs import OAuth2IntrospectionSettings, OAuth2ProxySettings from litellm.proxy.auth.core.python_core import PythonCore CORE: Final = PythonCore() diff --git a/tests/unit/proxy/auth/core/test_plan.py b/tests/unit/proxy/auth/core/test_plan.py index 7a25ec9f741..36cfabbe032 100644 --- a/tests/unit/proxy/auth/core/test_plan.py +++ b/tests/unit/proxy/auth/core/test_plan.py @@ -4,9 +4,12 @@ from typing import Final import pytest -from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RequestFacts, RoutingOverride -from litellm.proxy.auth.core.outcome import ( +from litellm.proxy.auth.core.identity import ( AnonymousReason, +) +from litellm.proxy.auth.core.inputs import PlanFacts, PlanSettings, RequestFacts, RoutingOverride +from litellm.proxy.auth.core.python_core import PythonCore +from litellm.proxy.auth.core.validators import ( AnonymousValidator, CustomValidator, JwtValidator, @@ -16,7 +19,6 @@ from litellm.proxy.auth.core.outcome import ( PremiumFeature, PremiumRequired, ) -from litellm.proxy.auth.core.python_core import PythonCore PLANNER: Final = PythonCore().plan() JWT: Final = "header.payload.signature" diff --git a/tests/unit/proxy/auth/core/test_python_core.py b/tests/unit/proxy/auth/core/test_python_core.py index 1431e3457e4..ed2a4cea0bf 100644 --- a/tests/unit/proxy/auth/core/test_python_core.py +++ b/tests/unit/proxy/auth/core/test_python_core.py @@ -1,5 +1,4 @@ import math -from collections.abc import Awaitable from dataclasses import dataclass, field, replace from datetime import datetime, timedelta, timezone from typing import Final @@ -16,43 +15,49 @@ from litellm.proxy.auth.core.answers import ( TeamMember, TeamMemberBudget, ) -from litellm.proxy.auth.core.inputs import RequestFacts, KeySettings -from litellm.proxy.auth.core.outcome import ( - CustomIdentity, - NoCustomIdentity, - ReplacedKey, - Allow, +from litellm.proxy.auth.core.identity import ( CachedAdminKey, - Deny, + CustomIdentity, ExpiredKey, + KeyAuthentication, KeyBlocked, - KeyBudgetExceeded, + KeyIdentity, KeyWithoutToken, MalformedKey, MasterKey, - MasterKeyOnlyRoute, MissingKey, - ModelBudgetExceeded, + NoCustomIdentity, NoDatabase, NoMasterKey, NotAVirtualKey, - Outcome, OwnerDeactivated, - TeamMemberBudgetExceeded, - UserAuthDisabled, + ReplacedKey, UserAuthRoute, VirtualKey, ) +from litellm.proxy.auth.core.inputs import AuthorizeSettings, KeySettings, RequestFacts from litellm.proxy.auth.core.key import hash_token +from litellm.proxy.auth.core.outcome import ( + Allow, + Deny, + KeyBudgetExceeded, + MasterKeyOnlyRoute, + ModelBudgetExceeded, + Outcome, + TeamMemberBudgetExceeded, + UserAuthDisabled, +) from litellm.proxy.auth.core.python_core import PythonCore MASTER: Final = "sk-master" +MASTER_ONLY_ROUTE: Final = "/global/spend/reset" NOW: Final = datetime(2026, 11, 6, 12, tzinfo=timezone.utc) -SETTINGS: Final = KeySettings( - master_key=MASTER, +KEY_SETTINGS: Final = KeySettings(master_key=MASTER, has_db=True) +NO_DB: Final = KeySettings(master_key=MASTER, has_db=False) +AUTHORIZE_SETTINGS: Final = AuthorizeSettings( allow_user_auth=False, has_db=True, - master_key_only_routes=frozenset({"/global/spend/reset"}), + master_key_only_routes=frozenset({MASTER_ONLY_ROUTE}), proxy_max_budget=0.0, ) PLAIN_KEY: Final = KeyFacts( @@ -66,15 +71,25 @@ PLAIN_KEY: Final = KeyFacts( team_member=None, budget_throttle_pct=None, ) -CHECKED: Final = ( - "prepare_key", - "inspect_key", +AUTHORIZED: Final = ( "enforce_model_access", "check_budget_exemption", "send_key_budget_alerts", "read_model_budget_scope", "load_key_context", ) +AUTHORIZE_CALLS: Final = frozenset( + { + *AUTHORIZED, + "read_team_member_budget", + "alert_team_member_budget", + "read_key_spend", + "alert_key_budget", + "alert_key_soft_budget", + "read_proxy_spend", + "alert_proxy_budget", + } +) NO_MODEL_BUDGETS: Final = ModelBudgetScope( models=(), @@ -190,167 +205,241 @@ def _facts(token: str | None, route: str = "/chat/completions", is_llm_api_route return RequestFacts(route=route, token=token, is_llm_api_route=is_llm_api_route) -async def _run(decision: Awaitable[Outcome], handler: RecordingHandler) -> tuple[Outcome, tuple[str, ...]]: - outcome: Final = await decision - return outcome, tuple(name for name, _ in handler.calls) +def _calls(handler: RecordingHandler) -> tuple[str, ...]: + return tuple(name for name, _ in handler.calls) -async def _decide( +async def _authenticate( token: str | None, *, route: str = "/chat/completions", - settings: KeySettings = SETTINGS, + settings: KeySettings = KEY_SETTINGS, + handler: RecordingHandler | None = None, +) -> tuple[KeyAuthentication, tuple[str, ...]]: + host: Final = handler or RecordingHandler() + authentication: Final = await PythonCore().keys().authenticate(host, _facts(token, route), settings, NOW) + return authentication, _calls(host) + + +async def _authorize( + identity: KeyIdentity, + *, + route: str = "/chat/completions", + is_llm_api_route: bool = True, + settings: AuthorizeSettings = AUTHORIZE_SETTINGS, handler: RecordingHandler | None = None, ) -> tuple[Outcome, tuple[str, ...]]: host: Final = handler or RecordingHandler() - return await _run(PythonCore().keys().authenticate(host, _facts(token, route), settings, NOW), host) + facts: Final = _facts("hash", route, is_llm_api_route) + outcome: Final = await PythonCore().authorize().authorize(host, identity, facts, settings) + return outcome, _calls(host) -async def _check( +async def _authorize_key( key: KeyFacts, *, - is_llm_api_route: bool = True, - settings: KeySettings = SETTINGS, - owner: OwnerFacts | None = None, exempt: bool = False, team_member_budget: TeamMemberBudget | None = None, key_spend: float = 0.0, + model_budget_scope: ModelBudgetScope = NO_MODEL_BUDGETS, + user_model_budget_exceeded: ModelBudgetExceeded | None = None, ) -> tuple[Outcome, tuple[str, ...]]: handler: Final = RecordingHandler( - key=key, owner=owner, exempt=exempt, team_member_budget=team_member_budget, key_spend=key_spend + key=key, + exempt=exempt, + team_member_budget=team_member_budget, + key_spend=key_spend, + model_budget_scope=model_budget_scope, + user_model_budget_exceeded=user_model_budget_exceeded, ) - facts: Final = _facts("hash", is_llm_api_route=is_llm_api_route) - return await _run(PythonCore().keys().check(handler, facts, settings, NOW), handler) + return await _authorize(VirtualKey(key=key), handler=handler) def _allowed(budget_throttle_pct: float | None = None) -> Outcome: - return Allow(VirtualKey(budget_throttle_pct=budget_throttle_pct)) + return Allow(budget_throttle_pct=budget_throttle_pct) @pytest.mark.asyncio @pytest.mark.parametrize( ("route", "token", "settings", "expected"), [ - ("/chat/completions", None, KeySettings(None, False, True, frozenset(), 0.0), Allow(NoMasterKey(token=None))), - ( - "/chat/completions", - "sk-x", - KeySettings(None, False, False, frozenset(), 0.0), - Allow(NoMasterKey(token="sk-x")), - ), - ("/chat/completions", None, SETTINGS, Deny(MissingKey())), - ("/chat/completions", "", SETTINGS, Deny(MalformedKey())), - ("/user/auth", "", KeySettings(MASTER, True, True, frozenset(), 0.0), Deny(MalformedKey())), - ("/user/auth", "sk-x", KeySettings(MASTER, True, True, frozenset(), 0.0), Allow(UserAuthRoute())), - ("/user/auth", "sk-x", SETTINGS, Deny(UserAuthDisabled())), - ("/user/auth", MASTER, SETTINGS, Deny(UserAuthDisabled())), + ("/chat/completions", None, KeySettings(None, True), NoMasterKey(token=None)), + ("/chat/completions", "sk-x", KeySettings(None, False), NoMasterKey(token="sk-x")), + ("/chat/completions", None, KEY_SETTINGS, MissingKey()), + ("/chat/completions", "", KEY_SETTINGS, MalformedKey()), + ("/user/auth", "", KEY_SETTINGS, MalformedKey()), + ("/user/auth", "sk-x", KEY_SETTINGS, UserAuthRoute()), + ("/user/auth", MASTER, KEY_SETTINGS, UserAuthRoute()), ], ) -async def test_credential_decisions_do_no_io(route: str, token: str | None, settings: KeySettings, expected: Outcome): - assert await _decide(token, route=route, settings=settings) == (expected, ()) +async def test_credential_decisions_do_no_io( + route: str, token: str | None, settings: KeySettings, expected: KeyAuthentication +): + assert await _authenticate(token, route=route, settings=settings) == (expected, ()) @pytest.mark.asyncio -@pytest.mark.parametrize("route", ["/chat/completions", "/global/spend/reset"]) +@pytest.mark.parametrize("route", ["/chat/completions", MASTER_ONLY_ROUTE]) async def test_master_key_resolves_end_user_without_key_prefetch(route: str): handler: Final = RecordingHandler() - assert await _decide(MASTER, route=route, handler=handler) == (Allow(MasterKey()), ("resolve_end_user",)) + assert await _authenticate(MASTER, route=route, handler=handler) == (MasterKey(), ("resolve_end_user",)) assert handler.calls == (("resolve_end_user", False),) @pytest.mark.asyncio -async def test_unknown_virtual_key_is_loaded_from_the_db_then_checked(): +async def test_unknown_virtual_key_is_loaded_from_the_db_then_verified(): handler: Final = RecordingHandler() - outcome, calls = await _decide("sk-x", handler=handler) - assert outcome == _allowed() - assert calls == ("resolve_end_user", "resolve_presented_key", "load_key", *CHECKED) + authentication, calls = await _authenticate("sk-x", handler=handler) + assert authentication == VirtualKey(key=PLAIN_KEY) + assert calls == ("resolve_end_user", "resolve_presented_key", "load_key", "prepare_key", "inspect_key") assert handler.calls[2] == ("load_key", hash_token("sk-x")) assert handler.calls[3] == ("prepare_key", True) @pytest.mark.asyncio -async def test_cached_virtual_key_is_checked_without_a_db_load(): +async def test_cached_virtual_key_is_verified_without_a_db_load(): handler: Final = RecordingHandler(presented=NonAdminKey()) - outcome, calls = await _decide("sk-x", handler=handler) - assert outcome == _allowed() - assert calls == ("resolve_end_user", "resolve_presented_key", *CHECKED) + authentication, calls = await _authenticate("sk-x", handler=handler) + assert authentication == VirtualKey(key=PLAIN_KEY) + assert calls == ("resolve_end_user", "resolve_presented_key", "prepare_key", "inspect_key") assert handler.calls[2] == ("prepare_key", False) @pytest.mark.asyncio @pytest.mark.parametrize("expires", [None, NOW + timedelta(seconds=1), NOW]) -async def test_unexpired_cached_admin_key_skips_route_db_and_key_checks(expires: datetime | None): - no_db: Final = KeySettings(MASTER, False, False, frozenset({"/global/spend/reset"}), 0.0) - outcome, calls = await _decide( - "sk-admin", - route="/global/spend/reset", - settings=no_db, - handler=RecordingHandler(presented=AdminKey(expires=expires)), +async def test_unexpired_cached_admin_key_skips_the_db_and_key_checks(expires: datetime | None): + authentication, calls = await _authenticate( + "sk-admin", route=MASTER_ONLY_ROUTE, settings=NO_DB, handler=RecordingHandler(presented=AdminKey(expires)) ) - assert (outcome, calls) == (Allow(CachedAdminKey()), ("resolve_end_user", "resolve_presented_key")) + assert (authentication, calls) == (CachedAdminKey(), ("resolve_end_user", "resolve_presented_key")) @pytest.mark.asyncio async def test_expired_cached_admin_key_is_evicted_then_denied_with_the_presented_key(): expires: Final = NOW - timedelta(seconds=1) handler: Final = RecordingHandler(presented=AdminKey(expires=expires)) - outcome, calls = await _decide("sk-admin", handler=handler) - assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label="sk-admin")) + authentication, calls = await _authenticate("sk-admin", handler=handler) + assert authentication == ExpiredKey(expires=expires, now=NOW, key_label="sk-admin") assert calls == ("resolve_end_user", "resolve_presented_key", "evict_presented_key") assert handler.calls[-1] == ("evict_presented_key", hash_token("sk-admin")) -@pytest.mark.asyncio -@pytest.mark.parametrize("presented", [None, NonAdminKey()]) -async def test_non_admin_key_cannot_use_master_key_only_routes(presented: PresentedKey | None): - outcome, calls = await _decide("sk-x", route="/global/spend/reset", handler=RecordingHandler(presented=presented)) - assert (outcome, calls) == (Deny(MasterKeyOnlyRoute()), ("resolve_end_user", "resolve_presented_key")) - - @pytest.mark.asyncio @pytest.mark.parametrize("presented", [None, NonAdminKey()]) async def test_virtual_key_needs_a_database(presented: PresentedKey | None): - no_db: Final = KeySettings(MASTER, False, False, frozenset(), 0.0) - outcome, _ = await _decide("sk-x", settings=no_db, handler=RecordingHandler(presented=presented)) - assert outcome == Deny(NoDatabase()) + authentication, _ = await _authenticate("sk-x", settings=NO_DB, handler=RecordingHandler(presented=presented)) + assert authentication == NoDatabase() @pytest.mark.asyncio @pytest.mark.parametrize("token", ["eyJhbGciOiJIUzI1NiJ9.e30.sig", "a" * 64, "sk"]) async def test_uncached_token_without_virtual_key_prefix_is_rejected_before_the_db(token: str): - assert await _decide(token) == (Deny(NotAVirtualKey()), ("resolve_end_user", "resolve_presented_key")) + assert await _authenticate(token) == (NotAVirtualKey(), ("resolve_end_user", "resolve_presented_key")) @pytest.mark.asyncio async def test_expired_virtual_key_is_labelled_by_the_hash_it_was_loaded_with(): expires: Final = NOW - timedelta(seconds=1) - outcome, _ = await _decide("sk-x", handler=RecordingHandler(key=replace(PLAIN_KEY, expires=expires))) - assert outcome == Deny(ExpiredKey(expires=expires, now=NOW, key_label=hash_token("sk-x"))) + authentication, _ = await _authenticate("sk-x", handler=RecordingHandler(key=replace(PLAIN_KEY, expires=expires))) + assert authentication == ExpiredKey(expires=expires, now=NOW, key_label=hash_token("sk-x")) @pytest.mark.asyncio -async def test_blocked_key_is_denied_before_model_access(): - assert await _check(replace(PLAIN_KEY, blocked=True)) == (Deny(KeyBlocked()), ("inspect_key",)) +@pytest.mark.parametrize( + ("key", "owner", "expected", "calls"), + [ + (replace(PLAIN_KEY, blocked=True), None, KeyBlocked(), ("inspect_key",)), + ( + replace(PLAIN_KEY, user_id="u1"), + OwnerFacts(deactivated=True), + OwnerDeactivated(user_id="u1"), + ("inspect_key", "load_key_owner"), + ), + (replace(PLAIN_KEY, has_token=False), None, KeyWithoutToken(), ("inspect_key",)), + (replace(PLAIN_KEY, user_id="u1"), OwnerFacts(deactivated=False), None, ("inspect_key", "load_key_owner")), + (replace(PLAIN_KEY, user_id="u1"), None, None, ("inspect_key", "load_key_owner")), + ], +) +async def test_verify_decides_whether_the_resolved_key_is_usable( + key: KeyFacts, owner: OwnerFacts | None, expected: KeyAuthentication | None, calls: tuple[str, ...] +): + handler: Final = RecordingHandler(key=key, owner=owner) + verified: Final = await PythonCore().keys().verify(handler, _facts("hash"), NOW) + assert verified == (VirtualKey(key=key) if expected is None else expected) + assert _calls(handler) == calls @pytest.mark.asyncio -async def test_deactivated_owner_is_denied_after_model_access(): - outcome, calls = await _check(replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True)) - assert (outcome, calls) == ( - Deny(OwnerDeactivated(user_id="u1")), - ("inspect_key", "enforce_model_access", "load_key_owner"), +async def test_verify_labels_an_expired_key_with_the_hash_it_was_given(): + expired: Final = replace(PLAIN_KEY, expires=NOW - timedelta(seconds=1)) + verified: Final = await PythonCore().keys().verify(RecordingHandler(key=expired), _facts("hash"), NOW) + assert verified == ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label="hash") + + +@pytest.mark.asyncio +async def test_expired_key_over_its_team_member_budget_is_denied_as_expired_without_a_budget_alert(): + key: Final = replace( + PLAIN_KEY, user_id="u1", team_member=TeamMember(user_id="u1", team_id="t1"), expires=NOW - timedelta(seconds=1) ) + handler: Final = RecordingHandler(key=key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + authentication, calls = await _authenticate("sk-x", handler=handler) + assert authentication == ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label=hash_token("sk-x")) + assert AUTHORIZE_CALLS.isdisjoint(calls) @pytest.mark.asyncio -@pytest.mark.parametrize("owner", [None, OwnerFacts(deactivated=False)]) -async def test_active_or_missing_owner_continues(owner: OwnerFacts | None): - outcome, _ = await _check(replace(PLAIN_KEY, user_id="u1"), owner=owner) - assert outcome == _allowed() +async def test_deactivated_owner_is_denied_before_model_access_is_checked(): + handler: Final = RecordingHandler(key=replace(PLAIN_KEY, user_id="u1"), owner=OwnerFacts(deactivated=True)) + authentication, calls = await _authenticate("sk-x", handler=handler) + assert authentication == OwnerDeactivated(user_id="u1") + assert "enforce_model_access" not in calls @pytest.mark.asyncio -async def test_exempt_request_skips_every_budget_check_but_not_expiry(): +@pytest.mark.parametrize( + ("handler", "settings", "expected"), + [ + (RecordingHandler(), NO_DB, NoDatabase()), + (RecordingHandler(key=replace(PLAIN_KEY, blocked=True)), KEY_SETTINGS, KeyBlocked()), + ], +) +async def test_master_key_only_route_no_longer_hides_why_a_key_is_unusable( + handler: RecordingHandler, settings: KeySettings, expected: KeyAuthentication +): + authentication, _ = await _authenticate("sk-x", route=MASTER_ONLY_ROUTE, settings=settings, handler=handler) + assert authentication == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("identity", "allow_user_auth", "expected"), + [ + (UserAuthRoute(), True, _allowed()), + (UserAuthRoute(), False, Deny(UserAuthDisabled())), + (NoMasterKey(token="sk-x"), False, _allowed()), + (MasterKey(), False, _allowed()), + (CachedAdminKey(), False, _allowed()), + ], +) +async def test_non_virtual_key_identities_are_authorized_without_io( + identity: KeyIdentity, allow_user_auth: bool, expected: Outcome +): + settings: Final = replace(AUTHORIZE_SETTINGS, allow_user_auth=allow_user_auth) + assert await _authorize(identity, route=MASTER_ONLY_ROUTE, settings=settings) == (expected, ()) + + +@pytest.mark.asyncio +async def test_virtual_key_on_a_master_key_only_route_is_denied_before_any_io(): + assert await _authorize(VirtualKey(key=PLAIN_KEY), route=MASTER_ONLY_ROUTE) == (Deny(MasterKeyOnlyRoute()), ()) + + +@pytest.mark.asyncio +async def test_plain_virtual_key_runs_every_authorization_step_in_order(): + assert await _authorize_key(PLAIN_KEY) == (_allowed(), AUTHORIZED) + + +@pytest.mark.asyncio +async def test_exempt_request_skips_every_budget_check(): key: Final = replace( PLAIN_KEY, user_id="u1", @@ -359,13 +448,10 @@ async def test_exempt_request_skips_every_budget_check_but_not_expiry(): spend=5.0, soft_budget=1.0, ) - assert await _check(key, exempt=True) == ( + assert await _authorize_key(key, exempt=True) == ( _allowed(), - ("inspect_key", "enforce_model_access", "load_key_owner", "check_budget_exemption", "load_key_context"), + ("enforce_model_access", "check_budget_exemption", "load_key_context"), ) - expired: Final = replace(key, expires=NOW - timedelta(seconds=1)) - outcome, _ = await _check(expired, exempt=True) - assert outcome == Deny(ExpiredKey(expires=NOW - timedelta(seconds=1), now=NOW, key_label="hash")) @pytest.mark.asyncio @@ -380,7 +466,7 @@ async def test_exempt_request_skips_every_budget_check_but_not_expiry(): ) async def test_team_member_budget_alerts_then_denies_at_the_limit(budget: TeamMemberBudget | None, denied: bool): key: Final = replace(PLAIN_KEY, user_id="u1", team_member=TeamMember(user_id="u1", team_id="t1")) - outcome, calls = await _check(key, team_member_budget=budget) + outcome, calls = await _authorize_key(key, team_member_budget=budget) if denied: assert budget is not None assert outcome == Deny( @@ -394,26 +480,20 @@ async def test_team_member_budget_alerts_then_denies_at_the_limit(budget: TeamMe @pytest.mark.asyncio async def test_team_member_budget_is_read_only_for_a_team_member(): - key: Final = replace(PLAIN_KEY, user_id="u1") - _, calls = await _check(key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + _, calls = await _authorize_key( + replace(PLAIN_KEY, user_id="u1"), team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5) + ) assert "read_team_member_budget" not in calls @pytest.mark.asyncio async def test_team_member_budget_needs_a_database(): key: Final = replace(PLAIN_KEY, user_id="u1", team_member=TeamMember(user_id="u1", team_id="t1")) - no_db: Final = KeySettings(MASTER, False, False, frozenset(), 0.0) - _, calls = await _check(key, settings=no_db, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) - assert "read_team_member_budget" not in calls - - -@pytest.mark.asyncio -async def test_team_member_budget_is_checked_before_expiry(): - key: Final = replace( - PLAIN_KEY, user_id="u1", team_member=TeamMember(user_id="u1", team_id="t1"), expires=NOW - timedelta(seconds=1) + handler: Final = RecordingHandler(key=key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) + _, calls = await _authorize( + VirtualKey(key=key), settings=replace(AUTHORIZE_SETTINGS, has_db=False), handler=handler ) - outcome, _ = await _check(key, team_member_budget=TeamMemberBudget(spend=1.0, max_budget=0.5)) - assert outcome == Deny(TeamMemberBudgetExceeded(user_id="u1", team_id="t1", spend=1.0, max_budget=0.5)) + assert "read_team_member_budget" not in calls @pytest.mark.asyncio @@ -425,18 +505,8 @@ async def test_team_member_budget_is_checked_before_expiry(): (4.9, 5.0, None, _allowed()), (6.0, 5.0, 0.5, _allowed(budget_throttle_pct=0.5)), (4.9, 5.0, 0.5, _allowed()), - ( - 1e9, - math.inf, - None, - _allowed(), - ), - ( - 1e9, - math.nan, - None, - _allowed(), - ), + (1e9, math.inf, None, _allowed()), + (1e9, math.nan, None, _allowed()), ], ) async def test_key_budget_denies_or_throttles_at_the_limit( @@ -444,17 +514,19 @@ async def test_key_budget_denies_or_throttles_at_the_limit( ): key: Final = replace(PLAIN_KEY, max_budget=max_budget, budget_throttle_pct=throttle) handler: Final = RecordingHandler(key=key, key_spend=spend) - outcome, calls = await _run(PythonCore().keys().check(handler, _facts("hash"), SETTINGS, NOW), handler) + outcome, calls = await _authorize(VirtualKey(key=key), handler=handler) assert outcome == expected - assert calls[3:6] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget") - assert handler.calls[5] == ("alert_key_budget", spend) + assert calls[2:5] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget") + assert handler.calls[4] == ("alert_key_budget", spend) assert ("read_model_budget_scope" in calls) == isinstance(expected, Allow) @pytest.mark.asyncio async def test_key_budget_is_enforced_only_on_llm_routes(): key: Final = replace(PLAIN_KEY, max_budget=1.0) - outcome, calls = await _check(key, is_llm_api_route=False, key_spend=2.0) + outcome, calls = await _authorize( + VirtualKey(key=key), is_llm_api_route=False, handler=RecordingHandler(key=key, key_spend=2.0) + ) assert outcome == _allowed() assert "read_key_spend" not in calls assert "read_model_budget_scope" in calls @@ -475,16 +547,16 @@ async def test_key_budget_is_enforced_only_on_llm_routes(): async def test_soft_budget_alerts_at_the_threshold_before_model_budgets( spend: float, soft_budget: float | None, alerted: bool ): - outcome, calls = await _check(replace(PLAIN_KEY, spend=spend, soft_budget=soft_budget)) + outcome, calls = await _authorize_key(replace(PLAIN_KEY, spend=spend, soft_budget=soft_budget)) assert outcome == _allowed() soft_alert: Final = ("alert_key_soft_budget",) if alerted else () - assert calls[3:] == ("send_key_budget_alerts", *soft_alert, "read_model_budget_scope", "load_key_context") + assert calls[2:] == ("send_key_budget_alerts", *soft_alert, "read_model_budget_scope", "load_key_context") @pytest.mark.asyncio async def test_throttled_key_still_gets_soft_and_model_budget_checks(): key: Final = replace(PLAIN_KEY, max_budget=5.0, budget_throttle_pct=0.5, spend=6.0, soft_budget=1.0) - outcome, calls = await _check(key, key_spend=6.0) + outcome, calls = await _authorize_key(key, key_spend=6.0) assert outcome == _allowed(budget_throttle_pct=0.5) assert calls[-4:] == ("alert_key_budget", "alert_key_soft_budget", "read_model_budget_scope", "load_key_context") @@ -510,31 +582,20 @@ async def test_throttled_key_still_gets_soft_and_model_budget_checks(): async def test_proxy_budget_only_alerts_and_runs_for_exempt_requests_too( exempt: bool, proxy_max_budget: float, has_db: bool, proxy_spend: float | None, tail: tuple[Call, ...] ): - settings: Final = replace(SETTINGS, proxy_max_budget=proxy_max_budget, has_db=has_db) + settings: Final = replace(AUTHORIZE_SETTINGS, proxy_max_budget=proxy_max_budget, has_db=has_db) handler: Final = RecordingHandler(exempt=exempt, proxy_spend=proxy_spend) - outcome, _ = await _run(PythonCore().keys().check(handler, _facts("hash"), settings, NOW), handler) + outcome, _ = await _authorize(VirtualKey(key=PLAIN_KEY), settings=settings, handler=handler) assert outcome == _allowed() assert handler.calls[-len(tail) :] == tail @pytest.mark.asyncio -@pytest.mark.parametrize("exempt", [False, True]) -async def test_key_without_token_is_denied_after_its_context_loads(exempt: bool): - settings: Final = replace(SETTINGS, proxy_max_budget=10.0) - handler: Final = RecordingHandler(key=replace(PLAIN_KEY, has_token=False), exempt=exempt, proxy_spend=1.0) - outcome, calls = await _run(PythonCore().keys().check(handler, _facts("hash"), settings, NOW), handler) - assert outcome == Deny(KeyWithoutToken()) - assert calls[-3:] == ("load_key_context", "read_proxy_spend", "alert_proxy_budget") - - -@pytest.mark.asyncio -async def test_a_model_budget_denial_stops_the_key_chain_before_its_context_loads(): +async def test_a_model_budget_denial_stops_authorization_before_the_key_context_loads(): exceeded: Final = ModelBudgetExceeded( message="over", current_cost=2.0, max_budget=1.0, entity_type="user", entity_id="u-1" ) scope: Final = replace(NO_MODEL_BUDGETS, models=("gpt-4o",), user_has_model_budgets=True) - handler: Final = RecordingHandler(model_budget_scope=scope, user_model_budget_exceeded=exceeded) - outcome, calls = await _run(PythonCore().keys().check(handler, _facts("hash"), SETTINGS, NOW), handler) + outcome, calls = await _authorize_key(PLAIN_KEY, model_budget_scope=scope, user_model_budget_exceeded=exceeded) assert outcome == Deny(exceeded) assert calls[-2:] == ("read_model_budget_scope", "check_user_model_budget")