From 211264519286fc4a1f3bed0d83d938fd50a32d43 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Thu, 5 Nov 2026 12:50:23 -0800 Subject: [PATCH] feat(mcp): validate credential headers when an MCP server is saved --- .../mcp_management_endpoints.py | 59 +++++++ litellm/types/mcp.py | 48 ++++++ .../test_mcp_management_endpoints.py | 159 ++++++++++++++++++ tests/unit/types/test_mcp.py | 56 ++++++ 4 files changed, 322 insertions(+) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f6b836bc4be..e44f1c46772 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -273,6 +273,7 @@ if MCP_AVAILABLE: MCPGatewaySessionsResponse, MCPGatewaySessionsTerminateResponse, MCPUpstreamProtocol, + credential_headers_problem, normalize_upstream_header_name, validate_mcp_protocol_transport, ) @@ -326,10 +327,60 @@ if MCP_AVAILABLE: }, ) + def _credential_headers_error(payload: McpServerPayloadLike) -> str | None: + raw_credentials: Final[object] = getattr(payload, "credentials", None) + credentials: Final = cast(Mapping[str, object] | None, raw_credentials) # cast-ok: request credentials mapping + has_headers_key: Final = credentials is not None and "headers" in credentials + raw_headers: Final[object] = credentials.get("headers") if credentials is not None else None + headers: Final = cast(Mapping[str, str] | None, raw_headers) # cast-ok: typed by MCPCredentials + raw_auth_type: Final[object] = getattr(payload, "auth_type", None) + auth_type: Final = cast(MCPAuth | str | None, raw_auth_type) # cast-ok: request auth type + is_create: Final = isinstance(payload, NewMCPServerRequest) + if has_headers_key and headers is None: + return "credentials.headers cannot be null" + if headers is None: + return ( + "auth_type secret_headers requires credentials.headers" + if is_create and auth_type == MCPAuth.secret_headers + else None + ) + if auth_type != MCPAuth.secret_headers and (is_create or auth_type is not None): + return "credentials.headers is only used with auth_type secret_headers" + if auth_type == MCPAuth.secret_headers and getattr(payload, "spec_path", None): + return "auth_type secret_headers is not supported for OpenAPI servers" + raw_static: Final[object] = getattr(payload, "static_headers", None) + static_headers: Final = cast(Mapping[str, str] | None, raw_static) # cast-ok: request static headers + return credential_headers_problem(headers, static_headers) + + def credential_headers_update_error( + payload: UpdateMCPServerRequest, stored_auth_type: MCPAuth | str | None, stored_spec_path: str | None + ) -> str | None: + """What a partial update gets wrong about secret headers once the stored server is known.""" + headers: Final = (payload.credentials or {}).get("headers") + switching_in: Final = payload.auth_type == MCPAuth.secret_headers and stored_auth_type != MCPAuth.secret_headers + effective_auth_type: Final = payload.auth_type or stored_auth_type + if headers is not None and effective_auth_type != MCPAuth.secret_headers: + return "credentials.headers is only used with auth_type secret_headers" + if effective_auth_type != MCPAuth.secret_headers: + return None + if switching_in and headers is None: + return "auth_type secret_headers requires credentials.headers" + return ( + "auth_type secret_headers is not supported for OpenAPI servers" + if payload.spec_path or stored_spec_path + else None + ) + + def _validate_credential_headers(payload: McpServerPayloadLike) -> None: + problem: Final = _credential_headers_error(payload) + if problem: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail={"error": problem}) + def validate_and_normalize_mcp_server_payload(payload: McpServerPayloadLike) -> None: _base_validate_and_normalize_mcp_server_payload(payload) _validate_mcp_server_name_fields(payload) _validate_upstream_token_header(payload) + _validate_credential_headers(payload) def mcp_identifier_conflict_message(conflict: McpIdentifierConflict) -> str: return ( @@ -3260,6 +3311,14 @@ if MCP_AVAILABLE: _validate_mcp_protocol_update(payload, payload_fields_set, old_server_record, old_server_record_read_failed) + headers_problem: Final = ( + None + if old_server_record is None + else credential_headers_update_error(payload, old_server_record.auth_type, old_server_record.spec_path) + ) + if headers_problem: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail={"error": headers_problem}) + if payload.per_server_oauth_discovery and (old_server_record is not None or old_server_record_read_failed): relay_eligible: Final = old_server_record is not None and is_per_server_oauth_discovery_eligible( payload.auth_type if "auth_type" in payload_fields_set else old_server_record.auth_type, diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index d01d630ff20..67490526278 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -351,6 +351,54 @@ def same_header(name: str, other: str) -> bool: return name.lower() == other.lower() +_HOP_BY_HOP_HEADERS: Final = frozenset( + ( + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + ) +) +_CLIENT_MANAGED_HEADERS: Final = _HOP_BY_HOP_HEADERS | frozenset( + ("accept", "content-type", "content-length", "host", "cookie", "last-event-id") +) +_CLIENT_MANAGED_PREFIX: Final = "mcp-" + + +def credential_headers_problem(headers: Mapping[str, str], static_headers: Mapping[str, str] | None) -> str | None: + """The first reason ``headers`` cannot be sent upstream, or None.""" + if not headers: + return "credentials.headers must name at least one header" + for name, value in headers.items(): + if normalize_upstream_header_name(name) != name: + return f"Invalid header name {name!r}: must be a valid HTTP header name (RFC 7230 token)" + if ( + not value.strip() + or value != value.strip() + or any(char != "\t" and not " " <= char <= "~" for char in value) + ): + return f"Invalid value for header {name!r}: must be printable ASCII with no leading or trailing space" + lowered: Final = [name.lower() for name in headers] + if len(set(lowered)) != len(lowered): + return "credentials.headers names must be unique, ignoring case" + reserved: Final = next( + ( + name + for name in headers + if name.lower() in _CLIENT_MANAGED_HEADERS or name.lower().startswith(_CLIENT_MANAGED_PREFIX) + ), + None, + ) + if reserved is not None: + return f"Header {reserved!r} is set by the gateway and cannot be configured" + clash: Final = next((name for name in headers if has_header(static_headers, name)), None) + return None if clash is None else f"Header {clash!r} is in both static_headers and credentials.headers" + + def has_header(headers: Mapping[str, str] | None, name: str) -> bool: """Whether ``headers`` carries ``name`` under any casing.""" return bool(headers) and any(same_header(key, name) for key in headers or {}) diff --git a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py index 4b1415c8879..a35b7145b71 100644 --- a/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -7622,6 +7622,32 @@ async def test_edit_mcp_server_dcr_bridge_on_unknown_server_returns_404_not_400( update_mock.assert_called_once() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stored_auth_type,fields,status_code", + [("api_key", {"auth_type": "secret_headers"}, 400), ("secret_headers", {"description": "renamed"}, 404)], +) +async def test_edit_mcp_server_checks_secret_headers_against_the_stored_server( + stored_auth_type: str, fields: Mapping[str, str], status_code: int +) -> None: + from litellm.proxy._types import UpdateMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import edit_mcp_server + + old_record = MagicMock() + old_record.auth_type = stored_auth_type + old_record.spec_path = None + update_mock = AsyncMock(return_value=None) + p1, p2, p3, p4, p5 = _edit_endpoint_patches(old_record, update_mock) + with p1, p2, p3, p4, p5: + payload: Final = UpdateMCPServerRequest(server_id="srv-1", **fields) + user_auth = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as exc: + await edit_mcp_server(payload=payload, user_api_key_dict=user_auth) + + assert exc.value.status_code == status_code + assert update_mock.call_count == (0 if status_code == 400 else 1) + + class TestPerUserCredentialConfigServerResolution: """Per-user credential and env-var endpoints must resolve config-defined MCP servers, which live only in the in-memory registry and never get a DB row, so @@ -11791,3 +11817,136 @@ async def test_mcp_register_application_type_reaches_upstream_or_is_rejected( assert posted["application_type"] == application_type finally: mgmt_endpoints._temporary_mcp_servers.pop(server.server_id, None) + + +_DD_URL: Final = "https://upstream.example/mcp" +_DD_CREDENTIAL_HEADERS: Final = {"DD-API-KEY": "dd-api-secret", "DD-APPLICATION-KEY": "dd-app-secret"} +_DD_STATIC_HEADERS: Final = {"X-Datadog-Site": "datadoghq.com"} + + +@pytest.mark.parametrize( + "payload,error", + [ + ( + NewMCPServerRequest(server_name="dd", url=_DD_URL, auth_type=MCPAuth.secret_headers), + "auth_type secret_headers requires credentials.headers", + ), + ( + NewMCPServerRequest( + server_name="dd", url=_DD_URL, auth_type=MCPAuth.secret_headers, credentials={"headers": {}} + ), + "credentials.headers must name at least one header", + ), + ( + NewMCPServerRequest( + server_name="dd", url=_DD_URL, auth_type=MCPAuth.api_key, credentials={"headers": {"DD-API-KEY": "key"}} + ), + "credentials.headers is only used with auth_type secret_headers", + ), + ( + NewMCPServerRequest(server_name="dd", url=_DD_URL, credentials={"headers": dict(_DD_CREDENTIAL_HEADERS)}), + "credentials.headers is only used with auth_type secret_headers", + ), + ( + UpdateMCPServerRequest(server_id="srv-dd", credentials={"headers": None}), + "credentials.headers cannot be null", + ), + ( + NewMCPServerRequest( + server_name="dd", + spec_path="https://spec.example/openapi.json", + auth_type=MCPAuth.secret_headers, + credentials={"headers": dict(_DD_CREDENTIAL_HEADERS)}, + ), + "auth_type secret_headers is not supported for OpenAPI servers", + ), + ( + UpdateMCPServerRequest(server_id="srv-dd", credentials={"headers": {"DD-API-KEY": "a\nb"}}), + "Invalid value for header 'DD-API-KEY': must be printable ASCII with no leading or trailing space", + ), + ( + NewMCPServerRequest( + server_name="dd", + url=_DD_URL, + auth_type=MCPAuth.secret_headers, + credentials={"headers": {"DD-API-KEY": "key"}}, + static_headers={"dd-api-key": "old"}, + ), + "Header 'DD-API-KEY' is in both static_headers and credentials.headers", + ), + ], +) +def test_saving_unusable_secret_headers_is_rejected( + payload: NewMCPServerRequest | UpdateMCPServerRequest, error: str +) -> None: + with pytest.raises(HTTPException) as rejected: + assert mgmt_endpoints.validate_and_normalize_mcp_server_payload(payload) is None + assert (rejected.value.status_code, rejected.value.detail) == (400, {"error": error}) + + +@pytest.mark.parametrize( + "payload", + [ + NewMCPServerRequest( + server_name="dd", + url=_DD_URL, + auth_type=MCPAuth.secret_headers, + credentials={"headers": dict(_DD_CREDENTIAL_HEADERS)}, + static_headers=dict(_DD_STATIC_HEADERS), + ), + UpdateMCPServerRequest(server_id="srv-dd", auth_type=MCPAuth.secret_headers), + UpdateMCPServerRequest(server_id="srv-dd", description="renamed"), + UpdateMCPServerRequest(server_id="srv-dd", credentials={"headers": dict(_DD_CREDENTIAL_HEADERS)}), + ], +) +def test_saving_usable_or_unchanged_secret_headers_is_accepted( + payload: NewMCPServerRequest | UpdateMCPServerRequest, +) -> None: + assert mgmt_endpoints.validate_and_normalize_mcp_server_payload(payload) is None + + +@pytest.mark.parametrize( + "payload,stored_auth_type,stored_spec_path,error", + [ + ( + UpdateMCPServerRequest(server_id="srv-dd", auth_type=MCPAuth.secret_headers), + MCPAuth.api_key, + None, + "auth_type secret_headers requires credentials.headers", + ), + ( + UpdateMCPServerRequest(server_id="srv-dd", credentials={"headers": dict(_DD_CREDENTIAL_HEADERS)}), + MCPAuth.api_key, + None, + "credentials.headers is only used with auth_type secret_headers", + ), + ( + UpdateMCPServerRequest(server_id="srv-dd", description="renamed"), + MCPAuth.secret_headers, + "https://spec.example/openapi.json", + "auth_type secret_headers is not supported for OpenAPI servers", + ), + (UpdateMCPServerRequest(server_id="srv-dd", description="renamed"), MCPAuth.secret_headers, None, None), + ( + UpdateMCPServerRequest(server_id="srv-dd", auth_type=MCPAuth.secret_headers), + MCPAuth.secret_headers, + None, + None, + ), + ( + UpdateMCPServerRequest( + server_id="srv-dd", auth_type=MCPAuth.secret_headers, credentials={"headers": dict(_DD_CREDENTIAL_HEADERS)} + ), + MCPAuth.api_key, + None, + None, + ), + (UpdateMCPServerRequest(server_id="srv-dd", description="renamed"), MCPAuth.api_key, None, None), + ], +) +def test_updating_a_stored_server_checks_its_secret_headers_against_what_is_stored( + payload: UpdateMCPServerRequest, stored_auth_type: MCPAuth, stored_spec_path: str | None, error: str | None +) -> None: + assert ( + mgmt_endpoints.credential_headers_update_error(payload, stored_auth_type, stored_spec_path) == error + ) diff --git a/tests/unit/types/test_mcp.py b/tests/unit/types/test_mcp.py index 5450ec4aa48..127bd79d8a9 100644 --- a/tests/unit/types/test_mcp.py +++ b/tests/unit/types/test_mcp.py @@ -6,9 +6,12 @@ either stack because a second implementation is exactly how an injected header c resolved credential on one path and not the other. """ +from collections.abc import Mapping + import pytest from litellm.types.mcp import ( + credential_headers_problem, credential_redirect_hook, crosses_origin, has_header, @@ -85,3 +88,56 @@ async def test_the_hook_drops_the_slot_only_once_the_origin_changes() -> None: foreign = httpx.Request("GET", "https://attacker.example.com/x", headers={"esb-oauth": "Bearer x"}) await hook(foreign) assert "esb-oauth" not in foreign.headers + + +@pytest.mark.parametrize( + "headers,static_headers,problem", + [ + ({"DD-API-KEY": "key", "DD-APPLICATION-KEY": "app"}, {"X-Datadog-Site": "datadoghq.com"}, None), + ({}, None, "credentials.headers must name at least one header"), + ( + {"Bad Name": "v"}, + None, + "Invalid header name 'Bad Name': must be a valid HTTP header name (RFC 7230 token)", + ), + ( + {" DD-API-KEY": "v"}, + None, + "Invalid header name ' DD-API-KEY': must be a valid HTTP header name (RFC 7230 token)", + ), + ( + {"DD-API-KEY": ""}, + None, + "Invalid value for header 'DD-API-KEY': must be printable ASCII with no leading or trailing space", + ), + ({"DD-API-KEY": " "}, None, "Invalid value for header 'DD-API-KEY': must be printable ASCII with no leading or trailing space"), + ({"DD-API-KEY": "caf\u00e9"}, None, "Invalid value for header 'DD-API-KEY': must be printable ASCII with no leading or trailing space"), + ({"DD-API-KEY": "a\tb c"}, None, None), + ({"DD-API-KEY": " key"}, None, "Invalid value for header 'DD-API-KEY': must be printable ASCII with no leading or trailing space"), + ({"DD-API-KEY": "key\t"}, None, "Invalid value for header 'DD-API-KEY': must be printable ASCII with no leading or trailing space"), + ({"Accept": "text/plain"}, None, "Header 'Accept' is set by the gateway and cannot be configured"), + ({"content-type": "text/plain"}, None, "Header 'content-type' is set by the gateway and cannot be configured"), + ({"Mcp-Method": "x"}, None, "Header 'Mcp-Method' is set by the gateway and cannot be configured"), + ({"MCP-Param-Region": "x"}, None, "Header 'MCP-Param-Region' is set by the gateway and cannot be configured"), + ({"Last-Event-ID": "1"}, None, "Header 'Last-Event-ID' is set by the gateway and cannot be configured"), + ({"Cookie": "a=b"}, None, "Header 'Cookie' is set by the gateway and cannot be configured"), + ({"Upgrade": "h2c"}, None, "Header 'Upgrade' is set by the gateway and cannot be configured"), + ({"Proxy-Authorization": "x"}, None, "Header 'Proxy-Authorization' is set by the gateway and cannot be configured"), + ({"User-Agent": "acme"}, None, None), + ( + {"DD-API-KEY": "key\r\nX-Injected: 1"}, + None, + "Invalid value for header 'DD-API-KEY': must be printable ASCII with no leading or trailing space", + ), + ({"DD-API-KEY": "a", "dd-api-key": "b"}, None, "credentials.headers names must be unique, ignoring case"), + ( + {"DD-API-KEY": "a"}, + {"dd-api-key": "old"}, + "Header 'DD-API-KEY' is in both static_headers and credentials.headers", + ), + ], +) +def test_credential_headers_problem_names_the_first_unusable_header( + headers: Mapping[str, str], static_headers: Mapping[str, str] | None, problem: str | None +) -> None: + assert credential_headers_problem(headers, static_headers) == problem