diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index db4c30500a7..f14e24d3383 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2337,6 +2337,68 @@ 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 {} + + +_CALLER_OWNED_KEY_FIELDS: Final[frozenset[str]] = frozenset({"user_api_key_end_user_id"}) + +_WITHHELD_FIELDS: Final[frozenset[str]] = frozenset({"user_api_key_request_route", "agent_id"}) + +_RESERVED_IDENTITY_FIELDS: Final[frozenset[str]] = frozenset( + {"requester_ip_address", "user_agent", "agent_id", "actor_agent_id"} +) + + +def _is_key_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, except the end user the call + is made for, which the caller owns here as it does on every LLM route. + + Attribution also reaches outside that namespace, so the fields pre-call processing fills in + from the key and the connection are reserved by name as well. + """ + if key in _CALLER_OWNED_KEY_FIELDS: + return False + return key in _RESERVED_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 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. + """ + caller_fields: Final = tuple( + (key, value) for key, value in _metadata_fields(caller).items() if not _is_key_identity_field(key) + ) + key_identity: Final = tuple( + (key, value) + for key, value in _metadata_fields(proxy).items() + if _is_key_identity_field(key) and key not in _WITHHELD_FIELDS + ) + 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 +2565,16 @@ 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, + ) + if caller_metadata is not None or metadata: + request_data["metadata"] = metadata + else: + request_data.pop("metadata", None) 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/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index f5e5cc32357..4c09f9afb4e 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,27 +2170,29 @@ 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, processed_data: dict[str, object] | None = None +) -> 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) + return_value=(processed_data or {"guardrail_name": "test-guardrail"}, mock_logging_obj) ) mocker.patch( "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", 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,21 +2486,189 @@ 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_caller_owns_the_end_user() -> None: + metadata: Final = _guardrail_request_metadata({"user_api_key_end_user_id": "customer-42"}, {}) + assert metadata == {"user_api_key_end_user_id": "customer-42"} + + +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: + metadata: Final = _guardrail_request_metadata( + { + "requester_ip_address": "203.0.113.9", + "user_agent": "forged/1.0", + "agent_id": "someone-elses-agent", + "topic": "tax", + }, + {"requester_ip_address": "198.51.100.4", "user_agent": "litellm-cli/2.3"}, + agent_id="billing-agent", + ) + assert metadata == { + "topic": "tax", + "requester_ip_address": "198.51.100.4", + "user_agent": "litellm-cli/2.3", + "agent_id": "billing-agent", + } + + +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"} + + @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"]}) +async def test_apply_guardrail_replaces_caller_identity_with_the_authenticated_key(mocker): + """Identity fields come from the proxy's own sanitized metadata, never from + the caller, so a request cannot impersonate another key or probe its policy.""" + mock_guardrail = _patch_apply_guardrail_env( + mocker, + {"texts": ["ok"]}, + processed_data={ + "guardrail_name": "test-guardrail", + "metadata": { + "route": "/apply_guardrail", + "user_api_key_alias": "billing-app", + "user_api_key_user_id": "u-1", + }, + }, + ) + + request = ApplyGuardrailRequest( + guardrail_name="test-guardrail", + text="hello", + metadata={ + "forbidden_topics": ["tax"], + "user_api_key_alias": "someone-else", + "user_api_key_user_email": "victim@example.com", + }, + ) + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"), + ) + + forwarded = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"] + assert forwarded == { + "forbidden_topics": ["tax"], + "user_api_key_alias": "billing-app", + "user_api_key_user_id": "u-1", + } + + +@pytest.mark.asyncio +async def test_apply_guardrail_ignores_a_caller_supplied_user_api_key(mocker): + """The bare key field carries attribution as much as the prefixed ones, so a caller + cannot set it either.""" + mock_guardrail = _patch_apply_guardrail_env( + mocker, + {"texts": ["ok"]}, + processed_data={ + "guardrail_name": "test-guardrail", + "metadata": {"route": "/apply_guardrail", "user_api_key": "hash-of-the-real-key"}, + }, + ) + + request = ApplyGuardrailRequest( + guardrail_name="test-guardrail", + text="hello", + metadata={"forbidden_topics": ["tax"], "user_api_key": "hash-of-someone-else"}, + ) + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"), + ) + + forwarded = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"] + assert forwarded == {"forbidden_topics": ["tax"], "user_api_key": "hash-of-the-real-key"} + + +@pytest.mark.asyncio +async def test_apply_guardrail_attaches_key_identity_when_caller_sends_no_metadata(mocker): + """A caller that sends no metadata still gets the authenticated identity + forwarded, the same shape every LLM route gives a guardrail.""" + mock_guardrail = _patch_apply_guardrail_env( + mocker, + {"texts": ["ok"]}, + processed_data={"guardrail_name": "test-guardrail", "metadata": {"user_api_key_alias": "billing-app"}}, + ) request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello") await apply_guardrail( fastapi_request=mocker.Mock(), request=request, - user_api_key_dict=UserAPIKeyAuth(), + user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app"), + ) + + assert mock_guardrail.apply_guardrail.await_args.kwargs["request_data"] == { + "metadata": {"user_api_key_alias": "billing-app"} + } + + +@pytest.mark.asyncio +async def test_apply_guardrail_sends_the_key_identity_when_the_caller_sends_no_metadata(mocker: MockerFixture) -> None: + """A caller that sends no metadata still arrives with the key's identity attached, because + pre-call processing fills it in. An empty request_data is not a shape the pipeline produces.""" + mock_guardrail = _patch_apply_guardrail_env( + mocker, + {"texts": ["ok"]}, + processed_data={"guardrail_name": "test-guardrail", "metadata": {"user_api_key_alias": "billing-app"}}, + ) + + request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello") + await apply_guardrail( + fastapi_request=mocker.Mock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app"), ) mock_guardrail.apply_guardrail.assert_awaited_once_with( inputs={"texts": ["hello"]}, - request_data={}, + request_data={"metadata": {"user_api_key_alias": "billing-app"}}, input_type="request", )