Revert "[Bugfix] Use MediaConnector for image/video URL fetching to prevent SSRF" (#4751)
This commit is contained in:
Gao Han authored and GitHub committed
1 parent
908e7d32d9
commit
c582805e7a
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
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import asyncio
|
||||
import base64
|
||||
import dataclasses
|
||||
import io
|
||||
import json
|
||||
@@ -19,6 +20,7 @@ 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
|
||||
@@ -83,7 +85,6 @@ 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
|
||||
@@ -2036,23 +2037,17 @@ 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,
|
||||
engine_client.model_config,
|
||||
normalize_rgb=normalize_edit_images_rgb,
|
||||
)
|
||||
pil_images = await _load_input_images(input_images_list, 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], engine_client.model_config, normalize_rgb=False)
|
||||
loaded = await _load_input_images([mask_image], normalize_rgb=False)
|
||||
prompt["multi_modal_data"]["mask_image"] = loaded[0]
|
||||
|
||||
if reference_image is not None:
|
||||
loaded = await _load_input_images(
|
||||
[reference_image], engine_client.model_config, normalize_rgb=normalize_edit_images_rgb
|
||||
)
|
||||
loaded = await _load_input_images([reference_image], normalize_rgb=normalize_edit_images_rgb)
|
||||
prompt["multi_modal_data"]["reference_image"] = loaded[0]
|
||||
|
||||
# 3 Build sample params
|
||||
@@ -2539,41 +2534,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:
|
||||
if isinstance(inp, str):
|
||||
# http(s), data: and file: all handled by MediaConnector
|
||||
# 1. URL + base64
|
||||
if isinstance(inp, str) and inp.startswith("data:image"):
|
||||
try:
|
||||
# Use RGBA to preserve alpha for transparent PNG.
|
||||
result = await connector.fetch_image_async(
|
||||
inp,
|
||||
image_mode="RGBA",
|
||||
)
|
||||
images.append(result.media)
|
||||
_, b64_data = inp.split(",", 1)
|
||||
image_bytes = base64.b64decode(b64_data)
|
||||
img = Image.open(io.BytesIO(image_bytes))
|
||||
images.append(img)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to load image: {e}") from 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
|
||||
elif hasattr(inp, "file"):
|
||||
# UploadFile
|
||||
try:
|
||||
img_data = await inp.read()
|
||||
img = Image.open(io.BytesIO(img_data))
|
||||
@@ -3025,7 +3020,6 @@ 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,
|
||||
)
|
||||
@@ -3037,15 +3031,8 @@ 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_connector,
|
||||
)
|
||||
audio_path = await decode_audio_url(request.audio_reference.audio_url)
|
||||
except InvalidInputReferenceError as exc:
|
||||
raise HTTPException(400, detail=str(exc)) from exc
|
||||
reference_audio = ReferenceAudio(path=audio_path)
|
||||
|
||||
@@ -10,14 +10,12 @@ 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 (
|
||||
@@ -147,12 +145,22 @@ 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, 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
|
||||
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.")
|
||||
|
||||
|
||||
def _decode_base64_video(
|
||||
@@ -176,117 +184,96 @@ def _decode_base64_video(
|
||||
raise InvalidInputReferenceError(f"Invalid {source}: video data is empty.")
|
||||
|
||||
|
||||
class _VideoFramesMediaIO(MediaIO[VideoFrames]):
|
||||
"""MediaIO that decodes video bytes into VideoFrames with frame extraction."""
|
||||
|
||||
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(
|
||||
data,
|
||||
source=self._source,
|
||||
max_frames=self._max_frames,
|
||||
keep=self._keep,
|
||||
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,
|
||||
)
|
||||
|
||||
def load_base64(self, media_type: str, data: str) -> VideoFrames:
|
||||
return self.load_bytes(base64.b64decode(data))
|
||||
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
|
||||
return _decode_video_bytes(
|
||||
response.content,
|
||||
source="video_reference.video_url",
|
||||
max_frames=max_frames,
|
||||
keep=keep,
|
||||
)
|
||||
|
||||
def load_file(self, filepath: Path) -> VideoFrames:
|
||||
return self.load_bytes(filepath.read_bytes())
|
||||
raise InvalidInputReferenceError("Invalid video_reference.video_url: must be an http(s) URL or data URL.")
|
||||
|
||||
|
||||
class _AudioFileMediaIO(MediaIO[str]):
|
||||
"""MediaIO that writes audio bytes to a temp file, returns the path."""
|
||||
async def decode_audio_url(audio_url: str) -> str:
|
||||
"""Decode an audio URL or data-URL to a temporary file path."""
|
||||
import tempfile
|
||||
|
||||
def load_bytes(self, data: bytes) -> str:
|
||||
return self._write_temp(data, suffix=".wav")
|
||||
audio_bytes: bytes | None = None
|
||||
|
||||
def load_base64(self, media_type: str, data: str) -> str:
|
||||
suffix = ".wav"
|
||||
ext = media_type.split("/")[-1] if media_type else ""
|
||||
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]
|
||||
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)
|
||||
|
||||
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
|
||||
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
|
||||
tmp.write(audio_bytes)
|
||||
tmp.close()
|
||||
return tmp.name
|
||||
|
||||
|
||||
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.")
|
||||
@@ -300,22 +287,13 @@ async def decode_input_reference(
|
||||
)
|
||||
|
||||
if isinstance(image_reference, UrlImageReference):
|
||||
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)
|
||||
return await decode_image_url(image_reference.image_url)
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -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,8 +26,6 @@ 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.
|
||||
@@ -53,51 +51,29 @@ def normalize_prompts(prompts: list[Any]) -> list[dict[str, Any]]:
|
||||
return out
|
||||
|
||||
|
||||
class _TempFileMediaIO(MediaIO[str]):
|
||||
"""MediaIO that writes fetched bytes to a temp file, returns the path."""
|
||||
|
||||
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
|
||||
|
||||
|
||||
def _get_media_connector() -> MediaConnector:
|
||||
"""Build a MediaConnector from the current forward context's model config."""
|
||||
from vllm_omni.diffusion.forward_context import get_forward_context
|
||||
|
||||
ctx = get_forward_context()
|
||||
model_config = ctx.vllm_config.model_config if ctx.vllm_config else None
|
||||
return MediaConnector(
|
||||
allowed_local_media_path=getattr(model_config, "allowed_local_media_path", ""),
|
||||
allowed_media_domains=getattr(model_config, "allowed_media_domains", None),
|
||||
)
|
||||
|
||||
|
||||
def materialize_media_source(source: str) -> str:
|
||||
"""Return a local filesystem path for ``source``.
|
||||
|
||||
Accepts a local path, a ``data:<mime>;base64,...`` URI, or an ``http(s)://`` URL.
|
||||
All inputs are routed through vLLM's ``MediaConnector`` which respects
|
||||
``--allowed-media-domains`` and ``--allowed-local-media-path`` to prevent
|
||||
SSRF. Bare local paths are converted to ``file://`` URIs so that the
|
||||
same path restriction applies.
|
||||
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 not urlparse(source).scheme:
|
||||
source = Path(source).resolve().as_uri()
|
||||
connector = _get_media_connector()
|
||||
return connector.load_from_url(source, _TempFileMediaIO())
|
||||
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()
|
||||
f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False)
|
||||
f.write(data)
|
||||
f.close()
|
||||
return f.name
|
||||
return source
|
||||
|
||||
|
||||
def _load_video_path_pyav(
|
||||
|
||||
Reference in new issue
Block a user