fix(guardrails): cap the image fetch instead of buffering whatever the url serves

_build_image_content_item hands a caller-supplied url to
BedrockImageProcessor.process_image_async, which buffers the whole body through
async_safe_get before returning. The 4 MB check runs on bytes that are already
resident, so it rejects an oversized image without preventing the allocation --
and _build_input_content_items gathers these concurrently, so one request with
several urls multiplies it. An arbitrarily large or indefinitely chunked response
is enough to exhaust proxy memory.

async_safe_get takes an optional max_bytes and, when given one, streams the body
and aborts past the cap with PayloadTooLargeError. Omitted, it keeps the previous
buffering, so every existing caller -- including the model-call image paths in
factory.py -- is byte-for-byte unchanged. Only the guardrail passes it.

The rebuilt response drops content-encoding and content-length: aiter_bytes
yields decoded bytes, so carrying those over would describe the body wrongly.

PayloadTooLargeError subclasses ValueError, like SSRFError, so callers already
treating a bad remote response as a rejected fetch need no new except arm. The
guardrail names it explicitly anyway, so an operator reading the log sees "too
large" rather than "could not be read".

The two existing remote-url tests now stub `stream` rather than `get`, which is
the transport the capped path uses. The new test asserts on how many bytes were
pulled -- without that, an unstubbed transport would raise for the wrong reason
and the test would pass against the unfixed code.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -3407,14 +3407,14 @@ class BedrockImageProcessor:
         return base64_bytes, content_type
 
     @staticmethod
-    async def get_image_details_async(image_url) -> tuple[str, str]:
+    async def get_image_details_async(image_url, max_bytes: int | None = None) -> tuple[str, str]:
         try:
             client: Final = get_async_httpx_client(
                 llm_provider=httpxSpecialProvider.PromptFactory,
                 params={"concurrent_limit": 1},
             )
             # Send a GET request to the image URL
-            response: Final[httpx.Response] = await async_safe_get(client, image_url)
+            response: Final[httpx.Response] = await async_safe_get(client, image_url, max_bytes=max_bytes)
             response.raise_for_status()  # Raise an exception for HTTP errors
 
             return BedrockImageProcessor._post_call_image_processing(response, image_url)
@@ -3595,13 +3595,20 @@ class BedrockImageProcessor:
         return cls._create_bedrock_block(img_bytes, mime_type, image_format)
 
     @classmethod
-    async def process_image_async(cls, image_url: str, format: str | None) -> BedrockContentBlock:
-        """Asynchronous image processing."""
+    async def process_image_async(
+        cls, image_url: str, format: str | None, max_bytes: int | None = None
+    ) -> BedrockContentBlock:
+        """Asynchronous image processing.
+
+        ``max_bytes`` caps a remote fetch and is ignored for a base64 data URI,
+        whose size the caller already knows before calling. Omitting it keeps the
+        previous unbounded fetch.
+        """
 
         if "base64" in image_url:
             img_bytes, mime_type, image_format = cls._parse_base64_image(image_url)
         elif "http://" in image_url or "https://" in image_url:
-            img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url)
+            img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url, max_bytes=max_bytes)
             image_format = mime_type.split("/")[1]
         else:
             raise ValueError("Unsupported image type. Expected either image url or base64 encoded string")

litellm/litellm_core_utils/url_utils.py

@@ -44,6 +44,14 @@ class SSRFError(ValueError):
     """Raised when a URL targets a blocked network."""
 
 
+class PayloadTooLargeError(ValueError):
+    """Raised when a fetched body exceeds the caller's byte cap.
+
+    A ``ValueError`` subclass, like :class:`SSRFError`, so callers that already
+    treat a malformed remote response as a rejected fetch need no new except arm.
+    """
+
+
 def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str:
     """Percent-encode one user-controlled URL path segment.
 
@@ -412,21 +420,88 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
     raise SSRFError("Too many redirects")
 
 
-async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any:
-    """Async version of safe_get."""
+# Headers that describe the wire encoding of the body. `aiter_bytes` yields
+# decoded bytes, so carrying these onto the rebuilt response would describe it
+# wrongly (a gzip label on already-inflated bytes, a length from before decoding).
+_TRANSFER_ENCODING_HEADERS: Final = frozenset({"content-encoding", "content-length"})
+
+
+def _underlying_httpx_client(client: object) -> httpx.AsyncClient:
+    """Return the object exposing ``stream``.
+
+    ``AsyncHTTPHandler``/``HTTPHandler`` wrap an httpx client and forward only
+    ``get``/``post``/..., so streaming has to go through the wrapped ``.client``.
+    A raw httpx client is returned unchanged.
+    """
+    inner: Final = getattr(client, "client", client)
+    if not isinstance(inner, httpx.AsyncClient):
+        raise TypeError(f"cannot stream from {type(client).__name__}: no httpx client to stream with")
+    return inner
+
+
+async def _async_get_capped(
+    client: object,
+    url: str,
+    max_bytes: int,
+    request_kwargs: dict[str, object],  # mutable-ok: forwarded straight to httpx as **kwargs
+) -> httpx.Response:
+    """GET ``url``, aborting the transfer once the body exceeds ``max_bytes``.
+
+    ``client.get`` buffers the whole body before returning, so a caller-supplied
+    URL serving an arbitrarily large or indefinitely chunked response is an
+    unbounded allocation. Streaming makes the cap effective during the transfer
+    rather than after it.
+
+    Returns a fully-read response so callers keep using ``.content`` as before.
+    """
+    async with _underlying_httpx_client(client).stream("GET", url, **request_kwargs) as response:
+        chunks: list[bytes] = []  # mutable-ok: accumulator for the capped body
+        total = 0
+        async for chunk in response.aiter_bytes():
+            total += len(chunk)
+            if total > max_bytes:
+                raise PayloadTooLargeError(f"remote body exceeds {max_bytes} bytes")
+            chunks.append(chunk)
+    kept_headers: Final = [  # mutable-ok: httpx.Response takes the header pairs as a list
+        (k, v) for k, v in response.headers.multi_items() if k.lower() not in _TRANSFER_ENCODING_HEADERS
+    ]
+    return httpx.Response(
+        status_code=response.status_code,
+        headers=kept_headers,
+        content=b"".join(chunks),
+        request=response.request,
+    )
+
+
+async def async_safe_get(client: Any, url: str, max_bytes: int | None = None, **kwargs: Any) -> Any:
+    """Async version of safe_get.
+
+    ``max_bytes`` caps the response body, rejecting an oversized transfer with
+    :class:`PayloadTooLargeError` while it is still in flight. Omitting it keeps
+    the previous unbounded buffering, so existing callers are unaffected.
+    """
+
+    async def _issue(
+        target_url: str,
+        request_kwargs: dict[str, object],  # mutable-ok: forwarded straight to httpx as **kwargs
+    ) -> httpx.Response:
+        if max_bytes is None:
+            return await client.get(target_url, **request_kwargs)
+        return await _async_get_capped(client, target_url, max_bytes, request_kwargs)
+
     if not getattr(litellm, "user_url_validation", True):
         kwargs.setdefault("follow_redirects", True)
-        return await client.get(url, **kwargs)
+        return await _issue(url, kwargs)
     kwargs.pop("follow_redirects", None)
     caller_headers: Final = kwargs.pop("headers", {})
     for _ in range(_MAX_REDIRECTS):
         validated_url, original_host = validate_url(url)
-        response = await client.get(
-            validated_url,
-            headers={**caller_headers, "Host": original_host},
-            follow_redirects=False,
+        hop_kwargs: dict[str, object] = {  # mutable-ok: a fresh per-hop kwargs dict, consumed by this call
             **kwargs,
-        )
+            "headers": {**caller_headers, "Host": original_host},  # mutable-ok: httpx takes headers as a dict
+            "follow_redirects": False,
+        }
+        response = await _issue(validated_url, hop_kwargs)
         if not response.is_redirect:
             return response
         # Resolve the next hop against the ORIGINAL (pre-rewrite) URL so

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -34,7 +34,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
 from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost
 from litellm.litellm_core_utils.prompt_templates.factory import BedrockImageProcessor
-from litellm.litellm_core_utils.url_utils import SSRFError
+from litellm.litellm_core_utils.url_utils import PayloadTooLargeError, SSRFError
 from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
 from litellm.llms.base_llm.guardrail_translation.utils import (
     effective_scan_only_tool_results_for_guardrail,
@@ -484,7 +484,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             return None
 
         try:
-            block: Final = await BedrockImageProcessor.process_image_async(image_url=image_url, format=None)
+            # Cap the fetch itself. The decoded-size check below runs after the
+            # bytes are already resident, so on its own it does not stop a caller
+            # from pointing the proxy at an arbitrarily large or chunked response
+            # -- and `_build_input_content_items` gathers these concurrently, so
+            # one request with several URLs multiplies the allocation.
+            block: Final = await BedrockImageProcessor.process_image_async(
+                image_url=image_url, format=None, max_bytes=_MAX_IMAGE_BYTES
+            )
+        except PayloadTooLargeError as e:
+            # Named before the ValueError arm it subclasses, so the operator sees
+            # "too large" rather than "could not be read" for a size rejection.
+            self._handle_unscannable_image(reason=f"remote image over ApplyGuardrail's 4 MB limit: {e}")
+            return None
         except (httpx.HTTPError, SSRFError, ValueError, TypeError, KeyError, binascii.Error) as e:
             self._handle_unscannable_image(reason=f"image content could not be read: {e}")
             return None

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -4,6 +4,7 @@ Unit tests for Bedrock Guardrails
 
 import asyncio
 import base64
+import contextlib
 import json
 import sys
 from datetime import datetime, timezone
@@ -5297,6 +5298,32 @@ class TestBedrockGuardrailImageInput:
             request=httpx.Request("GET", url),
         )
 
+    @staticmethod
+    def _fake_stream(chunks: list[bytes], served: list[int] | None = None):
+        """Stand in for httpx.AsyncClient.stream, serving `chunks` one at a time.
+
+        The guardrail caps the transfer, so the fetch goes through `stream` rather
+        than `get`. `served` counts the chunks actually pulled, which is how a test
+        tells "stopped mid-transfer" apart from "read everything, then rejected".
+        """
+
+        @contextlib.asynccontextmanager
+        async def _stream(self, method: str, url, **kwargs):
+            async def _aiter_bytes():
+                for chunk in chunks:
+                    if served is not None:
+                        served.append(len(chunk))
+                    yield chunk
+
+            response = MagicMock()
+            response.status_code = 200
+            response.headers = httpx.Headers({"content-type": "image/jpeg"})
+            response.request = httpx.Request("GET", str(url))
+            response.aiter_bytes = _aiter_bytes
+            yield response
+
+        return _stream
+
     def _guardrail(self, **kwargs) -> BedrockGuardrail:
         return BedrockGuardrail(
             guardrail_name="bedrock-image",
@@ -5430,9 +5457,7 @@ class TestBedrockGuardrailImageInput:
                 "content": [{"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}],
             }
         ]
-        get = AsyncMock(return_value=self._jpeg_response(self._REMOTE_IMAGE_URL))
-
-        with patch.object(httpx.AsyncClient, "get", new=get):
+        with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream([self._JPEG_BYTES])):
             request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
 
         assert request["content"] == [
@@ -5456,8 +5481,12 @@ class TestBedrockGuardrailImageInput:
         url = "http://169.254.169.254/latest/meta-data/"
         messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}}]}]
         get = AsyncMock(return_value=self._jpeg_response(url))
+        served: list[int] = []
 
-        with patch.object(httpx.AsyncClient, "get", new=get):
+        with (
+            patch.object(httpx.AsyncClient, "get", new=get),
+            patch.object(httpx.AsyncClient, "stream", new=self._fake_stream([self._JPEG_BYTES], served)),
+        ):
             with pytest.raises(HTTPException):
                 await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
 
@@ -5465,7 +5494,9 @@ class TestBedrockGuardrailImageInput:
                 source="INPUT", messages=messages
             )
 
+        # Both transports are stubbed: neither the plain nor the capped fetch ran.
         get.assert_not_awaited()
+        assert served == []
         assert request["content"] == []
 
     @pytest.mark.asyncio
@@ -5531,6 +5562,35 @@ class TestBedrockGuardrailImageInput:
         kinds = [k for item in sent[0]["content"] for k in item]
         assert "image" in kinds, f"image never reached the payload: {kinds}"
 
+    @pytest.mark.asyncio
+    async def test_oversized_remote_image_is_cut_off_during_the_transfer(self):
+        """A caller-supplied url can serve an unbounded or indefinitely chunked body.
+
+        The decoded-size check runs once the bytes are already resident, so the cap
+        has to apply while the transfer is in flight. Asserting on how much was
+        pulled is what separates that from buffering it all and rejecting after.
+        """
+        messages = [
+            {
+                "role": "user",
+                "content": [{"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}],
+            }
+        ]
+        # 8 MB offered one MB at a time against a 4 MB cap.
+        chunks: list[bytes] = [b"\0" * (1024 * 1024) for _ in range(8)]
+        served: list[int] = []
+
+        with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream(chunks, served)):
+            with pytest.raises(HTTPException):
+                await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
+
+        # Without the cap the fetch buffers through `get` instead, so `served` stays
+        # empty and the HTTPException above would come from an unstubbed transport
+        # rather than from the size rejection. Assert the capped path actually ran.
+        assert served, "the fetch did not go through the capped stream path"
+        assert sum(served) <= 5 * 1024 * 1024, f"read {sum(served)} bytes past a 4 MB cap"
+        assert len(served) < len(chunks), "the whole body was pulled before rejecting it"
+
     @pytest.mark.asyncio
     async def test_apply_guardrail_scans_images_from_inputs(self):
         """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook.