mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(cookbook): scan structured content parts in ATR guardrail
Previously the hook only scanned string content and silently skipped
list-shaped content, so a caller could split a blocked payload across
multiple {"type": "text"} parts and bypass detection. Add a small
_extract_text helper that flattens both legacy strings and OpenAI
structured content parts into a single string before regex matching.
Non-text parts (image_url, input_audio, etc.) are ignored.
Smoke-tested locally on a split-payload case ("ignore previous " +
"instructions") which now matches ATR-PI-001 instead of bypassing.
Closes the medium-severity finding from veria-ai's review on this PR.
This commit is contained in:
parent
6f6d046d5c
commit
91d285b19e
2 changed files with 37 additions and 4 deletions
|
|
@ -20,6 +20,12 @@ common categories: instruction override, system prompt exfiltration,
|
|||
role-play jailbreak, base64-wrapped payloads, MCP tool override, and
|
||||
`file://` SSRF references. The full ruleset lives in the ATR repository.
|
||||
|
||||
The hook scans both legacy string `content` and OpenAI structured
|
||||
content parts (`[{"type": "text", "text": "..."}]`), concatenating
|
||||
the text fields before pattern matching. This keeps the screen
|
||||
intact when an attacker splits a payload across multiple text parts.
|
||||
Non-text parts (image_url, input_audio, etc.) are ignored.
|
||||
|
||||
## Wire it up
|
||||
|
||||
Add a guardrail entry to your proxy config:
|
||||
|
|
|
|||
|
|
@ -78,12 +78,39 @@ ATR_INSPIRED_PATTERNS = [
|
|||
|
||||
|
||||
class ATRDetectionGuardrail(CustomGuardrail):
|
||||
"""Block requests that hit any ATR-inspired threat pattern."""
|
||||
"""Block requests that hit any ATR-inspired threat pattern.
|
||||
|
||||
Scans both legacy string content and OpenAI structured content parts
|
||||
(list of {"type": "text", "text": "..."} entries), so a payload split
|
||||
across multiple text parts is concatenated before pattern matching.
|
||||
Non-text parts (image_url, input_audio, etc.) are ignored.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.optional_params = kwargs
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@staticmethod
|
||||
def _extract_text(content) -> str:
|
||||
"""Return scannable text from any OpenAI-shaped message content.
|
||||
|
||||
Supports: string, list of strings, list of structured parts
|
||||
(dicts with type == "text" and a "text" string field). Anything
|
||||
else returns "".
|
||||
"""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts = []
|
||||
for part in content:
|
||||
if isinstance(part, str):
|
||||
parts.append(part)
|
||||
elif isinstance(part, dict):
|
||||
if part.get("type") == "text" and isinstance(part.get("text"), str):
|
||||
parts.append(part["text"])
|
||||
return "\n".join(parts)
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def _scan(text: str):
|
||||
for rule_id, label, pattern in ATR_INSPIRED_PATTERNS:
|
||||
|
|
@ -99,10 +126,10 @@ class ATRDetectionGuardrail(CustomGuardrail):
|
|||
call_type: CallTypesLiteral,
|
||||
) -> Optional[Union[Exception, str, dict]]:
|
||||
for message in data.get("messages") or []:
|
||||
content = message.get("content")
|
||||
if not isinstance(content, str):
|
||||
text = self._extract_text(message.get("content"))
|
||||
if not text:
|
||||
continue
|
||||
hit = self._scan(content)
|
||||
hit = self._scan(text)
|
||||
if hit is None:
|
||||
continue
|
||||
rule_id, label = hit
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue