mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(guardrails): forward only what a vendor can encode, without the credential headers
Copying the whole key namespace brought the authenticated key object with it, which no vendor can encode, and the header bucket went as it stood, with anthropic-api-key and x-amz-security-token in clear. Forward only values that survive JSON encoding, drop the headers carrying a credential, and reserve the end user budget a caller could still set.
This commit is contained in:
parent
bf125c77c7
commit
48bcd9178d
5 changed files with 90 additions and 17 deletions
|
|
@ -2347,7 +2347,7 @@ def _metadata_fields(value: object) -> Mapping[str, object]:
|
|||
return {}
|
||||
|
||||
|
||||
_WITHHELD_FIELDS: Final[frozenset[str]] = frozenset({"user_api_key_request_route", "agent_id"})
|
||||
_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(
|
||||
{
|
||||
|
|
@ -2360,9 +2360,34 @@ _PROXY_RESOLVED_IDENTITY_FIELDS: Final[frozenset[str]] = frozenset(
|
|||
"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:
|
||||
"""Whether a header name carries a credential rather than request context."""
|
||||
lowered: Final = name.lower()
|
||||
return any(marker in lowered for marker in _SECRET_HEADER_MARKERS)
|
||||
|
||||
|
||||
def _is_json_safe(value: object) -> bool:
|
||||
"""Whether a value survives the JSON encoding a guardrail vendor puts it through."""
|
||||
try:
|
||||
json.dumps(value)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _forwarded_identity_value(key: str, value: object) -> object:
|
||||
"""The value to forward, with the credential headers taken out of the header bucket."""
|
||||
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.
|
||||
|
|
@ -2384,7 +2409,9 @@ def _guardrail_request_metadata(
|
|||
|
||||
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.
|
||||
resolved. A vendor receives this encoded as JSON, so only JSON-safe values are forwarded,
|
||||
which leaves the authenticated key object behind, and the headers go without the ones
|
||||
carrying a credential.
|
||||
|
||||
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
|
||||
|
|
@ -2401,9 +2428,9 @@ def _guardrail_request_metadata(
|
|||
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, value)
|
||||
(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
|
||||
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))
|
||||
|
|
@ -2581,10 +2608,7 @@ async def apply_guardrail(
|
|||
_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)
|
||||
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]},
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2171,7 +2171,7 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker):
|
|||
|
||||
|
||||
def _patch_apply_guardrail_env(
|
||||
mocker: MockerFixture, guardrail_result: object, processed_data: dict[str, object] | None = None
|
||||
mocker: MockerFixture, guardrail_result: object
|
||||
) -> Mock:
|
||||
mock_guardrail: Final = mocker.Mock()
|
||||
mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result)
|
||||
|
|
@ -2185,7 +2185,7 @@ def _patch_apply_guardrail_env(
|
|||
mock_logging_obj.model_call_details = {}
|
||||
mock_processor: Final = mocker.Mock()
|
||||
mock_processor.common_processing_pre_call_logic = AsyncMock(
|
||||
return_value=(processed_data or {"guardrail_name": "test-guardrail"}, mock_logging_obj)
|
||||
return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj)
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing",
|
||||
|
|
@ -2597,6 +2597,43 @@ def test_a_reserved_identity_field_the_proxy_left_empty_is_dropped() -> None:
|
|||
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"})
|
||||
|
|
@ -2614,7 +2651,7 @@ class _MetadataRecordingGuardrail(CustomGuardrail):
|
|||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[object] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
|
|
@ -2668,6 +2705,7 @@ async def test_apply_guardrail_replaces_caller_identity_with_the_authenticated_k
|
|||
},
|
||||
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()
|
||||
|
|
@ -2685,8 +2723,9 @@ async def test_apply_guardrail_attaches_key_identity_when_caller_sends_no_metada
|
|||
@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")
|
||||
auth.invoked_agent_id = "billing-agent"
|
||||
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"
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue