mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
fix(guardrails): scan list items beside a capped parent key
List items keep their parent key in the scanned JSON line, capped to its last 64 characters, so a credential key over a list is still detected while a long key over many items no longer multiplies the scan text
This commit is contained in:
parent
1b6d30fb95
commit
4020206941
4 changed files with 20 additions and 4 deletions
|
|
@ -454,6 +454,8 @@ _CONFIG_SECTION: Final = "litellm-prompt"
|
|||
|
||||
_JSON_STRING: Final = TypeAdapter(str)
|
||||
|
||||
_KEY_CONTEXT: Final = 64
|
||||
|
||||
_ASSIGNMENT_LINE: Final = re.compile(r"[^\s\[#;:=][^:=]*[:=]")
|
||||
|
||||
_SHELL_ASSIGNMENT: Final = re.compile(r"(?P<key>[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P<value>\S+)")
|
||||
|
|
@ -546,7 +548,7 @@ def _json_unescaped(value: str) -> str:
|
|||
|
||||
def _json_line(key: str | None, text: str) -> str:
|
||||
value: Final = json.dumps(text, ensure_ascii=False)
|
||||
return value if key is None else f"{json.dumps(key, ensure_ascii=False)}: {value}"
|
||||
return value if key is None else f"{json.dumps(key[-_KEY_CONTEXT:], ensure_ascii=False)}: {value}"
|
||||
|
||||
|
||||
def _value_in_texts(value: str, texts: Sequence[str]) -> str | None:
|
||||
|
|
|
|||
|
|
@ -98,12 +98,12 @@ def _tool_arguments(arguments: str) -> JsonValue:
|
|||
parse_constant=_finite_number,
|
||||
)
|
||||
)
|
||||
json.dumps(parsed, ensure_ascii=False).encode("utf-8")
|
||||
json.dumps(parsed, ensure_ascii=False).encode("utf-8") # raises for lone surrogates, which stay raw text
|
||||
return parsed
|
||||
|
||||
|
||||
def iter_tool_json_fields(value: object) -> Iterator[tuple[str | None, str]]:
|
||||
"""Yield ``(key, text)`` for every non-empty string leaf; the key is None for list items and a bare root."""
|
||||
"""Yield ``(nearest key, text)`` for every non-empty string leaf; list items take their parent's key."""
|
||||
pending: Final[list[tuple[object, str | None]]] = [(value, None)] # mutable-ok: owned DFS stack avoids recursion
|
||||
while pending:
|
||||
current, key = pending.pop()
|
||||
|
|
@ -120,7 +120,7 @@ def iter_tool_json_fields(value: object) -> Iterator[tuple[str | None, str]]:
|
|||
items = _object_list(current)
|
||||
if items is None:
|
||||
raise ValueError("Invalid tool JSON payload")
|
||||
pending.extend((child, None) for child in reversed(items))
|
||||
pending.extend((child, key) for child in reversed(items))
|
||||
|
||||
|
||||
def iter_tool_json_text(value: object) -> Iterator[str]:
|
||||
|
|
|
|||
|
|
@ -214,6 +214,8 @@ def test_read_only_vendors_receive_tool_and_system_text(
|
|||
|
||||
|
||||
def test_moderation_flag_in_tool_result_blocks(rig: Rig) -> None:
|
||||
rig.wire.drain()
|
||||
response: Final = rig.proxy.request("POST", "/v1/messages", _messages(FLAGGED))
|
||||
assert response.status_code == 403, response.text
|
||||
assert "Violated content safety policy" in response.text
|
||||
assert _provider_calls(rig.wire.drain()) == ()
|
||||
|
|
|
|||
|
|
@ -147,6 +147,18 @@ async def test_pre_call_hook_records_secrets_found_only_in_tool_payloads():
|
|||
assert recorded["masked_entity_count"] == {"AWS Access Key": 1}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("payload", "expected"),
|
||||
(
|
||||
({"password": ["hunter2brahms9x", "safe"]}, {"password": ["[REDACTED]", "safe"]}),
|
||||
(["a8f5f167f44f4964e6c998dee827110c", "safe"], ["[REDACTED]", "safe"]),
|
||||
),
|
||||
ids=("credential-key-over-list", "unkeyed-high-entropy"),
|
||||
)
|
||||
def test_list_items_are_scanned_beside_their_parent_key_or_quoted(payload, expected):
|
||||
assert _guardrail().redact_tool_payload(payload) == expected
|
||||
|
||||
|
||||
def test_long_key_over_a_list_is_scanned_in_bounded_time():
|
||||
payload = {"k" * 10_000: ["v"] * 1000 + [f"AWS_ACCESS_KEY_ID={AWS_KEY}"]}
|
||||
started = time.perf_counter()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue