diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 067de31228a..1528579770d 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -113,7 +113,6 @@ from litellm.proxy.common_request_processing import ( from litellm.proxy.common_utils.callback_utils import ( GUARDRAIL_SCAN_IDS_METADATA_KEY, GUARDRAIL_SCAN_METADATA_METADATA_KEY, - TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY, get_logging_caching_headers, # pyright: ignore[reportUnknownVariableType] # untyped request_data param ) from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id @@ -199,14 +198,12 @@ _GUARDRAIL_HEADER_METADATA_FIELDS: Final = frozenset( "applied_guardrails", "applied_policies", "policy_sources", - "semantic-similarity", GUARDRAIL_SCAN_IDS_METADATA_KEY, GUARDRAIL_SCAN_METADATA_METADATA_KEY, - TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY, - "pillar_response_headers", - "pillar_flagged", ) ) +# h11's header field-value grammar: no NUL, line breaks or other whitespace beyond spaces and tabs between characters. +_SENDABLE_HEADER_VALUE: Final = re.compile(rb"(?:[^\x00\s]+(?:[ \t]+[^\x00\s]+)*)?") def _is_mergeable_client_metadata(key: str, value: object) -> bool: @@ -240,8 +237,10 @@ def _merged_metadata_entry(prior: object, value: object) -> object: ) -def _encodes_as_latin1(value: str) -> bool: - return all(ord(character) < 256 for character in value) +def _is_sendable_header_value(value: str) -> bool: + return all(ord(character) < 256 for character in value) and bool( + _SENDABLE_HEADER_VALUE.fullmatch(value.encode("latin-1")) + ) def _guardrail_response_headers( @@ -270,7 +269,7 @@ def _guardrail_response_headers( for field, entries in groupby(sorted(recorded, key=lambda entry: entry[0]), key=lambda entry: entry[0]) } headers: Final = _RESPONSE_HEADERS.validate_python(get_logging_caching_headers({"metadata": merged}) or {}) - return MappingProxyType({name: value for name, value in headers.items() if _encodes_as_latin1(value)}) + return MappingProxyType({name: value for name, value in headers.items() if _is_sendable_header_value(value)}) class _RegisteredPassThroughRoute(TypedDict): diff --git a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py index b6b40547d8b..b4c05f6a1ef 100644 --- a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py +++ b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py @@ -21,6 +21,9 @@ POST_CALL: Final = "passthrough-header-post" NON_LATIN1: Final = "passthrough-header-pre\u2014em" BLOCKER: Final = "passthrough-header-block" MISSING_MODEL: Final = "passthrough-header-missing-model" +PILLAR: Final = "passthrough-header-pillar" +TAG_POLICY: Final = "passthrough-header-tag-policy" +POLICY_TAG_PREFIX: Final = "passthrough-header-tag-" HEADER: Final = "x-litellm-applied-guardrails" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" @@ -47,6 +50,10 @@ def _block(_request: Request) -> Reply: return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode()) +def _pillar_allow(_request: Request) -> Reply: + return Reply(body=json.dumps({"flagged": False, "session_id": "pillar-vendor-session"}).encode()) + + def _openai_reply(request: Request) -> Reply: payload: Final = _JSON.validate_json(request.body) if payload.get("model") == MISSING_MODEL: @@ -84,6 +91,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I with ExitStack() as stack: vendor: Final = stack.enter_context(wire_server(_allow)) blocking_vendor: Final = stack.enter_context(wire_server(_block)) + pillar_vendor: Final = stack.enter_context(wire_server(_pillar_allow)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -114,7 +122,21 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I (NON_LATIN1, "pre_call", vendor.url), (BLOCKER, "pre_call", blocking_vendor.url), ) + ] + + [ + { + "guardrail_name": PILLAR, + "litellm_params": { + "guardrail": "pillar", + "mode": "pre_call", + "default_on": False, + "api_base": pillar_vendor.url, + "api_key": "synthetic-pillar-key", + }, + } ], + "policies": {TAG_POLICY: {"guardrails": {"add": [PRE_CALL]}}}, + "policy_attachments": [{"policy": TAG_POLICY, "tags": [f"{POLICY_TAG_PREFIX}*"]}], } ) ) @@ -240,3 +262,33 @@ def test_non_list_client_claim_does_not_fail_a_guarded_request(rig: HeaderRig, c assert response.status_code == 200, response.text assert response.headers.get(HEADER) == PRE_CALL + + +def test_pillar_headers_stay_off_pass_through_responses(rig: HeaderRig) -> None: + gateway: Final = rig.owned.gateway + forged: Final[dict[str, JsonValue]] = {"pillar_session_id_response": "client-session", "pillar_scanners": {"x": 1}} + with gateway.scenario() as scenario: + response: Final = _send( + gateway, scenario, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE), "metadata": forged}, PILLAR + ) + + assert response.status_code == 200, response.text + assert response.headers.get(HEADER) == PILLAR + assert [name for name, _ in response.headers.multi_items() if name.startswith("x-pillar-")] == [] + + +@pytest.mark.parametrize("suffix", ["team\n", "x\r\nX-Injected: yes"]) +def test_header_value_with_a_line_break_is_left_out(rig: HeaderRig, suffix: str) -> None: + gateway: Final = rig.owned.gateway + body: Final[dict[str, JsonValue]] = { + **_route_body(OPENAI_ROUTE), + "metadata": {"tags": [f"{POLICY_TAG_PREFIX}{suffix}"]}, + } + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, body) + + assert response.status_code == 200, response.text + assert response.headers.get(HEADER) == PRE_CALL + assert response.headers.get("x-litellm-applied-policies") == TAG_POLICY + assert "x-litellm-policy-sources" not in response.headers + assert "x-injected" not in response.headers diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 104dd5beb47..d007efdc93e 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -17,6 +17,7 @@ from typing import Final, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch from uuid import UUID +import h11 import httpx import pytest import respx @@ -51,6 +52,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( SafeRouteAdder, _build_passthrough_failure_request_payload, _registered_pass_through_routes, + _guardrail_response_headers, _strip_client_guardrail_fields, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, @@ -12069,4 +12071,35 @@ class TestGuardrailResponseHeaders: _strip_client_guardrail_fields(body) - assert body == {"metadata": {"tags": ["team-a"]}, "litellm_metadata": {}} + assert body == {"metadata": {"pillar_flagged": True, "tags": ["team-a"]}, "litellm_metadata": {}} + + def test_pillar_fields_recorded_on_the_request_add_no_pillar_headers(self): + headers: Final = _guardrail_response_headers( + {"metadata": {"pillar_session_id_response": "client-session"}}, + { + "litellm_params": { + "metadata": { + "applied_guardrails": ["pillar"], + "pillar_flagged": False, + "pillar_response_headers": {"x-pillar-flagged": "false", "x-pillar-session-id": "vendor"}, + "_pillar_response_headers_trusted": True, + } + } + }, + ) + + assert headers == {"x-litellm-applied-guardrails": "pillar"} + + @pytest.mark.parametrize("tag", ["prod-team\n", "prod-x\r\nX-Injected: yes", "prod\x00"]) + def test_header_value_an_http_response_cannot_carry_is_left_out(self, tag: str): + headers: Final = _guardrail_response_headers( + {}, + { + "litellm_params": { + "metadata": {"applied_policies": ["tag-policy"], "policy_sources": {"tag-policy": f"tag:{tag}"}} + } + }, + ) + + assert headers == {"x-litellm-applied-policies": "tag-policy"} + assert h11.Response(status_code=200, headers=[*headers.items()]).status_code == 200