diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 4722699b7fa..c32a76a020d 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -24,7 +24,7 @@ from copy import deepcopy from datetime import datetime from inspect import isawaitable -from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NewType, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NamedTuple, NewType, cast, overload from typing_extensions import Required, TypedDict @@ -80,6 +80,44 @@ def _parse_content_list(contents_data: Sequence[Any]) -> list[Content]: # region Internal Helper functions for unified Content +class _ParsedDataUri(NamedTuple): + """The parts of an RFC 2397 data URI.""" + + media_type: str + is_base64: bool + payload: str + + +def _parse_data_uri(uri: str) -> _ParsedDataUri: + """Parse a ``data:`` URI once, so every caller reads it the same way. + + RFC 2397: ``data:[][;parameter[=value]]...[;base64],``. The media type is + everything between ``data:`` and the first ``;`` and defaults to ``text/plain`` when empty. + ``base64`` is the only encoding marker and is only valid once, as the last parameter; other + parameters such as ``charset=utf-8`` are metadata. Only the metadata before the first comma + is inspected, so a ``;base64,`` sequence inside the payload is never treated as a marker. + + Raises: + ContentError: If the URI has no comma or carries an unsupported encoding marker. + """ + if "," not in uri: + raise ContentError("Data URI must contain a comma separating metadata and data") + prefix, payload = uri.split(",", 1) + parts = prefix.split(";") + parameters = parts[1:] + for index, parameter in enumerate(parameters): + if parameter == "base64": + if index != len(parameters) - 1: + raise ContentError("Data URI 'base64' marker must be the last parameter") + elif parameter and "=" not in parameter: + raise ContentError(f"Unsupported data URI encoding: {parameter}") + return _ParsedDataUri( + media_type=parts[0][5:] or "text/plain", # Remove 'data:' + is_base64=bool(parameters) and parameters[-1] == "base64", + payload=payload, + ) + + def detect_media_type_from_base64( *, data_bytes: bytes | None = None, @@ -132,14 +170,18 @@ def detect_media_type_from_base64( # Remove data URI prefix if present if not data_uri.startswith("data:") or "," not in data_uri: raise ValueError("Invalid data URI format.") - prefix, data_str = data_uri.split(",", 1) - if not prefix.endswith(";base64"): + try: + parsed = _parse_data_uri(data_uri) + except ContentError as exc: + raise ValueError(str(exc)) from exc + if not parsed.is_base64: raise ValueError("Data URI must use base64 encoding.") + data_str = parsed.payload if data_str is not None: if data is not None: raise ValueError("Provide exactly one of data_bytes, data_str, or data_uri.") try: - data = base64.b64decode(data_str) + data = base64.b64decode(data_str, validate=True) except Exception as exc: raise ValueError("Invalid base64 data provided.") from exc if data is None: @@ -155,9 +197,16 @@ def detect_media_type_from_base64( return "image/gif" if data.startswith(b"RIFF") and len(data) > 11 and data[8:12] == b"WEBP": return "image/webp" - if data.startswith(b"BM"): + if ( + data.startswith(b"BM") + and len(data) >= 18 + and int.from_bytes(data[14:18], "little") in (12, 40, 52, 56, 64, 108, 124) + ): return "image/bmp" - if data.startswith(b"])", head) or (head.startswith(b"])", head)): return "image/svg+xml" # Documents @@ -167,7 +216,13 @@ def detect_media_type_from_base64( # Audio if data.startswith(b"RIFF") and len(data) > 11 and data[8:12] == b"WAVE": return "audio/wav" - if data.startswith(b"ID3") or data.startswith(b"\xff\xfb") or data.startswith(b"\xff\xf3"): + if data.startswith(b"ID3") or ( + len(data) >= 2 + and data[0] == 0xFF + and (data[1] & 0xE0) == 0xE0 + and (data[1] & 0x18) != 0x08 # MPEG version 01 is reserved. + and (data[1] & 0x06) != 0 # Layer 00 is reserved. + ): return "audio/mpeg" if data.startswith(b"OggS"): return "audio/ogg" @@ -199,11 +254,10 @@ def _get_data_bytes_as_str(content: Content) -> str | None: if not uri.startswith("data:"): return None - if ";base64," not in uri: + parsed = _parse_data_uri(uri) + if not parsed.is_base64: raise ContentError("Data URI must use base64 encoding") - - _, data = uri.split(";base64,", 1) - return data + return parsed.payload def _get_data_bytes(content: Content) -> bytes | None: # pyright: ignore[reportUnusedFunction] @@ -246,20 +300,9 @@ def _validate_uri(uri: str, media_type: str | None) -> dict[str, Any]: # Check for data URI if uri.startswith("data:"): - if "," not in uri: - raise ContentError("Data URI must contain a comma separating metadata and data") - prefix, _ = uri.split(",", 1) - if ";" in prefix: - parts = prefix.split(";") - if len(parts) < 2: - raise ContentError("Invalid data URI format") - # Check encoding - encoding = parts[-1] - if encoding not in ("base64", ""): - raise ContentError(f"Unsupported data URI encoding: {encoding}") - if media_type is None: - # attempt to extract: - media_type = parts[0][5:] # Remove 'data:' + parsed = _parse_data_uri(uri) + if media_type is None: + media_type = parsed.media_type return {"type": "data", "uri": uri, "media_type": media_type} # Check for common URI schemes diff --git a/python/packages/core/tests/core/test_data_uri_media_detection.py b/python/packages/core/tests/core/test_data_uri_media_detection.py new file mode 100644 index 00000000000..01eac48ea7d --- /dev/null +++ b/python/packages/core/tests/core/test_data_uri_media_detection.py @@ -0,0 +1,120 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Regression tests for data-URI parameter handling and media-type sniffing. + +Covers two reported defects in agent_framework/_types.py: +- _validate_uri rejected RFC 2397 parameters such as charset=... and failed to + take the media type from a data URI without parameters (#8905). +- detect_media_type_from_base64 misclassified plain "BM..." text as BMP, any + XML document as SVG, missed MPEG-2/2.5 Layer III frames, and silently + discarded invalid base64 characters (#8906). +""" + +from __future__ import annotations + +import base64 + +import pytest + +from agent_framework import Content, detect_media_type_from_base64 +from agent_framework.exceptions import ContentError + + +class TestDataUriParameters: + def test_charset_parameter_accepted(self) -> None: + content = Content.from_uri("data:text/plain;charset=utf-8,hello%20world") + assert content.media_type == "text/plain" + + def test_media_type_extracted_without_parameters(self) -> None: + content = Content.from_uri("data:text/plain,hello") + assert content.media_type == "text/plain" + + def test_empty_media_type_defaults_to_text_plain(self) -> None: + content = Content.from_uri("data:;base64,aGVsbG8=") + assert content.media_type == "text/plain" + + def test_charset_and_base64_accepted(self) -> None: + content = Content.from_uri("data:application/json;charset=utf-8;base64,e30=") + assert content.media_type == "application/json" + + @pytest.mark.parametrize( + "uri", + [ + "data:image/png;base64;charset=utf-8,aGk=", + "data:image/png;base64;base64,aGk=", + ], + ) + def test_base64_must_be_last_parameter(self, uri: str) -> None: + with pytest.raises(ContentError, match="must be the last parameter"): + Content.from_uri(uri) + + def test_base64_as_last_parameter_accepted(self) -> None: + content = Content.from_uri("data:image/png;charset=utf-8;base64,aGk=") + assert content.media_type == "image/png" + + def test_unknown_bare_encoding_still_rejected(self) -> None: + with pytest.raises(ContentError, match="Unsupported data URI encoding"): + Content.from_uri("data:text/plain;gzip,hello") + + def test_explicit_media_type_wins_over_uri(self) -> None: + content = Content.from_uri("data:text/plain;charset=utf-8,hello", media_type="text/csv") + assert content.media_type == "text/csv" + + +class TestMediaTypeDetection: + def test_plain_text_starting_with_bm_is_not_bmp(self) -> None: + assert detect_media_type_from_base64(data_bytes=b"BMW is a car") is None + + @pytest.mark.parametrize("dib_size", [12, 40, 52, 56, 64, 108, 124]) + def test_real_bmp_header_detected(self, dib_size: int) -> None: + bmp = b"BM" + b"\x00" * 12 + dib_size.to_bytes(4, "little") + b"\x00" * dib_size + assert detect_media_type_from_base64(data_bytes=bmp) == "image/bmp" + + def test_xml_document_is_not_svg(self) -> None: + assert detect_media_type_from_base64(data_bytes=b"") is None + + def test_svg_with_xml_prolog_detected(self) -> None: + svg = b"" + assert detect_media_type_from_base64(data_bytes=svg) == "image/svg+xml" + + def test_bare_svg_element_detected(self) -> None: + assert detect_media_type_from_base64(data_bytes=b"") == "image/svg+xml" + + @pytest.mark.parametrize("prefix", [b" \n", b"\xef\xbb\xbf", b" \xef\xbb\xbf\n"]) + def test_svg_with_leading_whitespace_or_bom_detected(self, prefix: bytes) -> None: + assert detect_media_type_from_base64(data_bytes=prefix + b"") == "image/svg+xml" + + @pytest.mark.parametrize("prefix", [b"", b""]) + def test_svg_named_prefix_is_not_svg(self, prefix: bytes) -> None: + assert detect_media_type_from_base64(data_bytes=prefix + b"") is None + + def test_mpeg2_layer3_frame_sync_detected(self) -> None: + assert detect_media_type_from_base64(data_bytes=b"\xff\xf2\x90\x00" + b"\0" * 20) == "audio/mpeg" + + def test_mpeg25_layer3_frame_sync_detected(self) -> None: + assert detect_media_type_from_base64(data_bytes=b"\xff\xe3\x90\x00" + b"\0" * 20) == "audio/mpeg" + + def test_reserved_layer_not_mpeg(self) -> None: + assert detect_media_type_from_base64(data_bytes=b"\xff\xf9\x90\x00" + b"\0" * 20) is None + + def test_reserved_version_not_mpeg(self) -> None: + assert detect_media_type_from_base64(data_bytes=b"\xff\xeb\x90\x00" + b"\0" * 20) is None + + def test_invalid_base64_characters_rejected(self) -> None: + with pytest.raises(ValueError, match="Invalid base64 data provided."): + detect_media_type_from_base64(data_str="iVBO!!!!Rw0KGgowMDAwMDAwMA==") + + def test_valid_signatures_unchanged(self) -> None: + assert detect_media_type_from_base64(data_bytes=b"\x89PNG\r\n\x1a\n" + b"x" * 8) == "image/png" + assert detect_media_type_from_base64(data_bytes=b"\xff\xd8\xff\xe0" + b"x" * 8) == "image/jpeg" + assert detect_media_type_from_base64(data_bytes=b"GIF89a" + b"x" * 8) == "image/gif" + assert detect_media_type_from_base64(data_bytes=b"RIFF\x00\x00\x00\x00WEBP") == "image/webp" + assert detect_media_type_from_base64(data_bytes=b"RIFF\x00\x00\x00\x00WAVE") == "audio/wav" + assert detect_media_type_from_base64(data_bytes=b"ID3\x04\x00" + b"x" * 8) == "audio/mpeg" + assert detect_media_type_from_base64(data_bytes=b"OggS" + b"x" * 8) == "audio/ogg" + assert detect_media_type_from_base64(data_bytes=b"fLaC" + b"x" * 8) == "audio/flac" + assert detect_media_type_from_base64(data_bytes=b"%PDF-1.7") == "application/pdf" + assert ( + detect_media_type_from_base64(data_str=base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"x" * 8).decode()) + == "image/png" + ) diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index ff4bfedf014..8c29138ef4b 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -5939,3 +5939,23 @@ def test_agent_response_update_serialization_includes_finish_reason() -> None: # endregion + + +def test_get_data_bytes_as_str_ignores_base64_marker_inside_the_payload(): + """A ';base64,' sequence inside the payload is data, not the encoding marker.""" + content = Content.from_uri(uri="data:text/plain,a;base64,QUJD") + with raises(ContentError, match="base64 encoding"): + _get_data_bytes_as_str(content) + + +def test_data_uri_readers_agree_on_the_base64_marker(): + """Validation, media type detection and payload extraction share one reading of a data URI.""" + png_b64 = base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"\x00" * 16).decode() + uri = f"data:image/png;charset=utf-8;base64,{png_b64}" + content = Content.from_uri(uri=uri) + assert content.media_type == "image/png" + assert _get_data_bytes_as_str(content) == png_b64 + assert detect_media_type_from_base64(data_uri=uri) == "image/png" + payload_marker = f"data:image/png,x;base64,{png_b64}" + with pytest.raises(ValueError, match="Data URI must use base64 encoding."): + detect_media_type_from_base64(data_uri=payload_marker)