diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index db4c30500a7..1871048577c 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2337,6 +2337,100 @@ async def test_custom_code_guardrail( ) +_GUARDRAIL_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _metadata_fields(value: object) -> Mapping[str, object]: + try: + return _GUARDRAIL_METADATA_ADAPTER.validate_python(value) + except ValidationError: + return {} + + +_WITHHELD_FIELDS: Final[frozenset[str]] = frozenset({"user_api_key_request_route", "agent_id", "user_api_key_auth"}) + +_PROXY_RESOLVED_IDENTITY_FIELDS: Final[frozenset[str]] = frozenset( + { + "requester_ip_address", + "user_agent", + "headers", + "agent_id", + "actor_agent_id", + "target_agent_id", + "billing_agent_id", + "agent_execution_mode", + "verified_human_user_id", + "user_api_end_user_max_budget", + } +) + +_SECRET_HEADER_MARKERS: Final = ("authorization", "cookie", "key", "token", "secret", "password", "credential") + + +def _is_secret_header(name: str) -> bool: + lowered: Final = name.lower() + return any(marker in lowered for marker in _SECRET_HEADER_MARKERS) + + +def _is_json_safe(value: object) -> bool: + try: + json.dumps(value) + except (TypeError, ValueError): + return False + return True + + +def _forwarded_identity_value(key: str, value: object) -> object: + if key != "headers": + return value + return {name: item for name, item in _metadata_fields(value).items() if not _is_secret_header(name)} + + +def _is_proxy_identity_field(key: str) -> bool: + """Whether a metadata field attributes the authenticated caller rather than the request. + + Pre-call processing owns the whole ``user_api_key_`` namespace, the end user included, and + fills the agent, connection and header fields from the key and the connection right after, + so a caller value never reaches an LLM route. The same fields are reserved here. + """ + return key in _PROXY_RESOLVED_IDENTITY_FIELDS or key == "user_api_key" or key.startswith("user_api_key_") + + +def _guardrail_request_metadata( + caller: object, proxy: object, agent_id: str | None = None +) -> dict[str, object]: # mutable-ok: guardrails annotate the metadata they are handed + """Caller metadata with the identity fields taken from the authenticated key. + + The result is handed to the guardrail as ``request_data["metadata"]``, which providers + write their own detection details into, so it cannot be a read-only mapping. + + The caller's bucket goes through the sanitizer the LLM routes use, so this endpoint drops + the same fields they do, and what that leaves is then overlaid with the identity the proxy + resolved. + + The request route is withheld from both sides. The tool policy guardrail reads the request + itself when a route is present, instead of the tools replayed to it, so a route here turns + its checks off: the caller must not be able to set one, and this endpoint's own route + describes the evaluation rather than the call being evaluated. + + The agent is taken from ``agent_id``, the one the key authenticated with, and omitted when the + key is bound to none. Pre-call processing keeps the request's own agent when the key has none, + so the processed metadata is not a trustworthy source for it either. + """ + from litellm.proxy.litellm_pre_call_utils import strip_caller_supplied_identity + + sanitized: Final = dict(_metadata_fields(caller)) + strip_caller_supplied_identity(sanitized) + caller_fields: Final = tuple((key, value) for key, value in sanitized.items() if not _is_proxy_identity_field(key)) + key_identity: Final = tuple( + (key, _forwarded_identity_value(key, value)) + for key, value in _metadata_fields(proxy).items() + if _is_proxy_identity_field(key) and key not in _WITHHELD_FIELDS and _is_json_safe(value) + ) + bound_agent: Final = (("agent_id", agent_id),) if agent_id is not None else () + return dict((*caller_fields, *key_identity, *bound_agent)) + + def _execution_timeout_response(timeout: float) -> TestCustomCodeGuardrailResponse: return TestCustomCodeGuardrailResponse( success=False, @@ -2503,6 +2597,13 @@ async def apply_guardrail( response_handler: Final = GUARDRAIL_RESPONSE_HANDLERS[request.response_format] request_data: Final = dict(response_handler.prepare_request_data(request)) + caller_metadata: Final = request_data.get("metadata") + metadata: Final = _guardrail_request_metadata( + caller_metadata, + _GUARDRAIL_METADATA_ADAPTER.validate_python(data).get("metadata"), + agent_id=user_api_key_dict.invoked_agent_id or user_api_key_dict.agent_id, + ) + request_data["metadata"] = metadata input_type: Final = response_handler.resolve_input_type(request, active_guardrail) guardrailed_inputs: Final = await active_guardrail.apply_guardrail( inputs={"texts": [request.text]}, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 546521f9048..2f7b4541f6c 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -865,6 +865,29 @@ def strip_client_guardrail_information(data: dict[str, object]) -> None: # muta ) +def strip_caller_supplied_identity( + metadata: dict[str, object], # mutable-ok: strips in place on the metadata bucket it is handed + *, + allow_client_message_redaction_opt_out: bool = False, +) -> None: + """Drop from a caller-supplied metadata bucket the fields pre-call processing owns. + + The proxy fills a family of ``user_api_key_*`` fields into one of the two metadata dicts, + so a caller pre-populating any of them on the other dict would have its forged values + surface in guardrails, spend tracking, audit logs and identity resolution. Stripping by + prefix covers ``user_api_key_*`` fields added later without per-key maintenance. + + Shared with ``/apply_guardrail``, which accepts caller metadata the same way and so has to + take identity from the authenticated key rather than from the body. + """ + _strip_untrusted_request_header_controls( + metadata.get("headers"), + allow_client_message_redaction_opt_out=allow_client_message_redaction_opt_out, + ) + for key in [k for k in metadata if k.startswith("user_api_key_") or k in _UNTRUSTED_METADATA_CONTROL_FIELDS]: + metadata.pop(key, None) + + def _strip_router_reserved_metadata( data: dict[str, object], # mutable-ok: strips in place on the request body the pre-call pipeline threads through ) -> None: @@ -2842,14 +2865,10 @@ async def add_litellm_data_to_request( for _meta_key in ("metadata", "litellm_metadata"): _user_meta = data.get(_meta_key) if isinstance(_user_meta, dict): - _strip_untrusted_request_header_controls( - _user_meta.get("headers"), - allow_client_message_redaction_opt_out=(_allow_client_message_redaction_opt_out), + strip_caller_supplied_identity( + _user_meta, + allow_client_message_redaction_opt_out=_allow_client_message_redaction_opt_out, ) - for _k in [ - k for k in _user_meta if k.startswith("user_api_key_") or k in _UNTRUSTED_METADATA_CONTROL_FIELDS - ]: - _user_meta.pop(_k, None) # Strip pricing overrides AFTER the litellm_metadata string-to-dict parse # above, for the same reason as the user_api_key_* strip — JSON-string diff --git a/tests/integration/authorization/test_guardrail_submissions.py b/tests/integration/authorization/test_guardrail_submissions.py index fb095010eb2..bbb891d6bed 100644 --- a/tests/integration/authorization/test_guardrail_submissions.py +++ b/tests/integration/authorization/test_guardrail_submissions.py @@ -233,7 +233,11 @@ def _assert_guardrail_request( assert payload["litellm_version"] is None or isinstance(payload["litellm_version"], str), payload request_data: Final = object_value(payload["request_data"]) if model is None: - assert request_data == {}, payload + assert set(request_data) == {"metadata"}, payload + direct_metadata: Final = object_value(request_data["metadata"]) + assert direct_metadata.get("user_api_key_alias") == key_alias, direct_metadata + assert "user_api_key_auth" not in direct_metadata, direct_metadata + assert "authorization" not in object_value(direct_metadata.get("headers") or {}), direct_metadata else: assert set(request_data).issubset( { diff --git a/tests/integration/observability/test_apply_guardrail_endpoint.py b/tests/integration/observability/test_apply_guardrail_endpoint.py index f50679bd54e..402a2224013 100644 --- a/tests/integration/observability/test_apply_guardrail_endpoint.py +++ b/tests/integration/observability/test_apply_guardrail_endpoint.py @@ -116,7 +116,13 @@ def _assert_guardrail_request(request: Request, text: str, response: httpx.Respo "tool_calls": None, "model": None, }, response.text - assert payload["request_data"] == {}, response.text + direct_request_data: Final = payload["request_data"] + assert isinstance(direct_request_data, dict), response.text + assert set(direct_request_data) == {"metadata"}, response.text + direct_metadata: Final = direct_request_data["metadata"] + assert isinstance(direct_metadata, dict), response.text + assert "user_api_key_auth" not in direct_metadata, response.text + assert "authorization" not in (direct_metadata.get("headers") or {}), response.text request_headers: Final = payload["request_headers"] assert request_headers is None or isinstance(request_headers, dict), response.text if isinstance(request_headers, dict): diff --git a/tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py b/tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py index 2921011d6eb..828062c1c54 100644 --- a/tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py +++ b/tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py @@ -63,7 +63,7 @@ async def test_apply_guardrail_endpoint_returns_correct_response( # Verify the guardrail was called with correct parameters mock_guardrail.apply_guardrail.assert_called_once_with( inputs={"texts": ["Test text with PII"]}, - request_data={}, + request_data={"metadata": {}}, input_type="request", ) @@ -198,5 +198,5 @@ async def test_apply_guardrail_endpoint_without_optional_params(mock_proxy_loggi # Verify the guardrail was called with correct parameters mock_guardrail.apply_guardrail.assert_called_once_with( - inputs={"texts": ["Test text"]}, request_data={}, input_type="request" + inputs={"texts": ["Test text"]}, request_data={"metadata": {}}, input_type="request" ) diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index f5e5cc32357..827d95bac23 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -4,7 +4,7 @@ from collections.abc import Callable from datetime import datetime from types import MappingProxyType from typing import Dict, Final, List, Literal, Optional -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import httpx import pytest @@ -25,6 +25,7 @@ from litellm.proxy.guardrails.guardrail_endpoints import ( TestCustomCodeGuardrailRequest, UpdateGuardrailRequest, _get_guardrails_list_response, + _guardrail_request_metadata, apply_guardrail, approve_guardrail_submission, create_guardrail, @@ -2169,18 +2170,20 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker): } -def _patch_apply_guardrail_env(mocker, guardrail_result): - mock_guardrail = mocker.Mock() +def _patch_apply_guardrail_env( + mocker: MockerFixture, guardrail_result: object +) -> Mock: + mock_guardrail: Final = mocker.Mock() mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result) - mock_registry = mocker.Mock() + mock_registry: Final = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) - mock_logging_obj = mocker.Mock() + mock_logging_obj: Final = mocker.Mock() mock_logging_obj.async_success_handler = AsyncMock() mock_logging_obj.model_call_details = {} - mock_processor = mocker.Mock() + mock_processor: Final = mocker.Mock() mock_processor.common_processing_pre_call_logic = AsyncMock( return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj) ) @@ -2189,7 +2192,7 @@ def _patch_apply_guardrail_env(mocker, guardrail_result): return_value=mock_processor, ) - mock_proxy_logging = mocker.Mock() + mock_proxy_logging: Final = mocker.Mock() mock_proxy_logging.post_call_success_hook = AsyncMock() mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) mocker.patch("litellm.proxy.proxy_server.general_settings", {}) @@ -2483,23 +2486,248 @@ async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker): ) +def test_the_authenticated_key_identity_replaces_the_caller_version() -> None: + metadata: Final = _guardrail_request_metadata( + {"forbidden_topics": ["tax"], "user_api_key_alias": "someone-else"}, + {"user_api_key_alias": "billing-app", "user_api_key_user_id": "u-1"}, + ) + assert metadata == { + "forbidden_topics": ["tax"], + "user_api_key_alias": "billing-app", + "user_api_key_user_id": "u-1", + } + + +def test_the_bare_key_hash_cannot_come_from_the_caller() -> None: + metadata: Final = _guardrail_request_metadata({"user_api_key": "spoof"}, {"user_api_key": "hashed"}) + assert metadata == {"user_api_key": "hashed"} + + +def test_the_end_user_comes_from_the_proxy_not_the_caller() -> None: + """LLM routes resolve the end user and drop whatever the body said, so this endpoint does too.""" + metadata: Final = _guardrail_request_metadata( + {"user_api_key_end_user_id": "someone-elses-customer"}, + {"user_api_key_end_user_id": "customer-42"}, + ) + assert metadata == {"user_api_key_end_user_id": "customer-42"} + + +def test_an_end_user_the_proxy_did_not_resolve_is_dropped() -> None: + metadata: Final = _guardrail_request_metadata({"user_api_key_end_user_id": "someone-elses-customer"}, {}) + assert metadata == {} + + +def test_no_request_route_reaches_the_guardrail() -> None: + """A route makes the tool policy guardrail read the request instead of the replayed tools.""" + metadata: Final = _guardrail_request_metadata( + {"user_api_key_request_route": "/v1/chat/completions", "topic": "tax"}, + {"user_api_key_request_route": "/apply_guardrail"}, + ) + assert metadata == {"topic": "tax"} + + +def test_an_unbound_agent_is_not_taken_from_the_request() -> None: + """Pre-call processing keeps the request's own agent_id when the key has none, so the proxy's + metadata already carries a caller-chosen value and reserving the field is not enough.""" + metadata: Final = _guardrail_request_metadata( + {"agent_id": "someone-elses-agent"}, {"agent_id": "someone-elses-agent"} + ) + assert metadata == {} + + +def test_a_bound_agent_comes_from_the_authenticated_key() -> None: + metadata: Final = _guardrail_request_metadata( + {"agent_id": "someone-elses-agent"}, {"agent_id": "someone-elses-agent"}, agent_id="billing-agent" + ) + assert metadata == {"agent_id": "billing-agent"} + + +def test_identity_fields_outside_the_key_namespace_are_reserved_too() -> None: + """Pre-call processing fills every one of these from the key and the connection, so an LLM + route never carries a caller value for them and neither may this endpoint.""" + metadata: Final = _guardrail_request_metadata( + { + "requester_ip_address": "203.0.113.9", + "user_agent": "forged/1.0", + "agent_id": "someone-elses-agent", + "actor_agent_id": "someone-elses-actor", + "target_agent_id": "someone-elses-target", + "billing_agent_id": "someone-elses-payer", + "agent_execution_mode": "forged", + "verified_human_user_id": "someone-else", + "topic": "tax", + }, + { + "requester_ip_address": "198.51.100.4", + "user_agent": "litellm-cli/2.3", + "actor_agent_id": "support-agent", + "target_agent_id": "billing-agent", + "billing_agent_id": "support-agent", + "agent_execution_mode": "managed", + "verified_human_user_id": "user-7", + }, + agent_id="billing-agent", + ) + assert metadata == { + "topic": "tax", + "requester_ip_address": "198.51.100.4", + "user_agent": "litellm-cli/2.3", + "actor_agent_id": "support-agent", + "target_agent_id": "billing-agent", + "billing_agent_id": "support-agent", + "agent_execution_mode": "managed", + "verified_human_user_id": "user-7", + "agent_id": "billing-agent", + } + + +def test_a_reserved_identity_field_the_proxy_left_empty_is_dropped() -> None: + """Reserving the names is what keeps a forged value out when the proxy resolved none.""" + metadata: Final = _guardrail_request_metadata( + { + "actor_agent_id": "someone-elses-actor", + "target_agent_id": "someone-elses-target", + "billing_agent_id": "someone-elses-payer", + "agent_execution_mode": "forged", + "verified_human_user_id": "someone-else", + "topic": "tax", + }, + {}, + ) + assert metadata == {"topic": "tax"} + + +def test_the_end_user_budget_is_not_a_caller_field() -> None: + metadata: Final = _guardrail_request_metadata( + {"user_api_end_user_max_budget": 999.0, "topic": "tax"}, + {"user_api_end_user_max_budget": 10.0}, + ) + assert metadata == {"topic": "tax", "user_api_end_user_max_budget": 10.0} + + +def test_the_authenticated_key_object_is_left_behind() -> None: + """A guardrail vendor receives this encoded as JSON, which the auth object does not survive.""" + metadata: Final = _guardrail_request_metadata( + {}, + { + "user_api_key_auth": UserAPIKeyAuth(key_alias="billing-app"), + "user_api_key_metadata": {"tags": {"a", "b"}}, + "user_api_key_alias": "billing-app", + }, + ) + assert metadata == {"user_api_key_alias": "billing-app"} + + +def test_credential_headers_do_not_reach_the_guardrail() -> None: + """The header bucket carries provider credentials that the logging redaction does not name.""" + metadata: Final = _guardrail_request_metadata( + {}, + { + "headers": { + "anthropic-api-key": "sk-ant-secret", + "x-amz-security-token": "aws-session-token", + "authorization": "Bearer sk-caller", + "user-agent": "curl/8.7", + } + }, + ) + assert metadata == {"headers": {"user-agent": "curl/8.7"}} + + +def test_a_guardrail_can_annotate_what_it_is_handed() -> None: + """OpenAI moderation and Presidio write their detection details into this mapping.""" + metadata: Final = _guardrail_request_metadata({}, {"user_api_key_alias": "billing-app"}) + metadata["moderation_model"] = "omni-moderation-latest" + assert metadata == {"user_api_key_alias": "billing-app", "moderation_model": "omni-moderation-latest"} + + +class _MetadataRecordingGuardrail(CustomGuardrail): + """A guardrail registered the way the proxy registers one, keeping what the endpoint hands it.""" + + def __init__(self) -> None: + super().__init__(guardrail_name="identity-probe") + self.calls: Final[list[dict[str, object]]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional[object] = None, + ) -> GenericGuardrailAPIInputs: + self.calls.append(dict(request_data)) + return inputs + + +async def _metadata_the_guardrail_was_handed( + caller_metadata: dict[str, object] | None, auth: UserAPIKeyAuth +) -> dict[str, object]: + """Drive the real endpoint with a real registered guardrail and report what reached it. + + A registered guardrail also runs as an ordinary callback on this request, so the endpoint's + own call is the first one it receives.""" + import litellm + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.guardrails.guardrail_endpoints import router + from litellm.proxy.proxy_server import openai_exception_handler + + guardrail: Final = _MetadataRecordingGuardrail() + litellm.logging_callback_manager.add_litellm_callback(guardrail) + body: Final[dict[str, object]] = {"guardrail_name": "identity-probe", "text": "hello"} + if caller_metadata is not None: + body["metadata"] = caller_metadata + app: Final = FastAPI() + app.include_router(router) + app.add_exception_handler(ProxyException, openai_exception_handler) + app.dependency_overrides[user_api_key_auth] = lambda: auth + try: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://guardrails.test" + ) as client: + response: Final = await client.post("/guardrails/apply_guardrail", json=body) + assert response.status_code == 200, response.text + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + assert guardrail.calls, "the endpoint never reached the registered guardrail" + return dict(guardrail.calls[0].get("metadata") or {}) + + @pytest.mark.asyncio -async def test_apply_guardrail_omits_metadata_when_not_sent(mocker): - """Without metadata, request_data stays empty (backward-compatible).""" - mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]}) - - request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello") - await apply_guardrail( - fastapi_request=mocker.Mock(), - request=request, - user_api_key_dict=UserAPIKeyAuth(), +async def test_apply_guardrail_replaces_caller_identity_with_the_authenticated_key() -> None: + """Identity fields come from the proxy's own metadata, never from the caller, so a request + cannot impersonate another key or probe its policy.""" + forwarded: Final = await _metadata_the_guardrail_was_handed( + { + "forbidden_topics": ["tax"], + "user_api_key_alias": "someone-else", + "user_api_key_user_email": "victim@example.com", + "user_api_key": "hash-of-someone-else", + }, + UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"), ) + json.dumps(forwarded) + assert forwarded["forbidden_topics"] == ["tax"] + assert forwarded.get("user_api_key_alias") == "billing-app" + assert "victim@example.com" not in forwarded.values() + assert forwarded.get("user_api_key") != "hash-of-someone-else" - mock_guardrail.apply_guardrail.assert_awaited_once_with( - inputs={"texts": ["hello"]}, - request_data={}, - input_type="request", + +@pytest.mark.asyncio +async def test_apply_guardrail_attaches_key_identity_when_caller_sends_no_metadata() -> None: + """A caller that sends no metadata still arrives with the key's identity attached, because + pre-call processing fills it in.""" + forwarded: Final = await _metadata_the_guardrail_was_handed(None, UserAPIKeyAuth(key_alias="billing-app")) + assert forwarded.get("user_api_key_alias") == "billing-app" + + +@pytest.mark.asyncio +async def test_apply_guardrail_prefers_the_agent_the_key_was_invoked_as() -> None: + """A key bound to an agent can be invoked as another, and the invoked one is who is acting.""" + auth: Final = UserAPIKeyAuth(agent_id="support-agent").model_copy( + update={"invoked_agent_id": "billing-agent"} ) + forwarded: Final = await _metadata_the_guardrail_was_handed(None, auth) + assert forwarded.get("agent_id") == "billing-agent" @pytest.mark.asyncio