fix(proxy): keep the caller's Google token on credential-less Vertex passthrough under custom auth

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
@@ -48,7 +49,7 @@
     create_websocket_passthrough_route,
     websocket_passthrough_request,
 )
-from litellm.proxy.utils import is_known_model
+from litellm.proxy.utils import hash_token, is_known_model
 from litellm.proxy.vector_store_endpoints.utils import (
     assert_proxy_admin_for_vector_store_index_management,
     assert_user_can_access_vector_store,
@@ -1791,28 +1792,17 @@ 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).
+def _operator_configured_caller_key_header_names() -> tuple[str, ...]:
+    """Operator-configured caller-key header names.
 
-    ``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.
+    ``user_api_key_auth`` also reads the caller's key from
+    ``general_settings.litellm_key_header_name`` and from each
+    ``general_settings.pass_through_endpoints`` entry's
+    ``headers.litellm_user_api_key``. Google never consumes either, so both are
+    dropped by name.
     """
     from litellm.proxy.proxy_server import general_settings
 
@@ -1829,42 +1819,34 @@ 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_secret(value: str, user_api_key_dict: UserAPIKeyAuth) -> bool:
+    """Whether a header value is a LiteLLM secret that authenticated this caller.
 
-    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.
+    That is the proxy master key, or a LiteLLM-issued key whose ``hash_token`` is
+    the ``api_key`` ``user_api_key_auth`` resolved (a DB virtual key, or a custom
+    auth returning the key it validated). A value auth only echoed unhashed (a
+    custom auth's own identifier, the no-master-key path returning the raw
+    header) or never consumed (JWT auth resolves ``api_key=None``) is not a
+    LiteLLM secret, so a bring-your-own Google credential in that position is
+    still forwarded.
     """
-    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()
+    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
+    authenticated_key: Final = user_api_key_dict.api_key
+    if authenticated_key is None:
+        return False
+    return hmac.compare_digest(hash_token(normalized).encode(), authenticated_key.encode())
 
 
-def _forwarded_headers_for_credentialless_vertex_passthrough(request: Request) -> Mapping[str, str]:
+def _forwarded_headers_for_credentialless_vertex_passthrough(
+    request: Request, user_api_key_dict: UserAPIKeyAuth
+) -> 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.
@@ -1878,31 +1860,25 @@ def _forwarded_headers_for_credentialless_vertex_passthrough(request: Request) -
     ``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.
+    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
+    LiteLLM secret that authenticated this caller (``_is_authenticated_caller_secret``),
+    which is dropped by value wherever it appears. Matching on what actually
+    authenticated, rather than on header precedence, keeps a caller's own Google
+    token in ``Authorization`` flowing when a custom auth or JWT authenticated
+    them without consuming it. When neither a surviving ``Authorization`` nor
+    ``x-goog-api-key`` remains the request is rejected so the virtual key cannot
+    leak upstream.
     """
     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 +1894,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 +1907,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 +1923,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 +2083,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:

tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py

@@ -16,6 +16,7 @@
 
 
 import litellm
+from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
 from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
     BaseOpenAIPassThroughHandler,
     RouteChecks,
@@ -592,7 +593,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 +618,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 +3343,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"
@@ -3466,6 +3467,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 +3482,12 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
     )
 
     async def _run(
-        self, monkeypatch, headers: list[tuple[bytes, bytes]]
+        self,
+        monkeypatch,
+        headers: list[tuple[bytes, bytes]],
+        authenticated: UserAPIKeyAuth | None = None,
     ) -> tuple[HTTPException | None, dict | None]:
+        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,15 +3520,15 @@ 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:
                 await vertex_proxy_route(
                     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 +3778,64 @@ 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"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "authenticated",
+        [
+            pytest.param(
+                UserAPIKeyAuth(api_key="best-api-key-ever", user_role=LitellmUserRoles.PROXY_ADMIN),
+                id="custom-auth-returning-its-own-identifier",
+            ),
+            pytest.param(UserAPIKeyAuth(api_key=None, user_id="jwt-subject"), id="jwt-auth"),
+            pytest.param(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, 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,
+        )
+        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
+    async def test_master_key_in_authorization_alone_is_rejected(self, monkeypatch):
+        monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-master-1234")
+        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):
+        monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-master-1234")
+        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``.

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