mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): strip the header used for proxy auth when forwarding provider auth headers
With forward_llm_provider_auth_headers on, clean_headers kept every provider auth header except x-api-key, so a virtual key sent as x-goog-api-key reached the provider. generate_content now forwards client headers, which made this overwrite the deployment key litellm_auth_header_name now follows get_api_key's header precedence, honors litellm_key_header_name, and skips an empty authorization header. clean_headers drops whichever header that resolves to
This commit is contained in:
parent
8ed52b127d
commit
a7c1ddb99e
3 changed files with 225 additions and 10 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue