diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 1bb2665237c..e458c19a41c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -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( diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py index 1f68acbf9b3..fd2e4a27ece 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix_extraction.py @@ -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( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py index 4abf2e92120..577ae974678 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix.py @@ -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.""" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py index d9f27f4f56f..8d9582aa546 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_ovalix_extraction.py @@ -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():