diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 850b50227f4e..9a3bc82c6fa4 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -6,6 +6,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. """ +import hmac import json import os import re @@ -28,6 +29,7 @@ from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * +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 ( _get_bearer_token, @@ -1791,29 +1793,11 @@ def _normalize_credential_value(value: str) -> str: ) -_VERTEX_CALLER_KEY_HEADER_PRECEDENCE: Final = ( - SpecialHeaders.custom_litellm_api_key.value.lower(), - SpecialHeaders.openai_authorization.value.lower(), - SpecialHeaders.azure_authorization.value.lower(), - SpecialHeaders.anthropic_authorization.value.lower(), - SpecialHeaders.google_ai_studio_authorization.value.lower(), - SpecialHeaders.azure_apim_authorization.value.lower(), -) - _MAPPED_ROUTE_CALLER_KEY_HEADER: Final = "litellm_user_api_key" -def _operator_configured_caller_key_header_names() -> tuple[tuple[str, ...], tuple[str, ...]]: - """Operator-configured caller-key header names, as (override, pass_through). - - ``user_api_key_auth`` accepts the caller's key from two runtime-configured - header sources beyond the built-in ones, at opposite ends of its precedence. - ``general_settings.litellm_key_header_name`` overrides every built-in source - (it replaces the resolved key after ``get_api_key`` runs), so it is highest - precedence. Each ``general_settings.pass_through_endpoints`` entry's - ``headers.litellm_user_api_key`` is checked last inside ``get_api_key``, so it - is lowest. Google never consumes either, so both are also dropped by name. - """ +def _operator_configured_caller_key_header_names() -> tuple[str, ...]: + """Operator-configured header names ``user_api_key_auth`` reads the caller's key from.""" from litellm.proxy.proxy_server import general_settings custom_key_header: Final = general_settings.get("litellm_key_header_name") @@ -1829,80 +1813,53 @@ def _operator_configured_caller_key_header_names() -> tuple[tuple[str, ...], tup if isinstance(headers, dict) and isinstance(headers.get("litellm_user_api_key"), str) ) ) - return override, pass_through + return override + pass_through -def _authenticated_caller_key_values(request: Request) -> frozenset[str]: - """The value ``user_api_key_auth`` would accept as this caller's LiteLLM key. +def _is_authenticated_caller_jwt(value: str, jwt_claims: Mapping[str, object]) -> bool: + """Whether a header value is the JWT whose claims ``user_api_key_auth`` stored as ``jwt_claims``.""" + presented_claims: Final = JWTHandler.get_unverified_claims(value) + if presented_claims is None: + return False + return all( + presented_claims.get(name) == claim + for name, claim in jwt_claims.items() + if name not in JWTHandler.LITELLM_INTERNAL_CLAIMS + ) - The Vertex route authenticates through ``Depends(user_api_key_auth)``, which - resolves the key by precedence, matched here exactly. The ``/vertex_ai`` route - is a mapped pass-through route, so a header literally named - ``litellm_user_api_key`` overrides every other source (``user_api_key_auth`` - applies it last), making it highest precedence. Then an operator - ``litellm_key_header_name``, then the built-in headers in ``get_api_key`` order, - then a ``pass_through_endpoints`` ``litellm_user_api_key`` header which - ``get_api_key`` checks last. Some of those headers (``Authorization``, - ``x-goog-api-key``) are also kept as genuine bring-your-own Google credentials, - so returning only the value that actually authenticated lets the filter strip - that value wherever it appears while leaving a real Google credential in place. - An empty set means no caller key was found, so nothing is value-stripped. - """ - incoming: Final = _safe_get_request_headers(request) - override_headers, pass_through_headers = _operator_configured_caller_key_header_names() - ordered_names: Final = ( - (_MAPPED_ROUTE_CALLER_KEY_HEADER,) - + override_headers - + _VERTEX_CALLER_KEY_HEADER_PRECEDENCE - + pass_through_headers - ) - present_values: Final = (incoming[name] for name in ordered_names if incoming.get(name)) - authenticated_key: Final = next( - (stripped for value in present_values if (stripped := _normalize_credential_value(value))), - "", - ) - return frozenset({authenticated_key}) if authenticated_key else frozenset() - - -def _forwarded_headers_for_credentialless_vertex_passthrough(request: Request) -> Mapping[str, str]: - """ - Header set to forward on the bring-your-own-credentials Vertex passthrough - branch, used when the proxy has no Vertex credential configured. - - No credential the proxy accepts for caller authentication is forwarded to - Google. ``user_api_key_auth`` reads the caller's key from every header in - ``SpecialHeaders.litellm_credential_header_names()``, and Vertex only ever - authenticates with an OAuth token in ``Authorization`` or an API key in - ``x-goog-api-key``. So the proxy-only auth headers Google never consumes - (everything in that set except those two, e.g. ``x-litellm-api-key`` / - ``api-key`` / ``x-api-key`` / ``Ocp-Apim-Subscription-Key``, plus the mapped - pass-through ``litellm_user_api_key`` header and any operator-configured - ``litellm_key_header_name`` / ``pass_through_endpoints`` key header) are dropped - by name. ``Authorization`` and ``x-goog-api-key`` may - instead carry a genuine bring-your-own Google credential, so they are kept - unless their value is the caller's authenticated LiteLLM key, which is dropped - by value (normalizing any ``Bearer`` / ``Basic`` / ``AWS4`` auth-scheme prefix - the same way authentication does). Because the value that authenticated is - resolved by the same precedence ``user_api_key_auth`` uses, a virtual key sent - only in ``x-goog-api-key`` (or in an operator-configured key header) is dropped - too, while a real Google key in ``x-goog-api-key`` alongside a virtual key in a - higher-precedence header is preserved. When neither a surviving - ``Authorization`` nor ``x-goog-api-key`` remains the request is rejected so the - virtual key cannot leak upstream. - """ + +def _is_authenticated_caller_secret(value: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + """Whether a header value is the master key, the JWT that authenticated, or the key stored as ``api_key``.""" + from litellm.proxy.proxy_server import master_key + + normalized: Final = _normalize_credential_value(value) + if master_key is not None and hmac.compare_digest(normalized.encode(), master_key.encode()): + return True + jwt_claims: Final = user_api_key_dict.jwt_claims + if jwt_claims and _is_authenticated_caller_jwt(normalized, jwt_claims): + return True + authenticated_key: Final = user_api_key_dict.api_key + if authenticated_key is None: + return False + if master_key is None and not normalized.startswith("sk-"): + return False + stored_representation: Final = UserAPIKeyAuth._safe_hash_litellm_api_key(normalized) # pyright: ignore[reportPrivateUsage] # the exact transform auth applied when it stored api_key + return hmac.compare_digest(stored_representation.encode(), authenticated_key.encode()) + + +def _forwarded_headers_for_credentialless_vertex_passthrough( + request: Request, user_api_key_dict: UserAPIKeyAuth +) -> Mapping[str, str]: + """Caller headers to forward on the bring-your-own-credentials Vertex branch, minus LiteLLM secrets.""" incoming: Final = _safe_get_request_headers(request) - caller_key_values: Final = _authenticated_caller_key_values(request) - override_headers, pass_through_headers = _operator_configured_caller_key_header_names() - never_forwarded: Final = ( - _HEADERS_NEVER_FORWARDED_TO_VERTEX.union((_MAPPED_ROUTE_CALLER_KEY_HEADER,)) - .union(override_headers) - .union(pass_through_headers) + never_forwarded: Final = _HEADERS_NEVER_FORWARDED_TO_VERTEX.union( + (_MAPPED_ROUTE_CALLER_KEY_HEADER, *_operator_configured_caller_key_header_names()) ) forwarded: Final = MappingProxyType( { name: value for name, value in incoming.items() - if name not in never_forwarded and _normalize_credential_value(value) not in caller_key_values + if name not in never_forwarded and not _is_authenticated_caller_secret(value, user_api_key_dict) } ) if "authorization" not in forwarded and "x-goog-api-key" not in forwarded: @@ -1918,6 +1875,7 @@ async def _prepare_vertex_auth_headers( vertex_location: str | None, base_target_url: str | None, get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, + user_api_key_dict: UserAPIKeyAuth, ) -> tuple[Mapping[str, str], str | None, bool, str | None, str | None]: """ Prepare authentication headers for Vertex AI pass-through requests. @@ -1930,6 +1888,8 @@ async def _prepare_vertex_auth_headers( vertex_location: Vertex location base_target_url: Base URL for the Vertex AI service get_vertex_pass_through_handler: Handler for the specific Vertex AI service + user_api_key_dict: The caller's resolved authentication, so only the secret that + authenticated them is stripped on the credential-less branch Returns: Tuple containing: @@ -1944,7 +1904,7 @@ async def _prepare_vertex_auth_headers( # Use headers from the incoming request if no vertex credentials are found if (vertex_credentials is None or vertex_credentials.vertex_project is None) and router_credentials is None: - headers = _forwarded_headers_for_credentialless_vertex_passthrough(request) + headers = _forwarded_headers_for_credentialless_vertex_passthrough(request, user_api_key_dict) headers_passed_through = True verbose_proxy_logger.debug( "default_vertex_config not set, forwarding caller-provided headers %s", tuple(headers.keys()) @@ -2104,6 +2064,7 @@ async def _base_vertex_proxy_route( vertex_location=vertex_location, base_target_url=base_target_url, get_vertex_pass_through_handler=get_vertex_pass_through_handler, + user_api_key_dict=user_api_key_dict, ) if base_target_url is None: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index a1725005dbfe..3b506324ad73 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1,3 +1,4 @@ +import base64 import contextlib import json import os @@ -16,6 +17,7 @@ import litellm +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, RouteChecks, @@ -37,6 +39,7 @@ vllm_proxy_route, ) from litellm.proxy._types import LitellmUserRoles, SpecialHeaders, UserAPIKeyAuth +from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials @@ -592,7 +595,7 @@ async def test_vertex_passthrough_with_no_default_credentials(self, monkeypatch) "method": "POST", "path": endpoint, "headers": [ - (b"authorization", b"Bearer test-creds"), + (b"authorization", b"Bearer sk-test-creds"), ], } ) @@ -617,7 +620,7 @@ async def test_vertex_passthrough_with_no_default_credentials(self, monkeypatch) ): mock_ensure_token.return_value = ("test-auth-header", test_project) mock_get_token.return_value = (test_token, "") - mock_auth.return_value = MagicMock() + mock_auth.return_value = UserAPIKeyAuth(api_key="sk-test-creds") with pytest.raises(HTTPException) as exc_info: await vertex_proxy_route( @@ -3342,14 +3345,14 @@ def fake_create_pass_through_route(**kwargs): ), mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route), mock.patch(f"{module}.get_litellm_virtual_key", return_value="Bearer test-key"), - mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value={"api_key": "test-key"})), + mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=UserAPIKeyAuth(api_key="test-key"))), mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler), ): await vertex_proxy_route( endpoint=endpoint, request=request, fastapi_response=Response(), - user_api_key_dict=UserAPIKeyAuth(token="test-key"), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), ) assert captured, "create_pass_through_route was never called" @@ -3447,6 +3450,13 @@ def test_is_passthrough_request_streaming_tolerates_non_object_bodies(request_bo assert is_passthrough_request_streaming(request_body) is expected +def _unsigned_jwt(claims: Mapping[str, str]) -> str: + def segment(payload: Mapping[str, str]) -> str: + return base64.urlsafe_b64encode(json.dumps(dict(payload)).encode()).rstrip(b"=").decode() + + return ".".join((segment({"alg": "RS256", "typ": "JWT"}), segment(claims), "c2lnbmF0dXJl")) + + class TestVertexCredentiallessPassthroughVirtualKeyLeak: """Regression coverage for LIT-5997. @@ -3466,6 +3476,12 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: genuine bring-your-own Google credential that must still pass through. The by-value strip also covers a virtual key sent in the operator-configured ``general_settings.litellm_key_header_name``, whatever that header is named. + + The by-value strip keys off what actually authenticated the caller (the + master key, or the LiteLLM key whose hash ``user_api_key_auth`` resolved as + ``api_key``), never off header precedence: a custom auth or JWT that + authenticated the caller without consuming ``Authorization`` leaves the + caller's own Google token there, and it must keep flowing. """ VKEY = "sk-litellm-victim-key" @@ -3475,8 +3491,14 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak: ) async def _run( - self, monkeypatch, headers: list[tuple[bytes, bytes]] + self, + monkeypatch, + headers: list[tuple[bytes, bytes]], + authenticated: UserAPIKeyAuth | None = None, + master_key: str | None = "sk-master-1234", ) -> tuple[HTTPException | None, dict | None]: + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", master_key) + caller: Final = authenticated if authenticated is not None else UserAPIKeyAuth(api_key=self.VKEY) from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import ( PassthroughEndpointRouter, ) @@ -3509,7 +3531,7 @@ def fake_create_pass_through_route(**kwargs): raised: HTTPException | None = None with ( mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route), - mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=UserAPIKeyAuth(token="hashed"))), + mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=caller)), mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler), ): try: @@ -3517,7 +3539,7 @@ def fake_create_pass_through_route(**kwargs): endpoint=self.ENDPOINT, request=request, fastapi_response=Response(), - user_api_key_dict=UserAPIKeyAuth(token="hashed"), + user_api_key_dict=caller, ) except HTTPException as exc: raised = exc @@ -3767,6 +3789,135 @@ async def test_virtual_key_in_mapped_route_litellm_user_api_key_header_alone_is_ assert forwarded is None, "a virtual key in the mapped-route litellm_user_api_key header must be dropped, not forwarded" assert raised is not None and raised.status_code == 401 + GOOGLE_OAUTH_TOKEN = "ya29.byo-google-oauth-token" + + LITELLM_JWT_CLAIMS = MappingProxyType({"sub": "jwt-subject", "iss": "https://idp.example.com"}) + LITELLM_JWT = _unsigned_jwt(LITELLM_JWT_CLAIMS) + GOOGLE_SERVICE_ACCOUNT_JWT = _unsigned_jwt( + { + "sub": "vertex-caller@my-proj.iam.gserviceaccount.com", + "iss": "vertex-caller@my-proj.iam.gserviceaccount.com", + "aud": "https://aiplatform.googleapis.com/", + } + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("master_key", "authenticated"), + [ + pytest.param( + "sk-master-1234", + UserAPIKeyAuth(api_key="best-api-key-ever", user_role=LitellmUserRoles.PROXY_ADMIN), + id="custom-auth-returning-its-own-identifier", + ), + pytest.param( + "sk-master-1234", + UserAPIKeyAuth(api_key=None, user_id="jwt-subject", jwt_claims=dict(LITELLM_JWT_CLAIMS)), + id="jwt-auth", + ), + pytest.param(None, UserAPIKeyAuth(api_key=GOOGLE_OAUTH_TOKEN), id="no-master-key-echoes-raw-header"), + ], + ) + async def test_google_token_in_authorization_is_forwarded_when_auth_did_not_consume_it( + self, monkeypatch, master_key: str | None, authenticated: UserAPIKeyAuth + ): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"authorization", f"Bearer {self.GOOGLE_OAUTH_TOKEN}".encode()), + (b"content-type", b"application/json"), + ], + authenticated=authenticated, + master_key=master_key, + ) + assert raised is None, f"the caller's own Google token must not be mistaken for a LiteLLM key: {raised}" + assert forwarded is not None + assert forwarded.get("authorization") == f"Bearer {self.GOOGLE_OAUTH_TOKEN}" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("credential", "authenticated"), + [ + pytest.param("modified_key", UserAPIKeyAuth(api_key="modified_key"), id="custom-auth-echoing-opaque-credential"), + pytest.param( + LITELLM_JWT, + UserAPIKeyAuth(api_key=LITELLM_JWT, user_id="jwt-subject"), + id="custom-auth-echoing-jwt", + ), + pytest.param( + LITELLM_JWT, + UserAPIKeyAuth(api_key=None, user_id="jwt-subject", jwt_claims=dict(LITELLM_JWT_CLAIMS)), + id="jwt-auth", + ), + pytest.param( + LITELLM_JWT, + UserAPIKeyAuth( + api_key=None, + user_id="jwt-subject", + jwt_claims={ + **LITELLM_JWT_CLAIMS, + JWTHandler.LITELLM_JWT_ISSUER_CLAIM: "https://idp.example.com", + JWTHandler.LITELLM_USER_ID_CLAIM: "jwt-subject", + }, + ), + id="multi-issuer-jwt-auth-normalized-claims", + ), + ], + ) + async def test_non_sk_litellm_credential_that_authenticated_is_rejected_not_forwarded( + self, monkeypatch, credential: str, authenticated: UserAPIKeyAuth + ): + raised, forwarded = await self._run( + monkeypatch, + [(b"authorization", f"Bearer {credential}".encode()), (b"content-type", b"application/json")], + authenticated=authenticated, + ) + assert forwarded is None, "the credential that authenticated the caller must never reach the upstream forwarder" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_jwt_authenticated_caller_keeps_a_different_byo_google_jwt(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-litellm-api-key", self.LITELLM_JWT.encode()), + (b"authorization", f"Bearer {self.GOOGLE_SERVICE_ACCOUNT_JWT}".encode()), + (b"content-type", b"application/json"), + ], + authenticated=UserAPIKeyAuth(api_key=None, user_id="jwt-subject", jwt_claims=dict(self.LITELLM_JWT_CLAIMS)), + ) + assert raised is None, f"a Google JWT that is not the one that authenticated must keep flowing: {raised}" + assert forwarded is not None + assert forwarded.get("authorization") == f"Bearer {self.GOOGLE_SERVICE_ACCOUNT_JWT}" + assert "x-litellm-api-key" not in forwarded + + @pytest.mark.asyncio + async def test_master_key_in_authorization_alone_is_rejected(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [(b"authorization", b"Bearer sk-master-1234"), (b"content-type", b"application/json")], + authenticated=UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert forwarded is None, "the master key must never reach the upstream forwarder" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_master_key_is_stripped_and_byo_x_goog_api_key_forwards(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"authorization", b"Bearer sk-master-1234"), + (b"x-goog-api-key", b"AIza-real-google-api-key"), + (b"content-type", b"application/json"), + ], + authenticated=UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert raised is None + assert forwarded is not None + assert forwarded.get("x-goog-api-key") == "AIza-real-google-api-key" + assert "authorization" not in forwarded + assert "sk-master-1234" not in " ".join(f"{name}:{value}" for name, value in forwarded.items()) + class TestGetAzureAISearchIndexFromEndpoint: """The operable index is only the segment right after ``indexes``. diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index a48e9e9e17f4..9d1975513a15 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -6,7 +6,6 @@ """ import json -import sys from contextlib import ExitStack from unittest.mock import AsyncMock, MagicMock, patch @@ -66,21 +65,6 @@ def _make_mock_request(): return mock_request -def _ensure_proxy_server_mock(): - """Insert a mock proxy_server module if the real one can't import.""" - key = "litellm.proxy.proxy_server" - if key not in sys.modules: - mock_mod = MagicMock() - mock_mod.proxy_logging_obj = MagicMock() - sys.modules[key] = mock_mod - import litellm.proxy - - if not hasattr(litellm.proxy, "proxy_server"): - litellm.proxy.proxy_server = sys.modules[key] - - -_ensure_proxy_server_mock() - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( pass_through_request, ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 8e973fc37719..961479c03931 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -2,6 +2,7 @@ import pytest +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _base_vertex_proxy_route, ) @@ -323,6 +324,7 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header(): vertex_location="us-central1", base_target_url="https://us-central1-aiplatform.googleapis.com", get_vertex_pass_through_handler=mock_handler, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"), ) # Verify that allowlisted headers are preserved @@ -417,6 +419,7 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): vertex_location="us-central1", base_target_url="https://us-central1-aiplatform.googleapis.com", get_vertex_pass_through_handler=mock_handler, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"), ) # The ONLY Authorization header should be the Vertex token