A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 16:00 UTC
(security) Enforce frame limit in VideoMediaIO (#38636)
Signed-off-by: jperezde <jperezde@redhat.com>
Juan Pérez de Algaba committed
Apr 1, 2026 at 12:23 UTC
58ee61422169ce17e08248f8efa1e9df434fe395
2 files changed
+69
-10
tests/multimodal/media/test_video.py
+61
-9
@@ -239,6 +239,17 @@ def test_video_media_io_backend_env_var_fallback(monkeypatch: pytest.MonkeyPatch
239
assert metadata_missing["video_backend"] == "test_video_backend_override_2"
240
241
242
+def _make_jpeg_b64_frames(n: int, width: int = 8, height: int = 8) -> list[str]:
243
+ """Return *n* tiny base64-encoded JPEG frames."""
244
+ frames: list[str] = []
245
+ for i in range(n):
246
+ img = Image.new("RGB", (width, height), color=(i % 256, 0, 0))
247
+ buf = io.BytesIO()
248
+ img.save(buf, format="JPEG")
249
+ frames.append(pybase64.b64encode(buf.getvalue()).decode("ascii"))
250
+ return frames
251
+
252
+
253
def test_load_base64_jpeg_returns_metadata():
254
"""Regression test: load_base64 with video/jpeg must return metadata.
255
@@ -248,16 +259,8 @@ def test_load_base64_jpeg_returns_metadata():
259
"""
260
261
num_test_frames = 3
251
- frame_width, frame_height = 8, 8
252
-
253
- # Build a few tiny JPEG frames and base64-encode them
254
- b64_frames = []
255
- for i in range(num_test_frames):
256
- img = Image.new("RGB", (frame_width, frame_height), color=(i * 80, 0, 0))
257
- buf = io.BytesIO()
258
- img.save(buf, format="JPEG")
259
- b64_frames.append(pybase64.b64encode(buf.getvalue()).decode("ascii"))
262
263
+ b64_frames = _make_jpeg_b64_frames(num_test_frames)
264
data = ",".join(b64_frames)
265
266
imageio = ImageMediaIO()
@@ -287,3 +290,52 @@ def test_load_base64_jpeg_returns_metadata():
290
# Default fps=1 → duration == num_frames
291
assert metadata["fps"] == 1.0
292
assert metadata["duration"] == float(num_test_frames)
293
+
294
+
295
+def test_load_base64_jpeg_enforces_num_frames_limit():
296
+ """Frames beyond num_frames must be truncated in the video/jpeg path.
297
+
298
+ Without the limit an attacker can send thousands of base64 JPEG frames
299
+ in a single request and exhaust server memory (OOM).
300
+ """
301
+ num_frames_limit = 4
302
+ sent_frames = 20
303
+
304
+ b64_frames = _make_jpeg_b64_frames(sent_frames)
305
+ data = ",".join(b64_frames)
306
+
307
+ imageio = ImageMediaIO()
308
+ videoio = VideoMediaIO(imageio, num_frames=num_frames_limit)
309
+ frames, metadata = videoio.load_base64("video/jpeg", data)
310
+
311
+ assert frames.shape[0] == num_frames_limit
312
+ assert metadata["total_num_frames"] == num_frames_limit
313
+ assert metadata["frames_indices"] == list(range(num_frames_limit))
314
+
315
+
316
+def test_load_base64_jpeg_no_limit_when_num_frames_negative():
317
+ """When num_frames is -1, all frames should be loaded without truncation."""
318
+ sent_frames = 10
319
+
320
+ b64_frames = _make_jpeg_b64_frames(sent_frames)
321
+ data = ",".join(b64_frames)
322
+
323
+ imageio = ImageMediaIO()
324
+ videoio = VideoMediaIO(imageio, num_frames=-1)
325
+ frames, metadata = videoio.load_base64("video/jpeg", data)
326
+
327
+ assert frames.shape[0] == sent_frames
328
+ assert metadata["total_num_frames"] == sent_frames
329
+ assert metadata["frames_indices"] == list(range(sent_frames))
330
+
331
+
332
+def test_load_base64_jpeg_raises_on_zero_num_frames():
333
+ """num_frames=0 is invalid and should raise ValueError."""
334
+ b64_frames = _make_jpeg_b64_frames(3)
335
+ data = ",".join(b64_frames)
336
+
337
+ imageio = ImageMediaIO()
338
+ videoio = VideoMediaIO(imageio, num_frames=0)
339
+
340
+ with pytest.raises(ValueError, match="num_frames must be greater than 0 or -1"):
341
+ videoio.load_base64("video/jpeg", data)
vllm/multimodal/media/video.py
+8
-1
@@ -80,8 +80,15 @@ class VideoMediaIO(MediaIO[tuple[npt.NDArray, dict[str, Any]]]):
80
"image/jpeg",
81
)
82
83
+ if self.num_frames > 0:
84
+ frame_parts = data.split(",", self.num_frames)[: self.num_frames]
85
+ elif self.num_frames == 0:
86
+ raise ValueError("num_frames must be greater than 0 or -1")
87
+ else:
88
+ frame_parts = data.split(",")
89
+
90
frames = np.stack(
84
- [np.asarray(load_frame(frame_data)) for frame_data in data.split(",")]
91
+ [np.asarray(load_frame(frame_data)) for frame_data in frame_parts]
92
)
93
total = int(frames.shape[0])
94
fps = float(self.kwargs.get("fps", 1))