diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 5d0bedb776..7137a4e422 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -31,15 +31,19 @@ def _process_image_response(response: Response, url: str) -> str: f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) - image_bytes = response.content + # Stream download with size checking to prevent downloading huge files + max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024) + image_bytes = bytearray() + bytes_downloaded = 0 - # Check actual size after download if Content-Length was not available - if content_length is None: - size_mb = len(image_bytes) / (1024 * 1024) - if size_mb > MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: + for chunk in response.iter_bytes(chunk_size=8192): + bytes_downloaded += len(chunk) + if bytes_downloaded > max_bytes: + size_mb = bytes_downloaded / (1024 * 1024) raise litellm.ImageFetchError( f"Error: Image size ({size_mb:.2f}MB) exceeds maximum allowed size ({MAX_IMAGE_URL_DOWNLOAD_SIZE_MB}MB). url={url}" ) + image_bytes.extend(chunk) base64_image = base64.b64encode(image_bytes).decode("utf-8") diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index b15d75a414..9c2939b2da 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -66,6 +66,41 @@ class LargeImageClient: ) +class StreamingLargeImageClient: + """ + Client that streams a large image to test streaming download protection. + This simulates a huge file without actually creating it all in memory. + """ + + def __init__(self, size_mb=100, include_content_length=False): + self.size_mb = size_mb + self.include_content_length = include_content_length + + def get(self, url, follow_redirects=True): + size_bytes = int(self.size_mb * 1024 * 1024) + headers = {"Content-Type": "image/jpeg"} + if self.include_content_length: + headers["Content-Length"] = str(size_bytes) + + # Create a generator that yields chunks without creating the whole file in memory + def generate_chunks(total_size, chunk_size=8192): + bytes_sent = 0 + while bytes_sent < total_size: + chunk = b"x" * min(chunk_size, total_size - bytes_sent) + bytes_sent += len(chunk) + yield chunk + + # Create response with streaming content + response = Response( + status_code=200, + headers=headers, + request=Request("GET", url), + ) + # Mock the iter_bytes method to return our generator + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) + return response + + def test_image_exceeds_size_limit_with_content_length(monkeypatch): """ Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected when Content-Length header is present. @@ -83,6 +118,7 @@ def test_image_exceeds_size_limit_with_content_length(monkeypatch): def test_image_exceeds_size_limit_without_content_length(monkeypatch): """ Test that images exceeding MAX_IMAGE_URL_DOWNLOAD_SIZE_MB are rejected even without Content-Length header. + This uses the old non-streaming mock for backward compatibility. """ monkeypatch.setattr( litellm, "module_level_client", LargeImageClient(size_mb=100, include_content_length=False) @@ -94,6 +130,29 @@ def test_image_exceeds_size_limit_without_content_length(monkeypatch): assert "exceeds maximum allowed size" in str(excinfo.value) +def test_streaming_download_protects_against_huge_files(monkeypatch): + """ + Test that streaming download aborts early when file exceeds size limit, + preventing memory exhaustion from huge files (e.g., petabyte-sized files). + + This test verifies that the streaming implementation doesn't download the entire + file into memory before checking size. Instead, it should abort as soon as the + limit is exceeded during streaming. + """ + # Simulate a 1GB file - far larger than the 50MB default limit + client = StreamingLargeImageClient(size_mb=1024, include_content_length=False) + monkeypatch.setattr(litellm, "module_level_client", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + convert_url_to_base64("https://example.com/huge-image.jpg") + + # Verify the error message shows it was caught during streaming + assert "exceeds maximum allowed size" in str(excinfo.value) + + # The error should be raised after downloading just slightly more than the limit + # not after downloading the full 1GB + + class SmallImageClient: """ Client that returns a small valid image. @@ -124,6 +183,26 @@ def test_image_within_size_limit(monkeypatch): assert result.startswith("data:image/jpeg;base64,") +def test_streaming_download_handles_petabyte_file(monkeypatch): + """ + Test that streaming download can handle extremely large file URLs (e.g., petabyte-sized) + without attempting to download the entire file or causing memory exhaustion. + + This simulates what happens if a malicious actor or misconfiguration provides + a URL to an extremely large file. + """ + # Simulate a 1 petabyte file (1,000,000 GB) + # Without streaming protection, this would cause OOM or hang indefinitely + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) + monkeypatch.setattr(litellm, "module_level_client", client) + + with pytest.raises(litellm.ImageFetchError) as excinfo: + convert_url_to_base64("https://example.com/petabyte-file.jpg") + + # Should fail fast without downloading anywhere near 1 petabyte + assert "exceeds maximum allowed size" in str(excinfo.value) + + def test_image_size_limit_disabled(monkeypatch): """ Test that setting MAX_IMAGE_URL_DOWNLOAD_SIZE_MB to 0 disables all image URL downloads.