diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 18797b0a880..9740b5801bf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -50,7 +50,8 @@ HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" # Headroom writes a retrieval hash as either `<>` or # `Retrieve more: hash=HASH`, and its own reader accepts 12 to 24 hex chars # (the row-drop path emits SHA-256[:12]). -_HASH_PATTERN: Final = re.compile(r"(?:< tuple[str, ...]: raw: Final = body.get("ccr_hashes") if not _is_object_list(raw): return () - return tuple(item for item in raw if isinstance(item, str) and item) + return tuple(item for item in raw if isinstance(item, str) and _HASH_VALUE_PATTERN.fullmatch(item)) def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 3f2df420540..f5638f2cbe8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -419,6 +419,38 @@ async def test_apply_guardrail_trusts_ccr_hashes_from_compress_response( assert tool_result["content"] == original_content +@pytest.mark.asyncio +async def test_apply_guardrail_ignores_malformed_ccr_hashes( + guardrail: HeadroomGuardrail, +): + """A hash goes straight into the retrieval URL, so anything that is not a + Headroom hash must not reach the authorized set.""" + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response( + COMPRESSED_MESSAGES, + ccr_hashes=["../../etc/passwd", "", "NOTAHASH", "f3b3d2ef"], + ) + request_data = {"model": "gpt-4o", "litellm_call_id": "call-bad"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert not has_headroom_retrieve_tool(result.get("tools") or []) + assert "call-bad" not in guardrail._issued_hashes_by_call_id + + @pytest.mark.asyncio async def test_apply_guardrail_no_tool_injected_when_no_hashes( guardrail: HeadroomGuardrail,