diff --git a/tests/entrypoints/openai_api/test_decode_audio_url.py b/tests/entrypoints/openai_api/test_decode_audio_url.py index 2a5c5be47..fbf56b70e 100644 --- a/tests/entrypoints/openai_api/test_decode_audio_url.py +++ b/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) diff --git a/tests/entrypoints/openai_api/test_decode_video_url.py b/tests/entrypoints/openai_api/test_decode_video_url.py new file mode 100644 index 000000000..64ff7e2e0 --- /dev/null +++ b/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 diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 31a82102a..49c43d7a9 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/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("") == 2, ( f"N=2 reference images must emit 2 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) diff --git a/tests/entrypoints/openai_api/test_materialize_media_source.py b/tests/entrypoints/openai_api/test_materialize_media_source.py new file mode 100644 index 000000000..a0ce824b3 --- /dev/null +++ b/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) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 94b17b2b8..da97c4b94 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 11186a0e7..fbfd3aafe 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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) diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index aa6479189..ee5791dee 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/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, ) diff --git a/vllm_omni/transformers_utils/processors/audiox.py b/vllm_omni/transformers_utils/processors/audiox.py index b0e53cfc7..8839ffac9 100644 --- a/vllm_omni/transformers_utils/processors/audiox.py +++ b/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:;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:;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(