diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 216b32c9b50..9138f68e930 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -614,6 +614,12 @@ def _headers_of_the_decoded_body(headers: httpx.Headers) -> httpx.Headers: ) +def _declared_decoded_length(headers: httpx.Headers) -> int: + if headers.get("content-encoding", "identity") != "identity": + return 0 + return int(headers.get("content-length", "0")) + + class AsyncHTTPHandler: def __init__( self, @@ -789,7 +795,7 @@ class AsyncHTTPHandler: ) if response.is_redirect or response.is_error: return httpx.Response(response.status_code, headers=response.headers, request=response.request) - if int(response.headers.get("content-length", "0")) > max_bytes: + if _declared_decoded_length(response.headers) > max_bytes: raise HTTPResponseLimitError("Response exceeds the configured size limit") with BytesIO() as body: async for chunk in response.aiter_bytes(chunk_size=65536): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 6b0211c3866..262602a658f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -1569,7 +1569,10 @@ class TestBoundedOpenAPISpecLoading: assert route.calls[0].request.headers["accept-encoding"] == "identity" @pytest.mark.asyncio - @pytest.mark.parametrize("headers", [{"content-length": "1000000"}, {"content-encoding": "gzip"}]) + @pytest.mark.parametrize( + "headers", + [{"content-length": "1000000"}, {"content-length": "1000000", "content-encoding": "identity"}], + ) async def test_unsafe_response_headers_reject_before_reading(self, respx_mock, monkeypatch, headers): import httpx from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError @@ -1593,6 +1596,36 @@ class TestBoundedOpenAPISpecLoading: await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12) assert closed == [True] + @pytest.mark.asyncio + async def test_compressed_spec_loads_when_its_decoded_size_fits(self, respx_mock, monkeypatch): + import gzip + + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + compressed = gzip.compress(b'{"paths":{}}') + assert len(compressed) > 12 + respx_mock.get("https://93.184.216.34/spec.json").respond( + 200, content=compressed, headers={"content-encoding": "gzip"} + ) + assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12) == {"paths": {}} + + @pytest.mark.asyncio + async def test_compressed_spec_is_capped_on_its_decoded_size(self, respx_mock, monkeypatch): + import gzip + + from litellm.llms.custom_httpx.http_handler import HTTPResponseLimitError + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import load_openapi_spec_async + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + compressed = gzip.compress(b'{"paths":{},"pad":"' + b"x" * 100_000 + b'"}') + assert len(compressed) < 12_000 + respx_mock.get("https://93.184.216.34/spec.json").respond( + 200, content=compressed, headers={"content-encoding": "gzip"} + ) + with pytest.raises(HTTPResponseLimitError): + await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=12_000) + @pytest.mark.asyncio async def test_chunked_response_is_bounded_and_closed(self, respx_mock, monkeypatch): import httpx diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 81e462bb46c..c27f1536325 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1793,6 +1793,23 @@ async def test_bounded_get_decodes_a_compressed_body_and_reports_the_decoded_len assert response.headers["content-length"] == str(len(document)) +@pytest.mark.asyncio +async def test_bounded_get_reads_a_compressed_body_whose_wire_length_exceeds_the_cap(respx_mock, monkeypatch): + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + document = bytes(range(256)) + compressed = gzip.compress(document) + assert len(compressed) > len(document) + respx_mock.get("https://cdn.example/blob.bin").respond( + 200, content=compressed, headers={"content-encoding": "gzip", "content-length": str(len(compressed))} + ) + handler = AsyncHTTPHandler() + try: + response = await handler.get("https://cdn.example/blob.bin", max_response_bytes=len(document)) + finally: + await handler.close() + assert response.content == document + + @pytest.mark.asyncio async def test_bounded_get_caps_the_decoded_size_of_a_compressed_body(respx_mock, monkeypatch): monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")