Python: Fix data URI validation and media type detection

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:[<media-type>][;parameter[=value]]...[;base64],<data>``. 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"<svg") or data.startswith(b"<?xml"):
+    head = data[:512].lstrip()
+    if head.startswith(b"\xef\xbb\xbf"):
+        head = head[3:].lstrip()
+    if re.match(rb"<svg(?=[\s/>])", head) or (head.startswith(b"<?xml") and re.search(rb"<svg(?=[\s/>])", 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

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"<?xml version='1.0'?><rss/>") is None
+
+    def test_svg_with_xml_prolog_detected(self) -> None:
+        svg = b"<?xml version='1.0'?><svg xmlns='http://www.w3.org/2000/svg'></svg>"
+        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"<svg></svg>") == "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"<svg/>") == "image/svg+xml"
+
+    @pytest.mark.parametrize("prefix", [b"", b"<?xml version='1.0'?>"])
+    def test_svg_named_prefix_is_not_svg(self, prefix: bytes) -> None:
+        assert detect_media_type_from_base64(data_bytes=prefix + b"<svgish/>") 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"
+        )

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)