From 84266bf924ebafde269625cece6772a887719dbe Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 20 Jun 2026 18:49:41 -0700 Subject: [PATCH] feat(auth): resolve caller identity once into a Principal at the auth seam (#30887) Introduce a single, typed caller identity that is resolved once at the auth boundary and read by reference downstream, instead of being re-derived from a 50-field key object or rebuilt from request metadata. What this adds (litellm/proxy/auth/resolvers/), organized by responsibility: - Principal: a small, frozen, identity-only value type (user / organization / teams / project / end-user / roles / scopes / network), with its sub-models and the role mapping. No budget or policy state; those stay on the key object. - DbIdentityStore: the auth flow's resolver, owning both halves of resolving a caller. resolve_key does the one combined_view lookup (cache, then DB via the shared lower-level helpers, then write-back) and returns the key object, which still flows for budget / rate-limit / policy unchanged. principal_from_key projects the identity slice of that key object into a Principal, issuing no lookup. user_api_key_auth resolves every key through the store rather than calling get_key_object directly; auth_checks.get_key_object stays as the legacy entrypoint for its other callers until they migrate. - network: the X-Forwarded-For / trusted-proxy CIDR primitives live here in one place. trusted_proxy_utils now imports them rather than keeping a second copy. At the seam, user_api_key_auth projects one per-request Principal off the resolved key object and stamps the request network context onto it once (X-Forwarded-For is trusted only when trusted_proxy_ranges is configured). It is attached to request.state.principal for the downstream consumers later phases add. The projection is additive and defensive: a failure never rejects an already-authenticated request, and a missing principal must be treated as deny by any future reader. The Principal is always identifiable (credential_ref and a stable subject off the token), never anonymous. This is additive and changes no behavior today; it is the identity foundation the spend-attribution and authorization phases build on. --- litellm/proxy/auth/auth_method.py | 13 ++ litellm/proxy/auth/network.py | 101 +++++++++ litellm/proxy/auth/resolvers/__init__.py | 33 +++ litellm/proxy/auth/resolvers/exceptions.py | 50 +++++ litellm/proxy/auth/resolvers/models.py | 82 +++++++ litellm/proxy/auth/resolvers/store.py | 206 ++++++++++++++++++ litellm/proxy/auth/roles.py | 35 +++ litellm/proxy/auth/trusted_proxy_utils.py | 72 ++---- litellm/proxy/auth/user_api_key_auth.py | 117 +++++++--- scripts/ruff_strict_gate.py | 94 +++----- .../proxy/auth/test_user_api_key_auth.py | 25 ++- .../proxy_unit_tests/test_jwt_key_mapping.py | 31 ++- .../proxy/auth/test_auth_checks.py | 48 +++- .../auth/test_custom_auth_end_user_budget.py | 4 +- .../proxy/auth/test_handle_jwt.py | 27 ++- .../proxy/auth/test_model_checks.py | 4 +- tests/test_litellm/proxy/auth/test_network.py | 104 +++++++++ .../proxy/auth/test_onboarding.py | 1 - .../proxy/auth/test_resolvers_exceptions.py | 35 +++ .../proxy/auth/test_resolvers_models.py | 95 ++++++++ .../proxy/auth/test_resolvers_seam.py | 96 ++++++++ .../proxy/auth/test_resolvers_store.py | 70 ++++++ .../proxy/auth/test_user_api_key_auth.py | 10 +- tests/test_litellm/test_ruff_strict_gate.py | 10 - 24 files changed, 1149 insertions(+), 214 deletions(-) create mode 100644 litellm/proxy/auth/auth_method.py create mode 100644 litellm/proxy/auth/network.py create mode 100644 litellm/proxy/auth/resolvers/__init__.py create mode 100644 litellm/proxy/auth/resolvers/exceptions.py create mode 100644 litellm/proxy/auth/resolvers/models.py create mode 100644 litellm/proxy/auth/resolvers/store.py create mode 100644 litellm/proxy/auth/roles.py create mode 100644 tests/test_litellm/proxy/auth/test_network.py create mode 100644 tests/test_litellm/proxy/auth/test_resolvers_exceptions.py create mode 100644 tests/test_litellm/proxy/auth/test_resolvers_models.py create mode 100644 tests/test_litellm/proxy/auth/test_resolvers_seam.py create mode 100644 tests/test_litellm/proxy/auth/test_resolvers_store.py diff --git a/litellm/proxy/auth/auth_method.py b/litellm/proxy/auth/auth_method.py new file mode 100644 index 00000000000..a604eb563e6 --- /dev/null +++ b/litellm/proxy/auth/auth_method.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +from enum import Enum + + +class AuthMethod(str, Enum): + API_KEY = "api_key" + HTTP_BASIC = "http_basic" + BEARER_JWT = "bearer_jwt" + OAUTH2_INTROSPECTION = "oauth2_introspection" + OIDC = "oidc" + SAML = "saml" + MUTUAL_TLS = "mutual_tls" diff --git a/litellm/proxy/auth/network.py b/litellm/proxy/auth/network.py new file mode 100644 index 00000000000..4eb6f1dcec2 --- /dev/null +++ b/litellm/proxy/auth/network.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import ipaddress +from typing import Any, Union + +from fastapi import Request +from pydantic import BaseModel, Field + +from litellm._logging import verbose_proxy_logger + +TrustedProxyNetwork = Union[ipaddress.IPv4Network, ipaddress.IPv6Network] + + +class NetworkContext(BaseModel): + client_ip: str | None = None + host: str | None = None + via_trusted_proxy: bool = False + + +class TrustedProxyConfig(BaseModel): + use_forwarded_for: bool = False + trusted_proxy_cidrs: list[str] = Field(default_factory=list) + + +def normalize_cidr_ranges( + configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs" +) -> list[str]: + if not configured_ranges: + return [] + if isinstance(configured_ranges, str): + return [r.strip() for r in configured_ranges.split(",") if r.strip()] + if isinstance(configured_ranges, (list, tuple, set)): + return [str(r).strip() for r in configured_ranges if str(r).strip()] + verbose_proxy_logger.warning( + "Invalid %s value: expected a list of CIDR ranges, got %s", + setting_name, + type(configured_ranges).__name__, + ) + return [] + + +def parse_trusted_proxy_ranges( + configured_ranges: Any, *, setting_name: str = "trusted_proxy_cidrs" +) -> list[TrustedProxyNetwork]: + networks: list[TrustedProxyNetwork] = [] + for cidr in normalize_cidr_ranges(configured_ranges, setting_name=setting_name): + try: + networks.append(ipaddress.ip_network(cidr, strict=False)) + except ValueError: + verbose_proxy_logger.warning( + "Invalid CIDR in %s: %s, skipping", setting_name, cidr + ) + return networks + + +def ip_in_networks(client_ip: str | None, networks: list[TrustedProxyNetwork]) -> bool: + if not client_ip or not networks: + return False + try: + addr = ipaddress.ip_address(client_ip.strip()) + except ValueError: + return False + return any(addr in network for network in networks) + + +def _is_valid_ip(value: str) -> bool: + try: + ipaddress.ip_address(value) + return True + except ValueError: + return False + + +def resolve_client_ip( + request: Request, config: TrustedProxyConfig +) -> tuple[str | None, bool]: + """Resolve the real client IP, trusting X-Forwarded-For only when the direct + peer is itself a configured trusted proxy. Walks the header right-to-left and + returns the first hop that is not a trusted proxy, so a forged left-most entry + cannot spoof the client.""" + peer = request.client.host if request.client else None + networks = parse_trusted_proxy_ranges(config.trusted_proxy_cidrs) + if not config.use_forwarded_for or not ip_in_networks(peer, networks): + return peer, False + forwarded = request.headers.get("x-forwarded-for", "") + hops = [h.strip() for h in forwarded.split(",") if h.strip()] + for hop in reversed(hops): + if _is_valid_ip(hop) and not ip_in_networks(hop, networks): + return hop, True + return peer, True + + +def resolve_network_context( + request: Request, config: TrustedProxyConfig +) -> NetworkContext: + ip, via_proxy = resolve_client_ip(request, config) + return NetworkContext( + client_ip=ip, + host=request.headers.get("host"), + via_trusted_proxy=via_proxy, + ) diff --git a/litellm/proxy/auth/resolvers/__init__.py b/litellm/proxy/auth/resolvers/__init__.py new file mode 100644 index 00000000000..d6bfb335c09 --- /dev/null +++ b/litellm/proxy/auth/resolvers/__init__.py @@ -0,0 +1,33 @@ +from litellm.proxy.auth.resolvers.exceptions import ( + IdentityResolutionError, + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.resolvers.models import ( + CredentialRef, + EndUserIdentity, + OrganizationIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) + +__all__ = [ + "CredentialRef", + "EndUserIdentity", + "IdentityResolutionError", + "KeyNotFoundError", + "KeyNotInCacheError", + "NoDatabaseConnectionError", + "OrganizationIdentity", + "Principal", + "PrincipalMissingSourceKeyError", + "PrincipalType", + "ProjectIdentity", + "TeamIdentity", + "UserIdentity", +] diff --git a/litellm/proxy/auth/resolvers/exceptions.py b/litellm/proxy/auth/resolvers/exceptions.py new file mode 100644 index 00000000000..dd953e66659 --- /dev/null +++ b/litellm/proxy/auth/resolvers/exceptions.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from fastapi import status + +from litellm.proxy._types import ProxyErrorTypes, ProxyException + + +class IdentityResolutionError(Exception): + """Base for every failure raised while resolving a caller's identity.""" + + +class NoDatabaseConnectionError(IdentityResolutionError): + def __init__(self) -> None: + super().__init__( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + + +class KeyNotInCacheError(IdentityResolutionError): + def __init__(self, hashed_token: str) -> None: + super().__init__( + f"Key doesn't exist in cache + check_cache_only=True. key={hashed_token}." + ) + + +class KeyNotFoundError(IdentityResolutionError, ProxyException): + """The token matched nothing in the cache or the verification token table. + + Also a ``ProxyException`` so the auth flow keeps mapping a missing key to the + OpenAI 401 contract unchanged while callers migrate onto the typed hierarchy. + """ + + def __init__(self, hashed_token: str) -> None: + ProxyException.__init__( + self, + message="Authentication Error, Invalid proxy server token passed. key={}, not found in db. Create key via `/key/generate` call.".format( + hashed_token + ), + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ) + + +class PrincipalMissingSourceKeyError(IdentityResolutionError): + def __init__(self) -> None: + super().__init__( + "Principal carries no source key; it was not produced by " + "IdentityStore.resolve" + ) diff --git a/litellm/proxy/auth/resolvers/models.py b/litellm/proxy/auth/resolvers/models.py new file mode 100644 index 00000000000..97e66b8fe67 --- /dev/null +++ b/litellm/proxy/auth/resolvers/models.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from enum import Enum + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import NetworkContext +from litellm.proxy.auth.roles import Role, TeamRole + + +class PrincipalType(str, Enum): + HUMAN = "human" + SERVICE_ACCOUNT = "service_account" + + +class UserIdentity(BaseModel): + id: str + external_id: str | None = None + user_name: str | None = None + email: str | None = None + display_name: str | None = None + + +class OrganizationIdentity(BaseModel): + id: str + name: str | None = None + + +class TeamIdentity(BaseModel): + id: str + name: str | None = None + role: TeamRole = TeamRole.MEMBER + + +class ProjectIdentity(BaseModel): + id: str + name: str | None = None + + +class EndUserIdentity(BaseModel): + id: str + + +class CredentialRef(BaseModel): + key_id: str | None = None + token_id: str | None = None + + +class Principal(BaseModel): + """Normalized caller identity, resolved once per request at the auth seam. + + Frozen and constructed fresh per request, never cached or shared. The identity + fields carry no policy, budget, or rate-limit state. ``source_key`` is a + transitional carrier for the resolved key object so ``key_from_principal`` can + hand it to the request flow that still consumes ``UserAPIKeyAuth``; it is + excluded from serialization and repr and goes away once those consumers read + identity off the Principal directly. + """ + + model_config = ConfigDict(frozen=True) + + principal_type: PrincipalType + subject: str + issuer: str | None = None + audience: list[str] = Field(default_factory=list) + + user: UserIdentity | None = None + organization: OrganizationIdentity | None = None + teams: list[TeamIdentity] = Field(default_factory=list) + project: ProjectIdentity | None = None + end_user: EndUserIdentity | None = None + + roles: list[Role] = Field(default_factory=list) + scopes: list[str] = Field(default_factory=list) + + auth_method: AuthMethod + credential_ref: CredentialRef = Field(default_factory=CredentialRef) + network: NetworkContext = Field(default_factory=NetworkContext) + + source_key: UserAPIKeyAuth | None = Field(default=None, exclude=True, repr=False) diff --git a/litellm/proxy/auth/resolvers/store.py b/litellm/proxy/auth/resolvers/store.py new file mode 100644 index 00000000000..c43d4c84ca4 --- /dev/null +++ b/litellm/proxy/auth/resolvers/store.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence + +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import ( + _cache_key_object, + _copy_user_api_key_auth_for_cache, + _fetch_key_object_from_db_with_reconnect, + get_object_permission, +) +from litellm.proxy.auth.resolvers.exceptions import ( + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import NetworkContext +from litellm.proxy.auth.resolvers.models import ( + CredentialRef, + EndUserIdentity, + OrganizationIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) +from litellm.proxy.auth.roles import TeamRole, map_role, team_role + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.integrations.opentelemetry import Span + from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class IdentityStore: + """The auth flow's resolver: one combined_view lookup, projected into a Principal. + + ``resolve`` does the lookup (cache, then DB via the shared lower-level helpers, + then write-back) and returns the per-caller Principal. The Principal carries the + source key object so ``key_from_principal`` can hand it back to the parts of the + request flow that still consume ``UserAPIKeyAuth`` (budget, rate limits, policy); + that carrier is a stopgap until those consumers read identity off the Principal. + The Prisma client, key cache, the request's tracing span / logging sink, and + whether this store may only read the cache are injected so the composition root + can build the store once the proxy DB is connected; the span and logging sink + are infra the DB call is instrumented with and ``check_cache_only`` is a store + mode, none of them inputs to resolving identity. ``auth_checks.get_key_object`` + stays as the legacy entrypoint for its other callers until they migrate onto + this store. + """ + + def __init__( + self, + prisma_client: PrismaClient | None, + cache: DualCache, + *, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + check_cache_only: bool = False, + ) -> None: + self._prisma = prisma_client + self._cache = cache + self._parent_otel_span = parent_otel_span + self._proxy_logging_obj = proxy_logging_obj + self._check_cache_only = check_cache_only + + async def resolve( + self, + hashed_token: str, + *, + auth_method: AuthMethod = AuthMethod.API_KEY, + network: NetworkContext | None = None, + ) -> Principal: + key = await self._resolve_key(hashed_token) + return self._principal_from_key( + key, + auth_method=auth_method, + network=network, + subject_fallback=key.token, + credential_ref=CredentialRef(token_id=key.token), + ) + + @staticmethod + def key_from_principal(principal: Principal) -> UserAPIKeyAuth: + """Hand back the resolved key object carried on the Principal. + + Stopgap for the request flow that still consumes ``UserAPIKeyAuth`` for + budget, rate-limit, and policy state. Only Principals produced by + ``resolve`` carry a source key. + """ + if principal.source_key is None: + raise PrincipalMissingSourceKeyError() + return principal.source_key + + async def _resolve_key(self, hashed_token: str) -> UserAPIKeyAuth: + if self._prisma is None: + raise NoDatabaseConnectionError() + + cached = await self._cache.async_get_cache( + key=hashed_token, model_type=UserAPIKeyAuth + ) + if cached is not None: + return _copy_user_api_key_auth_for_cache(user_api_key_obj=cached) + + if self._check_cache_only: + raise KeyNotInCacheError(hashed_token) + + from_db: BaseModel | None = await _fetch_key_object_from_db_with_reconnect( + hashed_token=hashed_token, + prisma_client=self._prisma, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + if from_db is None: + raise KeyNotFoundError(hashed_token) + + key = UserAPIKeyAuth(**from_db.model_dump(exclude_none=True)) + + if key.object_permission_id and not key.object_permission: + try: + key.object_permission = await get_object_permission( + object_permission_id=key.object_permission_id, + prisma_client=self._prisma, + user_api_key_cache=self._cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Failed to load object_permission for key with object_permission_id={key.object_permission_id}: {e}" + ) + + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=key, + user_api_key_cache=self._cache, + proxy_logging_obj=self._proxy_logging_obj, + ) + return key + + @staticmethod + def _principal_from_key( + key: UserAPIKeyAuth, + *, + auth_method: AuthMethod, + issuer: str | None = None, + subject_fallback: str | None = None, + scopes: Sequence[str] = (), + credential_ref: CredentialRef | None = None, + network: NetworkContext | None = None, + ) -> Principal: + """Project the identity slice off an already-resolved key object and carry + the key on the Principal so ``key_from_principal`` can recover it. + + Pure: issues no lookup. Both ``resolve`` and the auth seam call this so + identity is projected once off whichever key object they already hold. + """ + teams: list[TeamIdentity] = [] + if key.team_id is not None: + role = ( + team_role(key.team_member.role) if key.team_member else TeamRole.MEMBER + ) + teams.append(TeamIdentity(id=key.team_id, name=key.team_alias, role=role)) + organization = ( + OrganizationIdentity(id=key.org_id, name=key.organization_alias) + if key.org_id is not None + else None + ) + user = ( + UserIdentity(id=key.user_id, email=key.user_email) + if key.user_id is not None + else None + ) + project = ( + ProjectIdentity(id=key.project_id, name=key.project_alias) + if key.project_id is not None + else None + ) + end_user = ( + EndUserIdentity(id=key.end_user_id) if key.end_user_id is not None else None + ) + mapped = map_role(key.user_role) + return Principal( + principal_type=( + PrincipalType.HUMAN if key.user_id else PrincipalType.SERVICE_ACCOUNT + ), + subject=key.user_id or key.key_alias or subject_fallback or "", + issuer=issuer, + user=user, + organization=organization, + teams=teams, + project=project, + end_user=end_user, + roles=[mapped] if mapped else [], + scopes=list(scopes), + auth_method=auth_method, + credential_ref=credential_ref or CredentialRef(), + network=network or NetworkContext(), + source_key=key, + ) diff --git a/litellm/proxy/auth/roles.py b/litellm/proxy/auth/roles.py new file mode 100644 index 00000000000..efe56a8b6b2 --- /dev/null +++ b/litellm/proxy/auth/roles.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from enum import Enum + + +class Role(str, Enum): + PLATFORM_ADMIN = "platform_admin" + PLATFORM_VIEWER = "platform_viewer" + ORG_ADMIN = "org_admin" + ORG_VIEWER = "org_viewer" + TEAM_ADMIN = "team_admin" + TEAM_MEMBER = "team_member" + + +class TeamRole(str, Enum): + ADMIN = "admin" + MEMBER = "member" + + +_ROLE_MAP: dict[str, Role] = { + "proxy_admin": Role.PLATFORM_ADMIN, + "proxy_admin_viewer": Role.PLATFORM_VIEWER, + "org_admin": Role.ORG_ADMIN, +} + + +def map_role(value: str | None) -> Role | None: + """Map a LiteLLM ``user_role`` string to a platform Role.""" + if value is None: + return None + return _ROLE_MAP.get(value) + + +def team_role(role: str | None) -> TeamRole: + return TeamRole.ADMIN if role == "admin" else TeamRole.MEMBER diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py index df7b3080f28..35bb79e7efe 100644 --- a/litellm/proxy/auth/trusted_proxy_utils.py +++ b/litellm/proxy/auth/trusted_proxy_utils.py @@ -1,12 +1,15 @@ -import ipaddress -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Optional from fastapi import Request from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.network import ( + ip_in_networks, + normalize_cidr_ranges, + parse_trusted_proxy_ranges, +) TRUSTED_PROXY_RANGES_KEY = "trusted_proxy_ranges" -TrustedProxyNetwork = Union[ipaddress.IPv4Network, ipaddress.IPv6Network] def _get_proxy_general_settings() -> Dict[str, Any]: @@ -18,43 +21,20 @@ def _get_proxy_general_settings() -> Dict[str, Any]: return {} -def _normalize_cidr_ranges(configured_ranges: Any, *, setting_name: str) -> List[str]: - if not configured_ranges: - return [] - if isinstance(configured_ranges, str): - return [ - raw_range.strip() - for raw_range in configured_ranges.split(",") - if raw_range.strip() - ] - if isinstance(configured_ranges, (list, tuple, set)): - return [ - str(raw_range).strip() - for raw_range in configured_ranges - if str(raw_range).strip() - ] - verbose_proxy_logger.warning( - "Invalid %s value: expected a list of CIDR ranges, got %s", - setting_name, - type(configured_ranges).__name__, +def get_trusted_proxy_cidrs( + general_settings: dict[str, Any] | None = None, +) -> list[str]: + """Operator-configured trusted reverse-proxy CIDRs, normalized to strings. + + Empty when none are configured, in which case X-Forwarded-For must not be + trusted and only the direct peer is authoritative. + """ + if general_settings is None: + general_settings = _get_proxy_general_settings() + return normalize_cidr_ranges( + general_settings.get(TRUSTED_PROXY_RANGES_KEY), + setting_name=TRUSTED_PROXY_RANGES_KEY, ) - return [] - - -def parse_trusted_proxy_ranges( - configured_ranges: Any, - *, - setting_name: str = TRUSTED_PROXY_RANGES_KEY, -) -> List[TrustedProxyNetwork]: - networks: List[TrustedProxyNetwork] = [] - for cidr in _normalize_cidr_ranges(configured_ranges, setting_name=setting_name): - try: - networks.append(ipaddress.ip_network(cidr, strict=False)) - except ValueError: - verbose_proxy_logger.warning( - "Invalid CIDR in %s: %s, skipping", setting_name, cidr - ) - return networks def _get_direct_client_ip(request: Request) -> Optional[str]: @@ -65,18 +45,6 @@ def _get_direct_client_ip(request: Request) -> Optional[str]: return None -def _is_ip_in_networks( - client_ip: Optional[str], networks: List[TrustedProxyNetwork] -) -> bool: - if not client_ip or not networks: - return False - try: - addr = ipaddress.ip_address(client_ip.strip()) - except ValueError: - return False - return any(addr in network for network in networks) - - def require_trusted_proxy_request( *, request: Request, @@ -105,7 +73,7 @@ def require_trusted_proxy_request( ) direct_client_ip = _get_direct_client_ip(request) - if not _is_ip_in_networks(direct_client_ip, trusted_networks): + if not ip_in_networks(direct_client_ip, trusted_networks): verbose_proxy_logger.warning( "%s rejected identity headers from untrusted direct client IP %r", feature_name, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 00d98a04a78..4a2df18b93b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -42,7 +42,6 @@ from litellm.proxy.auth.auth_checks import ( common_checks, get_end_user_object, get_jwt_key_mapping_object, - get_key_object, get_project_object, get_team_object, get_user_object, @@ -63,7 +62,12 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_context +from litellm.proxy.auth.resolvers import CredentialRef, Principal +from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, @@ -758,12 +762,13 @@ async def _auto_register_jwt_mapping( claim_value, ) - auto_registered_key = await get_key_object( - hashed_token=token_hash, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + auto_registered_key = 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=token_hash) ) if auto_registered_key is not None: auto_registered_key.org_id = org_id @@ -865,12 +870,13 @@ async def _resolve_jwt_to_virtual_key( ) return None elif cached_mapping is not None: - return await get_key_object( - hashed_token=cached_mapping, - 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 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=cached_mapping) ) # Resolve the mapping from DB, or treat prisma_client=None as a definitive @@ -889,12 +895,13 @@ async def _resolve_jwt_to_virtual_key( value=token_hash, ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl, ) - return await get_key_object( - hashed_token=token_hash, - 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 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=token_hash) ) # No mapping found (DB miss or no DB) — apply no-match policy. @@ -1493,13 +1500,14 @@ async def _user_api_key_auth_builder( ## Check CACHE try: with tracer.trace("litellm.proxy.auth.get_key_object_check_cache"): - valid_token = await get_key_object( - hashed_token=hash_token(api_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, - check_cache_only=True, + 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, + check_cache_only=True, + ).resolve(hashed_token=hash_token(api_key)) ) except Exception: verbose_logger.debug("api key not found in cache.") @@ -1679,12 +1687,13 @@ async def _user_api_key_auth_builder( try: with tracer.trace("litellm.proxy.auth.get_key_object_from_db"): - valid_token = await get_key_object( - hashed_token=api_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, + 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) ) except ProxyException as e: if e.code == 401 or e.code == "401": @@ -2501,6 +2510,34 @@ def _should_skip_budget_checks( return False +def _resolve_request_principal( + request: Request, valid_token: UserAPIKeyAuth +) -> Principal: + """Project the resolved identity into one per-request Principal, off the key + object the builder already fetched, and stamp the request network context + onto it once. X-Forwarded-For is only trusted when the operator configured + ``trusted_proxy_ranges``; otherwise the direct peer is authoritative. + + credential_ref and a stable subject fallback are always set off the token so + the Principal can never be anonymous, even for a keyless service-account key + with no user or alias.""" + cidrs = get_trusted_proxy_cidrs() + network = resolve_network_context( + request, + TrustedProxyConfig(use_forwarded_for=bool(cidrs), trusted_proxy_cidrs=cidrs), + ) + auth_method = ( + AuthMethod.BEARER_JWT if valid_token.jwt_claims else AuthMethod.API_KEY + ) + return IdentityStore._principal_from_key( + valid_token, + auth_method=auth_method, + network=network, + subject_fallback=valid_token.token, + credential_ref=CredentialRef(token_id=valid_token.token), + ) + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -2615,6 +2652,22 @@ async def user_api_key_auth( model=request_data.get("model") if isinstance(request_data, dict) else None, ) user_api_key_auth_obj.request_route = normalize_request_route(route) + + # Resolve caller identity once, here at the seam, into a single per-request + # Principal projected off the key object the builder already fetched (no + # second lookup). Downstream consumers read identity off this instead of + # re-resolving it. Additive and defensive: a projection failure must never + # reject an already-authenticated request, so it is left unset on failure; + # any future consumer must treat a missing principal as deny, not allow. + try: + request.state.principal = _resolve_request_principal( + request, user_api_key_auth_obj + ) + except Exception as e: + verbose_proxy_logger.warning( + "Principal projection at auth seam failed (non-fatal): %s", e + ) + return user_api_key_auth_obj diff --git a/scripts/ruff_strict_gate.py b/scripts/ruff_strict_gate.py index 9c406b8482b..5951a1215ed 100644 --- a/scripts/ruff_strict_gate.py +++ b/scripts/ruff_strict_gate.py @@ -5,13 +5,9 @@ Each rule has a hard ceiling (baseline + slack) in ruff-strict-budget.json. The gate counts each rule across the whole tree and fails when a rule is both over its ceiling and higher than the base it merges into, so a change is blamed for the violations it adds, never for drift that already exists in the base. - -The base is the merge-base of the current branch with --base; this matches CI, -which checks out the PR head sha and runs the gate against the PR's base sha. """ import argparse -import contextlib import json import re import shutil @@ -19,7 +15,6 @@ import subprocess import sys import tempfile from collections import Counter -from collections.abc import Iterator from pathlib import Path from typing import NamedTuple @@ -45,12 +40,6 @@ class Breach(NamedTuple): added: int -class GateInputs(NamedTuple): - head: list[Violation] - base: dict[str, int] - changed: dict[str, set[int]] - - def _run(cmd: list, cwd: Path = REPO_ROOT) -> str: proc = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True) if proc.returncode not in (0, 1): @@ -67,14 +56,14 @@ def _ruff_json(cwd: Path, config: Path) -> list: return json.loads(raw or "[]") -def collect_violations(root: Path, config: Path) -> list: +def head_violations() -> list: out = [] - for item in _ruff_json(root, config): + for item in _ruff_json(REPO_ROOT, STRICT_CONFIG): name = Path(item["filename"]) rel = ( - (name if name.is_absolute() else root / name) + (name if name.is_absolute() else REPO_ROOT / name) .resolve() - .relative_to(root) + .relative_to(REPO_ROOT) .as_posix() ) out.append(Violation(rel, item["location"]["row"], item["code"])) @@ -85,29 +74,27 @@ def count_by_rule(violations: list) -> dict: return dict(Counter(v.code for v in violations)) -@contextlib.contextmanager -def _temp_worktree(ref: str) -> Iterator[Path]: - parent = Path(tempfile.mkdtemp(prefix="ruff_wt_")) +def base_counts(ref: str) -> dict: + parent = Path(tempfile.mkdtemp(prefix="ruff_base_")) worktree = parent / "wt" try: _run(["git", "worktree", "add", "--detach", str(worktree), ref]) - yield worktree + shutil.copy(STRICT_CONFIG, worktree / "ruff-strict.toml") + items = _ruff_json(worktree, worktree / "ruff-strict.toml") + return dict(Counter(item["code"] for item in items)) finally: - subprocess.run( - ["git", "worktree", "remove", "--force", str(worktree)], - cwd=REPO_ROOT, - capture_output=True, - text=True, - ) + _run(["git", "worktree", "remove", "--force", str(worktree)]) shutil.rmtree(parent, ignore_errors=True) -def base_counts(ref: str) -> dict: - with _temp_worktree(ref) as worktree: - shutil.copy(STRICT_CONFIG, worktree / "ruff-strict.toml") - return count_by_rule( - collect_violations(worktree, worktree / "ruff-strict.toml") - ) +def evaluate(head: dict, base: dict, budget: dict) -> list: + breaches = [] + for rule, spec in budget.items(): + cap = spec["baseline"] + spec["slack"] + total = head.get(rule, 0) + if total > cap and total > base.get(rule, 0): + breaches.append(Breach(rule, total, cap, total - base.get(rule, 0))) + return sorted(breaches) def parse_changed_lines(diff_text: str) -> dict: @@ -123,31 +110,24 @@ def parse_changed_lines(diff_text: str) -> dict: return changed -def evaluate(head: dict, base: dict, budget: dict) -> list: - breaches = [] - for rule, spec in budget.items(): - cap = spec["baseline"] + spec["slack"] - total = head.get(rule, 0) - if total > cap and total > base.get(rule, 0): - breaches.append(Breach(rule, total, cap, total - base.get(rule, 0))) - return sorted(breaches) - - def introduced(violations: list, changed: dict) -> list: return [v for v in violations if v.line in changed.get(v.file, set())] -def gather(base: str) -> GateInputs: +def cmd_check(base: str) -> None: + budget = json.loads(BUDGET_PATH.read_text()) + head = head_violations() base_point = _run(["git", "merge-base", base, "HEAD"]).strip() or base - diff = _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET]) - return GateInputs( - collect_violations(REPO_ROOT, STRICT_CONFIG), - base_counts(base_point), - parse_changed_lines(diff), + breaches = evaluate(count_by_rule(head), base_counts(base_point), budget) + if not breaches: + print(f"OK: every strict rule is within its codebase ceiling (base {base})") + return + new = introduced( + head, + parse_changed_lines( + _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET]) + ), ) - - -def report(breaches: list, new: list, base: str) -> None: print(f"FAIL: strict-rule totals exceed their ceiling (base {base}):") for breach in breaches: print( @@ -158,24 +138,12 @@ def report(breaches: list, new: list, base: str) -> None: print( "Reduce the new violations or remove an equal number elsewhere; the ceiling is baseline + slack in ruff-strict-budget.json." ) - summary = "; ".join(f"{b.rule} {b.total}/{b.cap} (+{b.added})" for b in breaches) - print(f"BREACHED RULES: {summary}") - - -def cmd_check(base: str) -> None: - budget = json.loads(BUDGET_PATH.read_text()) - inputs = gather(base) - breaches = evaluate(count_by_rule(inputs.head), inputs.base, budget) - if not breaches: - print(f"OK: every strict rule is within its codebase ceiling (base {base})") - return - report(breaches, introduced(inputs.head, inputs.changed), base) raise SystemExit(1) def cmd_update() -> None: budget = json.loads(BUDGET_PATH.read_text()) - head = count_by_rule(collect_violations(REPO_ROOT, STRICT_CONFIG)) + head = count_by_rule(head_violations()) for rule in budget: budget[rule]["baseline"] = head.get(rule, 0) BUDGET_PATH.write_text(json.dumps(budget, indent=2, sort_keys=True) + "\n") diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py b/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py index a45df5df008..fd674afd85a 100644 --- a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py +++ b/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py @@ -4,6 +4,8 @@ import pytest from fastapi import Request from litellm_enterprise.proxy.auth.user_api_key_auth import enterprise_custom_auth +from litellm.proxy._types import UserAPIKeyAuth + @pytest.mark.asyncio async def test_enterprise_custom_auth_none_user_auth(): @@ -49,16 +51,19 @@ async def test_enterprise_custom_auth_returns_string(): mock_user_auth = AsyncMock(return_value="sk-test-key") request = MagicMock(spec=Request) - with patch( - "litellm.proxy.auth.user_api_key_auth.enterprise_custom_auth", mock_user_auth - ), patch("litellm.proxy.proxy_server.master_key", "sk-1234"), patch( - "litellm.proxy.proxy_server.prisma_client", MagicMock() + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.enterprise_custom_auth", + mock_user_auth, + ), + patch("litellm.proxy.proxy_server.master_key", "sk-1234"), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), ): # Verify the key is correctly handled in _user_api_key_auth_builder with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object" + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key" ) as mock_get_key_object: - mock_get_key_object.return_value = MagicMock( + mock_get_key_object.return_value = UserAPIKeyAuth( token="sk-test-key", user_role="internal_user", team_id=None, @@ -82,9 +87,7 @@ async def test_enterprise_custom_auth_returns_string(): except Exception as e: print("error:", e) - # Verify get_key_object was called with the correct key + # Verify the key lookup was called with the correct hashed key mock_get_key_object.assert_called_once() - # The key should be hashed before being passed to get_key_object - assert mock_get_key_object.call_args[1]["hashed_token"] == hash_token( - "sk-test-key" - ) + # The key should be hashed before being passed to the resolver + assert mock_get_key_object.call_args[0][0] == hash_token("sk-test-key") diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 61c24183964..65d075b7f99 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -60,7 +60,8 @@ async def test_jwt_to_virtual_key_mapping_resolution(): # Use patch to mock get_key_object in the module where it's used with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ) as mock_get_key: mock_get_key.return_value = mock_key_obj @@ -105,7 +106,8 @@ async def test_jwt_to_virtual_key_mapping_no_mapping(): # Mock get_key_object just in case with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): user_api_key_cache = DualCache() @@ -481,7 +483,8 @@ async def test_reject_behavior_raises_403_on_no_mapping(): user_api_key_cache = DualCache() with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): with pytest.raises(HTTPException) as exc_info: await _resolve_jwt_to_virtual_key( @@ -519,7 +522,8 @@ async def test_reject_behavior_caches_sentinel_after_db_miss(): user_api_key_cache = DualCache() with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): # First call — DB miss, should raise 403 and write sentinel with pytest.raises(HTTPException) as exc_info: @@ -578,7 +582,8 @@ async def test_reject_behavior_raises_403_on_cached_no_mapping(): await user_api_key_cache.async_set_cache(cache_key, "__NO_MAPPING__") with patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, ): with pytest.raises(HTTPException) as exc_info: await _resolve_jwt_to_virtual_key( @@ -671,7 +676,7 @@ async def test_auto_register_creates_key_and_mapping_when_helper_invoked(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -803,7 +808,7 @@ async def test_auto_register_race_condition_unique_conflict(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -833,13 +838,7 @@ async def test_auto_register_race_condition_unique_conflict(): # Cache should hold the winner's token, not the loser's cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:user-42") assert cached == "winner_token_hash" - mock_get_key.assert_called_once_with( - hashed_token="winner_token_hash", - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) + mock_get_key.assert_called_once_with("winner_token_hash") # ────────────────────────────────────────────── @@ -1093,7 +1092,7 @@ async def test_auto_register_race_conflict_tolerates_delete_failure(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( @@ -1249,7 +1248,7 @@ async def test_auto_register_helper_stamps_validated_identity_context(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key, patch( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5ec5d12784f..52634cc25fe 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2365,7 +2365,9 @@ async def test_virtual_key_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:key:test-hashed-token": return 1.5 return fallback_spend @@ -2397,7 +2399,9 @@ async def test_virtual_key_budget_check_fallback_no_counter(): proxy_logging_obj.budget_alerts = AsyncMock() # get_current_spend returns fallback_spend when no counter exists - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): return fallback_spend with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): @@ -2424,7 +2428,9 @@ async def test_team_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) proxy_logging_obj.budget_alerts = AsyncMock() - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team:test-team": return 1.5 return fallback_spend @@ -2449,7 +2455,9 @@ async def test_end_user_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:end_user:customer-1": return 1.5 return fallback_spend @@ -2475,7 +2483,9 @@ async def test_tag_budget_check_reads_from_spend_counter(): litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:tag:paid-tag": return 1.5 return fallback_spend @@ -2523,7 +2533,9 @@ async def test_team_member_budget_check_reads_from_spend_counter(): proxy_logging_obj = ProxyLogging(user_api_key_cache=None) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 1.5 return fallback_spend @@ -2756,7 +2768,9 @@ async def test_team_member_budget_check_falls_back_to_team_default_budget_id(): return_value=fake_budget_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 70.0 return fallback_spend @@ -2853,7 +2867,9 @@ async def test_team_member_budget_check_per_member_override_wins_over_team_defau mocked_spend = 70.0 - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return mocked_spend return fallback_spend @@ -2943,7 +2959,9 @@ async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 500.0 return fallback_spend @@ -3010,7 +3028,9 @@ async def test_team_member_budget_check_null_clone_with_null_default_skips_enfor return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 1000.0 return fallback_spend @@ -3077,7 +3097,9 @@ async def test_team_member_budget_check_zero_team_default_treated_as_no_cap(): return_value=fake_default_row ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend @@ -3135,7 +3157,9 @@ async def test_team_member_budget_check_zero_per_member_row_still_blocks(): prisma_client = MagicMock() prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:team_member:test-user:test-team": return 0.0 return fallback_spend diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index e49f025df2e..3203878a1e0 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -106,7 +106,9 @@ async def test_custom_auth_enforces_end_user_budget_when_common_checks_skipped() litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), ) - async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + async def mock_get_current_spend( + counter_key, fallback_spend, max_budget=None, **kwargs + ): if counter_key == "spend:end_user:customer-1": return 5.0 return fallback_spend diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 63510086f95..b547ec877e2 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1206,7 +1206,13 @@ async def test_auth_builder_returns_team_membership_object(): JWTAuthManager, "get_objects", new_callable=AsyncMock, - return_value=(user_object, None, None, mock_team_membership, user_object.user_id), + return_value=( + user_object, + None, + None, + mock_team_membership, + user_object.user_id, + ), ) as mock_get_objects, patch.object( JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock @@ -3509,9 +3515,7 @@ def test_canonical_user_id_no_change_when_ids_match(): user_object = LiteLLM_UserTable(user_id=same, user_email=same) assert ( - JWTAuthManager._canonical_user_id_from_db( - user_id=same, user_object=user_object - ) + JWTAuthManager._canonical_user_id_from_db(user_id=same, user_object=user_object) == same ) @@ -3802,12 +3806,15 @@ async def test_get_objects_team_membership_uses_rebound_user_id(): user_id_jwt_field="email", user_id_upsert=True ) - with patch( - "litellm.proxy.auth.handle_jwt.get_user_object", - side_effect=fake_get_user_object, - ), patch( - "litellm.proxy.auth.handle_jwt.get_team_membership", - side_effect=fake_get_team_membership, + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_user_object", + side_effect=fake_get_user_object, + ), + patch( + "litellm.proxy.auth.handle_jwt.get_team_membership", + side_effect=fake_get_team_membership, + ), ): ( user_object, diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 8d686900ea6..02b1f698132 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -436,7 +436,9 @@ def test_wildcard_custom_prefix_keeps_org_segment_for_non_provider_first_segment result = get_known_models_from_wildcard( wildcard_model="my_hf/*", - litellm_params=LiteLLM_Params(model="huggingface/*", custom_llm_provider="huggingface"), + litellm_params=LiteLLM_Params( + model="huggingface/*", custom_llm_provider="huggingface" + ), ) assert result == ["my_hf/meta-llama/Llama-3-8B"] diff --git a/tests/test_litellm/proxy/auth/test_network.py b/tests/test_litellm/proxy/auth/test_network.py new file mode 100644 index 00000000000..b67723305e4 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_network.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from fastapi import Request + +from litellm.proxy.auth.network import ( + TrustedProxyConfig, + resolve_client_ip, + resolve_network_context, +) + +TRUSTED = TrustedProxyConfig(use_forwarded_for=True, trusted_proxy_cidrs=["10.0.0.0/8"]) + + +def make_request( + *, + headers: Optional[Dict[str, str]] = None, + client: Optional[Tuple[str, int]] = ("203.0.113.7", 5555), +) -> Request: + raw_headers: List[Tuple[bytes, bytes]] = [ + (key.lower().encode(), value.encode()) for key, value in (headers or {}).items() + ] + scope: Dict[str, Any] = { + "type": "http", + "http_version": "1.1", + "method": "GET", + "path": "/", + "raw_path": b"/", + "query_string": b"", + "headers": raw_headers, + "client": client, + "server": ("testserver", 80), + "scheme": "http", + } + return Request(scope) + + +def test_xff_ignored_when_forwarding_disabled(): + config = TrustedProxyConfig( + use_forwarded_for=False, trusted_proxy_cidrs=["10.0.0.0/8"] + ) + request = make_request( + headers={"x-forwarded-for": "203.0.113.9"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, config) + assert ip == "10.0.0.1" + assert via_proxy is False + + +def test_xff_honored_from_trusted_peer(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9, 10.0.0.5"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "203.0.113.9" + assert via_proxy is True + + +def test_spoofed_xff_from_untrusted_peer_is_ignored(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9"}, client=("8.8.8.8", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "8.8.8.8" + assert via_proxy is False + + +def test_right_to_left_parse_skips_chained_trusted_proxies(): + request = make_request( + headers={"x-forwarded-for": "198.51.100.4, 10.1.1.1, 10.0.0.9"}, + client=("10.0.0.1", 1), + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "198.51.100.4" + assert via_proxy is True + + +def test_all_trusted_hops_fall_back_to_peer(): + request = make_request( + headers={"x-forwarded-for": "10.1.1.1, 10.0.0.9"}, client=("10.0.0.1", 1) + ) + ip, via_proxy = resolve_client_ip(request, TRUSTED) + assert ip == "10.0.0.1" + assert via_proxy is True + + +def test_invalid_xff_token_is_skipped(): + request = make_request( + headers={"x-forwarded-for": "not-an-ip, 203.0.113.50"}, client=("10.0.0.1", 1) + ) + ip, _ = resolve_client_ip(request, TRUSTED) + assert ip == "203.0.113.50" + + +def test_network_context_captures_host_and_proxy_flag(): + request = make_request( + headers={"x-forwarded-for": "203.0.113.9", "host": "proxy.litellm.ai"}, + client=("10.0.0.1", 1), + ) + ctx = resolve_network_context(request, TRUSTED) + assert ctx.client_ip == "203.0.113.9" + assert ctx.host == "proxy.litellm.ai" + assert ctx.via_trusted_proxy is True diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index c81f4cb7d66..d55a5472af1 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -18,7 +18,6 @@ from fastapi import HTTPException import litellm from litellm.proxy._types import InvitationClaim - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py b/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py new file mode 100644 index 00000000000..f88b7e2be14 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_exceptions.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.auth.resolvers.exceptions import ( + IdentityResolutionError, + KeyNotFoundError, + KeyNotInCacheError, + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) + + +def test_all_resolution_errors_share_one_base(): + errors = [ + NoDatabaseConnectionError(), + KeyNotInCacheError("hashed"), + KeyNotFoundError("hashed"), + PrincipalMissingSourceKeyError(), + ] + assert all(isinstance(e, IdentityResolutionError) for e in errors) + + +def test_key_not_found_preserves_the_public_401_contract(): + # The auth seam catches ProxyException and rewrites the 401 message, so a + # missing key must keep mapping to that exact contract. + error = KeyNotFoundError("hashed-token") + + assert isinstance(error, ProxyException) + assert error.code == "401" + assert error.type == ProxyErrorTypes.token_not_found_in_db.value + assert error.param == "key" + + +def test_key_not_in_cache_names_the_token(): + assert "hashed-token" in str(KeyNotInCacheError("hashed-token")) diff --git a/tests/test_litellm/proxy/auth/test_resolvers_models.py b/tests/test_litellm/proxy/auth/test_resolvers_models.py new file mode 100644 index 00000000000..0fbe10b6f28 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_models.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import ( + EndUserIdentity, + Principal, + PrincipalType, + ProjectIdentity, + TeamIdentity, + UserIdentity, +) +from litellm.proxy.auth.roles import Role, TeamRole + + +def _principal() -> Principal: + return Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + ) + + +def test_principal_is_frozen(): + principal = _principal() + with pytest.raises(ValidationError): + principal.subject = "mutated" + + +def test_principal_defaults_are_independent_instances(): + a = _principal() + b = _principal() + assert a.teams == [] and a.scopes == [] and a.audience == [] + assert a.teams is not b.teams + + +def test_principal_requires_identity_core_fields(): + with pytest.raises(ValidationError): + Principal(subject="u1") # missing principal_type + auth_method + + +def test_principal_roles_are_validated_against_role_enum(): + principal = Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + roles=["org_admin"], + ) + assert principal.roles == [Role.ORG_ADMIN] + assert isinstance(principal.roles[0], Role) + + with pytest.raises(ValidationError): + Principal( + principal_type=PrincipalType.HUMAN, + subject="u1", + auth_method=AuthMethod.OIDC, + roles=["not_a_real_role"], + ) + + +def test_principal_default_network_and_collections(): + principal = Principal( + principal_type=PrincipalType.SERVICE_ACCOUNT, + subject="svc", + auth_method=AuthMethod.MUTUAL_TLS, + ) + assert principal.teams == [] + assert principal.scopes == [] + assert principal.project is None + assert principal.end_user is None + assert principal.network.client_ip is None + assert principal.network.via_trusted_proxy is False + + +def test_team_identity_defaults_to_member_role(): + team = TeamIdentity(id="g1") + assert team.role == TeamRole.MEMBER + + +def test_user_identity_optional_fields_default_none(): + user = UserIdentity(id="u1") + assert user.email is None + assert user.external_id is None + + +def test_project_identity_name_is_optional(): + assert ProjectIdentity(id="p1").name is None + assert ProjectIdentity(id="p1", name="Acme").name == "Acme" + + +def test_end_user_identity_requires_id(): + with pytest.raises(ValidationError): + EndUserIdentity() diff --git a/tests/test_litellm/proxy/auth/test_resolvers_seam.py b/tests/test_litellm/proxy/auth/test_resolvers_seam.py new file mode 100644 index 00000000000..cb42d4d3d2d --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_seam.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Tuple + +from fastapi import Request + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import PrincipalType +from litellm.proxy.auth.roles import Role +from litellm.proxy.auth.user_api_key_auth import _resolve_request_principal + + +def _request( + *, + headers: Optional[Dict[str, str]] = None, + client: Optional[Tuple[str, int]] = ("203.0.113.7", 5555), +) -> Request: + raw: List[Tuple[bytes, bytes]] = [ + (k.lower().encode(), v.encode()) for k, v in (headers or {}).items() + ] + scope: Dict[str, Any] = { + "type": "http", + "http_version": "1.1", + "method": "POST", + "path": "/v1/chat/completions", + "raw_path": b"/v1/chat/completions", + "query_string": b"", + "headers": raw, + "client": client, + "server": ("testserver", 80), + "scheme": "http", + } + return Request(scope) + + +def test_seam_projects_full_identity_from_key_object(): + token = UserAPIKeyAuth( + token="hashed-token", + user_id="u-1", + user_role="org_admin", + team_id="t-1", + team_alias="Eng", + org_id="o-1", + organization_alias="Acme", + end_user_id="cust-9", + ) + + principal = _resolve_request_principal(_request(), token) + + assert principal.principal_type == PrincipalType.HUMAN + assert principal.auth_method == AuthMethod.API_KEY + assert principal.user is not None and principal.user.id == "u-1" + assert principal.roles == [Role.ORG_ADMIN] + assert [t.id for t in principal.teams] == ["t-1"] + assert principal.teams[0].name == "Eng" + assert principal.organization is not None and principal.organization.id == "o-1" + assert principal.organization.name == "Acme" + assert principal.end_user is not None and principal.end_user.id == "cust-9" + # the key is always identifiable via credential_ref, even when other ids exist + assert principal.credential_ref.token_id == "hashed-token" + + +def test_seam_principal_is_never_anonymous_for_keyless_service_account(): + # no user_id and no key_alias -> subject must still identify the key + token = UserAPIKeyAuth(token="hashed-token") + + principal = _resolve_request_principal(_request(), token) + + assert principal.principal_type == PrincipalType.SERVICE_ACCOUNT + assert principal.user is None + assert principal.subject == "hashed-token" + assert principal.credential_ref.token_id == "hashed-token" + + +def test_seam_stamps_direct_peer_when_no_trusted_proxy_configured(): + token = UserAPIKeyAuth(token="hashed-token", user_id="u-1") + + # No trusted_proxy_ranges configured -> XFF is not trusted, direct peer wins. + principal = _resolve_request_principal( + _request(headers={"x-forwarded-for": "10.9.9.9"}, client=("203.0.113.7", 5555)), + token, + ) + + assert principal.network.client_ip == "203.0.113.7" + assert principal.network.via_trusted_proxy is False + + +def test_seam_detects_jwt_auth_method(): + token = UserAPIKeyAuth( + token="hashed-token", user_id="u-2", jwt_claims={"sub": "u-2"} + ) + + principal = _resolve_request_principal(_request(), token) + + assert principal.auth_method == AuthMethod.BEARER_JWT diff --git a/tests/test_litellm/proxy/auth/test_resolvers_store.py b/tests/test_litellm/proxy/auth/test_resolvers_store.py new file mode 100644 index 00000000000..5e644a084a5 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_store.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from typing import Dict, Optional + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth, hash_token +from litellm.proxy.auth.resolvers.exceptions import ( + NoDatabaseConnectionError, + PrincipalMissingSourceKeyError, +) +from litellm.proxy.auth.auth_method import AuthMethod +from litellm.proxy.auth.resolvers.models import Principal, PrincipalType +from litellm.proxy.auth.resolvers.store import IdentityStore + + +class _FakeCache: + """Stands in for the DualCache get_key_object reads. It returns a cache hit + before the DB is touched, so seeding it exercises resolve without a database + (a non-None prisma client is still required; it is never reached on a hit).""" + + def __init__(self, entries: Optional[Dict[str, object]] = None) -> None: + self._entries = entries or {} + + async def async_get_cache(self, key, *args, **kwargs): + return self._entries.get(key) + + async def async_set_cache(self, *args, **kwargs): + return None + + +async def test_resolve_returns_a_principal_projected_from_the_looked_up_key(): + raw = "sk-live-abc" + key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1", team_id="t-1") + store = IdentityStore(object(), _FakeCache({hash_token(raw): key})) + + principal = await store.resolve(hashed_token=hash_token(raw)) + + assert isinstance(principal, Principal) + assert principal.principal_type == PrincipalType.HUMAN + assert principal.user is not None and principal.user.id == "u-1" + assert [t.id for t in principal.teams] == ["t-1"] + + +async def test_resolve_carries_the_key_for_key_from_principal(): + raw = "sk-live-abc" + key = UserAPIKeyAuth(token=hash_token(raw), user_id="u-1", team_id="t-1") + store = IdentityStore(object(), _FakeCache({hash_token(raw): key})) + + principal = await store.resolve(hashed_token=hash_token(raw)) + recovered = IdentityStore.key_from_principal(principal) + + assert recovered.user_id == "u-1" + assert recovered.team_id == "t-1" + + +def test_key_from_principal_raises_when_no_source_key_is_carried(): + bare = Principal( + principal_type=PrincipalType.SERVICE_ACCOUNT, + subject="svc", + auth_method=AuthMethod.API_KEY, + ) + with pytest.raises(PrincipalMissingSourceKeyError): + IdentityStore.key_from_principal(bare) + + +async def test_resolve_raises_without_a_db_connection(): + store = IdentityStore(None, _FakeCache()) + with pytest.raises(NoDatabaseConnectionError): + await store.resolve(hashed_token="missing") diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 80f12d4459f..27c1b04fbd9 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1106,7 +1106,7 @@ async def test_proxy_admin_expired_key_from_cache(): # Mock get_key_object to return expired token from cache with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, ) as mock_get_key_object, patch( @@ -1261,7 +1261,7 @@ async def test_scim_deactivated_user_key_is_rejected(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ), @@ -2484,7 +2484,7 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): stack.enter_context(p) stack.enter_context( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ) @@ -2598,7 +2598,7 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", new_callable=AsyncMock, return_value=valid_token, ), @@ -3652,7 +3652,7 @@ async def _run_builder_with_key_lookup(get_key_object_mock): request._url = URL(url="/chat/completions") with ( patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", get_key_object_mock, ), patch( diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/test_litellm/test_ruff_strict_gate.py index 96852e3a8a5..22255f0555e 100644 --- a/tests/test_litellm/test_ruff_strict_gate.py +++ b/tests/test_litellm/test_ruff_strict_gate.py @@ -82,13 +82,3 @@ def test_introduced_keeps_only_violations_on_changed_lines(): @pytest.mark.parametrize("hunk", ["@@ -1 +1 @@", "@@ -1,0 +1,2 @@"]) def test_parse_changed_lines_handles_single_and_ranged_hunks(hunk): assert gate.parse_changed_lines(f"+++ b/litellm/a.py\n{hunk}\n")["litellm/a.py"] - - -def test_report_emits_breached_rules_as_final_line(capsys): - # CI surfaces only the tail of the log, so the breached-rule summary (rule, - # total/cap, added) must be the last line or it gets truncated away. - breaches = sorted([gate.Breach("UP045", 530, 529, 1), gate.Breach("ANN401", 12, 10, 2)]) - new = [gate.Violation("litellm/types/llms/bedrock.py", 16, "UP045")] - gate.report(breaches, new, "origin/litellm_internal_staging") - last = capsys.readouterr().out.strip().splitlines()[-1] - assert last == "BREACHED RULES: ANN401 12/10 (+2); UP045 530/529 (+1)"