mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(ovalix): block requests whose inline attachments can't be decoded
Inline images, files and audio whose base64 or data URL failed to decode were dropped from inspection, while some providers decode them permissively and still forward the content. Flag them as undecodable and block the request before any file checkpoint call
This commit is contained in:
parent
f184f9d73f
commit
5ffc268585
4 changed files with 84 additions and 52 deletions
|
|
@ -64,6 +64,9 @@ _NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
|
|||
_FILE_BLOCK_ESCALATION_REASON: Final = (
|
||||
"This message was blocked by Ovalix because file content anonymization isn't possible via LiteLLM"
|
||||
)
|
||||
_UNDECODABLE_FILE_BLOCK_REASON: Final = (
|
||||
"This message was blocked by Ovalix because an attached file couldn't be decoded for inspection"
|
||||
)
|
||||
_TOOL_BLOCK_ESCALATION_REASON: Final = (
|
||||
"This message was blocked by Ovalix because tool call anonymization isn't possible via LiteLLM"
|
||||
)
|
||||
|
|
@ -380,6 +383,8 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
session_id: str,
|
||||
target: CheckpointTarget,
|
||||
) -> str | None:
|
||||
if any(part.undecodable for part in file_parts):
|
||||
return _UNDECODABLE_FILE_BLOCK_REASON
|
||||
for part in sorted(file_parts, key=lambda p: p.message_index, reverse=True):
|
||||
data = await self._file_part_to_data(part)
|
||||
reason = await self._block_reason_for_item(
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ class FilePart(NamedTuple):
|
|||
inline: bool
|
||||
oversize: bool
|
||||
message_index: int
|
||||
undecodable: bool = False
|
||||
|
||||
|
||||
def _split_data_url(value: str) -> tuple[str | None, str | None]:
|
||||
|
|
@ -60,24 +61,28 @@ def _b64decode_or_none(payload: str) -> bytes | None:
|
|||
return None
|
||||
|
||||
|
||||
def _decode_base64_with_limit(b64_payload: str, size_limit: int | None) -> tuple[bytes | None, bool]:
|
||||
cleaned: Final = "".join(b64_payload.split())
|
||||
if not cleaned:
|
||||
return None, False
|
||||
if size_limit is not None and (len(cleaned) * 3) // 4 - 2 > size_limit:
|
||||
return None, True
|
||||
def _decoded_base64(cleaned: str) -> bytes | None:
|
||||
strict: Final = _b64decode_or_none(cleaned)
|
||||
urlsafe: Final = (
|
||||
_b64decode_or_none(cleaned.translate(_URLSAFE_TO_STANDARD_B64))
|
||||
if strict is None and ("-" in cleaned or "_" in cleaned)
|
||||
else None
|
||||
)
|
||||
data: Final = strict if strict is not None else urlsafe
|
||||
if strict is not None or ("-" not in cleaned and "_" not in cleaned):
|
||||
return strict
|
||||
return _b64decode_or_none(cleaned.translate(_URLSAFE_TO_STANDARD_B64))
|
||||
|
||||
|
||||
def _inline_part(name: str | None, value: str, size_limit: int | None, message_index: int) -> FilePart | None:
|
||||
mime_hint, payload = _split_data_url(value)
|
||||
if payload is None:
|
||||
return FilePart(name, None, mime_hint, True, False, message_index, undecodable=True)
|
||||
cleaned: Final = "".join(payload.split())
|
||||
if not cleaned:
|
||||
return None
|
||||
if size_limit is not None and (len(cleaned) * 3) // 4 - 2 > size_limit:
|
||||
return FilePart(name, None, mime_hint, True, True, message_index)
|
||||
data: Final = _decoded_base64(cleaned)
|
||||
if data is None:
|
||||
return None, False
|
||||
return FilePart(name, None, mime_hint, True, False, message_index, undecodable=True)
|
||||
if size_limit is not None and len(data) > size_limit:
|
||||
return None, True
|
||||
return (data, False) if data else (None, False)
|
||||
return FilePart(name, None, mime_hint, True, True, message_index)
|
||||
return FilePart(name, data, mime_hint, True, False, message_index)
|
||||
|
||||
|
||||
def _name_from_url(url: str) -> str | None:
|
||||
|
|
@ -92,13 +97,9 @@ def _part_from_file_block(block: Mapping[str, object], size_limit: int | None, m
|
|||
if file_obj is None:
|
||||
return None
|
||||
name: Final = str_or_none(file_obj.get("filename")) or str_or_none(file_obj.get("file_id"))
|
||||
file_data: Final = file_obj.get("file_data")
|
||||
if isinstance(file_data, str) and file_data:
|
||||
mime_hint, payload = _split_data_url(file_data)
|
||||
data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False)
|
||||
if data is not None or oversize:
|
||||
return FilePart(name, data, mime_hint, True, oversize, message_index)
|
||||
return FilePart(name, None, None, False, False, message_index)
|
||||
file_data: Final = str_or_none(file_obj.get("file_data"))
|
||||
inline: Final = _inline_part(name, file_data, size_limit, message_index) if file_data else None
|
||||
return inline or FilePart(name, None, None, False, False, message_index)
|
||||
|
||||
|
||||
def _part_from_image_url_block(
|
||||
|
|
@ -110,11 +111,7 @@ def _part_from_image_url_block(
|
|||
if not isinstance(url, str) or not url:
|
||||
return None
|
||||
if url.startswith("data:"):
|
||||
mime_hint, payload = _split_data_url(url)
|
||||
data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False)
|
||||
if data is not None or oversize:
|
||||
return FilePart(None, data, mime_hint, True, oversize, message_index)
|
||||
return None
|
||||
return _inline_part(None, url, size_limit, message_index)
|
||||
return FilePart(_name_from_url(url), None, None, False, False, message_index)
|
||||
|
||||
|
||||
|
|
@ -122,12 +119,10 @@ def _part_from_input_file_block(
|
|||
block: Mapping[str, object], size_limit: int | None, message_index: int
|
||||
) -> FilePart | None:
|
||||
declared_name: Final = str_or_none(block.get("filename")) or str_or_none(block.get("file_id"))
|
||||
file_data: Final = block.get("file_data")
|
||||
if isinstance(file_data, str) and file_data:
|
||||
mime_hint, payload = _split_data_url(file_data)
|
||||
data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False)
|
||||
if data is not None or oversize:
|
||||
return FilePart(declared_name, data, mime_hint, True, oversize, message_index)
|
||||
file_data: Final = str_or_none(block.get("file_data"))
|
||||
inline: Final = _inline_part(declared_name, file_data, size_limit, message_index) if file_data else None
|
||||
if inline is not None:
|
||||
return inline
|
||||
file_url: Final = block.get("file_url")
|
||||
name: Final = (
|
||||
_name_from_url(file_url) if isinstance(file_url, str) and file_url and not declared_name else declared_name
|
||||
|
|
@ -145,10 +140,9 @@ def _part_from_input_audio_block(
|
|||
if not isinstance(data_b64, str) or not data_b64:
|
||||
return None
|
||||
name: Final = f"audio.{audio.get('format') or 'bin'}"
|
||||
data, oversize = _decode_base64_with_limit(data_b64, size_limit)
|
||||
if data is None and not oversize:
|
||||
return FilePart(name, None, None, False, False, message_index)
|
||||
return FilePart(name, data, None, True, oversize, message_index)
|
||||
return _inline_part(name, data_b64, size_limit, message_index) or FilePart(
|
||||
name, None, None, False, False, message_index
|
||||
)
|
||||
|
||||
|
||||
_BLOCK_PARSERS: Final[Mapping[str, Callable[[Mapping[str, object], int | None, int], FilePart | None]]] = (
|
||||
|
|
@ -200,11 +194,7 @@ def _file_part_of_image(value: str, size_limit: int | None, index: int) -> FileP
|
|||
if value.startswith(("http://", "https://")):
|
||||
name: Final = _name_from_url(value)
|
||||
return FilePart(name, None, None, False, False, index) if name else None
|
||||
mime_hint, payload = _split_data_url(value)
|
||||
data, oversize = _decode_base64_with_limit(payload, size_limit) if payload else (None, False)
|
||||
if data is None and not oversize:
|
||||
return None
|
||||
return FilePart(None, data, mime_hint, True, oversize, index)
|
||||
return _inline_part(None, value, size_limit, index)
|
||||
|
||||
|
||||
def extract_file_parts_from_images(
|
||||
|
|
|
|||
|
|
@ -1301,6 +1301,42 @@ async def test_file_checkpoint_call_failure_fails_closed():
|
|||
await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||
|
||||
|
||||
_GIF_B64 = base64.b64encode(b"GIF89a\x01\x00\x01\x00\x80\x00\x00\xff\xff\xff\x00\x00\x00;").decode()
|
||||
_CORRUPT_GIF_B64 = f"{_GIF_B64[:8]}!{_GIF_B64[8:]}"
|
||||
|
||||
|
||||
def _user_turn_with(block):
|
||||
return GenericGuardrailAPIInputs(
|
||||
texts=["what is this?"],
|
||||
structured_messages=[{"role": "user", "content": [{"type": "text", "text": "what is this?"}, block]}],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("inputs", "input_type"),
|
||||
[
|
||||
(
|
||||
_user_turn_with({"type": "image_url", "image_url": {"url": f"data:image/gif;base64,{_CORRUPT_GIF_B64}"}}),
|
||||
"request",
|
||||
),
|
||||
(_user_turn_with({"type": "file", "file": {"filename": "a.gif", "file_data": _CORRUPT_GIF_B64}}), "request"),
|
||||
(_user_turn_with({"type": "input_file", "filename": "a.gif", "file_data": _CORRUPT_GIF_B64}), "request"),
|
||||
(
|
||||
_user_turn_with({"type": "input_audio", "input_audio": {"data": _CORRUPT_GIF_B64, "format": "wav"}}),
|
||||
"request",
|
||||
),
|
||||
(GenericGuardrailAPIInputs(texts=[], images=[f"data:image/gif;base64,{_CORRUPT_GIF_B64}"]), "response"),
|
||||
],
|
||||
)
|
||||
async def test_inline_content_that_cannot_be_decoded_blocks_instead_of_skipping_inspection(inputs, input_type):
|
||||
g = _static_guardrail()
|
||||
with patch.object(g._async_handler, "post", new=_post_returning(lambda body: _ALLOW)):
|
||||
with pytest.raises(OvalixGuardrailBlockedException) as exc:
|
||||
await g.apply_guardrail(inputs=inputs, request_data={}, input_type=input_type, logging_obj=None)
|
||||
assert "couldn't be decoded" in str(exc.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discovery_resolved_without_any_checkpoint_raises():
|
||||
"""An application with no checkpoints in any direction is a tracker misconfiguration, so fail closed."""
|
||||
|
|
|
|||
|
|
@ -185,10 +185,10 @@ def test_file_block_urlsafe_base64_decoded_via_fallback():
|
|||
assert len(parts) == 1 and parts[0].data == raw
|
||||
|
||||
|
||||
def test_file_block_data_url_without_base64_marker_has_no_bytes():
|
||||
def test_file_block_data_url_without_base64_marker_is_undecodable():
|
||||
block = {"type": "file", "file": {"filename": "n.txt", "file_data": "data:text/plain,hello"}}
|
||||
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
|
||||
assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "n.txt"
|
||||
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None and parts[0].name == "n.txt"
|
||||
|
||||
|
||||
def test_input_file_block_data_url_decoded():
|
||||
|
|
@ -209,10 +209,10 @@ def test_input_audio_block_decoded():
|
|||
assert len(parts) == 1 and parts[0].data == b"wav" and parts[0].name == "audio.wav" and parts[0].inline
|
||||
|
||||
|
||||
def test_input_audio_block_undecodable_is_name_only():
|
||||
def test_input_audio_block_undecodable_is_flagged():
|
||||
block = {"type": "input_audio", "input_audio": {"data": "!!!not-base64!!!"}}
|
||||
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
|
||||
assert len(parts) == 1 and parts[0].name == "audio.bin" and parts[0].data is None and parts[0].inline is False
|
||||
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None and parts[0].name == "audio.bin"
|
||||
|
||||
|
||||
def test_input_audio_block_non_dict_skipped():
|
||||
|
|
@ -267,9 +267,10 @@ def test_messages_skip_non_dict_and_unknown_blocks():
|
|||
assert len(parts) == 1 and parts[0].name == "a.txt" and parts[0].data == b"x"
|
||||
|
||||
|
||||
def test_image_url_block_invalid_data_url_returns_no_part():
|
||||
def test_image_url_block_invalid_data_url_is_flagged():
|
||||
block = {"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%invalid%%%"}}
|
||||
assert extract_file_parts_from_messages(_msgs(block), size_limit=1000) == ()
|
||||
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
|
||||
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None
|
||||
|
||||
|
||||
def test_tool_message_with_empty_content_is_skipped():
|
||||
|
|
@ -286,10 +287,10 @@ def test_extract_tool_results_skips_non_dict_messages_and_tool_calls():
|
|||
assert extract_tool_results(msgs) == (("f", "ok", "c1"),)
|
||||
|
||||
|
||||
def test_malformed_data_url_yields_no_bytes():
|
||||
def test_malformed_data_url_is_undecodable():
|
||||
block = {"type": "file", "file": {"filename": "x.bin", "file_data": "data:garbage-no-comma"}}
|
||||
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
|
||||
assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "x.bin"
|
||||
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None and parts[0].name == "x.bin"
|
||||
|
||||
|
||||
def test_input_audio_block_without_data_skipped():
|
||||
|
|
@ -324,10 +325,10 @@ def test_whitespace_only_base64_payload_yields_no_bytes():
|
|||
assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "n.txt"
|
||||
|
||||
|
||||
def test_base64_invalid_even_after_urlsafe_translate_yields_no_bytes():
|
||||
def test_base64_invalid_even_after_urlsafe_translate_is_undecodable():
|
||||
block = {"type": "file", "file": {"filename": "b.bin", "file_data": "_@@@"}}
|
||||
parts = extract_file_parts_from_messages(_msgs(block), size_limit=1000)
|
||||
assert len(parts) == 1 and parts[0].data is None and parts[0].inline is False and parts[0].name == "b.bin"
|
||||
assert len(parts) == 1 and parts[0].undecodable and parts[0].data is None and parts[0].name == "b.bin"
|
||||
|
||||
|
||||
def test_file_oversize_detected_after_decode_when_estimate_passes():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue