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.