fix(hosted_vllm): cap remote media downloads on video create

Stream image, video, and audio URL bodies with Content-Length and byte caps so an authenticated caller cannot exhaust worker memory

litellm/constants.py

@@ -69,6 +69,8 @@ DEFAULT_IMAGE_HEIGHT: Final = int(os.getenv("DEFAULT_IMAGE_HEIGHT", 300))
 # Maps to OpenAI's 50 MB payload limit - requests with images exceeding this size will be rejected
 # Set MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0 to disable image URL handling entirely
 MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: Final = float(os.getenv("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 50))
+MAX_VIDEO_MEDIA_URLS_PER_REQUEST: Final = get_env_int("MAX_VIDEO_MEDIA_URLS_PER_REQUEST", 10)
+MAX_VIDEO_MEDIA_URL_TOTAL_DOWNLOAD_SIZE_MB: Final = float(os.getenv("MAX_VIDEO_MEDIA_URL_TOTAL_DOWNLOAD_SIZE_MB", 100))
 MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB: Final = int(
     os.getenv("MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB", 1024)
 )  # 1MB = 1024KB

litellm/llms/hosted_vllm/videos/transformation.py

@@ -4,13 +4,20 @@ import base64
 import json
 import mimetypes
 from collections.abc import Mapping
+from dataclasses import dataclass
 from io import BufferedReader
 from types import MappingProxyType
 from typing import Final
 from urllib.parse import urlparse
 
+from httpx import Response
 from httpx._types import FileTypes, RequestFiles
 
+from litellm.constants import (
+    MAX_IMAGE_URL_DOWNLOAD_SIZE_MB,
+    MAX_VIDEO_MEDIA_URL_TOTAL_DOWNLOAD_SIZE_MB,
+    MAX_VIDEO_MEDIA_URLS_PER_REQUEST,
+)
 from litellm.images.utils import ImageEditRequestUtils
 from litellm.litellm_core_utils.url_utils import SSRFError, safe_get
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
@@ -19,6 +26,9 @@ from litellm.secret_managers.main import get_secret_str
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.videos.main import VideoCreateOptionalRequestParams
 
+_BYTES_PER_MB: Final = 1024 * 1024
+_STREAM_CHUNK_SIZE: Final = 8192
+
 _EXCLUDED_FORM_KEYS: Final = frozenset(
     {
         "model",
@@ -69,6 +79,28 @@ _REFERENCE_URL_KEYS: Final = MappingProxyType(
 )
 
 
+@dataclass(slots=True)
+class _MediaDownloadBudget:
+    remaining_urls: int
+    remaining_bytes: int
+    max_bytes_per_url: int
+
+    def consume_url_slot(self) -> None:
+        if self.max_bytes_per_url <= 0:
+            raise ValueError("remote media URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0)")
+        if self.remaining_urls < 1:
+            raise ValueError("too many remote media URL references on one video request")
+        self.remaining_urls -= 1
+
+    def max_read_bytes(self) -> int:
+        return min(self.max_bytes_per_url, self.remaining_bytes)
+
+    def consume_bytes(self, nbytes: int) -> None:
+        if nbytes > self.remaining_bytes:
+            raise ValueError("remote media download exceeded the per-request size limit")
+        self.remaining_bytes -= nbytes
+
+
 def _serialize_form_value(value: object) -> str:
     if isinstance(value, str):
         return value
@@ -95,14 +127,41 @@ def _content_type_from_headers(headers: Mapping[str, object]) -> str:
     return raw.split(";", 1)[0].strip() or "application/octet-stream"
 
 
-def _fetch_url_as_data_url(url: str, client: HTTPHandler) -> str:
+def _declared_content_length(headers: Mapping[str, object]) -> int | None:
+    raw: Final = headers.get("content-length")
+    if not isinstance(raw, str):
+        return None
+    stripped: Final = raw.strip()
+    if not stripped.isdigit():
+        return None
+    return int(stripped)
+
+
+def _read_capped_body(response: Response, max_bytes: int) -> bytes:
+    declared: Final = _declared_content_length(response.headers)
+    if declared is not None and declared > max_bytes:
+        response.close()
+        raise ValueError("remote media URL Content-Length exceeds the maximum allowed size")
+    body: Final = bytearray()  # mutable-ok: streaming accumulator with a hard cap
+    for chunk in response.iter_bytes(chunk_size=_STREAM_CHUNK_SIZE):
+        body.extend(chunk)
+        if len(body) > max_bytes:
+            response.close()
+            raise ValueError("remote media download exceeded the maximum allowed size")
+    return bytes(body)
+
+
+def _fetch_url_as_data_url(url: str, client: HTTPHandler, budget: _MediaDownloadBudget) -> str:
+    budget.consume_url_slot()
     response: Final = safe_get(client, url)
     response.raise_for_status()
-    encoded: Final = base64.b64encode(response.content).decode("ascii")
+    content: Final = _read_capped_body(response, budget.max_read_bytes())
+    budget.consume_bytes(len(content))
+    encoded: Final = base64.b64encode(content).decode("ascii")
     return f"data:{_content_type_from_headers(response.headers)};base64,{encoded}"
 
 
-def _inline_reference_item(url_key: str, item: object, client: HTTPHandler) -> object:
+def _inline_reference_item(url_key: str, item: object, client: HTTPHandler, budget: _MediaDownloadBudget) -> object:
     if not isinstance(item, Mapping):
         return item
     url: Final = item.get(url_key)
@@ -113,18 +172,23 @@ def _inline_reference_item(url_key: str, item: object, client: HTTPHandler) -> o
         return item
     if scheme not in ("http", "https"):
         raise SSRFError(f"URL scheme '{scheme}' is not allowed")
-    return {**item, url_key: _fetch_url_as_data_url(url, client)}  # mutable-ok: JSON form field serialized immediately
+    return {  # mutable-ok: JSON form field serialized immediately
+        **item,
+        url_key: _fetch_url_as_data_url(url, client, budget),
+    }
 
 
-def _inline_media_reference(field_name: str, value: object, client: HTTPHandler) -> object:
+def _inline_media_reference(
+    field_name: str, value: object, client: HTTPHandler, budget: _MediaDownloadBudget
+) -> object:
     url_key: Final = _REFERENCE_URL_KEYS.get(field_name)
     if url_key is None:
         return value
     parsed: Final = _maybe_json(value)
     if isinstance(parsed, list):
-        return tuple(_inline_reference_item(url_key, item, client) for item in parsed)
+        return tuple(_inline_reference_item(url_key, item, client, budget) for item in parsed)
     if isinstance(parsed, Mapping):
-        return _inline_reference_item(url_key, parsed, client)
+        return _inline_reference_item(url_key, parsed, client, budget)
     return value
 
 
@@ -150,9 +214,41 @@ class HostedVLLMVideoConfig(OpenAIVideoConfig):
     https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/videos_api/
     """
 
-    def __init__(self, media_http_client: HTTPHandler | None = None) -> None:
+    def __init__(
+        self,
+        media_http_client: HTTPHandler | None = None,
+        *,
+        max_media_bytes_per_url: int | None = None,
+        max_media_bytes_per_request: int | None = None,
+        max_media_urls_per_request: int | None = None,
+    ) -> None:
         super().__init__()
         self._media_http_client = media_http_client
+        self._max_media_bytes_per_url = max_media_bytes_per_url
+        self._max_media_bytes_per_request = max_media_bytes_per_request
+        self._max_media_urls_per_request = max_media_urls_per_request
+
+    def _media_budget(self) -> _MediaDownloadBudget:
+        per_url: Final = (
+            self._max_media_bytes_per_url
+            if self._max_media_bytes_per_url is not None
+            else int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * _BYTES_PER_MB)
+        )
+        per_request: Final = (
+            self._max_media_bytes_per_request
+            if self._max_media_bytes_per_request is not None
+            else int(MAX_VIDEO_MEDIA_URL_TOTAL_DOWNLOAD_SIZE_MB * _BYTES_PER_MB)
+        )
+        url_slots: Final = (
+            self._max_media_urls_per_request
+            if self._max_media_urls_per_request is not None
+            else MAX_VIDEO_MEDIA_URLS_PER_REQUEST
+        )
+        return _MediaDownloadBudget(
+            remaining_urls=url_slots,
+            remaining_bytes=per_request,
+            max_bytes_per_url=per_url,
+        )
 
     def get_supported_openai_params(self, model: str) -> list:  # mutable-ok: BaseVideoConfig contract
         return [  # mutable-ok: BaseVideoConfig returns list
@@ -212,9 +308,13 @@ class HostedVLLMVideoConfig(OpenAIVideoConfig):
         headers: dict,  # mutable-ok: BaseVideoConfig contract
     ) -> tuple[dict, RequestFiles, str]:  # mutable-ok: BaseVideoConfig contract
         media_client: Final = self._media_http_client or HTTPHandler(concurrent_limit=1)
+        media_budget: Final = self._media_budget()
         input_reference: Final = video_create_optional_request_params.get("input_reference")
         form_files: Final = tuple(
-            (key, (None, _serialize_form_value(_inline_media_reference(key, value, media_client))))
+            (
+                key,
+                (None, _serialize_form_value(_inline_media_reference(key, value, media_client, media_budget))),
+            )
             for key, value in video_create_optional_request_params.items()
             if key not in _EXCLUDED_FORM_KEYS and value is not None
         )

tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py

@@ -349,3 +349,60 @@ def test_file_scheme_media_reference_is_rejected():
             litellm_params=GenericLiteLLMParams(),
             headers={},
         )
+
+
+def _transform_with_references(config: HostedVLLMVideoConfig, **references: object):
+    return config.transform_video_create_request(
+        model="MiniMax-H3",
+        prompt="a person singing",
+        api_base="http://localhost:8091/v1/videos",
+        video_create_optional_request_params=references,
+        litellm_params=GenericLiteLLMParams(),
+        headers={},
+    )
+
+
+def test_oversized_content_length_is_rejected_before_body():
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            content=b"",
+            headers={"content-type": "video/mp4", "content-length": str(51 * 1024 * 1024)},
+        )
+
+    config = HostedVLLMVideoConfig(media_http_client=_http_handler_for(handler))
+    with pytest.raises(ValueError, match="Content-Length"):
+        _transform_with_references(config, video_reference={"video_url": "http://1.1.1.1/clip.mp4"})
+
+
+def test_streamed_body_over_per_url_cap_is_rejected():
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            content=iter((b"1234", b"5678")),
+            headers={"content-type": "image/png"},
+        )
+
+    config = HostedVLLMVideoConfig(
+        media_http_client=_http_handler_for(handler),
+        max_media_bytes_per_url=4,
+        max_media_bytes_per_request=4,
+    )
+    with pytest.raises(ValueError, match="exceeded the maximum allowed size"):
+        _transform_with_references(config, image_reference={"image_url": "http://1.1.1.1/face.png"})
+
+
+def test_too_many_remote_media_urls_are_rejected():
+    def handler(request: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, content=b"ok", headers={"content-type": "image/png"})
+
+    config = HostedVLLMVideoConfig(
+        media_http_client=_http_handler_for(handler),
+        max_media_urls_per_request=1,
+    )
+    with pytest.raises(ValueError, match="too many remote media URL references"):
+        _transform_with_references(
+            config,
+            image_reference={"image_url": "http://1.1.1.1/a.png"},
+            audio_reference={"audio_url": "http://1.1.1.1/b.mp3"},
+        )