[Bugfix] Use MediaConnector for image/video URL fetching to prevent SSRF (#2565)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/entrypoints/openai_api/test_decode_audio_url.py

@@ -1,16 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for decode_audio_url helper."""
+"""Unit tests for decode_audio_url helper and _AudioFileMediaIO."""
 
 import asyncio
 import base64
 import os
 
-import httpx
 import pytest
+from vllm.multimodal.media import MediaConnector
 
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
-from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url
+from vllm_omni.entrypoints.openai.video_api_utils import (
+    _AudioFileMediaIO,
+    decode_audio_url,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -24,15 +27,127 @@ def _cleanup(path):
         os.unlink(path)
 
 
+def _permissive_connector():
+    return MediaConnector(
+        allowed_local_media_path="",
+        allowed_media_domains=None,
+    )
+
+
+def _restricted_connector(allowed_domains=None):
+    return MediaConnector(
+        allowed_local_media_path="",
+        allowed_media_domains=allowed_domains or [],
+    )
+
+
+# ---------------------------------------------------------------------------
+# _AudioFileMediaIO direct tests
+# ---------------------------------------------------------------------------
+
+
+class TestAudioFileMediaIO:
+    """Direct tests for _AudioFileMediaIO methods."""
+
+    def test_load_bytes_writes_temp_file(self):
+        audio_bytes = b"\xff\xfb\x90\x00" * 50
+        io = _AudioFileMediaIO()
+        path = io.load_bytes(audio_bytes)
+        try:
+            assert os.path.isfile(path)
+            assert path.endswith(".wav")
+            with open(path, "rb") as f:
+                assert f.read() == audio_bytes
+        finally:
+            _cleanup(path)
+
+    def test_load_bytes_empty_raises(self):
+        io = _AudioFileMediaIO()
+        with pytest.raises(InvalidInputReferenceError, match="audio data is empty"):
+            io.load_bytes(b"")
+
+    def test_load_base64_mp3_suffix(self):
+        audio_bytes = b"\xff\xfb\x90\x00" * 50
+        b64 = base64.b64encode(audio_bytes).decode()
+        io = _AudioFileMediaIO()
+        path = io.load_base64("audio/mp3", b64)
+        try:
+            assert path.endswith(".mp3")
+            with open(path, "rb") as f:
+                assert f.read() == audio_bytes
+        finally:
+            _cleanup(path)
+
+    def test_load_base64_mpeg_uses_mp3(self):
+        io = _AudioFileMediaIO()
+        path = io.load_base64("audio/mpeg", base64.b64encode(b"\x00" * 10).decode())
+        try:
+            assert path.endswith(".mp3")
+        finally:
+            _cleanup(path)
+
+    def test_load_base64_wav_suffix(self):
+        io = _AudioFileMediaIO()
+        path = io.load_base64("audio/wav", base64.b64encode(b"\x00" * 10).decode())
+        try:
+            assert path.endswith(".wav")
+        finally:
+            _cleanup(path)
+
+    def test_load_base64_flac_suffix(self):
+        io = _AudioFileMediaIO()
+        path = io.load_base64("audio/flac", base64.b64encode(b"\x00" * 10).decode())
+        try:
+            assert path.endswith(".flac")
+        finally:
+            _cleanup(path)
+
+    def test_load_base64_empty_media_type_uses_wav(self):
+        io = _AudioFileMediaIO()
+        path = io.load_base64("", base64.b64encode(b"\x00" * 10).decode())
+        try:
+            assert path.endswith(".wav")
+        finally:
+            _cleanup(path)
+
+    def test_load_base64_non_alnum_falls_back_to_wav(self):
+        io = _AudioFileMediaIO()
+        path = io.load_base64("audio/foo-bar", base64.b64encode(b"\x00" * 10).decode())
+        try:
+            assert path.endswith(".wav")
+        finally:
+            _cleanup(path)
+
+    def test_load_base64_too_long_extension_falls_back_to_wav(self):
+        io = _AudioFileMediaIO()
+        path = io.load_base64("audio/verylongextension", base64.b64encode(b"\x00" * 10).decode())
+        try:
+            assert path.endswith(".wav")
+        finally:
+            _cleanup(path)
+
+    def test_load_file_returns_path_string(self, tmp_path):
+        audio_file = tmp_path / "test.wav"
+        audio_file.write_bytes(b"\x00" * 100)
+        io = _AudioFileMediaIO()
+        result = io.load_file(audio_file)
+        assert result == str(audio_file)
+
+
+# ---------------------------------------------------------------------------
+# decode_audio_url with data URLs (exercises real MediaConnector + MediaIO)
+# ---------------------------------------------------------------------------
+
+
 class TestDecodeAudioUrlDataUrl:
-    """Tests for base64 data URL decoding."""
+    """Tests for base64 data URL decoding — no mocks, real code path."""
 
     def test_valid_mp3_data_url(self):
         audio_bytes = b"\xff\xfb\x90\x00" * 100
         b64 = base64.b64encode(audio_bytes).decode()
         url = f"data:audio/mp3;base64,{b64}"
 
-        path = _run(decode_audio_url(url))
+        path = _run(decode_audio_url(url, _permissive_connector()))
         try:
             assert os.path.isfile(path)
             assert path.endswith(".mp3")
@@ -46,7 +161,7 @@ class TestDecodeAudioUrlDataUrl:
         b64 = base64.b64encode(audio_bytes).decode()
         url = f"data:audio/wav;base64,{b64}"
 
-        path = _run(decode_audio_url(url))
+        path = _run(decode_audio_url(url, _permissive_connector()))
         try:
             assert path.endswith(".wav")
             with open(path, "rb") as f:
@@ -54,134 +169,103 @@ class TestDecodeAudioUrlDataUrl:
         finally:
             _cleanup(path)
 
-    def test_valid_mpeg_data_url_uses_mp3_suffix(self):
-        audio_bytes = b"\xff\xfb\x90\x00" * 50
-        b64 = base64.b64encode(audio_bytes).decode()
-        url = f"data:audio/mpeg;base64,{b64}"
-
-        path = _run(decode_audio_url(url))
-        try:
-            assert path.endswith(".mp3")
-        finally:
-            _cleanup(path)
-
-    def test_valid_ogg_data_url(self):
-        audio_bytes = b"OggS" + b"\x00" * 100
-        b64 = base64.b64encode(audio_bytes).decode()
-        url = f"data:audio/ogg;base64,{b64}"
-
-        path = _run(decode_audio_url(url))
-        try:
-            assert path.endswith(".ogg")
-        finally:
-            _cleanup(path)
-
     def test_invalid_base64_raises(self):
         url = "data:audio/mp3;base64,!!!not-valid-base64!!!"
-
-        with pytest.raises(InvalidInputReferenceError, match="not valid base64"):
-            _run(decode_audio_url(url))
+        with pytest.raises(InvalidInputReferenceError, match="Incorrect padding"):
+            _run(decode_audio_url(url, _permissive_connector()))
 
     def test_empty_audio_data_raises(self):
         url = "data:audio/mp3;base64,"
-
         with pytest.raises(InvalidInputReferenceError, match="audio data is empty"):
-            _run(decode_audio_url(url))
+            _run(decode_audio_url(url, _permissive_connector()))
 
 
-class TestDecodeAudioUrlSuffixSanitization:
-    """Tests for MIME extension sanitization (path traversal prevention)."""
-
-    def test_path_traversal_in_mime_is_neutralized(self):
-        audio_bytes = b"\x00" * 100
-        b64 = base64.b64encode(audio_bytes).decode()
-        url = f"data:audio/../..;base64,{b64}"
-
-        path = _run(decode_audio_url(url))
-        try:
-            # ".." contains non-alnum dots, so falls back to .wav
-            assert path.endswith(".wav")
-            assert os.path.dirname(path) == "/tmp"
-        finally:
-            _cleanup(path)
-
-    def test_non_alnum_extension_falls_back_to_wav(self):
-        audio_bytes = b"\x00" * 100
-        b64 = base64.b64encode(audio_bytes).decode()
-        url = f"data:audio/foo-bar;base64,{b64}"
-
-        path = _run(decode_audio_url(url))
-        try:
-            assert path.endswith(".wav")
-        finally:
-            _cleanup(path)
-
-    def test_too_long_extension_falls_back_to_wav(self):
-        audio_bytes = b"\x00" * 100
-        b64 = base64.b64encode(audio_bytes).decode()
-        url = f"data:audio/verylongextension;base64,{b64}"
-
-        path = _run(decode_audio_url(url))
-        try:
-            assert path.endswith(".wav")
-        finally:
-            _cleanup(path)
-
-    def test_valid_short_alnum_extension_is_used(self):
-        audio_bytes = b"\x00" * 100
-        b64 = base64.b64encode(audio_bytes).decode()
-        url = f"data:audio/flac;base64,{b64}"
-
-        path = _run(decode_audio_url(url))
-        try:
-            assert path.endswith(".flac")
-        finally:
-            _cleanup(path)
+# ---------------------------------------------------------------------------
+# decode_audio_url with file:// URIs (exercises load_file through connector)
+# ---------------------------------------------------------------------------
 
 
-class TestDecodeAudioUrlHttp:
-    """Tests for HTTP URL decoding."""
+class TestDecodeAudioUrlFileUri:
+    """Tests using file:// URIs — exercises MediaConnector + load_file."""
 
-    def test_valid_http_url(self, monkeypatch):
-        audio_bytes = b"\xff\xfb\x90\x00" * 50
-        fake_resp = httpx.Response(200, content=audio_bytes)
+    def test_file_uri_with_allowed_path(self, tmp_path):
+        audio_file = tmp_path / "test.wav"
+        audio_file.write_bytes(b"\xff\xfb\x90\x00" * 50)
+        file_uri = audio_file.as_uri()
+        connector = MediaConnector(
+            allowed_local_media_path=str(tmp_path),
+            allowed_media_domains=None,
+        )
+        path = _run(decode_audio_url(file_uri, connector))
+        assert path == str(audio_file)
 
-        async def _mock_get(self, url):
-            return fake_resp
+    def test_file_uri_without_allowed_path_raises(self, tmp_path):
+        audio_file = tmp_path / "test.wav"
+        audio_file.write_bytes(b"\x00" * 50)
+        file_uri = audio_file.as_uri()
+        connector = MediaConnector(
+            allowed_local_media_path="",
+            allowed_media_domains=None,
+        )
+        with pytest.raises(InvalidInputReferenceError):
+            _run(decode_audio_url(file_uri, connector))
 
-        monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
-        monkeypatch.setattr(httpx.Response, "raise_for_status", lambda self: None)
+    def test_file_uri_outside_allowed_path_raises(self, tmp_path):
+        audio_file = tmp_path / "test.wav"
+        audio_file.write_bytes(b"\x00" * 50)
+        file_uri = audio_file.as_uri()
+        other_dir = tmp_path / "other"
+        other_dir.mkdir()
+        connector = MediaConnector(
+            allowed_local_media_path=str(other_dir),
+            allowed_media_domains=None,
+        )
+        with pytest.raises(InvalidInputReferenceError):
+            _run(decode_audio_url(file_uri, connector))
 
-        path = _run(decode_audio_url("https://example.com/audio.mp3"))
-        try:
-            assert os.path.isfile(path)
-            assert path.endswith(".wav")
-            with open(path, "rb") as f:
-                assert f.read() == audio_bytes
-        finally:
-            _cleanup(path)
 
-    def test_http_error_raises(self, monkeypatch):
-        async def _mock_get(self, url):
-            raise httpx.HTTPStatusError("Not Found", request=httpx.Request("GET", url), response=httpx.Response(404))
-
-        monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
-
-        with pytest.raises(InvalidInputReferenceError, match="failed to download audio"):
-            _run(decode_audio_url("https://example.com/missing.mp3"))
+# ---------------------------------------------------------------------------
+# decode_audio_url — invalid schemes and SSRF
+# ---------------------------------------------------------------------------
 
 
 class TestDecodeAudioUrlInvalid:
     """Tests for invalid URL schemes."""
 
     def test_local_file_path_raises(self):
-        with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
-            _run(decode_audio_url("/path/to/local/file.mp3"))
+        with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+            _run(decode_audio_url("/path/to/local/file.mp3", _permissive_connector()))
 
     def test_ftp_url_raises(self):
-        with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
-            _run(decode_audio_url("ftp://example.com/file.mp3"))
+        with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+            _run(decode_audio_url("ftp://example.com/file.mp3", _permissive_connector()))
 
     def test_empty_string_raises(self):
-        with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
-            _run(decode_audio_url(""))
+        with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+            _run(decode_audio_url("", _permissive_connector()))
+
+
+class TestDecodeAudioUrlSSRF:
+    """Tests for SSRF protection via MediaConnector domain restrictions."""
+
+    def test_restricted_domain_rejects_disallowed_url(self):
+        connector = _restricted_connector(allowed_domains=["trusted.example.com"])
+        with pytest.raises(InvalidInputReferenceError, match="allowed domains"):
+            _run(
+                decode_audio_url(
+                    "https://evil.internal/admin/audio.mp3",
+                    connector,
+                )
+            )
+
+    def test_data_url_bypasses_domain_restriction(self):
+        audio_bytes = b"\xff\xfb\x90\x00" * 50
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/wav;base64,{b64}"
+
+        connector = _restricted_connector(allowed_domains=["trusted.example.com"])
+        path = _run(decode_audio_url(url, connector))
+        try:
+            assert os.path.isfile(path)
+        finally:
+            _cleanup(path)

tests/entrypoints/openai_api/test_decode_video_url.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for decode_video_url helper and _VideoFramesMediaIO."""
+
+import asyncio
+import base64
+import io
+
+import pytest
+from PIL import Image
+from vllm.multimodal.media import MediaConnector
+
+from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
+from vllm_omni.entrypoints.openai.video_api_utils import (
+    VideoFrames,
+    _VideoFramesMediaIO,
+    decode_video_url,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _run(coro):
+    return asyncio.run(coro)
+
+
+def _permissive_connector():
+    return MediaConnector(
+        allowed_local_media_path="",
+        allowed_media_domains=None,
+    )
+
+
+def _restricted_connector(allowed_domains=None):
+    return MediaConnector(
+        allowed_local_media_path="",
+        allowed_media_domains=allowed_domains or [],
+    )
+
+
+def _make_mp4_bytes(width=4, height=4, num_frames=2):
+    """Create a minimal valid MP4 video using PyAV."""
+    import av
+
+    buf = io.BytesIO()
+    with av.open(buf, mode="w", format="mp4") as container:
+        stream = container.add_stream("libx264", rate=24)
+        stream.width = width
+        stream.height = height
+        stream.pix_fmt = "yuv420p"
+        stream.options = {"preset": "ultrafast"}
+        for i in range(num_frames):
+            img = Image.new("RGB", (width, height), color=(i * 40, 128, 200))
+            frame = av.VideoFrame.from_image(img)
+            for packet in stream.encode(frame):
+                container.mux(packet)
+        for packet in stream.encode():
+            container.mux(packet)
+    return buf.getvalue()
+
+
+# ---------------------------------------------------------------------------
+# _VideoFramesMediaIO direct tests
+# ---------------------------------------------------------------------------
+
+
+class TestVideoFramesMediaIO:
+    """Direct tests for _VideoFramesMediaIO methods."""
+
+    def test_load_bytes_decodes_video(self):
+        video_bytes = _make_mp4_bytes(num_frames=3)
+        media_io = _VideoFramesMediaIO()
+        result = media_io.load_bytes(video_bytes)
+        assert isinstance(result, VideoFrames)
+        assert len(result) >= 1
+        assert all(isinstance(f, Image.Image) for f in result)
+
+    def test_load_bytes_respects_max_frames(self):
+        video_bytes = _make_mp4_bytes(num_frames=5)
+        media_io = _VideoFramesMediaIO(max_frames=2)
+        result = media_io.load_bytes(video_bytes)
+        assert len(result) <= 2
+
+    def test_load_bytes_invalid_data_raises(self):
+        media_io = _VideoFramesMediaIO()
+        with pytest.raises(InvalidInputReferenceError, match="not a valid video"):
+            media_io.load_bytes(b"not a video")
+
+    def test_load_base64_decodes_video(self):
+        video_bytes = _make_mp4_bytes()
+        b64 = base64.b64encode(video_bytes).decode()
+        media_io = _VideoFramesMediaIO()
+        result = media_io.load_base64("video/mp4", b64)
+        assert isinstance(result, VideoFrames)
+        assert len(result) >= 1
+
+    def test_load_file_decodes_video(self, tmp_path):
+        video_bytes = _make_mp4_bytes()
+        video_file = tmp_path / "test.mp4"
+        video_file.write_bytes(video_bytes)
+        media_io = _VideoFramesMediaIO()
+        result = media_io.load_file(video_file)
+        assert isinstance(result, VideoFrames)
+        assert len(result) >= 1
+
+
+# ---------------------------------------------------------------------------
+# decode_video_url with data URLs (real MediaConnector + MediaIO)
+# ---------------------------------------------------------------------------
+
+
+class TestDecodeVideoUrlDataUrl:
+    """Tests for data URL decoding — no mocks, real code path."""
+
+    def test_valid_video_data_url(self):
+        video_bytes = _make_mp4_bytes()
+        b64 = base64.b64encode(video_bytes).decode()
+        url = f"data:video/mp4;base64,{b64}"
+
+        result = _run(decode_video_url(url, _permissive_connector()))
+        assert isinstance(result, VideoFrames)
+        assert len(result) >= 1
+        assert all(isinstance(f, Image.Image) for f in result)
+
+    def test_invalid_video_data_url_raises(self):
+        b64 = base64.b64encode(b"not a video").decode()
+        url = f"data:video/mp4;base64,{b64}"
+        with pytest.raises(InvalidInputReferenceError, match="not a valid video"):
+            _run(decode_video_url(url, _permissive_connector()))
+
+
+# ---------------------------------------------------------------------------
+# decode_video_url with file:// URIs (exercises load_file through connector)
+# ---------------------------------------------------------------------------
+
+
+class TestDecodeVideoUrlFileUri:
+    """Tests using file:// URIs — exercises MediaConnector + load_file."""
+
+    def test_file_uri_with_allowed_path(self, tmp_path):
+        video_file = tmp_path / "test.mp4"
+        video_file.write_bytes(_make_mp4_bytes())
+        file_uri = video_file.as_uri()
+        connector = MediaConnector(
+            allowed_local_media_path=str(tmp_path),
+            allowed_media_domains=None,
+        )
+        result = _run(decode_video_url(file_uri, connector))
+        assert isinstance(result, VideoFrames)
+        assert len(result) >= 1
+
+    def test_file_uri_without_allowed_path_raises(self, tmp_path):
+        video_file = tmp_path / "test.mp4"
+        video_file.write_bytes(_make_mp4_bytes())
+        file_uri = video_file.as_uri()
+        connector = MediaConnector(
+            allowed_local_media_path="",
+            allowed_media_domains=None,
+        )
+        with pytest.raises(InvalidInputReferenceError):
+            _run(decode_video_url(file_uri, connector))
+
+
+# ---------------------------------------------------------------------------
+# decode_video_url — invalid schemes and SSRF
+# ---------------------------------------------------------------------------
+
+
+class TestDecodeVideoUrlInvalid:
+    """Tests for invalid URL schemes."""
+
+    def test_local_file_path_raises(self):
+        with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+            _run(decode_video_url("/path/to/local/file.mp4", _permissive_connector()))
+
+    def test_ftp_url_raises(self):
+        with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+            _run(decode_video_url("ftp://example.com/file.mp4", _permissive_connector()))
+
+    def test_empty_string_raises(self):
+        with pytest.raises(InvalidInputReferenceError, match="must be either a HTTP, data or file URL"):
+            _run(decode_video_url("", _permissive_connector()))
+
+
+class TestDecodeVideoUrlSSRF:
+    """Tests for SSRF protection via MediaConnector domain restrictions."""
+
+    def test_restricted_domain_rejects_disallowed_url(self):
+        connector = _restricted_connector(
+            allowed_domains=["trusted.example.com"],
+        )
+        with pytest.raises(InvalidInputReferenceError, match="allowed domains"):
+            _run(
+                decode_video_url(
+                    "https://evil.internal/admin/video.mp4",
+                    connector,
+                )
+            )
+
+    def test_data_url_bypasses_domain_restriction(self):
+        video_bytes = _make_mp4_bytes()
+        b64 = base64.b64encode(video_bytes).decode()
+        url = f"data:video/mp4;base64,{b64}"
+
+        connector = _restricted_connector(
+            allowed_domains=["trusted.example.com"],
+        )
+        result = _run(decode_video_url(url, connector))
+        assert isinstance(result, VideoFrames)
+        assert len(result) >= 1

tests/entrypoints/openai_api/test_image_server.py

@@ -13,6 +13,7 @@ import json
 from argparse import Namespace
 from http import HTTPStatus
 from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from fastapi import FastAPI, HTTPException
@@ -24,12 +25,19 @@ from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.sampling_params import RequestOutputKind
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
-from vllm_omni.entrypoints.openai.api_server import _check_max_generated_image_size, _DiffusionServingModels, router
+from vllm_omni.entrypoints.openai.api_server import (
+    _check_max_generated_image_size,
+    _DiffusionServingModels,
+    _load_input_images,
+    router,
+)
 from vllm_omni.entrypoints.openai.image_api_utils import (
     encode_image_base64,
     parse_size,
 )
+from vllm_omni.entrypoints.openai.protocol.videos import UrlImageReference
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference
 from vllm_omni.errors import GuardrailViolationError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -182,6 +190,7 @@ def mock_async_diffusion(mocker: MockerFixture):
             self.captured_sampling_params_list = None
             self.captured_prompt = None
             self.generate_calls = 0
+            self.model_config = SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
 
         async def generate(self, **kwargs):
             self.generate_calls += 1
@@ -271,6 +280,10 @@ def async_omni_test_client():
         def get_diffusion_od_config(self):
             return self.od_config
 
+        @property
+        def model_config(self):
+            return SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
+
     app = FastAPI()
     app.include_router(router)
 
@@ -335,6 +348,10 @@ def async_omni_rgba_test_client():
         def get_diffusion_od_config(self):
             return self.od_config
 
+        @property
+        def model_config(self):
+            return SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
+
     app = FastAPI()
     app.include_router(router)
 
@@ -399,6 +416,10 @@ def async_omni_stage_configs_only_client():
         def get_diffusion_od_config(self):
             return self.od_config
 
+        @property
+        def model_config(self):
+            return SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
+
     app = FastAPI()
     app.include_router(router)
 
@@ -435,9 +456,24 @@ def streaming_image_edit_client():
                 SamplingParams(temperature=0.1),
                 OmniDiffusionSamplingParams(),
             ]
+            stage_clients = [
+                SimpleNamespace(stage_type="llm", is_comprehension=True),
+                SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+            ]
+            stage_vllm_configs = [
+                SimpleNamespace(
+                    model_config=SimpleNamespace(
+                        allowed_local_media_path="",
+                        allowed_media_domains=None,
+                    ),
+                ),
+                None,
+            ]
             self.engine = SimpleNamespace(
                 stage_configs=stage_configs,
                 default_sampling_params_list=default_sampling_params_list,
+                stage_clients=stage_clients,
+                stage_vllm_configs=stage_vllm_configs,
             )
             self.default_sampling_params_list = default_sampling_params_list
             self.captured_sampling_params_list = None
@@ -729,6 +765,10 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
         def get_diffusion_od_config(self):
             return self.od_config
 
+        @property
+        def model_config(self):
+            return SimpleNamespace(allowed_local_media_path="", allowed_media_domains=None)
+
     app = FastAPI()
     app.include_router(router)
     engine = FakeAsyncOmniClass()
@@ -2136,3 +2176,107 @@ def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_on
         assert captured_prompt["prompt"].count("<img>") == 2, (
             f"N=2 reference images must emit 2 <img> placeholders in AR prompt; got {captured_prompt[KEY].count(IMG)} -- prompt: {captured_prompt[KEY]!r}"
         )
+
+
+# ---------------------------------------------------------------------------
+# MediaConnector / SSRF protection tests
+# ---------------------------------------------------------------------------
+
+
+_DATA_IMAGE_URL = f"data:image/png;base64,{base64.b64encode(make_test_image_bytes((8, 8))).decode()}"
+
+
+def _make_upload_image():
+    upload = MagicMock()
+    upload.file = True
+    upload.read = AsyncMock(return_value=make_test_image_bytes((8, 8)))
+    return upload
+
+
+class TestLoadInputImagesMediaConnector:
+    """SSRF protection via MediaConnector for _load_input_images and
+    decode_input_reference."""
+
+    @staticmethod
+    def _config(
+        domains=None,  # allow all domains
+        local_path="",  # disallow all local paths
+    ):
+        return SimpleNamespace(
+            allowed_local_media_path=local_path,
+            allowed_media_domains=domains,
+        )
+
+    @pytest.mark.parametrize(
+        "domains, url",
+        [
+            # restricted + data url from allowlisted domain → accept (data: bypasses domain check)
+            (["example.com"], _DATA_IMAGE_URL),
+            # unrestricted + data url → accept
+            (None, _DATA_IMAGE_URL),
+            # restricted + file upload → accept (UploadFile bypasses MediaConnector)
+            (["example.com"], _make_upload_image()),
+        ],
+        ids=[
+            "restricted-data-url",
+            "unrestricted-data-url",
+            "restricted-file-upload",
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_accepted_inputs(self, domains, url):
+        config = self._config(domains=domains)
+
+        images = await _load_input_images([url], config)
+        assert len(images) == 1
+        assert isinstance(images[0], Image.Image)
+
+        if isinstance(url, str):
+            ref = UrlImageReference(image_url=url)
+            image = await decode_input_reference(ref, None, None, config)
+            assert isinstance(image, Image.Image)
+
+    @pytest.mark.asyncio
+    async def test_http_url_rejected_by_domain_filter(self):
+        """An http(s) URL whose domain is not in the allowlist must be rejected."""
+        config = self._config(domains=["example.com"])
+        http_url = "https://blocked.example.org/image.png"
+
+        with pytest.raises(ValueError):
+            await _load_input_images([http_url], config)
+
+        ref = UrlImageReference(image_url=http_url)
+        with pytest.raises(ValueError):
+            await decode_input_reference(ref, None, None, config)
+
+    @pytest.mark.parametrize(
+        "allow_path",
+        [True, False],
+        ids=["allowed-path", "disallowed-path"],
+    )
+    @pytest.mark.asyncio
+    async def test_file_uri_path_restriction(self, tmp_path, allow_path):
+        """file:// URIs are gated by allowed_local_media_path."""
+
+        img_path = tmp_path / "test.png"
+        img_path.write_bytes(make_test_image_bytes((8, 8)))
+        file_uri = img_path.as_uri()
+        config = self._config(local_path=str(tmp_path) if allow_path else "")
+
+        # --- _load_input_images ---
+        if allow_path:
+            images = await _load_input_images([file_uri], config)
+            assert len(images) == 1
+            assert isinstance(images[0], Image.Image)
+        else:
+            with pytest.raises(ValueError):
+                await _load_input_images([file_uri], config)
+
+        # --- decode_input_reference ---
+        ref = UrlImageReference(image_url=file_uri)
+        if allow_path:
+            image = await decode_input_reference(ref, None, None, config)
+            assert isinstance(image, Image.Image)
+        else:
+            with pytest.raises(ValueError):
+                await decode_input_reference(ref, None, None, config)

tests/entrypoints/openai_api/test_materialize_media_source.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for audiox materialize_media_source and _TempFileMediaIO."""
+
+import base64
+import os
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from vllm_omni.transformers_utils.processors.audiox import (
+    _TempFileMediaIO,
+    materialize_media_source,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _cleanup(path):
+    if path and os.path.exists(path):
+        os.unlink(path)
+
+
+def _fake_forward_context(allowed_domains=None, allowed_local_media_path=""):
+    """Return a fake ForwardContext with model_config for SSRF testing."""
+    return SimpleNamespace(
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                allowed_local_media_path=allowed_local_media_path,
+                allowed_media_domains=allowed_domains,
+            ),
+        ),
+    )
+
+
+# ---------------------------------------------------------------------------
+# _TempFileMediaIO direct tests
+# ---------------------------------------------------------------------------
+
+
+class TestTempFileMediaIO:
+    """Direct tests for _TempFileMediaIO methods."""
+
+    def test_load_bytes_writes_temp_file(self):
+        data = b"\x00\x01\x02" * 50
+        io = _TempFileMediaIO()
+        path = io.load_bytes(data)
+        try:
+            assert os.path.isfile(path)
+            assert path.endswith(".bin")
+            with open(path, "rb") as f:
+                assert f.read() == data
+        finally:
+            _cleanup(path)
+
+    def test_load_base64_writes_temp_file(self):
+        data = b"\xff" * 100
+        b64 = base64.b64encode(data).decode()
+        io = _TempFileMediaIO()
+        path = io.load_base64("video/mp4", b64)
+        try:
+            assert os.path.isfile(path)
+            with open(path, "rb") as f:
+                assert f.read() == data
+        finally:
+            _cleanup(path)
+
+    def test_load_file_returns_path_string(self, tmp_path):
+        media_file = tmp_path / "video.mp4"
+        media_file.write_bytes(b"\x00" * 50)
+        io = _TempFileMediaIO()
+        result = io.load_file(media_file)
+        assert result == str(media_file)
+
+
+# ---------------------------------------------------------------------------
+# materialize_media_source
+# ---------------------------------------------------------------------------
+
+
+class TestMaterializeMediaSource:
+    """Tests for materialize_media_source — exercises real MediaConnector."""
+
+    def test_data_url_with_no_vllm_config(self):
+        data = b"\x00" * 50
+        b64 = base64.b64encode(data).decode()
+        url = f"data:audio/wav;base64,{b64}"
+        ctx = SimpleNamespace(vllm_config=None)
+
+        with patch(
+            "vllm_omni.diffusion.forward_context.get_forward_context",
+            return_value=ctx,
+        ):
+            path = materialize_media_source(url)
+        try:
+            assert os.path.isfile(path)
+        finally:
+            _cleanup(path)
+
+    def test_data_url_materializes_to_temp_file(self):
+        data = b"\x00\x01\x02" * 50
+        b64 = base64.b64encode(data).decode()
+        url = f"data:audio/wav;base64,{b64}"
+
+        with patch(
+            "vllm_omni.diffusion.forward_context.get_forward_context",
+            return_value=_fake_forward_context(),
+        ):
+            path = materialize_media_source(url)
+        try:
+            assert os.path.isfile(path)
+            with open(path, "rb") as f:
+                assert f.read() == data
+        finally:
+            _cleanup(path)
+
+    def test_local_path_with_allowed_path(self, tmp_path):
+        local_file = tmp_path / "media.bin"
+        local_file.write_bytes(b"\x00" * 50)
+        ctx = _fake_forward_context(allowed_local_media_path=str(tmp_path))
+        with patch(
+            "vllm_omni.diffusion.forward_context.get_forward_context",
+            return_value=ctx,
+        ):
+            path = materialize_media_source(str(local_file))
+        assert path == str(local_file)
+
+    def test_local_path_without_allowed_path_raises(self, tmp_path):
+        local_file = tmp_path / "media.bin"
+        local_file.write_bytes(b"\x00" * 50)
+        ctx = _fake_forward_context(allowed_local_media_path="")
+        with patch(
+            "vllm_omni.diffusion.forward_context.get_forward_context",
+            return_value=ctx,
+        ):
+            with pytest.raises(RuntimeError, match="allowed-local-media-path"):
+                materialize_media_source(str(local_file))
+
+    def test_http_url_rejected_by_domain_filter(self):
+        ctx = _fake_forward_context(allowed_domains=["trusted.example.com"])
+        with patch(
+            "vllm_omni.diffusion.forward_context.get_forward_context",
+            return_value=ctx,
+        ):
+            with pytest.raises(ValueError, match="allowed domains"):
+                materialize_media_source("https://evil.internal/media.bin")
+
+    def test_file_uri_with_allowed_path(self, tmp_path):
+        media_file = tmp_path / "audio.wav"
+        media_file.write_bytes(b"\x00" * 50)
+        file_uri = media_file.as_uri()
+
+        ctx = _fake_forward_context(allowed_local_media_path=str(tmp_path))
+        with patch(
+            "vllm_omni.diffusion.forward_context.get_forward_context",
+            return_value=ctx,
+        ):
+            path = materialize_media_source(file_uri)
+        assert path == str(media_file)
+
+    def test_file_uri_without_allowed_path_raises(self, tmp_path):
+        media_file = tmp_path / "audio.wav"
+        media_file.write_bytes(b"\x00" * 50)
+        file_uri = media_file.as_uri()
+
+        ctx = _fake_forward_context(allowed_local_media_path="")
+        with patch(
+            "vllm_omni.diffusion.forward_context.get_forward_context",
+            return_value=ctx,
+        ):
+            with pytest.raises(RuntimeError, match="allowed-local-media-path"):
+                materialize_media_source(file_uri)

tests/entrypoints/openai_api/test_video_server.py

@@ -67,6 +67,10 @@ class FakeAsyncOmni:
         self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
         self.captured_prompt = None
         self.captured_sampling_params_list = None
+        self.model_config = SimpleNamespace(
+            allowed_local_media_path="",
+            allowed_media_domains=None,
+        )
 
     async def generate(self, prompt, request_id, sampling_params_list):
         self.captured_prompt = prompt
@@ -82,6 +86,12 @@ class BlockingVideoHandler:
         self.stage_configs = None
         self.started = threading.Event()
         self.cancelled = threading.Event()
+        self._engine_client = SimpleNamespace(
+            model_config=SimpleNamespace(
+                allowed_local_media_path="",
+                allowed_media_domains=None,
+            ),
+        )
 
     def set_stage_configs_if_missing(self, stage_configs):
         if self.stage_configs is None:
@@ -1765,3 +1775,77 @@ def test_worker_fps_multiplier_is_applied_to_sync_encoding(test_client, mocker:
     assert response.status_code == 200
     assert response.content == b"fps-multiplied"
     assert fps_values == [16]
+
+
+# ---------------------------------------------------------------------------
+# MediaConnector / SSRF protection (end-to-end through /v1/videos)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def restricted_test_client():
+    """Test client whose engine restricts media to trusted.example.com."""
+    engine = FakeAsyncOmni()
+    engine.model_config = SimpleNamespace(
+        allowed_local_media_path="",
+        allowed_media_domains=["trusted.example.com"],
+    )
+    app = FastAPI()
+    app.include_router(router)
+    app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion(
+        diffusion_engine=engine,
+        model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    )
+    with TestClient(app) as client:
+        yield client
+
+
+class TestVideoEndpointSSRF:
+    """SSRF protection via /v1/videos endpoint with domain-restricted engine."""
+
+    def test_blocked_image_reference_url_returns_400(self, restricted_test_client):
+        response = restricted_test_client.post(
+            "/v1/videos",
+            data={
+                "prompt": "A fox.",
+                "image_reference": json.dumps({"image_url": "https://evil.internal/image.png"}),
+            },
+        )
+        assert response.status_code == 400
+        assert "allowed domains" in response.json()["detail"].lower()
+
+    def test_blocked_video_reference_url_returns_400(self, restricted_test_client):
+        response = restricted_test_client.post(
+            "/v1/videos",
+            data={
+                "prompt": "Continue.",
+                "video_reference": json.dumps({"video_url": "https://evil.internal/video.mp4"}),
+            },
+        )
+        assert response.status_code == 400
+        assert "allowed domains" in response.json()["detail"].lower()
+
+    def test_blocked_audio_reference_url_returns_400(self, restricted_test_client):
+        response = restricted_test_client.post(
+            "/v1/videos",
+            data={
+                "prompt": "Narrate.",
+                "audio_reference": json.dumps({"audio_url": "https://evil.internal/audio.mp3"}),
+            },
+        )
+        assert response.status_code == 400
+        assert "allowed domains" in response.json()["detail"].lower()
+
+    def test_data_url_bypasses_domain_restriction(self, restricted_test_client, mocker: MockerFixture):
+        mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+            return_value=b"fake-video",
+        )
+        response = restricted_test_client.post(
+            "/v1/videos",
+            data={
+                "prompt": "A fox.",
+                "image_reference": json.dumps({"image_url": _make_test_image_data_url((40, 24))}),
+            },
+        )
+        assert response.status_code == 200

vllm_omni/entrypoints/openai/api_server.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
-import base64
 import dataclasses
 import io
 import json
@@ -20,7 +19,6 @@ from numbers import Integral
 from pathlib import Path
 from typing import Annotated, Any, Literal, cast
 
-import httpx
 import numpy as np
 import vllm.envs as envs
 from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket
@@ -85,6 +83,7 @@ from vllm.entrypoints.speech_to_text.translation.serving import (
     OpenAIServingTranslation,
 )
 from vllm.logger import init_logger
+from vllm.multimodal.media import MediaConnector
 from vllm.tasks import POOLING_TASKS
 from vllm.tool_parsers import ToolParserManager
 from vllm.utils import random_uuid
@@ -2038,17 +2037,23 @@ async def edit_images(
         # Match the offline path: RGB normalize when the caller opts into
         # Hunyuan-aware behavior. RGBA/P uploads otherwise diverge from offline.
         normalize_edit_images_rgb = bot_task is not None or sys_type is not None
-        pil_images = await _load_input_images(input_images_list, normalize_rgb=normalize_edit_images_rgb)
+        pil_images = await _load_input_images(
+            input_images_list,
+            engine_client.model_config,
+            normalize_rgb=normalize_edit_images_rgb,
+        )
         prompt["multi_modal_data"] = {}
         prompt["multi_modal_data"]["image"] = pil_images
 
         if mask_image is not None:
             # Mask role is different (alpha channel matters); never normalize.
-            loaded = await _load_input_images([mask_image], normalize_rgb=False)
+            loaded = await _load_input_images([mask_image], engine_client.model_config, normalize_rgb=False)
             prompt["multi_modal_data"]["mask_image"] = loaded[0]
 
         if reference_image is not None:
-            loaded = await _load_input_images([reference_image], normalize_rgb=normalize_edit_images_rgb)
+            loaded = await _load_input_images(
+                [reference_image], engine_client.model_config, normalize_rgb=normalize_edit_images_rgb
+            )
             prompt["multi_modal_data"]["reference_image"] = loaded[0]
 
         # 3 Build sample params
@@ -2535,41 +2540,41 @@ def _extract_images_from_result(result: Any) -> list[Any]:
 
 async def _load_input_images(
     inputs: list[str],
+    model_config: Any,
     *,
     normalize_rgb: bool = True,
 ) -> list[Image.Image]:
+    """Load images from data URIs, http(s) URLs, file:// URIs, or UploadFiles.
+
+    http(s) and ``file:`` URLs are fetched through vLLM's ``MediaConnector`` which
+    respects ``--allowed-media-domains`` and ``--allowed-local-media-path``
+    to prevent SSRF.
     """
-    convert to PIL.Image.Image list
-    """
+
     if isinstance(inputs, str):
         inputs = [inputs]
 
+    connector = MediaConnector(
+        allowed_local_media_path=model_config.allowed_local_media_path,
+        allowed_media_domains=model_config.allowed_media_domains,
+    )
+
     images: list[Image.Image] = []
 
     for inp in inputs:
-        # 1. URL + base64
-        if isinstance(inp, str) and inp.startswith("data:image"):
+        if isinstance(inp, str):
+            # http(s), data: and file: all handled by MediaConnector
             try:
-                _, b64_data = inp.split(",", 1)
-                image_bytes = base64.b64decode(b64_data)
-                img = Image.open(io.BytesIO(image_bytes))
-                images.append(img)
+                # Use RGBA to preserve alpha for transparent PNG.
+                result = await connector.fetch_image_async(
+                    inp,
+                    image_mode="RGBA",
+                )
+                images.append(result.media)
             except Exception as e:
-                raise ValueError(f"Invalid base64 image: {e}")
-
-        # 2. URL
-        elif isinstance(inp, str) and inp.startswith("http"):
-            async with httpx.AsyncClient(timeout=60) as client:
-                try:
-                    resp = await client.get(inp)
-                    resp.raise_for_status()
-                    img = Image.open(io.BytesIO(resp.content))
-                    images.append(img)
-                except Exception as e:
-                    raise ValueError(f"Failed to download image from URL {inp}: {e}")
-
-        # 3. UploadFile
+                raise ValueError(f"Failed to load image: {e}") from e
         elif hasattr(inp, "file"):
+            # UploadFile
             try:
                 img_data = await inp.read()
                 img = Image.open(io.BytesIO(img_data))
@@ -3021,6 +3026,7 @@ async def _parse_video_form(
             request.image_reference,
             request.video_reference,
             input_reference_bytes,
+            model_config=handler._engine_client.model_config,
             max_video_frames=decode_spec.max_frames,
             video_keep=decode_spec.keep,
         )
@@ -3032,8 +3038,15 @@ async def _parse_video_form(
 
     reference_audio: ReferenceAudio | None = None
     if request.audio_reference is not None:
+        audio_connector = MediaConnector(
+            allowed_local_media_path=handler._engine_client.model_config.allowed_local_media_path,
+            allowed_media_domains=handler._engine_client.model_config.allowed_media_domains,
+        )
         try:
-            audio_path = await decode_audio_url(request.audio_reference.audio_url)
+            audio_path = await decode_audio_url(
+                request.audio_reference.audio_url,
+                audio_connector,
+            )
         except InvalidInputReferenceError as exc:
             raise HTTPException(400, detail=str(exc)) from exc
         reference_audio = ReferenceAudio(path=audio_path)

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -10,12 +10,14 @@ import base64
 import binascii
 from collections import deque
 from io import BytesIO
+from pathlib import Path
 from typing import Any, Literal
 
-import httpx
 import numpy as np
 import torch
 from PIL import Image, UnidentifiedImageError
+from vllm.multimodal.media import MediaConnector
+from vllm.multimodal.media.base import MediaIO
 
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.protocol.videos import (
@@ -145,22 +147,12 @@ def _decode_base64_image(input_reference: str, *, source: str) -> Image.Image:
     raise InvalidInputReferenceError(f"Invalid {source}: image data is empty.")
 
 
-async def decode_image_url(image_url: str) -> Image.Image:
-    if image_url.startswith("data:image"):
-        return _decode_base64_image(image_url, source="image_reference.image_url")
-
-    if image_url.startswith(("http://", "https://")):
-        async with httpx.AsyncClient(timeout=60) as client:
-            try:
-                response = await client.get(image_url)
-                response.raise_for_status()
-            except httpx.HTTPError as exc:
-                raise InvalidInputReferenceError(
-                    "Invalid image_reference.image_url: failed to download image."
-                ) from exc
-        return _decode_image_bytes(response.content, source="image_reference.image_url")
-
-    raise InvalidInputReferenceError("Invalid image_reference.image_url: must be an http(s) URL or data URL.")
+async def decode_image_url(image_url: str, connector: MediaConnector) -> Image.Image:
+    """Fetch and decode an image URL using MediaConnector for SSRF protection."""
+    try:
+        return (await connector.fetch_image_async(image_url)).media
+    except Exception as exc:
+        raise InvalidInputReferenceError(f"Invalid image_reference.image_url: {exc}") from exc
 
 
 def _decode_base64_video(
@@ -184,96 +176,117 @@ def _decode_base64_video(
     raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.")
 
 
-async def decode_video_url(
-    video_url: str,
-    *,
-    max_frames: int | None = None,
-    keep: Literal["first", "last"] = "first",
-) -> VideoFrames:
-    if video_url.startswith("data:video"):
-        return _decode_base64_video(
-            video_url,
-            source="video_reference.video_url",
-            max_frames=max_frames,
-            keep=keep,
-        )
+class _VideoFramesMediaIO(MediaIO[VideoFrames]):
+    """MediaIO that decodes video bytes into VideoFrames with frame extraction."""
 
-    if video_url.startswith(("http://", "https://")):
-        async with httpx.AsyncClient(timeout=60) as client:
-            try:
-                response = await client.get(video_url)
-                response.raise_for_status()
-            except httpx.HTTPError as exc:
-                raise InvalidInputReferenceError(
-                    "Invalid video_reference.video_url: failed to download video."
-                ) from exc
+    def __init__(
+        self,
+        *,
+        source: str = "video_reference.video_url",
+        max_frames: int | None = None,
+        keep: Literal["first", "last"] = "first",
+    ) -> None:
+        self._source = source
+        self._max_frames = max_frames
+        self._keep = keep
+
+    def load_bytes(self, data: bytes) -> VideoFrames:
         return _decode_video_bytes(
-            response.content,
-            source="video_reference.video_url",
-            max_frames=max_frames,
-            keep=keep,
+            data,
+            source=self._source,
+            max_frames=self._max_frames,
+            keep=self._keep,
         )
 
-    raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.")
+    def load_base64(self, media_type: str, data: str) -> VideoFrames:
+        return self.load_bytes(base64.b64decode(data))
+
+    def load_file(self, filepath: Path) -> VideoFrames:
+        return self.load_bytes(filepath.read_bytes())
 
 
-async def decode_audio_url(audio_url: str) -> str:
-    """Decode an audio URL or data-URL to a temporary file path."""
-    import tempfile
+class _AudioFileMediaIO(MediaIO[str]):
+    """MediaIO that writes audio bytes to a temp file, returns the path."""
 
-    audio_bytes: bytes | None = None
+    def load_bytes(self, data: bytes) -> str:
+        return self._write_temp(data, suffix=".wav")
 
-    if audio_url.startswith("data:audio"):
-        _, b64_data = audio_url.split(",", 1)
-        try:
-            audio_bytes = base64.b64decode(b64_data)
-        except (binascii.Error, ValueError) as exc:
-            raise InvalidInputReferenceError(
-                "Invalid audio_reference.audio_url: audio data is not valid base64."
-            ) from exc
-    elif audio_url.startswith(("http://", "https://")):
-        async with httpx.AsyncClient(timeout=60) as client:
-            try:
-                response = await client.get(audio_url)
-                response.raise_for_status()
-            except httpx.HTTPError as exc:
-                raise InvalidInputReferenceError(
-                    "Invalid audio_reference.audio_url: failed to download audio."
-                ) from exc
-        audio_bytes = response.content
-    else:
-        raise InvalidInputReferenceError("Invalid audio_reference.audio_url: must be an http(s) URL or data URL.")
-
-    if not audio_bytes:
-        raise InvalidInputReferenceError("Invalid audio_reference: audio data is empty.")
-
-    suffix = ".wav"
-    if audio_url.startswith("data:audio/"):
-        mime = audio_url.split(";")[0].removeprefix("data:")
-        ext = mime.split("/")[-1]
+    def load_base64(self, media_type: str, data: str) -> str:
+        suffix = ".wav"
+        ext = media_type.split("/")[-1] if media_type else ""
         if ext in ("mpeg", "mp3"):
             suffix = ".mp3"
         elif ext == "wav":
             suffix = ".wav"
         elif ext.isalnum() and len(ext) <= 8:
             suffix = f".{ext}"
+        return self._write_temp(base64.b64decode(data), suffix=suffix)
 
-    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
-    tmp.write(audio_bytes)
-    tmp.close()
-    return tmp.name
+    def load_file(self, filepath: Path) -> str:
+        return str(filepath)
+
+    @staticmethod
+    def _write_temp(data: bytes, *, suffix: str) -> str:
+        import tempfile
+
+        if not data:
+            raise InvalidInputReferenceError("Invalid audio_reference: audio data is empty.")
+        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
+        tmp.write(data)
+        tmp.close()
+        return tmp.name
+
+
+async def decode_video_url(
+    video_url: str,
+    connector: MediaConnector,
+    *,
+    max_frames: int | None = None,
+    keep: Literal["first", "last"] = "first",
+) -> VideoFrames:
+    """Fetch and decode a video URL using MediaConnector for SSRF protection."""
+    try:
+        return await connector.load_from_url_async(
+            video_url,
+            _VideoFramesMediaIO(max_frames=max_frames, keep=keep),
+        )
+    except InvalidInputReferenceError:
+        raise
+    except Exception as exc:
+        raise InvalidInputReferenceError(f"Invalid video_reference.video_url: {exc}") from exc
+
+
+async def decode_audio_url(
+    audio_url: str,
+    connector: MediaConnector,
+) -> str:
+    """Fetch and decode an audio URL using MediaConnector for SSRF protection."""
+    try:
+        return await connector.load_from_url_async(
+            audio_url,
+            _AudioFileMediaIO(),
+        )
+    except InvalidInputReferenceError:
+        raise
+    except Exception as exc:
+        raise InvalidInputReferenceError(f"Invalid audio_reference.audio_url: {exc}") from exc
 
 
 async def decode_input_reference(
     image_reference: ImageReference | None,
     video_reference: VideoReference | None,
     input_reference_bytes: bytes | None,
+    model_config: Any,
     *,
     max_video_frames: int | None = None,
     video_keep: Literal["first", "last"] = "first",
 ) -> Image.Image | VideoFrames | None:
-    """Decode media input from multipart bytes, data URLs, or typed references."""
+    """Decode media input from multipart bytes, data URLs, or typed references.
 
+    http(s) image and video URLs are fetched through vLLM's
+    ``MediaConnector`` which respects ``--allowed-media-domains`` and
+    ``--allowed-local-media-path`` to prevent SSRF.
+    """
     provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference))
     if provided > 1:
         raise InvalidInputReferenceError("Provide only one of input_reference, image_reference, or video_reference.")
@@ -287,13 +300,22 @@ async def decode_input_reference(
         )
 
     if isinstance(image_reference, UrlImageReference):
-        return await decode_image_url(image_reference.image_url)
+        connector = MediaConnector(
+            allowed_local_media_path=model_config.allowed_local_media_path,
+            allowed_media_domains=model_config.allowed_media_domains,
+        )
+        return await decode_image_url(image_reference.image_url, connector)
     elif isinstance(image_reference, FileImageReference):
         raise InvalidInputReferenceError("Invalid image_reference: file_id is not supported yet.")
 
     if isinstance(video_reference, UrlVideoReference):
+        connector = MediaConnector(
+            allowed_local_media_path=model_config.allowed_local_media_path,
+            allowed_media_domains=model_config.allowed_media_domains,
+        )
         return await decode_video_url(
             video_reference.video_url,
+            connector,
             max_frames=max_video_frames,
             keep=video_keep,
         )

vllm_omni/transformers_utils/processors/audiox.py

@@ -14,9 +14,9 @@ from __future__ import annotations
 import base64
 import os
 import tempfile
+from pathlib import Path
 from typing import Any
 from urllib.parse import urlparse
-from urllib.request import urlopen
 
 import av
 import numpy as np
@@ -26,6 +26,8 @@ import torch.nn.functional as F
 import torchaudio.functional as taF
 from einops import rearrange
 from torchvision.io import read_image
+from vllm.multimodal.media import MediaConnector
+from vllm.multimodal.media.base import MediaIO
 
 # AudioX task taxonomy. Tasks beginning with "v" require a video input; tasks containing
 # "t" carry a text prompt. tv2*/v2* share the same conditioning pathways.
@@ -51,29 +53,51 @@ def normalize_prompts(prompts: list[Any]) -> list[dict[str, Any]]:
     return out
 
 
-def materialize_media_source(source: str) -> str:
-    """Return a local filesystem path for ``source``.
+class _TempFileMediaIO(MediaIO[str]):
+    """MediaIO that writes fetched bytes to a temp file, returns the path."""
 
-    Accepts a local path, a ``data:<mime>;base64,...`` URI, or an ``http(s)://`` URL.
-    Anything non-local is fetched into a NamedTemporaryFile and that path is returned;
-    callers don't need to clean the tempfile up (the OS does on exit).
-    """
-    if source.startswith("data:"):
-        _, _, payload = source.partition(",")
-        raw = base64.b64decode(payload)
-        f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
-        f.write(raw)
-        f.close()
-        return f.name
-    parsed = urlparse(source)
-    if parsed.scheme in ("http", "https"):
-        with urlopen(source) as resp:
-            data = resp.read()
+    def load_bytes(self, data: bytes) -> str:
+        return self._write_temp(data)
+
+    def load_base64(self, media_type: str, data: str) -> str:
+        return self._write_temp(base64.b64decode(data))
+
+    def load_file(self, filepath: Path) -> str:
+        return str(filepath)
+
+    @staticmethod
+    def _write_temp(data: bytes) -> str:
         f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
         f.write(data)
         f.close()
         return f.name
-    return source
+
+
+def _get_media_connector() -> MediaConnector:
+    """Build a MediaConnector from the current forward context's model config."""
+    from vllm_omni.diffusion.forward_context import get_forward_context
+
+    ctx = get_forward_context()
+    model_config = ctx.vllm_config.model_config if ctx.vllm_config else None
+    return MediaConnector(
+        allowed_local_media_path=getattr(model_config, "allowed_local_media_path", ""),
+        allowed_media_domains=getattr(model_config, "allowed_media_domains", None),
+    )
+
+
+def materialize_media_source(source: str) -> str:
+    """Return a local filesystem path for ``source``.
+
+    Accepts a local path, a ``data:<mime>;base64,...`` URI, or an ``http(s)://`` URL.
+    All inputs are routed through vLLM's ``MediaConnector`` which respects
+    ``--allowed-media-domains`` and ``--allowed-local-media-path`` to prevent
+    SSRF.  Bare local paths are converted to ``file://`` URIs so that the
+    same path restriction applies.
+    """
+    if not urlparse(source).scheme:
+        source = Path(source).resolve().as_uri()
+    connector = _get_media_connector()
+    return connector.load_from_url(source, _TempFileMediaIO())
 
 
 def _load_video_path_pyav(