From 657dbce543d50591b00a10b1e4db35ded6af7465 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:13:09 +0800 Subject: [PATCH] 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) --- .../prompt_templates/factory.py | 17 +++- litellm/litellm_core_utils/url_utils.py | 91 +++++++++++++++++-- .../guardrail_hooks/bedrock_guardrails.py | 16 +++- .../test_bedrock_guardrails.py | 68 +++++++++++++- 4 files changed, 173 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b676077ab0e..dafdf986c33 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/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") diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0a59eaa75d3..0b26eaa869e 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 61f3003a0ea..139674b8aff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 3b4064f4385..a849cfbfd71 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/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.