diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 679ae3daa78..c64f3a9db2a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast import httpx import openai -from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket +from fastapi import APIRouter, Depends, HTTPException, Request, Response, Security, WebSocket from fastapi.responses import StreamingResponse from pydantic import ConfigDict, TypeAdapter, ValidationError from starlette.websockets import WebSocketState @@ -77,6 +77,8 @@ 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 + api_key_header, + custom_litellm_key_header, get_bearer_token, is_no_auth_dev_mode, user_api_key_auth, @@ -3137,6 +3139,16 @@ async def vertex_discovery_proxy_route( ) +async def _authenticate_vertex_caller( + request: Request, + authorization: Annotated[str | None, Security(api_key_header)], + litellm_api_key: Annotated[str | None, Security(custom_litellm_key_header)], +) -> UserAPIKeyAuth: + return await user_api_key_auth( + request=request, api_key=authorization or "", custom_litellm_key_header=litellm_api_key + ) + + @router.api_route( "/vertex-ai/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -3152,7 +3164,7 @@ async def vertex_proxy_route( endpoint: str, request: Request, fastapi_response: Response, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)], ): """ Call LiteLLM proxy via Vertex AI SDK. diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 376b2f18d3b..54874b51478 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -8128,11 +8128,14 @@ class TestAzureSpeechRawBodyThroughRealAuth: class _CountingCustomAuth: - def __init__(self) -> None: + def __init__(self, only_accepts: str | None = None) -> None: + self.only_accepts: Final = only_accepts self.api_keys_seen: list[str] = [] async def __call__(self, request: Request, api_key: str) -> UserAPIKeyAuth: self.api_keys_seen.append(api_key) + if self.only_accepts is not None and api_key != self.only_accepts: + raise ValueError("unknown caller") return UserAPIKeyAuth(user_id=f"custom-auth-caller-{len(self.api_keys_seen)}") @@ -8274,6 +8277,23 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] assert not google.called + @pytest.mark.parametrize("header", ["x-goog-api-key", "x-api-key", "API-Key", "Ocp-Apim-Subscription-Key"]) + def test_key_outside_litellm_headers_is_rejected_and_never_reaches_google_without_a_vertex_credential( + self, monkeypatch: pytest.MonkeyPatch, header: str + ) -> None: + custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY) + with respx.mock(assert_all_called=False) as upstream: + google: Final = upstream.route().mock(return_value=httpx.Response(200)) + + response, recorder = self._post( + monkeypatch, f"/vertex_ai/{self.GENERATE_CONTENT}", {header: self.VIRTUAL_KEY}, custom_auth, () + ) + + assert response.status_code == 401, response.text + assert self.VIRTUAL_KEY not in custom_auth.api_keys_seen + assert recorder.served_user_ids == [] + assert not google.called + class TestTypeSafePassthroughRoute: @staticmethod