From 6b3ee118104640cf6ed6fdbb2ddf7ffc1bbf09ea Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 26 Oct 2026 15:33:56 -0700 Subject: [PATCH] refactor(auth): extract credential selection into litellm._auth Move scheme stripping, key-source precedence, websocket key selection, hash_token and is_master_key into litellm/_auth as the interface the Rust auth port will mirror. Authorization schemes now match case-insensitively like FastAPI's HTTPBearer. Pass-through endpoint keys join the precedence table instead of a separate fallback, and the dead custom-header helper plus the duplicated hash_token and is_master_key bodies are gone --- litellm/_auth/AGENTS.md | 7 + litellm/_auth/__init__.py | 0 litellm/_auth/credential.py | 141 ++++++++++ litellm/_auth/key.py | 12 + litellm/proxy/auth/user_api_key_auth.py | 266 ++++++------------ .../mcp_server/auth/user_api_key_auth_mcp.py | 8 +- .../llm_passthrough_endpoints.py | 15 +- .../spend_tracking/spend_tracking_utils.py | 12 +- litellm/proxy/utils.py | 17 +- litellm/types/proxy/auth/user_api_key_auth.py | 10 +- tests/unit/_auth/__init__.py | 0 tests/unit/_auth/test_credential.py | 153 ++++++++++ tests/unit/_auth/test_key.py | 35 +++ .../unit/proxy/auth/test_user_api_key_auth.py | 45 --- .../test_spend_tracking_utils.py | 32 --- 15 files changed, 453 insertions(+), 300 deletions(-) create mode 100644 litellm/_auth/AGENTS.md create mode 100644 litellm/_auth/__init__.py create mode 100644 litellm/_auth/credential.py create mode 100644 litellm/_auth/key.py create mode 100644 tests/unit/_auth/__init__.py create mode 100644 tests/unit/_auth/test_credential.py create mode 100644 tests/unit/_auth/test_key.py diff --git a/litellm/_auth/AGENTS.md b/litellm/_auth/AGENTS.md new file mode 100644 index 00000000000..630cb8259aa --- /dev/null +++ b/litellm/_auth/AGENTS.md @@ -0,0 +1,7 @@ +# litellm/_auth + +## Why this exists + +- Lots of Python code like `litellm/proxy/auth/user_api_key_auth.py` is too complex +- Soon, Rust will implement the same auth stack, and expose common functionality as Python API through pyo3, to avoid maintaining 2 copy of code. +- So the goal of `_auth` here is to define that interface first, not writing really good python code. Should be behavior preserving refactor though diff --git a/litellm/_auth/__init__.py b/litellm/_auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/_auth/credential.py b/litellm/_auth/credential.py new file mode 100644 index 00000000000..105eb4afe73 --- /dev/null +++ b/litellm/_auth/credential.py @@ -0,0 +1,141 @@ +import re +from dataclasses import dataclass +from enum import Enum +from typing import Final + +_AWS4_PREFIX: Final = "AWS4-HMAC-SHA256" +_AWS4_BEARER_CREDENTIAL: Final = re.compile(r"Credential=Bearer\s+([^/\s,]+)") +_AWS4_CREDENTIAL: Final = re.compile(r"Credential=([^/\s,]+)") +_TOKEN_SCHEMES: Final = frozenset({"bearer", "basic"}) +_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX: Final = "openai-insecure-api-key." + + +class KeySource(Enum): + CUSTOM_HEADER = "custom_header" + LITELLM_HEADER = "x-litellm-api-key" + AUTHORIZATION = "authorization" + API_KEY = "api-key" + X_API_KEY = "x-api-key" + GOOGLE = "x-goog-api-key" + APIM = "ocp-apim-subscription-key" + QUERY = "key" + PASS_THROUGH = "pass_through_endpoint" + + +@dataclass(frozen=True, slots=True) +class PresentedKeys: + custom_header: str | None = None + litellm_header: str | None = None + authorization: str | None = None + api_key: str | None = None + x_api_key: str | None = None + google: str | None = None + apim: str | None = None + query: str | None = None + pass_through: str | None = None + + +@dataclass(frozen=True, slots=True) +class SelectedKey: + source: KeySource + presented: str + token: str + + +class WebSocketDenial(Enum): + NO_API_KEY = "no_api_key" + INVALID_AUTHORIZATION = "invalid_authorization" + + +@dataclass(frozen=True, slots=True) +class WebSocketKeys: + authorization: str | None = None + api_key: str | None = None + subprotocols: str | None = None + + +def _aws4_access_key(value: str) -> str | None: + match: Final = _AWS4_BEARER_CREDENTIAL.search(value) or _AWS4_CREDENTIAL.search(value) + return match.group(1) if match else None + + +def _split_scheme(value: str) -> tuple[str, str] | None: + scheme, separator, credential = value.partition(" ") + return (scheme.lower(), credential) if separator else None + + +def scheme_credential(value: str) -> str | None: + if value.startswith(_AWS4_PREFIX): + return _aws4_access_key(value) + split: Final = _split_scheme(value) + if split is None: + return None + scheme, credential = split + return credential if scheme in _TOKEN_SCHEMES else None + + +def litellm_header_token(value: str) -> str: + token: Final = scheme_credential(value) + return value if token is None else token + + +def _raw_key(keys: PresentedKeys) -> SelectedKey | None: + raw_sources: Final = ( + (KeySource.API_KEY, keys.api_key), + (KeySource.X_API_KEY, keys.x_api_key), + (KeySource.GOOGLE, keys.google), + (KeySource.APIM, keys.apim), + (KeySource.QUERY, keys.query or None), + (KeySource.PASS_THROUGH, keys.pass_through), + ) + return next( + ( + SelectedKey(source=source, presented=value, token=value) + for source, value in raw_sources + if value is not None + ), + None, + ) + + +def select_api_key(keys: PresentedKeys) -> SelectedKey | None: + if keys.custom_header is not None: + return SelectedKey( + source=KeySource.CUSTOM_HEADER, + presented=keys.custom_header, + token=scheme_credential(keys.custom_header) or "", + ) + if keys.litellm_header is not None: + return SelectedKey( + source=KeySource.LITELLM_HEADER, + presented=keys.litellm_header, + token=litellm_header_token(keys.litellm_header), + ) + if keys.authorization: + return SelectedKey( + source=KeySource.AUTHORIZATION, + presented=keys.authorization, + token=scheme_credential(keys.authorization) or "", + ) + return _raw_key(keys) + + +def _subprotocol_api_key(subprotocols: str | None) -> str | None: + protocols: Final = tuple(protocol.strip() for protocol in (subprotocols or "").split(",")) + return next( + ( + protocol.removeprefix(_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX) + for protocol in protocols + if protocol.startswith(_WEBSOCKET_SUBPROTOCOL_KEY_PREFIX) + ), + None, + ) + + +def select_websocket_api_key(keys: WebSocketKeys) -> str | WebSocketDenial: + if keys.authorization: + split: Final = _split_scheme(keys.authorization) + if split is None or split[0] != "bearer": + return WebSocketDenial.INVALID_AUTHORIZATION + return split[1].strip() + return keys.api_key or _subprotocol_api_key(keys.subprotocols) or WebSocketDenial.NO_API_KEY diff --git a/litellm/_auth/key.py b/litellm/_auth/key.py new file mode 100644 index 00000000000..af63778e3cf --- /dev/null +++ b/litellm/_auth/key.py @@ -0,0 +1,12 @@ +import hashlib +import secrets + + +def hash_token(token: str) -> str: + return hashlib.sha256(token.encode()).hexdigest() + + +def is_master_key(token: str | None, master_key: str | None) -> bool: + if token is None or master_key is None: + return False + return secrets.compare_digest(token.encode(), master_key.encode()) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 82135cb1b74..62de6b6088d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -9,8 +9,6 @@ Returns a UserAPIKeyAuth object if the API key is valid import asyncio import fnmatch -import re -import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone from enum import Enum @@ -24,6 +22,15 @@ from pydantic import TypeAdapter from starlette.exceptions import WebSocketException import litellm +from litellm._auth.credential import ( + PresentedKeys, + WebSocketDenial, + WebSocketKeys, + scheme_credential, + select_api_key, + select_websocket_api_key, +) +from litellm._auth.key import is_master_key from litellm._internal_context import service_target from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging @@ -1073,29 +1080,6 @@ async def _check_key_model_budget_with_fallback( request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model -def _get_bearer_token_or_received_api_key(api_key: str) -> str: - if api_key.startswith("Bearer "): # ensure Bearer token passed in - api_key = api_key.replace("Bearer ", "") # extract the token - elif api_key.startswith("Basic "): - api_key = api_key.replace("Basic ", "") # handle langfuse input - elif api_key.startswith("bearer "): - api_key = api_key.replace("bearer ", "") - elif api_key.startswith("AWS4-HMAC-SHA256"): - # Handle AWS Signature V4 format from LangChain - # Format: AWS4-HMAC-SHA256 Credential=Bearer $LITELLM_MASTER_KEY/date/region/service/aws4_request, SignedHeaders=..., Signature=... - # Extract the Bearer token from the Credential field - match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r"Credential=([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - - return api_key - - def _routing_selector_matches_claim( selector_value: object, claim_value: object, @@ -1188,29 +1172,7 @@ def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> def get_bearer_token( api_key: str, ) -> str: - if api_key.startswith("Bearer "): # ensure Bearer token passed in - api_key = api_key.replace("Bearer ", "") # extract the token - elif api_key.startswith("Basic "): - api_key = api_key.replace("Basic ", "") # handle langfuse input - elif api_key.startswith("bearer "): - api_key = api_key.replace("bearer ", "") - elif api_key.startswith("AWS4-HMAC-SHA256"): - # Handle AWS Signature V4 format from LangChain - # Format: AWS4-HMAC-SHA256 Credential=Bearer $LITELLM_MASTER_KEY/date/region/service/aws4_request, SignedHeaders=..., Signature=... - # Extract the Bearer token from the Credential field - match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r"Credential=([^/\s,]+)", api_key) - if match: - api_key = match.group(1) - else: - api_key = "" - else: - api_key = "" - return api_key + return scheme_credential(api_key) or "" _get_bearer_token: Final = get_bearer_token @@ -1247,6 +1209,14 @@ def _apply_budget_limits_to_end_user_params( verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) +_WEBSOCKET_DENIAL_DETAIL: Final = MappingProxyType( + { + WebSocketDenial.NO_API_KEY: "No API key provided", + WebSocketDenial.INVALID_AUTHORIZATION: "Invalid Authorization header format", + } +) + + async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth: return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model", "")) @@ -1307,27 +1277,17 @@ async def user_api_key_auth_websocket_for_model( request.body = return_body - authorization: Final = websocket.headers.get("authorization") - # If no Authorization header, try the api-key header - if not authorization: - api_key = websocket.headers.get("api-key") - if not api_key: - # Try extracting from WebSocket subprotocol (browser clients) - for protocol in websocket.headers.get("sec-websocket-protocol", "").split(","): - protocol = protocol.strip() - if protocol.startswith("openai-insecure-api-key."): - api_key = protocol[len("openai-insecure-api-key.") :] - break - if not api_key: - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="No API key provided") - else: - # Extract the API key from the Bearer token - if not authorization.startswith("Bearer "): - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="Invalid Authorization header format") - - api_key = authorization[len("Bearer ") :].strip() + selected: Final = select_websocket_api_key( + WebSocketKeys( + authorization=websocket.headers.get("authorization"), + api_key=websocket.headers.get("api-key"), + subprotocols=websocket.headers.get("sec-websocket-protocol"), + ) + ) + if isinstance(selected, WebSocketDenial): + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + raise HTTPException(status_code=403, detail=_WEBSOCKET_DENIAL_DETAIL[selected]) + api_key: Final = selected # Call user_api_key_auth with the extracted API key # Note: You'll need to modify this to work with WebSocket context if needed @@ -1455,9 +1415,43 @@ def get_rbac_role(jwt_handler: JWTHandler, scopes: list[str]) -> str: return LitellmUserRoles.TEAM +def _query_api_key(route: str, request: Request) -> str | None: + from litellm.proxy.auth.route_checks import RouteChecks + from litellm.proxy.common_utils.http_parsing_utils import ( + safe_get_request_query_params, + ) + + if not RouteChecks.is_generate_content_route(route=route): + return None + return safe_get_request_query_params(request).get("key") + + +def _pass_through_header_api_key(endpoint: dict, route: str, request: Request) -> str | None: + headers: Final[dict | None] = endpoint.get("headers") if endpoint.get("path", "") == route else None + if headers is None: + return None + return request.headers.get(headers.get("litellm_user_api_key", "")) + + +def _pass_through_api_key(pass_through_endpoints: list[dict] | None, route: str, request: Request) -> str | None: + header_values: Final = tuple( + value + for value in ( + _pass_through_header_api_key(endpoint, route, request) for endpoint in pass_through_endpoints or () + ) + if value is not None + ) + return header_values[-1] if header_values else None + + +def _header_or_none(value: object) -> str | None: + """Direct callers of ``user_api_key_auth`` leave unpassed headers as their ``fastapi.Security`` defaults.""" + return value if isinstance(value, str) else None + + def get_api_key( custom_litellm_key_header: str | None, - api_key: str, + api_key: str | None, azure_api_key_header: str | None, anthropic_api_key_header: str | None, google_ai_studio_api_key_header: str | None, @@ -1465,54 +1459,26 @@ def get_api_key( pass_through_endpoints: list[dict] | None, route: str, request: Request, -) -> tuple[str, str | None]: - """ - Returns: - Tuple[Optional[str], Optional[str]]: Tuple of the api_key and the passed_in_key - """ - from litellm.proxy.auth.route_checks import RouteChecks - from litellm.proxy.common_utils.http_parsing_utils import ( - safe_get_request_query_params, + *, + custom_header: str | None = None, +) -> tuple[str | None, str | None]: + selected: Final = select_api_key( + PresentedKeys( + custom_header=custom_header, + litellm_header=_header_or_none(custom_litellm_key_header), + authorization=_header_or_none(api_key), + api_key=_header_or_none(azure_api_key_header), + x_api_key=_header_or_none(anthropic_api_key_header), + google=_header_or_none(google_ai_studio_api_key_header), + apim=_header_or_none(azure_apim_header), + query=_query_api_key(route, request), + pass_through=_pass_through_api_key(pass_through_endpoints, route, request), + ) ) - - api_key = api_key - passed_in_key: str | None = None - if isinstance(custom_litellm_key_header, str): - passed_in_key = custom_litellm_key_header - api_key = _get_bearer_token_or_received_api_key(custom_litellm_key_header) - elif isinstance(api_key, str) and len(api_key) > 0: - passed_in_key = api_key - api_key = get_bearer_token(api_key=api_key) - elif isinstance(azure_api_key_header, str): - passed_in_key = azure_api_key_header - api_key = azure_api_key_header - elif isinstance(anthropic_api_key_header, str): - passed_in_key = anthropic_api_key_header - api_key = anthropic_api_key_header - elif isinstance(google_ai_studio_api_key_header, str): - passed_in_key = google_ai_studio_api_key_header - api_key = google_ai_studio_api_key_header - elif isinstance(azure_apim_header, str): - passed_in_key = azure_apim_header - api_key = azure_apim_header - elif ( - RouteChecks.is_generate_content_route(route=route) - and request is not None - and safe_get_request_query_params(request).get("key") - ): - google_auth_key: Final[str] = safe_get_request_query_params(request).get("key") or "" - passed_in_key = google_auth_key - api_key = google_auth_key - elif pass_through_endpoints is not None: - for endpoint in pass_through_endpoints: - if endpoint.get("path", "") == route: - headers: dict | None = endpoint.get("headers", None) - if headers is not None: - header_key: str = headers.get("litellm_user_api_key", "") - if request.headers.get(header_key) is not None: - api_key = request.headers.get(header_key) or "" - passed_in_key = api_key - return api_key, passed_in_key + if selected is None: + return api_key, None + verbose_proxy_logger.debug("Selected API key from %s", selected.source.value) + return selected.token, selected.presented async def check_api_key_for_custom_headers_or_pass_through_endpoints( @@ -2189,7 +2155,7 @@ def _custom_auth_token_for_this_request(response: object) -> UserAPIKeyAuth: async def user_api_key_auth_builder( request: Request, api_key: str, - azure_api_key_header: str, + azure_api_key_header: str | None, anthropic_api_key_header: str | None, google_ai_studio_api_key_header: str | None, azure_apim_header: str | None, @@ -2231,8 +2197,13 @@ async def user_api_key_auth_builder( route=route, ) pass_through_endpoints: Final[list[dict] | None] = general_settings.get("pass_through_endpoints", None) - ## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header - api_key, passed_in_key = get_api_key( + custom_litellm_key_header_name: Final = general_settings.get("litellm_key_header_name") + api_key, _ = get_api_key( + custom_header=( + request.headers.get(custom_litellm_key_header_name) + if custom_litellm_key_header_name is not None + else None + ), custom_litellm_key_header=custom_litellm_key_header, api_key=api_key, azure_api_key_header=azure_api_key_header, @@ -2243,12 +2214,6 @@ async def user_api_key_auth_builder( route=route, request=request, ) - custom_litellm_key_header_name: Final = general_settings.get("litellm_key_header_name") - if custom_litellm_key_header_name is not None and custom_litellm_key_header_name in request.headers: - api_key = get_api_key_from_custom_header( - request=request, - custom_litellm_key_header_name=custom_litellm_key_header_name, - ) if open_telemetry_logger is not None: # Reuse the span created by user_api_key_auth (before body parse) @@ -2692,10 +2657,6 @@ async def user_api_key_auth_builder( raise e verbose_proxy_logger.debug("Unable to find user in db. Error - %s", e) - ### CHECK IF ADMIN ### - # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead - ### CHECK IF ADMIN ### - # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead if valid_token is None: ## Check CACHE try: @@ -2806,19 +2767,13 @@ async def user_api_key_auth_builder( except Exception as e: verbose_logger.debug(e) # moving from .warning to .debug as it spams logs when team missing from cache. - try: - is_master_key_valid = secrets.compare_digest(api_key, master_key) - except Exception: - is_master_key_valid = False - - ## VALIDATE MASTER KEY ## if not isinstance(master_key, str): raise HTTPException( status_code=500, detail={f"Master key must be a valid string. Current type={type(master_key)}"}, ) - if is_master_key_valid: + if is_master_key(api_key, master_key): # Substitute a stable alias for the raw master key so neither the # master key nor its hash propagates into spend logs, Prometheus # /metrics labels, audit trails, rate-limit buckets, or any other @@ -4344,7 +4299,7 @@ def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Reque async def user_api_key_auth( request: Request, api_key: str = fastapi.Security(api_key_header), - azure_api_key_header: str = fastapi.Security(azure_api_key_header), + azure_api_key_header: str | None = fastapi.Security(azure_api_key_header), anthropic_api_key_header: str | None = fastapi.Security(anthropic_api_key_header), google_ai_studio_api_key_header: str | None = fastapi.Security(google_ai_studio_api_key_header), azure_apim_header: str | None = fastapi.Security(azure_apim_header), @@ -4494,41 +4449,6 @@ async def _return_user_api_key_auth_obj( return UserAPIKeyAuth.model_validate(user_api_key_kwargs) -def get_api_key_from_custom_header(request: Request, custom_litellm_key_header_name: str) -> str: - """ - Get API key from custom header - - Args: - request (Request): Request object - custom_litellm_key_header_name (str): Custom header name - - Returns: - Optional[str]: API key - """ - api_key: str = "" - # use this as the virtual key passed to litellm proxy - custom_litellm_key_header_name = custom_litellm_key_header_name.lower() - _headers: Final = {k.lower(): v for k, v in request.headers.items()} - verbose_proxy_logger.debug( - "searching for custom_litellm_key_header_name= %s, in headers=%s", - custom_litellm_key_header_name, - _headers, - ) - custom_api_key: Final = _headers.get(custom_litellm_key_header_name) - if custom_api_key: - api_key = get_bearer_token(api_key=custom_api_key) # rebind-ok: pre-existing rebinding on a rename-only line - verbose_proxy_logger.debug( - "Found custom API key using header: %s, setting api_key=%s", - custom_litellm_key_header_name, - abbreviate_api_key(api_key), - ) - else: - verbose_proxy_logger.exception( - "No LiteLLM Virtual Key pass. Please set header=%s: Bearer ", custom_litellm_key_header_name - ) - return api_key - - def _get_temp_budget_increase(valid_token: UserAPIKeyAuth): valid_token_metadata: Final = valid_token.metadata if "temp_budget_increase" in valid_token_metadata and "temp_budget_expiry" in valid_token_metadata: diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index be24b1630b0..ba5f25403d7 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -19,6 +19,7 @@ from starlette.types import Scope from typing_extensions import assert_never import litellm +from litellm._auth.credential import litellm_header_token from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.constants import MCP_ALL_TOOLS_WILDCARD, MCP_PEEKED_BODY_SCOPE_KEY @@ -48,7 +49,6 @@ from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_age from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.oauth2_check import OAuth2TokenValidationError from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports - _get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth _run_centralized_common_checks, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export oauth2_introspection_applies, run_centralized_common_checks, @@ -110,7 +110,7 @@ _CACHED_VALUE: Final = TypeAdapter(object) def _normalize_caller_admission_credential(value: str) -> str: - return _get_bearer_token_or_received_api_key(strip_auth_scheme(value, "Bearer")).strip() + return litellm_header_token(strip_auth_scheme(value, "Bearer")).strip() def _admission_request(scope: Scope) -> Request: @@ -797,7 +797,7 @@ class MCPRequestHandler: # for a delegated server, so validate it: identity / spend / rate # limits resolve and any stored upstream token can be forwarded. validated_user_api_key_auth = await user_api_key_auth( - api_key=f"Bearer {_get_bearer_token_or_received_api_key(litellm_api_key)}", + api_key=f"Bearer {litellm_header_token(litellm_api_key)}", request=request, ) elif MCPRequestHandler._target_servers_are_true_passthrough( @@ -1278,7 +1278,7 @@ class MCPRequestHandler: permissions, budgets, rate limits) and the sealed upstream token is injected under the server's per-server auth-header key, while the leak-defense chokepoint strips the envelope ``Authorization`` itself from egress.""" - presented_token: Final = _get_bearer_token_or_received_api_key(litellm_api_key) + presented_token: Final = litellm_header_token(litellm_api_key) explicit_auth: Final = await user_api_key_auth(api_key=f"Bearer {presented_token}", request=request) result: Final = await MCPRequestHandler._open_dcr_bridge_envelope( server=server, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 627164abb93..701559c12fa 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -31,6 +31,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm import get_llm_provider +from litellm._auth.credential import scheme_credential from litellm._logging import redact_internal_details_from_client_message, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( @@ -97,7 +98,7 @@ from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports _get_bearer_token, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - get_bearer_token, + get_bearer_token, # pyright: ignore[reportUnusedImport] # backwards-compatible package export is_no_auth_dev_mode, user_api_key_auth, user_api_key_auth_accepted_websocket_for_model, @@ -2851,17 +2852,7 @@ _CREDENTIALLESS_VERTEX_MISSING_CREDENTIAL_DETAIL: Final = ( def _normalize_credential_value(value: str) -> str: - """Reduce a header value to the bare token, matching how ``user_api_key_auth`` - reads a caller's key. - - Reuses the auth module's ``_get_bearer_token`` so the caller-key comparison - strips exactly the schemes authentication accepts (``Bearer`` / ``bearer`` / - ``Basic`` / ``AWS4-HMAC-SHA256`` credential), rather than re-deriving a - narrower normalization here. ``_get_bearer_token`` returns ``""`` for a value - with no recognized scheme prefix, so a bare token (or a real Google - credential that carries no scheme) falls back to its own value. - """ - return get_bearer_token(value) or value + return scheme_credential(value) or value _VERTEX_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-goog-api-key"}) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index fc54d1da0b2..591b70b53c1 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -2,7 +2,6 @@ import ipaddress import json import os import re -import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone from datetime import datetime as dt @@ -14,6 +13,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm +from litellm._auth.key import is_master_key as is_master_key # noqa: PLC0414 # public re-export from litellm._logging import verbose_proxy_logger from litellm.constants import ( BILLED_RETRY_ATTEMPT_METADATA_KEY, @@ -125,16 +125,6 @@ def _get_max_string_length_prompt_in_db() -> int: return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB -def is_master_key(api_key: str | None, _master_key: str | None) -> bool: - """ - Raw-only constant-time master-key comparison. The hashed form is never - considered equivalent — only the raw master-key string matches. - """ - if _master_key is None or api_key is None: - return False - return secrets.compare_digest(api_key, _master_key) - - _is_master_key: Final = is_master_key diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 00187790901..cc5ced5c3bd 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -136,6 +136,7 @@ from litellm import ( ModelResponseStream, Router, ) +from litellm._auth.key import hash_token as hash_token # noqa: PLC0414 # public re-export from litellm._internal_context import service_target from litellm._logging import redact_string, verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes @@ -5791,11 +5792,8 @@ class PrismaClient: # Default to success if metadata parsing fails return "success" - def hash_token(self, token: str): - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token + def hash_token(self, token: str) -> str: + return hash_token(token) def jsonify_object(self, data: Mapping[str, object]) -> dict[str, object]: db_data: Final[dict[str, object]] = copy.deepcopy(dict(data)) @@ -8361,15 +8359,6 @@ async def send_email( return True -def hash_token(token: str): - import hashlib - - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token - - def hash_password(password: str) -> str: """Hash a password using scrypt with a random salt.""" import base64 diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 1183c36eda7..92067be067a 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -9,6 +9,7 @@ from typing import ( from pydantic import ConfigDict, Field, PrivateAttr, model_validator from typing_extensions import TypeIs # noqa: TID251 # narrows the before-validator input without copying it +from litellm._auth.key import hash_token as hash_token # noqa: PLC0414 # public re-export from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.team import Member from litellm.models.verification_token import LiteLLM_VerificationToken @@ -37,15 +38,6 @@ def is_jwt(token: str | None) -> bool: return len(parts) == 3 -def hash_token(token: str) -> str: - import hashlib - - # Hash the string using SHA-256 - hashed_token: Final = hashlib.sha256(token.encode()).hexdigest() - - return hashed_token - - _UntypedDict: TypeAlias = dict[Any, Any] diff --git a/tests/unit/_auth/__init__.py b/tests/unit/_auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/_auth/test_credential.py b/tests/unit/_auth/test_credential.py new file mode 100644 index 00000000000..f0b73313b03 --- /dev/null +++ b/tests/unit/_auth/test_credential.py @@ -0,0 +1,153 @@ +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.credential import ( + KeySource, + PresentedKeys, + SelectedKey, + WebSocketDenial, + WebSocketKeys, + litellm_header_token, + scheme_credential, + select_api_key, + select_websocket_api_key, +) + +AWS4_BEARER: Final = ( + "AWS4-HMAC-SHA256 Credential=Bearer token/date/region/service/aws4_request, SignedHeaders=host, Signature=sig" +) +AWS4_RAW: Final = ( + "AWS4-HMAC-SHA256 Credential=token/date/region/service/aws4_request, SignedHeaders=host, Signature=sig" +) + +PRECEDENCE: Final = ( + ("custom_header", KeySource.CUSTOM_HEADER, "Bearer {}"), + ("litellm_header", KeySource.LITELLM_HEADER, "Bearer {}"), + ("authorization", KeySource.AUTHORIZATION, "Bearer {}"), + ("api_key", KeySource.API_KEY, "{}"), + ("x_api_key", KeySource.X_API_KEY, "{}"), + ("google", KeySource.GOOGLE, "{}"), + ("apim", KeySource.APIM, "{}"), + ("query", KeySource.QUERY, "{}"), + ("pass_through", KeySource.PASS_THROUGH, "{}"), +) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("Bearer token", "token"), + ("bearer token", "token"), + ("Basic dXNlcjpwYXNz", "dXNlcjpwYXNz"), + (AWS4_BEARER, "token"), + (AWS4_RAW, "token"), + ("Bearer ", ""), + ("Bearer token", " token"), + ("BEARER token", "token"), + ("BASIC dXNlcjpwYXNz", "dXNlcjpwYXNz"), + ("AWS4-HMAC-SHA256 SignedHeaders=host", None), + ("token", None), + ("Bearer", None), + ("Token token", None), + ("Bearertoken", None), + ], +) +def test_scheme_credential_strips_a_known_scheme_in_any_case(value: str, expected: str | None): + assert scheme_credential(value) == expected + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("token", "token"), + ("Bearer token", "token"), + (AWS4_BEARER, "token"), + ("AWS4-HMAC-SHA256 malformed", "AWS4-HMAC-SHA256 malformed"), + ("Bearer ", ""), + ], +) +def test_litellm_header_strips_a_known_scheme_or_keeps_the_raw_value(value: str, expected: str): + assert litellm_header_token(value) == expected + + +@pytest.mark.parametrize("rank", range(len(PRECEDENCE))) +def test_higher_priority_source_wins_over_every_lower_one(rank: int): + higher, source, template = PRECEDENCE[rank] + winner: Final = template.format(higher) + losers: Final = {field: lower_template.format("loser") for field, _, lower_template in PRECEDENCE[rank + 1 :]} + assert select_api_key(PresentedKeys(**losers, **{higher: winner})) == SelectedKey( + source=source, presented=winner, token=higher + ) + + +@pytest.mark.parametrize( + ("keys", "expected"), + [ + (PresentedKeys(custom_header="custom", x_api_key="lower"), ""), + (PresentedKeys(custom_header="", x_api_key="lower"), ""), + (PresentedKeys(authorization="token", x_api_key="lower"), ""), + (PresentedKeys(authorization="Token x", x_api_key="lower"), ""), + ], +) +def test_selected_source_without_a_known_scheme_yields_an_empty_token(keys: PresentedKeys, expected: str): + selected: Final = select_api_key(keys) + assert selected is not None + assert selected.token == expected + + +@pytest.mark.parametrize( + ("keys", "expected_source"), + [ + (PresentedKeys(authorization="", x_api_key="lower"), KeySource.X_API_KEY), + (PresentedKeys(query="", apim="lower"), KeySource.APIM), + (PresentedKeys(litellm_header=""), KeySource.LITELLM_HEADER), + (PresentedKeys(api_key=""), KeySource.API_KEY), + ], +) +def test_only_empty_authorization_and_query_fall_through(keys: PresentedKeys, expected_source: KeySource): + selected: Final = select_api_key(keys) + assert selected is not None + assert selected.source is expected_source + + +@pytest.mark.parametrize("keys", [PresentedKeys(), PresentedKeys(authorization=""), PresentedKeys(query="")]) +def test_nothing_presented_selects_nothing(keys: PresentedKeys): + assert select_api_key(keys) is None + + +def test_presented_keeps_the_raw_value_of_the_selected_source(): + keys: Final = PresentedKeys(authorization=AWS4_BEARER) + assert select_api_key(keys) == SelectedKey(KeySource.AUTHORIZATION, AWS4_BEARER, "token") + assert select_api_key(replace(keys, custom_header="Basic abc")) == SelectedKey( + KeySource.CUSTOM_HEADER, "Basic abc", "abc" + ) + + +@pytest.mark.parametrize( + ("keys", "expected"), + [ + (WebSocketKeys(authorization="Bearer token"), "token"), + (WebSocketKeys(authorization="bearer token "), "token"), + ( + WebSocketKeys(authorization="Bearer token", api_key="other", subprotocols="openai-insecure-api-key.x"), + "token", + ), + (WebSocketKeys(authorization="Basic token"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(authorization="token"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(authorization="Basic token", api_key="other"), WebSocketDenial.INVALID_AUTHORIZATION), + (WebSocketKeys(api_key="token"), "token"), + (WebSocketKeys(authorization="", api_key="token"), "token"), + (WebSocketKeys(api_key="token", subprotocols="openai-insecure-api-key.other"), "token"), + (WebSocketKeys(subprotocols="realtime, openai-insecure-api-key.token, openai-beta.realtime-v1"), "token"), + (WebSocketKeys(api_key="", subprotocols="openai-insecure-api-key.token"), "token"), + (WebSocketKeys(subprotocols="realtime"), WebSocketDenial.NO_API_KEY), + (WebSocketKeys(subprotocols="openai-insecure-api-key."), WebSocketDenial.NO_API_KEY), + (WebSocketKeys(), WebSocketDenial.NO_API_KEY), + ], +) +def test_websocket_takes_a_bearer_authorization_then_api_key_then_subprotocol( + keys: WebSocketKeys, expected: str | WebSocketDenial +): + assert select_websocket_api_key(keys) == expected diff --git a/tests/unit/_auth/test_key.py b/tests/unit/_auth/test_key.py new file mode 100644 index 00000000000..1d41920e4dd --- /dev/null +++ b/tests/unit/_auth/test_key.py @@ -0,0 +1,35 @@ +import hashlib + +import pytest + +from litellm._auth.key import hash_token, is_master_key + + +@pytest.mark.parametrize("token", ["sk-1234", "", "sk-ünïcode"]) +def test_hash_token_is_sha256_hex_of_utf8(token: str): + assert hash_token(token) == hashlib.sha256(token.encode()).hexdigest() + + +def test_hash_token_differs_per_token(): + assert hash_token("sk-a") != hash_token("sk-b") + + +@pytest.mark.parametrize( + ("token", "master_key", "expected"), + [ + ("sk-master", "sk-master", True), + ("sk-master", "sk-other", False), + ("sk-mast", "sk-master", False), + (None, "sk-master", False), + ("sk-master", None, False), + (None, None, False), + ("sk-mästér", "sk-mästér", True), + ("sk-mästér", "sk-master", False), + ], +) +def test_is_master_key_matches_only_the_exact_raw_key(token: str | None, master_key: str | None, expected: bool): + assert is_master_key(token, master_key) is expected + + +def test_hashed_master_key_is_not_the_master_key(): + assert is_master_key(hash_token("sk-master"), "sk-master") is False diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 19f9d0dcc73..cea01a7cf69 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -18,7 +18,6 @@ import litellm from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth, UserAPIKeyAuth, - get_api_key_from_custom_header, ) from fastapi import WebSocket, HTTPException, status @@ -506,50 +505,6 @@ async def test_auth_not_connected_to_db(): assert valid_token.token == "failed-to-connect-to-db" -def _assert_api_key_from_custom_header(headers, custom_header_name, expected_api_key): - verbose_proxy_logger.setLevel(logging.DEBUG) - request = MagicMock(spec=Request) - request.headers = headers - api_key = get_api_key_from_custom_header(request=request, custom_litellm_key_header_name=custom_header_name) - assert api_key == expected_api_key - - -def test_get_api_key_from_custom_header_bearer_token(): - token = "sk-" + "1" * 8 - _assert_api_key_from_custom_header( - headers={"x-custom-api-key": f"Bearer {token}"}, - custom_header_name="x-custom-api-key", - expected_api_key=token, - ) - - -def test_get_api_key_from_custom_header_empty_value(): - _assert_api_key_from_custom_header( - headers={"x-custom-api-key": ""}, - custom_header_name="x-custom-api-key", - expected_api_key="", - ) - - -def test_get_api_key_from_custom_header_missing_header(): - _assert_api_key_from_custom_header( - headers={}, - custom_header_name="X-Custom-API-Key", - expected_api_key="", - ) - - -def test_get_api_key_from_custom_header_different_casing(): - token = "sk-" + "1" * 8 - _assert_api_key_from_custom_header( - headers={"X-CUSTOM-API-KEY": f"Bearer {token}"}, - custom_header_name="X-Custom-API-Key", - expected_api_key=token, - ) - - - - @pytest.mark.parametrize( "user_role, auth_user_id, requested_user_id, expected_result", [ diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 45ded91b61e..8b68377ddd8 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -46,7 +46,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_session_id_for_spend_log, _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, - is_master_key, redact_logged_api_key, _redact_prompt_leaks_in_error_string, sanitize_error_information_for_spend_logs, @@ -3305,37 +3304,6 @@ def test_get_logging_payload_includes_request_duration_ms(): assert payload["request_duration_ms"] == 3000 -class TestIsMasterKey: - """Tests for _is_master_key handling None inputs without raising TypeError.""" - - def test_none_api_key_returns_false(self): - """Regression: _is_master_key(None, 'sk-master') should return False, not raise TypeError.""" - assert is_master_key(api_key=None, _master_key="sk-master-key") is False - - def test_none_master_key_returns_false(self): - assert is_master_key(api_key="sk-some-key", _master_key=None) is False - - def test_both_none_returns_false(self): - assert is_master_key(api_key=None, _master_key=None) is False - - def test_matching_key_returns_true(self): - assert is_master_key(api_key="sk-master", _master_key="sk-master") is True - - def test_non_matching_key_returns_false(self): - assert is_master_key(api_key="sk-other", _master_key="sk-master") is False - - def test_master_key_hash_is_rejected(self): - """ - ``_is_master_key`` must not accept ``hash_token(master_key)`` as - equivalent to the raw master key — only the raw value matches. - """ - from litellm.proxy.utils import hash_token - - master = "sk-master-key-123" - hashed = hash_token(master) - assert is_master_key(api_key=hashed, _master_key=master) is False - - def test_sanitize_request_body_strips_secret_fields(): """ secret_fields contains raw HTTP headers (including Authorization Bearer