diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..19c41f2cfc1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1498,12 +1498,30 @@ def _tenant_service_name(user_api_key_dict: UserAPIKeyAuth) -> str | None: ) -def litellm_auth_header_name(headers: Headers) -> str: - if "x-litellm-api-key" in headers: - return "x-litellm-api-key" - if "authorization" in headers: - return "authorization" - return "x-api-key" +_PROXY_AUTH_HEADER_PRECEDENCE: Final = tuple( + header.value.lower() + for header in ( + SpecialHeaders.custom_litellm_api_key, + SpecialHeaders.openai_authorization, + SpecialHeaders.azure_authorization, + SpecialHeaders.anthropic_authorization, + SpecialHeaders.google_ai_studio_authorization, + SpecialHeaders.azure_apim_authorization, + ) +) + + +def litellm_auth_header_name(headers: Headers, litellm_key_header_name: str | None = None) -> str: + if litellm_key_header_name is not None and litellm_key_header_name in headers: + return litellm_key_header_name.lower() + return next( + ( + name + for name in _PROXY_AUTH_HEADER_PRECEDENCE + if name in headers and (name != "authorization" or headers[name]) + ), + "x-api-key", + ) def forwarded_anthropic_oauth_token(headers: Headers) -> str | None: @@ -1560,6 +1578,8 @@ def clean_headers( elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE: if litellm_key_lower and header_lower == litellm_key_lower: continue + if authenticated_with_header is not None and header_lower == authenticated_with_header.lower(): + continue if header_lower == "authorization": continue # Never forward x-litellm-api-key (it's for proxy auth only) @@ -2669,13 +2689,14 @@ async def add_litellm_data_to_request( forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) if not forward_llm_auth: forward_llm_auth = getattr(litellm, "forward_llm_provider_auth_headers", False) + litellm_key_header_name: Final = ( + general_settings.get("litellm_key_header_name") if general_settings is not None else None + ) _headers: Final[dict[str, str]] = clean_headers( request.headers, - litellm_key_header_name=( - general_settings.get("litellm_key_header_name") if general_settings is not None else None - ), + litellm_key_header_name=litellm_key_header_name, forward_llm_provider_auth_headers=forward_llm_auth, - authenticated_with_header=litellm_auth_header_name(request.headers), + authenticated_with_header=litellm_auth_header_name(request.headers, litellm_key_header_name), ) _strip_untrusted_request_header_controls( _headers, diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index a30dbd6e501..0f59daee1b6 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -727,3 +727,63 @@ async def test_proxy_forwarded_headers_reach_generate_content_without_leaking_pr assert wire_headers["x-goog-api-key"] == _UPSTREAM_KEY assert wire_headers["x-cost-center"] == "core-platform" assert "authorization" not in wire_headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("incoming", "expected_upstream_key"), + [ + ({"x-goog-api-key": "sk-litellm-virtual-key", "x-cost-center": "core-platform"}, _UPSTREAM_KEY), + ( + { + "x-litellm-api-key": "sk-litellm-virtual-key", + "x-goog-api-key": "AIza-caller-key", + "x-cost-center": "core-platform", + }, + "AIza-caller-key", + ), + ], + ids=["proxy-key-in-goog-header", "caller-key-beside-proxy-key"], +) +async def test_provider_auth_forwarding_never_sends_proxy_key_to_generate_content( + header_capture_wire: _WireCapture, + incoming: Mapping[str, str], + expected_upstream_key: str, +) -> None: + from starlette.datastructures import Headers + + from litellm.google_genai.main import agenerate_content + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + clean_headers, + litellm_auth_header_name, + ) + + incoming_headers: Final = Headers(incoming) + cleaned: Final = clean_headers( + incoming_headers, + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(incoming_headers), + ) + backend_data: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers=cleaned, + request_data={"model": "gemini/gemini-2.0-flash"}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"), + general_settings={"forward_client_headers_to_llm_api": True}, + ) + + resp: Final = await agenerate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "ping"}]}], + api_base=header_capture_wire.url, + api_key=_UPSTREAM_KEY, + headers=backend_data.get("headers"), + ) + + assert _extract_chunk_texts(resp) == ("pong",) + assert len(header_capture_wire.requests) == 1 + wire_headers: Final = header_capture_wire.requests[0][1] + assert wire_headers["x-goog-api-key"] == expected_upstream_key + assert wire_headers["x-cost-center"] == "core-platform" + assert "x-litellm-api-key" not in wire_headers diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..074aa0c34e8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -4,6 +4,7 @@ import copy import json import os import time +from collections.abc import Mapping from datetime import datetime, timezone from types import SimpleNamespace from typing import Final @@ -58,6 +59,7 @@ from litellm.proxy.litellm_pre_call_utils import ( forwarded_anthropic_oauth_token, get_dynamic_logging_metadata, get_metadata_variable_name, + litellm_auth_header_name, move_guardrails_to_metadata, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -6447,6 +6449,84 @@ def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth assert "x-api-key" not in result +@pytest.mark.parametrize( + ("headers", "litellm_key_header_name", "expected"), + [ + ({"x-litellm-api-key": "sk-proxy", "authorization": "Bearer sk-other"}, None, "x-litellm-api-key"), + ({"authorization": "Bearer sk-proxy", "x-goog-api-key": "AIza-caller"}, None, "authorization"), + ({"authorization": "", "x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"), + ({"api-key": "sk-proxy", "x-api-key": "sk-ant-caller"}, None, "api-key"), + ({"x-api-key": "sk-proxy", "x-goog-api-key": "AIza-caller"}, None, "x-api-key"), + ({"x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"), + ({"ocp-apim-subscription-key": "sk-proxy"}, None, "ocp-apim-subscription-key"), + ({"x-custom-key": "sk-proxy", "authorization": "Bearer sk-other"}, "X-Custom-Key", "x-custom-key"), + ({"authorization": "Bearer sk-proxy"}, "x-custom-key", "authorization"), + ({"content-type": "application/json"}, None, "x-api-key"), + ], + ids=[ + "litellm-key-beats-authorization", + "authorization-beats-goog", + "empty-authorization-skipped", + "azure-beats-anthropic", + "anthropic-beats-goog", + "goog", + "azure-apim", + "custom-header-wins-when-present", + "custom-header-absent", + "no-auth-header", + ], +) +def test_litellm_auth_header_name_follows_user_api_key_auth_precedence( + headers: Mapping[str, str], litellm_key_header_name: str | None, expected: str +) -> None: + assert litellm_auth_header_name(Headers(headers), litellm_key_header_name) == expected + + +@pytest.mark.parametrize( + ("headers", "expected"), + [ + ( + {"x-goog-api-key": "sk-proxy-virtual-key", "x-cost-center": "core-platform"}, + {"x-cost-center": "core-platform"}, + ), + ( + {"x-litellm-api-key": "sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"}, + {"x-goog-api-key": "AIza-caller-key"}, + ), + ( + {"api-key": "sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"}, + {"x-goog-api-key": "AIza-caller-key"}, + ), + ], + ids=["goog-was-proxy-auth", "goog-beside-litellm-key", "goog-beside-azure-proxy-key"], +) +def test_clean_headers_byok_strips_only_the_header_used_for_proxy_auth( + headers: Mapping[str, str], expected: Mapping[str, str] +) -> None: + incoming: Final = Headers(headers) + + result: Final = clean_headers( + incoming, + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(incoming), + ) + + assert result == expected + + +def test_clean_headers_byok_strips_custom_litellm_key_header_named_x_api_key() -> None: + incoming: Final = Headers({"x-api-key": "sk-proxy-virtual-key", "authorization": "Bearer sk-ant-api03-caller"}) + + result: Final = clean_headers( + incoming, + litellm_key_header_name="x-api-key", + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(incoming, "x-api-key"), + ) + + assert "x-api-key" not in result + + # --------------------------------------------------------------------------- # Team guardrail + global policy regression tests # --------------------------------------------------------------------------- @@ -7965,6 +8045,60 @@ async def test_add_litellm_data_to_request_keeps_every_forwarded_credential_out_ assert value not in logged +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_headers", "general_settings", "expected_forwarded_auth"), + [ + ( + {"x-goog-api-key": "sk-virtual-key", "x-cost-center": "core-platform"}, + FORWARD_PROVIDER_AUTH_HEADERS, + {}, + ), + ( + { + "x-litellm-api-key": "Bearer sk-virtual-key", + "x-goog-api-key": "AIza-caller-key", + "x-cost-center": "core-platform", + }, + FORWARD_PROVIDER_AUTH_HEADERS, + {"x-goog-api-key": "AIza-caller-key"}, + ), + ( + {"x-api-key": "sk-virtual-key", "authorization": "Bearer sk-other", "x-cost-center": "core-platform"}, + {**FORWARD_PROVIDER_AUTH_HEADERS, "litellm_key_header_name": "x-api-key"}, + {}, + ), + ], + ids=["goog-was-proxy-auth", "goog-beside-litellm-key", "custom-litellm-key-header"], +) +async def test_add_litellm_data_to_request_never_forwards_the_header_used_for_proxy_auth( + request_headers: Mapping[str, str], + general_settings: Mapping[str, object], + expected_forwarded_auth: Mapping[str, str], +) -> None: + request_mock: Final = _make_request_mock( + "/v1/chat/completions", {"Content-Type": "application/json", **request_headers} + ) + + updated: Final = await add_litellm_data_to_request( + data={"model": "gemini-2.0-flash", "messages": [{"role": "user", "content": "hello"}]}, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings=dict(general_settings), + version="test-version", + ) + + forwarded_auth: Final = { + name: value + for name, value in updated["headers"].items() + if name.lower() in {"x-goog-api-key", "x-api-key", "x-litellm-api-key"} + } + assert forwarded_auth == expected_forwarded_auth + assert updated["headers"]["x-cost-center"] == "core-platform" + assert "sk-virtual-key" not in str(updated.get("api_key")) + + _SERVICE_ACCOUNT_TOKEN_REFERENCE: Final = "oidc/file//var/run/secrets/kubernetes.io/serviceaccount/token"