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")