diff --git a/tests/multimodal/media/test_connector.py b/tests/multimodal/media/test_connector.py index 5cd22a8b18..ddf6dbcef8 100644 --- a/tests/multimodal/media/test_connector.py +++ b/tests/multimodal/media/test_connector.py @@ -225,6 +225,44 @@ async def test_fetch_image_local_files_with_space_in_name(image_url: str): assert not ImageChops.difference(image_sync, image_async).getbbox() +@pytest.mark.asyncio +async def test_fetch_image_data_url_with_params(): + """RFC 2397 allows parameters between the mediatype and the base64 + marker; they must not be rejected or leak into the media type.""" + connector = MediaConnector() + + image = Image.new("RGB", (4, 4), color=(255, 0, 0)) + with NamedTemporaryFile(suffix=".png") as f: + image.save(f.name) + base64_image = base64.b64encode(f.read()).decode("utf-8") + + data_url = f"data:image/png;charset=utf-8;base64,{base64_image}" + image_sync = connector.fetch_image(data_url) + image_async = await connector.fetch_image_async(data_url) + assert _image_equals(image_sync, image_async) + + +def test_fetch_image_data_url_malformed(): + connector = MediaConnector() + + with pytest.raises(ValueError, match="missing ','"): + connector.fetch_image("data:image/png;base64") + + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:text/plain,hello") + + # ";base64" requires the ";"; here "base64" is a (bogus) media type. + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:base64,aGVsbG8=") + + # Strict RFC 2397 grammar: lowercase "base64", no whitespace. + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:image/png;BASE64,aGVsbG8=") + + with pytest.raises(NotImplementedError, match="base64"): + connector.fetch_image("data:image/png; base64,aGVsbG8=") + + @pytest.mark.asyncio async def test_fetch_image_error_conversion(): connector = MediaConnector() diff --git a/vllm/multimodal/media/connector.py b/vllm/multimodal/media/connector.py index bf9b7345ca..656f537c5b 100644 --- a/vllm/multimodal/media/connector.py +++ b/vllm/multimodal/media/connector.py @@ -302,14 +302,18 @@ class MediaConnector: media_io: MediaIO[_M], ) -> _M: # type: ignore[type-var] # Format per RFC 2397: - # data:[][;base64], - data_spec, data = url[5:].split(",", 1) - media_type, data_type = data_spec.split(";", 1) + # data:[][;=]*[;base64], + data_spec, sep, data = url[5:].partition(",") + if not sep: + msg = f"Invalid data URL {url[:32]!r}: missing ',' separator." + raise ValueError(msg) - if data_type != "base64": + media_type, sep, encoding = data_spec.rpartition(";") + if not sep or encoding != "base64": msg = "Only base64 data URLs are supported for now." raise NotImplementedError(msg) + media_type = media_type.partition(";")[0] return media_io.load_base64(media_type, data) def _load_file_url(