From 42b57b78fef49db563d54d6011d1e4d0f7d5393b Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 17 Oct 2026 14:48:18 +0300 Subject: [PATCH 1/2] perf: count capped bytes on the stream for unencoded bodies For a response with no Content-Encoding, the body cap now wraps response.stream in a stream that counts bytes and raises ResponseTooLargeError once the count passes the cap, and httpx2 reads the response with its own read()/aread(). The capped response is the object httpx2 built, so its headers, extensions and num_bytes_downloaded match an uncapped one. Encoded bodies keep the rebuild path. Closes #170 --- .../0009-the-body-cap-counts-decoded-bytes.md | 12 +- docs/errors.md | 2 +- src/httpware/_internal/body_cap.py | 85 +++++++++++- tests/test_body_cap.py | 129 +++++++++++++++++- tests/test_client_body_cap.py | 14 ++ 5 files changed, 232 insertions(+), 10 deletions(-) diff --git a/docs/adr/0009-the-body-cap-counts-decoded-bytes.md b/docs/adr/0009-the-body-cap-counts-decoded-bytes.md index 86ec7fe..d5c573b 100644 --- a/docs/adr/0009-the-body-cap-counts-decoded-bytes.md +++ b/docs/adr/0009-the-body-cap-counts-decoded-bytes.md @@ -4,8 +4,10 @@ status. `Content-Length` may reject early but never admit early, because it reports the compressed size and a 1000:1 gzip bomb passes a header check by construction. It replaced an error-only pre-check, since memory exhaustion has no status code, and it never applies to caller-driven -`stream()` iteration, whose purpose is bodies too large to buffer. The response is rebuilt -through `httpx2.Response(content=...)` with public API only: `elapsed`, `next_request` and -`default_encoding` are copied over, and `num_bytes_downloaded`, which has no public setter, reads 0. -Counting bytes in a wrapper around the response stream would avoid the rebuild; that is deferred to -[#170](https://github.com/modern-python/httpware/issues/170). +`stream()` iteration, whose purpose is bodies too large to buffer. For a body with no +`Content-Encoding` the decoded bytes are the wire bytes, so the cap wraps the public +`Response.stream` in a stream that counts bytes and raises once the count passes the cap, and httpx2 +reads the response with its own `read()`. The response stays the one httpx2 built. An encoded body +is read through `iter_bytes()` and rebuilt with `httpx2.Response(content=...)`, using public API +only: `elapsed`, `next_request` and `default_encoding` are copied over, and `num_bytes_downloaded`, +which has no public setter, reads 0. diff --git a/docs/errors.md b/docs/errors.md index 2c3e365..1f2eae9 100644 --- a/docs/errors.md +++ b/docs/errors.md @@ -202,7 +202,7 @@ Unlike `DecodeError`, this error is raised before the request is sent. Both clients accept `max_response_body_bytes: int | None = None`. By default there is no limit. When it is set, a response body larger than the cap raises `ResponseTooLargeError` instead of being returned, whatever the status: a `200` trips it as easily as a `500`. The cap counts decoded bytes, after decompression. It applies to `send()` and the verb methods, and to the error body that `stream()` reads before raising a `StatusError`. Bytes you read yourself while iterating a `stream()` are never capped, except with an `auth` that sets `requires_response_body`: that auth needs every body, so `stream()` reads the final one under the cap before yielding it. With a cap set, httpware runs the client's `auth` flow and, with `follow_redirects=True`, follows the redirects itself, and caps only the final response. It reads the body of each intermediate response, such as a redirect or a `DigestAuth` challenge, and throws it away, so the connection can be reused and the responses in `response.history` have no content. An intermediate body declared or read past the cap is not an error: httpware stops reading it and closes that connection. An `auth` that sets `requires_response_body`, like a token refresh that parses the token endpoint's JSON, gets each response read under the cap instead. Client `auth` is sent to the first URL only. `httpx2` keeps its `Authorization` header on a redirect within the same origin or from `http` to `https` on the same host, and drops it otherwise. -A response read under the cap is rebuilt from its decoded body. It keeps `elapsed`, `next_request` and the client's `default_encoding`. It loses the `Content-Encoding` and `Transfer-Encoding` headers, its `Content-Length` is the decoded size, and `num_bytes_downloaded` is `0`. +A response read under the cap keeps `elapsed`, `next_request` and the client's `default_encoding`. A response without `Content-Encoding` is the one `httpx2` built, with the same headers and `num_bytes_downloaded` it has without a cap. A response with `Content-Encoding` is rebuilt from its decoded body: it loses the `Content-Encoding` and `Transfer-Encoding` headers, its `Content-Length` is the decoded size, and `num_bytes_downloaded` is `0`. `ResponseTooLargeError` carries: diff --git a/src/httpware/_internal/body_cap.py b/src/httpware/_internal/body_cap.py index 41b4825..ea1e3b9 100644 --- a/src/httpware/_internal/body_cap.py +++ b/src/httpware/_internal/body_cap.py @@ -1,8 +1,9 @@ """Response-body cap enforcement: validate, read-capped (sync + async).""" import contextlib +import datetime as dt import typing -from collections.abc import Mapping +from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping from http import HTTPStatus import httpx2 @@ -98,6 +99,73 @@ def _reject_if_declared_length_exceeds_cap(response: httpx2.Response, cap: int) return content_length +class _CappedSyncStream(httpx2.SyncByteStream): + """Wrap a response stream, raising ResponseTooLargeError(reason="streamed") once it yields more than `cap` bytes.""" + + def __init__(self, stream: httpx2.SyncByteStream, cap: int, status_code: int, content_length: int | None) -> None: + self._stream = stream + self._cap = cap + self._status_code = status_code + self._content_length = content_length + + @property + def elapsed(self) -> dt.timedelta | None: + """The wrapped stream's `elapsed`, which httpx2 reads back as `Response.elapsed`.""" + return getattr(self._stream, "elapsed", None) + + def __iter__(self) -> Iterator[bytes]: + read = 0 + for chunk in self._stream: + read += len(chunk) + if read > self._cap: + raise ResponseTooLargeError( + status_code=self._status_code, + limit=self._cap, + content_length=self._content_length, + reason="streamed", + ) + yield chunk + + def close(self) -> None: + self._stream.close() + + +class _CappedAsyncStream(httpx2.AsyncByteStream): + """Async mirror of `_CappedSyncStream`.""" + + def __init__(self, stream: httpx2.AsyncByteStream, cap: int, status_code: int, content_length: int | None) -> None: + self._stream = stream + self._cap = cap + self._status_code = status_code + self._content_length = content_length + + @property + def elapsed(self) -> dt.timedelta | None: + """The wrapped stream's `elapsed`, which httpx2 reads back as `Response.elapsed`.""" + return getattr(self._stream, "elapsed", None) + + async def __aiter__(self) -> AsyncIterator[bytes]: + chunks = self._stream.__aiter__() + try: + read = 0 + async for chunk in chunks: + read += len(chunk) + if read > self._cap: + raise ResponseTooLargeError( + status_code=self._status_code, + limit=self._cap, + content_length=self._content_length, + reason="streamed", + ) + yield chunk + finally: + if isinstance(chunks, AsyncGenerator): + await chunks.aclose() + + async def aclose(self) -> None: + await self._stream.aclose() + + def _build_buffered_response(response: httpx2.Response, content: bytes, request: httpx2.Request) -> httpx2.Response: """Rebuild a buffered Response from the closed streaming `response` and its final decoded `content`.""" buffered = httpx2.Response( @@ -120,13 +188,19 @@ def _read_capped(response: httpx2.Response, cap: int, request: httpx2.Request) - Raises `ResponseTooLargeError` (reason="declared") if the declared Content-Length already exceeds `cap` — before any byte is read — and - (reason="streamed") if the decoded body crosses `cap` mid-read. Closes - `response` once its whole body is read; on any other path the caller closes it. + (reason="streamed") if the decoded body crosses `cap` mid-read. An unencoded + body is read into `response` itself; an encoded one into a rebuilt Response. + Closes `response` once its whole body is read; on any other path the caller closes it. """ if not _response_has_body(request.method, response.status_code): response.read() # empty body; preserve the original response (and its headers) return response content_length = _reject_if_declared_length_exceeds_cap(response, cap) + if "content-encoding" not in response.headers: + stream = typing.cast("httpx2.SyncByteStream", response.stream) + response.stream = _CappedSyncStream(stream, cap, response.status_code, content_length) + response.read() + return response try: content = _accumulate_capped(response.iter_bytes(), cap) except _CapExceeded: @@ -143,6 +217,11 @@ async def _read_capped_async(response: httpx2.Response, cap: int, request: httpx await response.aread() # empty body; preserve the original response (and its headers) return response content_length = _reject_if_declared_length_exceeds_cap(response, cap) + if "content-encoding" not in response.headers: + stream = typing.cast("httpx2.AsyncByteStream", response.stream) + response.stream = _CappedAsyncStream(stream, cap, response.status_code, content_length) + await response.aread() + return response buf = bytearray() async for chunk in response.aiter_bytes(): buf += chunk diff --git a/tests/test_body_cap.py b/tests/test_body_cap.py index 2b61c51..5fddf60 100644 --- a/tests/test_body_cap.py +++ b/tests/test_body_cap.py @@ -7,7 +7,7 @@ """ import gzip -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Iterator import httpx2 import pytest @@ -157,6 +157,69 @@ def handler(request: httpx2.Request) -> httpx2.Response: # noqa: ARG001 client.close() +def test_read_capped_returns_the_same_response_for_an_unencoded_body() -> None: + body = b"hello world" + + def handler(request: httpx2.Request) -> httpx2.Response: # noqa: ARG001 + return httpx2.Response(200, content=iter([body])) + + client, resp = _sync_stream(handler) + try: + out = _read_capped(resp, 1000, resp.request) + assert out is resp + assert out.content == body + assert out.is_closed + assert out.num_bytes_downloaded == len(body) + assert out.headers["transfer-encoding"] == "chunked" + finally: + resp.close() + client.close() + + +def test_read_capped_rebuilds_the_response_for_an_encoded_body() -> None: + raw = gzip.compress(b"A" * 500) + + def handler(request: httpx2.Request) -> httpx2.Response: # noqa: ARG001 + return httpx2.Response(200, headers={"content-encoding": "gzip"}, content=raw) + + client, resp = _sync_stream(handler) + try: + out = _read_capped(resp, 1_000_000, resp.request) + assert out is not resp + assert out.num_bytes_downloaded == 0 + finally: + resp.close() + client.close() + + +def test_read_capped_unencoded_over_cap_closes_the_stream() -> None: + closed: list[bool] = [] + + class _Body(httpx2.SyncByteStream): + def __iter__(self) -> Iterator[bytes]: + yield b"a" * 50 + yield b"b" * 50 + + def close(self) -> None: + closed.append(True) + + def handler(request: httpx2.Request) -> httpx2.Response: # noqa: ARG001 + return httpx2.Response(200, stream=_Body()) + + client, resp = _sync_stream(handler) + try: + with pytest.raises(ResponseTooLargeError) as caught: + _read_capped(resp, 70, resp.request) + assert caught.value.reason == "streamed" + assert caught.value.limit == 70 # noqa: PLR2004 — cap above + assert caught.value.status_code == 200 # noqa: PLR2004 — mirrors handler + assert resp.is_closed + assert closed == [True] + finally: + resp.close() + client.close() + + # ---- async ---- @@ -259,6 +322,70 @@ def handler(request: httpx2.Request) -> httpx2.Response: # noqa: ARG001 await client.aclose() +async def test_read_capped_async_returns_the_same_response_for_an_unencoded_body() -> None: + body = b"hello world" + + async def chunks() -> AsyncIterator[bytes]: + yield body + + def handler(request: httpx2.Request) -> httpx2.Response: # noqa: ARG001 + return httpx2.Response(200, content=chunks()) + + client, resp = await _async_stream(handler) + try: + out = await _read_capped_async(resp, 1000, resp.request) + assert out is resp + assert out.content == body + assert out.is_closed + assert out.num_bytes_downloaded == len(body) + finally: + await resp.aclose() + await client.aclose() + + +async def test_read_capped_async_rebuilds_the_response_for_an_encoded_body() -> None: + raw = gzip.compress(b"A" * 500) + + def handler(request: httpx2.Request) -> httpx2.Response: # noqa: ARG001 + return httpx2.Response(200, headers={"content-encoding": "gzip"}, content=raw) + + client, resp = await _async_stream(handler) + try: + out = await _read_capped_async(resp, 1_000_000, resp.request) + assert out is not resp + assert out.num_bytes_downloaded == 0 + finally: + await resp.aclose() + await client.aclose() + + +async def test_read_capped_async_unencoded_over_cap_closes_the_stream() -> None: + closed: list[bool] = [] + + class _Body(httpx2.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"a" * 50 + yield b"b" * 50 + + async def aclose(self) -> None: + closed.append(True) + + def handler(request: httpx2.Request) -> httpx2.Response: # noqa: ARG001 + return httpx2.Response(200, stream=_Body()) + + client, resp = await _async_stream(handler) + try: + with pytest.raises(ResponseTooLargeError) as caught: + await _read_capped_async(resp, 70, resp.request) + assert caught.value.reason == "streamed" + assert caught.value.limit == 70 # noqa: PLR2004 — cap above + assert resp.is_closed + assert closed == [True] + finally: + await resp.aclose() + await client.aclose() + + # ---- content-length parsing ---- diff --git a/tests/test_client_body_cap.py b/tests/test_client_body_cap.py index a10aeeb..c9768b1 100644 --- a/tests/test_client_body_cap.py +++ b/tests/test_client_body_cap.py @@ -277,3 +277,17 @@ async def test_async_response_keeps_default_encoding_with_or_without_a_cap(cap: ) as client: response = await client.get("https://example.test/x") assert response.text == "café" + + +@pytest.mark.parametrize("cap", [None, 1024]) +def test_sync_response_keeps_num_bytes_downloaded_with_or_without_a_cap(cap: int | None) -> None: + with Client(transport=httpx2.MockTransport(_streamed_body), max_response_body_bytes=cap) as client: + response = client.get("https://example.test/x") + assert response.num_bytes_downloaded == len(b"hello") + + +@pytest.mark.parametrize("cap", [None, 1024]) +async def test_async_response_keeps_num_bytes_downloaded_with_or_without_a_cap(cap: int | None) -> None: + async with AsyncClient(transport=httpx2.MockTransport(_streamed_body_async), max_response_body_bytes=cap) as client: + response = await client.get("https://example.test/x") + assert response.num_bytes_downloaded == len(b"hello") From 0ea99388c7093de6941a9fdd0db6888910402a4b Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 17 Oct 2026 14:56:31 +0300 Subject: [PATCH 2/2] fix: close the response when a capped stream read raises httpx2 2.0.0 closes a response only when its body is read to the end, so an over-cap read left it open on the floor version. --- src/httpware/_internal/body_cap.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/httpware/_internal/body_cap.py b/src/httpware/_internal/body_cap.py index ea1e3b9..4b04c93 100644 --- a/src/httpware/_internal/body_cap.py +++ b/src/httpware/_internal/body_cap.py @@ -189,8 +189,9 @@ def _read_capped(response: httpx2.Response, cap: int, request: httpx2.Request) - Raises `ResponseTooLargeError` (reason="declared") if the declared Content-Length already exceeds `cap` — before any byte is read — and (reason="streamed") if the decoded body crosses `cap` mid-read. An unencoded - body is read into `response` itself; an encoded one into a rebuilt Response. - Closes `response` once its whole body is read; on any other path the caller closes it. + body is read into `response` itself, which is closed on any error; an encoded + one is read into a rebuilt Response. Closes `response` once its whole body is read; + on any other path the caller closes it. """ if not _response_has_body(request.method, response.status_code): response.read() # empty body; preserve the original response (and its headers) @@ -199,7 +200,11 @@ def _read_capped(response: httpx2.Response, cap: int, request: httpx2.Request) - if "content-encoding" not in response.headers: stream = typing.cast("httpx2.SyncByteStream", response.stream) response.stream = _CappedSyncStream(stream, cap, response.status_code, content_length) - response.read() + try: + response.read() + except BaseException: + response.close() + raise return response try: content = _accumulate_capped(response.iter_bytes(), cap) @@ -220,7 +225,11 @@ async def _read_capped_async(response: httpx2.Response, cap: int, request: httpx if "content-encoding" not in response.headers: stream = typing.cast("httpx2.AsyncByteStream", response.stream) response.stream = _CappedAsyncStream(stream, cap, response.status_code, content_length) - await response.aread() + try: + await response.aread() + except BaseException: + await response.aclose() + raise return response buf = bytearray() async for chunk in response.aiter_bytes():