From 8683d4a77ebfef9039ab4e23da9b3a13afe7a801 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 07:58:45 +0900 Subject: [PATCH] refactor(auth): resolve virtual keys through AuthCore effects The Python core now decides the whole virtual key resolution path. It asks the proxy to resolve the end user, the presented key (preresolved, cached or UI login), evict an expired cached admin key and load an uncached key from the DB, each as a typed effect with its own continuation. It decides the admin key expiry, master-key-only routes, the no-DB case and the sk- prefix check itself, so the builder no longer has a legacy fallback. The budget and team checks in validate_resolved_virtual_key still run after Allow(VirtualKey) and move into the core in later passes Co-Authored-By: Claude Opus 5.5 --- litellm/_auth/AGENTS.md | 10 +- litellm/_auth/decision.py | 104 +++++- litellm/_auth/driver.py | 42 +++ litellm/_auth/python_core.py | 80 +++- litellm/proxy/auth/auth_denial.py | 71 +++- litellm/proxy/auth/user_api_key_auth.py | 467 ++++++++++++------------ tests/unit/_auth/test_python_core.py | 163 +++++++-- 7 files changed, 647 insertions(+), 290 deletions(-) create mode 100644 litellm/_auth/driver.py diff --git a/litellm/_auth/AGENTS.md b/litellm/_auth/AGENTS.md index b961828dd37..dba12c1c412 100644 --- a/litellm/_auth/AGENTS.md +++ b/litellm/_auth/AGENTS.md @@ -8,8 +8,8 @@ ## The swap point -- `decision.py` is the interface. `AuthCore.decide(RequestFacts, Settings) -> Outcome` is the only thing a core implements, and `python_core.PythonCore` is the Python one. A Rust core will implement the same protocol through pyo3, and the builder picks one without changing anything else -- A core does no IO and reads no globals. The same facts and settings always give the same outcome -- `Defer` means the core does not handle this request yet, and the legacy builder in `litellm/proxy/auth/user_api_key_auth.py` takes it from there. Move one decision at a time out of the builder into the core, and delete the builder code it replaces in the same change -- Projection stays in the proxy. `Allow` grants become `UserAPIKeyAuth` in the builder, and `litellm/proxy/auth/auth_denial.py` maps every `DenyReason` to the existing error contract -- When a decision needs a lookup, add it as an effect the core yields and the proxy answers, so both cores share the same IO code +- `decision.py` is the interface. `AuthCore.start(RequestFacts, Settings, now) -> Step` is the only thing a core implements, and `python_core.PythonCore` is the Python one. A Rust core will implement the same protocol through pyo3, and the builder picks one without changing anything else +- A core does no IO and reads no globals. The same facts, settings, time and effect answers always give the same outcome +- A `Step` is either an `Outcome` or an effect. Each effect carries its typed `resume` continuation, and `driver.drive` answers effects through an `EffectHandler` until an outcome comes back. The proxy's handler (`_BuilderEffects` in `litellm/proxy/auth/user_api_key_auth.py`) owns all cache and DB IO, so both cores share it +- Projection stays in the proxy. `Allow` grants become `UserAPIKeyAuth` in the builder, and `litellm/proxy/auth/auth_denial.py` maps every `Denial` to the existing error contract +- The checks in `validate_resolved_virtual_key` and the centralized common checks still run as projection after `Allow(VirtualKey(...))`. Move them into the core one at a time, as effects plus pure decisions, and delete the builder code each one replaces in the same change diff --git a/litellm/_auth/decision.py b/litellm/_auth/decision.py index d2eb53891ac..648877f18ed 100644 --- a/litellm/_auth/decision.py +++ b/litellm/_auth/decision.py @@ -1,5 +1,6 @@ +from collections.abc import Callable from dataclasses import dataclass -from enum import Enum +from datetime import datetime from typing import Protocol, TypeAlias @@ -13,6 +14,8 @@ class RequestFacts: class Settings: master_key: str | None allow_user_auth: bool + has_db: bool + master_key_only_routes: frozenset[str] @dataclass(frozen=True, slots=True) @@ -30,13 +33,58 @@ class MasterKey: pass -Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey +@dataclass(frozen=True, slots=True) +class CachedAdminKey: + pass -class DenyReason(Enum): - MISSING_KEY = "missing_key" - MALFORMED_KEY = "malformed_key" - USER_AUTH_DISABLED = "user_auth_disabled" +@dataclass(frozen=True, slots=True) +class VirtualKey: + loaded_from_db: bool + + +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 + + +@dataclass(frozen=True, slots=True) +class MasterKeyOnlyRoute: + pass + + +@dataclass(frozen=True, slots=True) +class NoDatabase: + pass + + +@dataclass(frozen=True, slots=True) +class NotAVirtualKey: + pass + + +Denial: TypeAlias = ( + MissingKey | MalformedKey | UserAuthDisabled | ExpiredKey | MasterKeyOnlyRoute | NoDatabase | NotAVirtualKey +) @dataclass(frozen=True, slots=True) @@ -46,16 +94,52 @@ class Allow: @dataclass(frozen=True, slots=True) class Deny: - reason: DenyReason + denial: Denial + + +Outcome: TypeAlias = Allow | Deny @dataclass(frozen=True, slots=True) -class Defer: +class AdminKey: + expires: datetime | None + + +@dataclass(frozen=True, slots=True) +class NonAdminKey: pass -Outcome: TypeAlias = Allow | Deny | Defer +PresentedKey: TypeAlias = AdminKey | NonAdminKey + + +@dataclass(frozen=True, slots=True) +class ResolveEndUser: + prefetch_key: bool + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class ResolvePresentedKey: + hashed_token: str + resume: "Callable[[PresentedKey | None], Step]" + + +@dataclass(frozen=True, slots=True) +class EvictPresentedKey: + hashed_token: str + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class LoadKey: + hashed_token: str + resume: "Callable[[], Step]" + + +Effect: TypeAlias = ResolveEndUser | ResolvePresentedKey | EvictPresentedKey | LoadKey +Step: TypeAlias = Outcome | Effect class AuthCore(Protocol): - def decide(self, facts: RequestFacts, settings: Settings) -> Outcome: ... + def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... diff --git a/litellm/_auth/driver.py b/litellm/_auth/driver.py new file mode 100644 index 00000000000..1e8b0804f98 --- /dev/null +++ b/litellm/_auth/driver.py @@ -0,0 +1,42 @@ +from typing import Protocol, assert_never + +from litellm._auth.decision import ( + Allow, + Deny, + EvictPresentedKey, + LoadKey, + Outcome, + PresentedKey, + ResolveEndUser, + ResolvePresentedKey, + Step, +) + + +class EffectHandler(Protocol): + async def resolve_end_user(self, prefetch_key: bool) -> None: ... + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: ... + + async def evict_presented_key(self, hashed_token: str) -> None: ... + + async def load_key(self, hashed_token: str) -> None: ... + + +async def drive(step: Step, handler: EffectHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case ResolveEndUser(prefetch_key=prefetch_key, resume=resume): + await handler.resolve_end_user(prefetch_key) + return await drive(resume(), handler) + case ResolvePresentedKey(hashed_token=hashed_token, resume=resume): + return await drive(resume(await handler.resolve_presented_key(hashed_token)), handler) + case EvictPresentedKey(hashed_token=hashed_token, resume=resume): + await handler.evict_presented_key(hashed_token) + return await drive(resume(), handler) + case LoadKey(hashed_token=hashed_token, resume=resume): + await handler.load_key(hashed_token) + return await drive(resume(), handler) + case _: + assert_never(step) diff --git a/litellm/_auth/python_core.py b/litellm/_auth/python_core.py index 1bf07265cd2..bb160e214e5 100644 --- a/litellm/_auth/python_core.py +++ b/litellm/_auth/python_core.py @@ -1,32 +1,84 @@ +from datetime import datetime from typing import Final from litellm._auth.decision import ( + AdminKey, Allow, - Defer, + CachedAdminKey, Deny, - DenyReason, + EvictPresentedKey, + ExpiredKey, + LoadKey, + MalformedKey, MasterKey, + MasterKeyOnlyRoute, + MissingKey, + NoDatabase, NoMasterKey, - Outcome, + NonAdminKey, + NotAVirtualKey, + PresentedKey, RequestFacts, + ResolveEndUser, + ResolvePresentedKey, Settings, + Step, + UserAuthDisabled, UserAuthRoute, + VirtualKey, ) -from litellm._auth.key import is_master_key +from litellm._auth.key import hash_token, is_master_key _USER_AUTH_ROUTE: Final = "/user/auth" +_VIRTUAL_KEY_PREFIX: Final = "sk-" class PythonCore: - def decide(self, facts: RequestFacts, settings: Settings) -> Outcome: + def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: + token: Final = facts.token if settings.master_key is None: - return Allow(NoMasterKey(token=facts.token)) - if facts.token is None: - return Deny(DenyReason.MISSING_KEY) - if facts.token == "": - return Deny(DenyReason.MALFORMED_KEY) + return Allow(NoMasterKey(token=token)) + if token is None: + return Deny(MissingKey()) + if token == "": + return Deny(MalformedKey()) if facts.route == _USER_AUTH_ROUTE: - return Allow(UserAuthRoute()) if settings.allow_user_auth else Deny(DenyReason.USER_AUTH_DISABLED) - if is_master_key(facts.token, settings.master_key): - return Allow(MasterKey()) - return Defer() + return Allow(UserAuthRoute()) if settings.allow_user_auth else Deny(UserAuthDisabled()) + if is_master_key(token, settings.master_key): + return ResolveEndUser(prefetch_key=False, resume=lambda: Allow(MasterKey())) + hashed_token: Final = hash_token(token) + return ResolveEndUser( + prefetch_key=True, + resume=lambda: ResolvePresentedKey( + hashed_token=hashed_token, + resume=lambda presented: _after_presented_key(presented, token, hashed_token, facts, settings, now), + ), + ) + + +def _after_presented_key( + presented: PresentedKey | None, + token: str, + hashed_token: str, + facts: RequestFacts, + settings: Settings, + now: datetime, +) -> Step: + match presented: + case AdminKey(expires=expires) if expires is not None and expires < now: + return EvictPresentedKey( + hashed_token=hashed_token, resume=lambda: Deny(ExpiredKey(expires=expires, now=now)) + ) + case AdminKey(): + return Allow(CachedAdminKey()) + case NonAdminKey() | None: + pass + if facts.route in settings.master_key_only_routes: + return Deny(MasterKeyOnlyRoute()) + if not settings.has_db: + return Deny(NoDatabase()) + if presented is not None: + return Allow(VirtualKey(loaded_from_db=False)) + if not token.startswith(_VIRTUAL_KEY_PREFIX): + return Deny(NotAVirtualKey()) + return LoadKey(hashed_token=hashed_token, resume=lambda: Allow(VirtualKey(loaded_from_db=True))) diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py index 6ab2b56dcda..69dece93d99 100644 --- a/litellm/proxy/auth/auth_denial.py +++ b/litellm/proxy/auth/auth_denial.py @@ -1,31 +1,86 @@ -from typing import assert_never +from typing import Final, assert_never from fastapi import HTTPException, status -from litellm._auth.decision import DenyReason +from litellm._auth.decision import ( + Denial, + ExpiredKey, + MalformedKey, + MasterKeyOnlyRoute, + MissingKey, + NoDatabase, + NotAVirtualKey, + UserAuthDisabled, +) +from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER, INVALID_VIRTUAL_KEY_ERROR_MESSAGE from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.auth.auth_utils import abbreviate_api_key +from litellm.types.proxy.auth.user_api_key_auth import is_jwt + +_JWT_AUTH_DISABLED_HINT: Final = ( + " This key has the structure of a JWT, but JWT auth is not enabled on this proxy, so it was treated as a" + " virtual key. Set `enable_jwt_auth: true` under `general_settings` in your proxy config to authenticate" + " with JWTs." +) -def denial_exception(reason: DenyReason) -> ProxyException | HTTPException: - match reason: - case DenyReason.MISSING_KEY: +def _not_a_virtual_key(token: str, jwt_auth_enabled: bool) -> HTTPException: + masked: Final = f"{token[:4]}****{token[-4:]}" if len(token) > 8 else "****" + hint: Final = _JWT_AUTH_DISABLED_HINT if not jwt_auth_enabled and is_jwt(token=token) else "" + error: Final = HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={masked}, expected to start with 'sk-'.{hint}", + ) + # Log routing classifies this 401 by where it was raised, never by its message text. + setattr(error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) + return error + + +def denial_exception( + denial: Denial, token: str | None, route: str, jwt_auth_enabled: bool +) -> ProxyException | HTTPException: + match denial: + case MissingKey(): return ProxyException( message="Authentication Error, No api key passed in.", type=ProxyErrorTypes.auth_error, param=None, code=status.HTTP_401_UNAUTHORIZED, ) - case DenyReason.MALFORMED_KEY: + case MalformedKey(): return ProxyException( message="Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix.", type=ProxyErrorTypes.auth_error, param=None, code=status.HTTP_401_UNAUTHORIZED, ) - case DenyReason.USER_AUTH_DISABLED: + case UserAuthDisabled(): return HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="'allow_user_auth' not set or set to False", ) + case ExpiredKey(expires=expires, now=now): + return ProxyException( + message=f"Authentication Error - Expired Key. Key Expiry time {expires} and current time {now}", + type=ProxyErrorTypes.expired_key, + code=status.HTTP_401_UNAUTHORIZED, + param=abbreviate_api_key(api_key=token or ""), + ) + case MasterKeyOnlyRoute(): + return ProxyException( + message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ) + case NoDatabase(): + return ProxyException( + message="No connected db.", + type=ProxyErrorTypes.no_db_connection, + code=400, + param=None, + ) + case NotAVirtualKey(): + return _not_a_virtual_key(token or "", jwt_auth_enabled) case _: - assert_never(reason) + assert_never(denial) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ed3427136a9..e61c501ed4b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -31,16 +31,21 @@ from litellm._auth.credential import ( select_websocket_api_key, ) from litellm._auth.decision import ( + AdminKey, Allow, AuthCore, - Defer, + CachedAdminKey, Deny, MasterKey, NoMasterKey, + NonAdminKey, + PresentedKey, RequestFacts, Settings, UserAuthRoute, + VirtualKey, ) +from litellm._auth.driver import drive from litellm._auth.python_core import PythonCore from litellm._internal_context import service_target from litellm._logging import verbose_logger, verbose_proxy_logger @@ -49,8 +54,6 @@ from litellm.caching.redis_cache import RedisCache from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, GLOBAL_PROXY_SPEND_CACHE_KEY, - INVALID_VIRTUAL_KEY_ERROR_MARKER, - INVALID_VIRTUAL_KEY_ERROR_MESSAGE, LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_MASTER_KEY_ALIAS, MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY, @@ -1555,12 +1558,6 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( # the lookup and return None (caller proceeds to auth_builder). _JWT_PROXY_ADMIN_SENTINEL: Final = "__JWT_PROXY_ADMIN__" -_JWT_AUTH_DISABLED_HINT = ( - " This key has the structure of a JWT, but JWT auth is not enabled on this proxy, so it was treated as a" - " virtual key. Set `enable_jwt_auth: true` under `general_settings` in your proxy config to authenticate" - " with JWTs." -) - class _JWTKeyMappingFailure(Enum): NO_ELIGIBLE_KEY = "JWT Key Mapping: No eligible existing virtual key for this identity. Access denied." @@ -2196,7 +2193,7 @@ class _ResolvedEndUser(NamedTuple): async def _resolve_end_user( request: Request, - request_data: dict[str, object], + request_data: Mapping[str, object], token_end_user_id: str | None, api_key: str, key_is_resolved: bool, @@ -2291,6 +2288,182 @@ async def _master_key_auth( return with_end_user +_MASTER_KEY_ONLY_ROUTES: Final = frozenset(LiteLLMRoutes.master_key_only_routes.value) + + +def _aware_expiry(expires: datetime | str) -> datetime: + expiry: Final = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) + return ( + expiry if expiry.tzinfo is not None and expiry.utcoffset() is not None else expiry.replace(tzinfo=timezone.utc) + ) + + +def _presented_key(key: UserAPIKeyAuth) -> PresentedKey: + if key.user_role != LitellmUserRoles.PROXY_ADMIN: + return NonAdminKey() + return AdminKey(expires=None if key.expires is None else _aware_expiry(key.expires)) + + +def _with_end_user(key: UserAPIKeyAuth, end_user: _ResolvedEndUser, parent_otel_span: Span | None) -> UserAPIKeyAuth: + with_end_user: Final = update_valid_token_with_end_user_params(valid_token=key, end_user_params=end_user.params) + with_end_user.parent_otel_span = parent_otel_span + if end_user.end_user_object is not None: + with_end_user.end_user_object_permission = end_user.end_user_object.object_permission + return with_end_user + + +def _reset_end_user_fields(key: UserAPIKeyAuth, end_user: _ResolvedEndUser) -> UserAPIKeyAuth: + key.end_user_id = end_user.end_user_id or None + key.end_user_tpm_limit = end_user.budget.tpm_limit if end_user.budget is not None else None + key.end_user_rpm_limit = end_user.budget.rpm_limit if end_user.budget is not None else None + key.end_user_tpd_limit = end_user.budget.tpd_limit if end_user.budget is not None else None + key.allowed_model_region = ( + end_user.end_user_object.allowed_model_region if end_user.end_user_object is not None else None + ) + return key + + +async def _refresh_team_fields_from_cache( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> None: + if valid_token.team_id is None or valid_token.team_id == UI_TEAM_ID: + return + ## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token + try: + team_obj: Final[LiteLLM_TeamTableCachedObj] = await get_team_object( + team_id=valid_token.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=True, + ) + + if ( + team_obj.last_refreshed_at is not None + and valid_token.last_refreshed_at is not None + and team_obj.last_refreshed_at > valid_token.last_refreshed_at + ): + team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__) + + for k, v in team_obj_dict.items(): + field_name = f"team_{k}" + if field_name in valid_token.__fields__: + setattr(valid_token, field_name, v) + except Exception as e: + verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. + + +class _BuilderEffects: + def __init__( + self, + *, + request: Request, + request_data: Mapping[str, object], + token: str, + preresolved_key: UserAPIKeyAuth | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, + route: str, + ) -> None: + self._request: Final = request + self._request_data: Final = request_data + self._token: Final = token + self._preresolved_key: Final = preresolved_key + self._prisma_client: Final = prisma_client + self._user_api_key_cache: Final = user_api_key_cache + self._parent_otel_span: Final = parent_otel_span + self._proxy_logging_obj: Final = proxy_logging_obj + self._route: Final = route + self.end_user: _ResolvedEndUser = _ResolvedEndUser(end_user_id=None, end_user_object=None, budget=None) + self.resolved_key: UserAPIKeyAuth | None = preresolved_key + + @property + def key(self) -> UserAPIKeyAuth: + if self.resolved_key is None: + raise RuntimeError("the auth core granted a key it never resolved") + return self.resolved_key + + def _identity_store(self, *, check_cache_only: bool) -> IdentityStore: + return IdentityStore( + self._prisma_client, + self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + check_cache_only=check_cache_only, + ) + + async def resolve_end_user(self, prefetch_key: bool) -> None: + self.end_user = await _resolve_end_user( + request=self._request, + request_data=self._request_data, + token_end_user_id=self._preresolved_key.end_user_id if self._preresolved_key is not None else None, + api_key=self._token, + key_is_resolved=not prefetch_key or self._preresolved_key is not None, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + route=self._route, + ) + + async def _cached_key(self, hashed_token: str) -> UserAPIKeyAuth | None: + try: + with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): + cached: Final = IdentityStore.key_from_principal( + await self._identity_store(check_cache_only=True).resolve(hashed_token=hashed_token) + ) + except Exception: + verbose_logger.debug("api key not found in cache.") + return None + # Key-cache entries are written only after the proxy validated a virtual key or the master key, but + # via_virtual_key is exclude=True so serialization drops it. Restore it at this trusted boundary. + cached.via_virtual_key = True + return cached + + def _ui_login_key(self) -> UserAPIKeyAuth | None: + # Decryption fails closed for anything that is not a genuine UI login blob. + if self._token.startswith("sk-") or get_secret_bool("EXPERIMENTAL_UI_LOGIN") is False: + return None + return ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(self._token) + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: + found: Final = self._preresolved_key or await self._cached_key(hashed_token) or self._ui_login_key() + if found is None: + return None + self.resolved_key = ( + await _refresh_session_token_grants( + valid_token=found, + prisma_client=self._prisma_client, + user_api_key_cache=self._user_api_key_cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + if found.is_session_token and self._prisma_client is not None + else found + ) + return _presented_key(self.resolved_key) + + async def evict_presented_key(self, hashed_token: str) -> None: + await delete_cache_key_object( + hashed_token=hashed_token, + user_api_key_cache=self._user_api_key_cache, + proxy_logging_obj=self._proxy_logging_obj, + ) + + async def load_key(self, hashed_token: str) -> None: + with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): + self.resolved_key = IdentityStore.key_from_principal( + await self._identity_store(check_cache_only=False).resolve(hashed_token=hashed_token) + ) + + async def user_api_key_auth_builder( request: Request, api_key: str, @@ -2701,10 +2874,34 @@ async def user_api_key_auth_builder( status_code=500, detail={f"Master key must be a valid string. Current type={type(master_key)}"}, ) - match _AUTH_CORE.decide( - RequestFacts(route=route, token=api_key), - Settings(master_key=master_key, allow_user_auth=general_settings.get("allow_user_auth", False) is True), - ): + effects: Final = _BuilderEffects( + request=request, + request_data=request_data, + token=api_key, + preresolved_key=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + outcome: Final = await drive( + _AUTH_CORE.start( + RequestFacts(route=route, token=api_key), + Settings( + master_key=master_key, + allow_user_auth=general_settings.get("allow_user_auth", False) is True, + has_db=prisma_client is not None, + master_key_only_routes=_MASTER_KEY_ONLY_ROUTES, + ), + datetime.now(timezone.utc), + ), + effects, + ) + valid_token = effects.resolved_key # rebind-ok: the exception handler below reports the resolved identity + match outcome: + case Deny(denial=denial): + raise denial_exception(denial, token=api_key, route=route, jwt_auth_enabled=enable_jwt_auth) case Allow(grant=grant): match grant: case NoMasterKey(token=token): @@ -2718,18 +2915,7 @@ async def user_api_key_auth_builder( case MasterKey(): return await _master_key_auth( token=api_key, - end_user=await _resolve_end_user( - request=request, - request_data=request_data, - token_end_user_id=valid_token.end_user_id if valid_token is not None else None, - api_key=api_key, - key_is_resolved=True, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - ), + end_user=effects.end_user, admin_user_id=litellm_proxy_admin_name, parent_otel_span=parent_otel_span, route=route, @@ -2737,212 +2923,31 @@ async def user_api_key_auth_builder( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - case Deny(reason=reason): - raise denial_exception(reason) - case Defer(): - pass - - end_user: Final = await _resolve_end_user( - request=request, - request_data=request_data, - token_end_user_id=valid_token.end_user_id if valid_token is not None else None, - api_key=api_key, - key_is_resolved=valid_token is not None, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - ) - - if valid_token is None: - ## Check CACHE - try: - with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): - valid_token = IdentityStore.key_from_principal( - await IdentityStore( - prisma_client, - user_api_key_cache, + case CachedAdminKey(): + return _with_end_user(effects.key, effects.end_user, parent_otel_span) + case VirtualKey(loaded_from_db=loaded_from_db): + if not loaded_from_db: + await _refresh_team_fields_from_cache( + valid_token=effects.key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return await validate_resolved_virtual_key( + request=request, + request_data=request_data, + valid_token=( + _reset_end_user_fields(effects.key, effects.end_user) if loaded_from_db else effects.key + ), + api_key=hash_token(api_key) if loaded_from_db else api_key, + route=route, + start_time=start_time, parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, - ).resolve(hashed_token=hash_token(api_key)) - ) - # Key-cache entries are written only after the proxy validated a - # virtual key or the master key, but via_virtual_key is exclude=True - # so serialization drops it; restore it at this trusted boundary. - # The UI-login JWT fallback below constructs its token from a - # decrypted blob, not this cache, and stays unmarked. - if isinstance(valid_token, UserAPIKeyAuth): - valid_token.via_virtual_key = True - except Exception: - verbose_logger.debug("api key not found in cache.") - valid_token = None - - ## Check UI/CLI Hash Key - # Attempt decryption for non-sk- tokens unless the operator has - # explicitly set EXPERIMENTAL_UI_LOGIN=false to disable it. - # Unset (None) keeps the new default of always attempting decryption; - # decryption fails closed for anything that is not a genuine blob. - if ( - valid_token is None - and not api_key.startswith("sk-") - and get_secret_bool("EXPERIMENTAL_UI_LOGIN") is not False - ): - valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key) - - if valid_token is not None and valid_token.is_session_token and prisma_client is not None: - valid_token = await _refresh_session_token_grants( # rebind-ok: later checks read this name - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - - if ( - valid_token is not None - and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.user_role == LitellmUserRoles.PROXY_ADMIN - ): - if valid_token.expires is not None: - current_time = datetime.now(timezone.utc) - if isinstance(valid_token.expires, datetime): - expiry_time = valid_token.expires - else: - expiry_time = datetime.fromisoformat(valid_token.expires) - if expiry_time.tzinfo is None or expiry_time.tzinfo.utcoffset(expiry_time) is None: - expiry_time = expiry_time.replace(tzinfo=timezone.utc) - if expiry_time < current_time: - await delete_cache_key_object( - hashed_token=hash_token(api_key), - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - raise ProxyException( - message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", - type=ProxyErrorTypes.expired_key, - code=status.HTTP_401_UNAUTHORIZED, - param=abbreviate_api_key(api_key=api_key), - ) - valid_token = update_valid_token_with_end_user_params( - valid_token=valid_token, end_user_params=end_user.params - ) - valid_token.parent_otel_span = parent_otel_span - if end_user.end_user_object is not None: - valid_token.end_user_object_permission = end_user.end_user_object.object_permission - - return valid_token - - if ( - valid_token is not None - and isinstance(valid_token, UserAPIKeyAuth) - and valid_token.team_id is not None - and valid_token.team_id != UI_TEAM_ID - ): - ## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token - try: - team_obj: Final[LiteLLM_TeamTableCachedObj] = await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=True, - ) - - if ( - team_obj.last_refreshed_at is not None - and valid_token.last_refreshed_at is not None - and team_obj.last_refreshed_at > valid_token.last_refreshed_at - ): - team_obj_dict: Final[Mapping[str, object]] = MappingProxyType(team_obj.__dict__) - - for k, v in team_obj_dict.items(): - field_name = f"team_{k}" - if field_name in valid_token.__fields__: - setattr(valid_token, field_name, v) - except Exception as e: - verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. - - ## IF it's not a master key - ## Route should not be in master_key_only_routes - if route in LiteLLMRoutes.master_key_only_routes.value: - raise ProxyException( - message=f"Authentication Error, Tried to access route={route}, which is only for MASTER KEY", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, - ) - - ## Check DB - - if ( - prisma_client is None - ): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error - raise ProxyException( - message="No connected db.", - type=ProxyErrorTypes.no_db_connection, - code=400, - param=None, - ) - - if valid_token is None: - if isinstance(api_key, str): # if generated token, make sure it starts with sk-. - _masked_key: Final = f"{api_key[:4]}****{api_key[-4:]}" if len(api_key) > 8 else "****" - if not api_key.startswith("sk-"): - _hint = _JWT_AUTH_DISABLED_HINT if not enable_jwt_auth and JWTHandler.is_jwt(token=api_key) else "" - _malformed_key_error = HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail=( - f"{INVALID_VIRTUAL_KEY_ERROR_MESSAGE}. Received={_masked_key}, " - f"expected to start with 'sk-'.{_hint}" - ), - ) # prevent token hashes from being used - # Stamp provenance here so log routing classifies this 401 by - # where it was raised, never by its message text. - setattr(_malformed_key_error, INVALID_VIRTUAL_KEY_ERROR_MARKER, True) - raise _malformed_key_error - else: - verbose_logger.warning( - "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format( - type(api_key) if api_key is not None else "None" - ) - ) - if api_key.startswith("sk-"): - api_key = hash_token(token=api_key) - - with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): - valid_token = IdentityStore.key_from_principal( - await IdentityStore( - prisma_client, - user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ).resolve(hashed_token=api_key) - ) - # update end-user params on valid token - # These can change per request - it's important to update them here - valid_token.end_user_id = end_user.end_user_id or None - valid_token.end_user_tpm_limit = end_user.budget.tpm_limit if end_user.budget is not None else None - valid_token.end_user_rpm_limit = end_user.budget.rpm_limit if end_user.budget is not None else None - valid_token.end_user_tpd_limit = end_user.budget.tpd_limit if end_user.budget is not None else None - valid_token.allowed_model_region = ( - end_user.end_user_object.allowed_model_region if end_user.end_user_object is not None else None - ) - - return await validate_resolved_virtual_key( - request=request, - request_data=request_data, - valid_token=valid_token, - api_key=api_key, - route=route, - start_time=start_time, - parent_otel_span=parent_otel_span, - end_user_id=end_user.end_user_id, - end_user_params=end_user.params, - _end_user_object=end_user.end_user_object, - ) + end_user_id=effects.end_user.end_user_id, + end_user_params=effects.end_user.params, + _end_user_object=effects.end_user.end_user_object, + ) except Exception as e: return await UserAPIKeyAuthExceptionHandler.handle_authentication_error( e=e, diff --git a/tests/unit/_auth/test_python_core.py b/tests/unit/_auth/test_python_core.py index 1fe0a9862a4..753e8ef19ca 100644 --- a/tests/unit/_auth/test_python_core.py +++ b/tests/unit/_auth/test_python_core.py @@ -1,43 +1,162 @@ +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final + import pytest from litellm._auth.decision import ( + AdminKey, Allow, - Defer, + CachedAdminKey, Deny, - DenyReason, + ExpiredKey, + MalformedKey, MasterKey, + MasterKeyOnlyRoute, + MissingKey, + NoDatabase, NoMasterKey, + NonAdminKey, + NotAVirtualKey, Outcome, + PresentedKey, RequestFacts, Settings, + UserAuthDisabled, UserAuthRoute, + VirtualKey, ) +from litellm._auth.driver import drive +from litellm._auth.key import hash_token from litellm._auth.python_core import PythonCore -MASTER: str = "sk-master" +MASTER: Final = "sk-master" +NOW: Final = datetime(2026, 11, 6, 12, tzinfo=timezone.utc) +SETTINGS: Final = Settings( + master_key=MASTER, + allow_user_auth=False, + has_db=True, + master_key_only_routes=frozenset({"/global/spend/reset"}), +) +@dataclass +class RecordingHandler: + presented: PresentedKey | None = None + calls: tuple[tuple[str, object], ...] = () + + async def resolve_end_user(self, prefetch_key: bool) -> None: + self.calls = (*self.calls, ("resolve_end_user", prefetch_key)) + + async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: + self.calls = (*self.calls, ("resolve_presented_key", hashed_token)) + return self.presented + + async def evict_presented_key(self, hashed_token: str) -> None: + self.calls = (*self.calls, ("evict_presented_key", hashed_token)) + + async def load_key(self, hashed_token: str) -> None: + self.calls = (*self.calls, ("load_key", hashed_token)) + + +async def _decide( + token: str | None, + *, + route: str = "/chat/completions", + settings: Settings = SETTINGS, + presented: PresentedKey | None = None, +) -> tuple[Outcome, tuple[tuple[str, object], ...]]: + handler: Final = RecordingHandler(presented=presented) + outcome: Final = await drive(PythonCore().start(RequestFacts(route=route, token=token), settings, NOW), handler) + return outcome, handler.calls + + +@pytest.mark.asyncio @pytest.mark.parametrize( ("route", "token", "settings", "expected"), [ - ("/chat/completions", None, Settings(master_key=None, allow_user_auth=False), Allow(NoMasterKey(token=None))), - ( - "/chat/completions", - "sk-x", - Settings(master_key=None, allow_user_auth=False), - Allow(NoMasterKey(token="sk-x")), - ), - ("/chat/completions", None, Settings(master_key=MASTER, allow_user_auth=False), Deny(DenyReason.MISSING_KEY)), - ("/chat/completions", "", Settings(master_key=MASTER, allow_user_auth=False), Deny(DenyReason.MALFORMED_KEY)), - ("/user/auth", "", Settings(master_key=MASTER, allow_user_auth=True), Deny(DenyReason.MALFORMED_KEY)), - ("/user/auth", "sk-x", Settings(master_key=MASTER, allow_user_auth=True), Allow(UserAuthRoute())), - ("/user/auth", "sk-x", Settings(master_key=MASTER, allow_user_auth=False), Deny(DenyReason.USER_AUTH_DISABLED)), - ("/user/auth", MASTER, Settings(master_key=MASTER, allow_user_auth=False), Deny(DenyReason.USER_AUTH_DISABLED)), - ("/chat/completions", MASTER, Settings(master_key=MASTER, allow_user_auth=False), Allow(MasterKey())), - ("/key/generate", MASTER, Settings(master_key=MASTER, allow_user_auth=False), Allow(MasterKey())), - ("/chat/completions", "sk-master2", Settings(master_key=MASTER, allow_user_auth=False), Defer()), - ("/chat/completions", "sk-x", Settings(master_key=MASTER, allow_user_auth=False), Defer()), + ("/chat/completions", None, Settings(None, False, True, frozenset()), Allow(NoMasterKey(token=None))), + ("/chat/completions", "sk-x", Settings(None, False, False, frozenset()), Allow(NoMasterKey(token="sk-x"))), + ("/chat/completions", None, SETTINGS, Deny(MissingKey())), + ("/chat/completions", "", SETTINGS, Deny(MalformedKey())), + ("/user/auth", "", Settings(MASTER, True, True, frozenset()), Deny(MalformedKey())), + ("/user/auth", "sk-x", Settings(MASTER, True, True, frozenset()), Allow(UserAuthRoute())), + ("/user/auth", "sk-x", SETTINGS, Deny(UserAuthDisabled())), + ("/user/auth", MASTER, SETTINGS, Deny(UserAuthDisabled())), ], ) -def test_decide(route: str, token: str | None, settings: Settings, expected: Outcome): - assert PythonCore().decide(RequestFacts(route=route, token=token), settings) == expected +async def test_credential_decisions_do_no_io(route: str, token: str | None, settings: Settings, expected: Outcome): + assert await _decide(token, route=route, settings=settings) == (expected, ()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/chat/completions", "/global/spend/reset"]) +async def test_master_key_resolves_end_user_without_key_prefetch(route: str): + assert await _decide(MASTER, route=route) == (Allow(MasterKey()), (("resolve_end_user", False),)) + + +@pytest.mark.asyncio +async def test_unknown_virtual_key_is_loaded_from_the_db_after_a_cache_miss(): + hashed: Final = hash_token("sk-x") + assert await _decide("sk-x") == ( + Allow(VirtualKey(loaded_from_db=True)), + (("resolve_end_user", True), ("resolve_presented_key", hashed), ("load_key", hashed)), + ) + + +@pytest.mark.asyncio +async def test_cached_virtual_key_is_not_loaded_again(): + assert await _decide("sk-x", presented=NonAdminKey()) == ( + Allow(VirtualKey(loaded_from_db=False)), + (("resolve_end_user", True), ("resolve_presented_key", hash_token("sk-x"))), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("expires", [None, NOW + timedelta(seconds=1), NOW]) +async def test_unexpired_cached_admin_key_skips_route_and_db_checks(expires: datetime | None): + no_db: Final = Settings(MASTER, False, False, frozenset({"/global/spend/reset"})) + outcome, _ = await _decide( + "sk-admin", route="/global/spend/reset", settings=no_db, presented=AdminKey(expires=expires) + ) + assert outcome == Allow(CachedAdminKey()) + + +@pytest.mark.asyncio +async def test_expired_cached_admin_key_is_evicted_then_denied(): + expires: Final = NOW - timedelta(seconds=1) + hashed: Final = hash_token("sk-admin") + assert await _decide("sk-admin", presented=AdminKey(expires=expires)) == ( + Deny(ExpiredKey(expires=expires, now=NOW)), + (("resolve_end_user", True), ("resolve_presented_key", hashed), ("evict_presented_key", hashed)), + ) + + +@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", presented=presented) + assert outcome == Deny(MasterKeyOnlyRoute()) + assert "load_key" not in {name for name, _ in calls} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("presented", [None, NonAdminKey()]) +async def test_virtual_key_needs_a_database(presented: PresentedKey | None): + no_db: Final = Settings(MASTER, False, False, frozenset()) + outcome, _ = await _decide("sk-x", settings=no_db, presented=presented) + assert outcome == Deny(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): + outcome, calls = await _decide(token) + assert outcome == Deny(NotAVirtualKey()) + assert [name for name, _ in calls] == ["resolve_end_user", "resolve_presented_key"] + + +@pytest.mark.asyncio +async def test_cached_token_without_virtual_key_prefix_is_accepted(): + outcome, _ = await _decide("ui-login-blob", presented=NonAdminKey()) + assert outcome == Allow(VirtualKey(loaded_from_db=False))