scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

feat(auth): resolve caller identity once into a Principal at the auth seam (#30887)

Browse source 
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.
This commit is contained in:
Yassin Kortam 2026-06-20 18:49:41 -07:00 • committed by GitHub
parent df5619f9b2
commit 84266bf924
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 1149 additions and 214 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/auth/auth_method.py Normal file
View file

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

101
litellm/proxy/auth/network.py Normal file
View file

@ -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,
)

33
litellm/proxy/auth/resolvers/__init__.py Normal file
View file

@ -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",
]

50
litellm/proxy/auth/resolvers/exceptions.py Normal file
View file

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

82
litellm/proxy/auth/resolvers/models.py Normal file
View file

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

206
litellm/proxy/auth/resolvers/store.py Normal file
View file

@ -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,
)

35
litellm/proxy/auth/roles.py Normal file
View file

@ -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

72
litellm/proxy/auth/trusted_proxy_utils.py
View file

@ -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,

117
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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

94
scripts/ruff_strict_gate.py
View file

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

25
tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py
View file

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

31
tests/proxy_unit_tests/test_jwt_key_mapping.py
View file

@ -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(

48
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -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

4
tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py
View file

@ -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

27
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -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,

4
tests/test_litellm/proxy/auth/test_model_checks.py
View file

@ -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"]

104
tests/test_litellm/proxy/auth/test_network.py Normal file
View file

@ -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

1
tests/test_litellm/proxy/auth/test_onboarding.py
View file

@ -18,7 +18,6 @@ from fastapi import HTTPException
import litellm
from litellm.proxy._types import InvitationClaim
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

35
tests/test_litellm/proxy/auth/test_resolvers_exceptions.py Normal file
View file

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

95
tests/test_litellm/proxy/auth/test_resolvers_models.py Normal file
View file

@ -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()

96
tests/test_litellm/proxy/auth/test_resolvers_seam.py Normal file
View file

@ -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

70
tests/test_litellm/proxy/auth/test_resolvers_store.py Normal file
View file

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

10
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -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(

10
tests/test_litellm/test_ruff_strict_gate.py
View file

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