Revert "[Bugfix] Use MediaConnector for image/video URL fetching to prevent S…"
This reverts commit d8edf0c9d5.
This commit is contained in:
Gao Han authored and GitHub committed
1 parent
908e7d32d9
commit
d98f94fa29
8 files changed
+243
-997
No files matched your search
@@ -1,19 +1,16 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for decode_audio_url helper and _AudioFileMediaIO."""
|
||||
"""Unit tests for decode_audio_url helper."""
|
||||
|
||||
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 (
|
||||
_AudioFileMediaIO,
|
||||
decode_audio_url,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -27,127 +24,15 @@ 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 — no mocks, real code path."""
|
||||
"""Tests for base64 data URL decoding."""
|
||||
|
||||
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, _permissive_connector()))
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert os.path.isfile(path)
|
||||
assert path.endswith(".mp3")
|
||||
@@ -161,7 +46,7 @@ class TestDecodeAudioUrlDataUrl:
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/wav;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url, _permissive_connector()))
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert path.endswith(".wav")
|
||||
with open(path, "rb") as f:
|
||||
@@ -169,103 +54,134 @@ 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="Incorrect padding"):
|
||||
_run(decode_audio_url(url, _permissive_connector()))
|
||||
|
||||
with pytest.raises(InvalidInputReferenceError, match="not valid base64"):
|
||||
_run(decode_audio_url(url))
|
||||
|
||||
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, _permissive_connector()))
|
||||
_run(decode_audio_url(url))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# decode_audio_url with file:// URIs (exercises load_file through 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)
|
||||
|
||||
|
||||
class TestDecodeAudioUrlFileUri:
|
||||
"""Tests using file:// URIs — exercises MediaConnector + load_file."""
|
||||
class TestDecodeAudioUrlHttp:
|
||||
"""Tests for HTTP URL decoding."""
|
||||
|
||||
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)
|
||||
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_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))
|
||||
async def _mock_get(self, url):
|
||||
return fake_resp
|
||||
|
||||
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))
|
||||
monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
|
||||
monkeypatch.setattr(httpx.Response, "raise_for_status", lambda self: None)
|
||||
|
||||
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)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# decode_audio_url — invalid schemes and SSRF
|
||||
# ---------------------------------------------------------------------------
|
||||
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"))
|
||||
|
||||
|
||||
class TestDecodeAudioUrlInvalid:
|
||||
"""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_audio_url("/path/to/local/file.mp3", _permissive_connector()))
|
||||
with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
|
||||
_run(decode_audio_url("/path/to/local/file.mp3"))
|
||||
|
||||
def test_ftp_url_raises(self):
|
||||
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()))
|
||||
with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
|
||||
_run(decode_audio_url("ftp://example.com/file.mp3"))
|
||||
|
||||
def test_empty_string_raises(self):
|
||||
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)
|
||||
with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
|
||||
_run(decode_audio_url(""))
|
||||
@@ -1,210 +0,0 @@
|
||||
# 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
|
||||
@@ -13,7 +13,6 @@ 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
|
||||
@@ -25,19 +24,12 @@ 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,
|
||||
_load_input_images,
|
||||
router,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.api_server import _check_max_generated_image_size, _DiffusionServingModels, 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
|
||||
|
||||
@@ -190,7 +182,6 @@ 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
|
||||
@@ -280,10 +271,6 @@ 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)
|
||||
|
||||
@@ -348,10 +335,6 @@ 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)
|
||||
|
||||
@@ -416,10 +399,6 @@ 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)
|
||||
|
||||
@@ -456,24 +435,9 @@ 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
|
||||
@@ -765,10 +729,6 @@ 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()
|
||||
@@ -2176,107 +2136,3 @@ 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)
|
||||
@@ -1,173 +0,0 @@
|
||||
# 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)
|
||||
@@ -68,10 +68,6 @@ 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
|
||||
@@ -87,12 +83,6 @@ 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:
|
||||
@@ -1869,77 +1859,3 @@ 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
|
||||
Reference in new issue
Block a user