perf: count capped bytes on the stream for unencoded bodies (#192)

* 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

* 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.

Author: lesnik512

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.

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:
 

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,24 @@ 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, 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)
         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)
+        try:
+            response.read()
+        except BaseException:
+            response.close()
+            raise
+        return response
     try:
         content = _accumulate_capped(response.iter_bytes(), cap)
     except _CapExceeded:
@@ -143,6 +222,15 @@ 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)
+        try:
+            await response.aread()
+        except BaseException:
+            await response.aclose()
+            raise
+        return response
     buf = bytearray()
     async for chunk in response.aiter_bytes():
         buf += chunk

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 ----
 
 

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")