diff --git a/src/httpx2/httpx2/_client.py b/src/httpx2/httpx2/_client.py index d1916e4e..4246a4c9 100644 --- a/src/httpx2/httpx2/_client.py +++ b/src/httpx2/httpx2/_client.py @@ -160,8 +160,13 @@ def __init__(self, stream: AsyncByteStream, start: float) -> None: self.elapsed: datetime.timedelta | None = None async def __aiter__(self) -> typing.AsyncIterator[bytes]: - async for chunk in self._stream: - yield chunk + stream = self._stream.__aiter__() + try: + async for chunk in stream: + yield chunk + finally: + if isinstance(stream, AsyncGenerator): + await stream.aclose() async def aclose(self) -> None: self.elapsed = datetime.timedelta(seconds=time.perf_counter() - self._start) diff --git a/src/httpx2/httpx2/_content.py b/src/httpx2/httpx2/_content.py index c75b6121..f5ba32ef 100644 --- a/src/httpx2/httpx2/_content.py +++ b/src/httpx2/httpx2/_content.py @@ -2,7 +2,7 @@ import inspect import warnings -from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Mapping +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterable, Iterator, Mapping from json import dumps as json_dumps from typing import ( Any, @@ -79,9 +79,15 @@ async def __aiter__(self) -> AsyncIterator[bytes]: yield chunk chunk = await self._stream.aread(self.CHUNK_SIZE) else: - # Otherwise iterate. - async for part in self._stream: - yield part + # Otherwise iterate, making sure the wrapped stream is closed even if the + # consumer stops early (e.g. an exception is raised part-way through decoding). + stream = self._stream.__aiter__() + try: + async for part in stream: + yield part + finally: + if isinstance(stream, AsyncGenerator): + await stream.aclose() class UnattachedStream(AsyncByteStream, SyncByteStream): diff --git a/src/httpx2/httpx2/_decoders.py b/src/httpx2/httpx2/_decoders.py index a05f5fd4..73b302db 100644 --- a/src/httpx2/httpx2/_decoders.py +++ b/src/httpx2/httpx2/_decoders.py @@ -8,6 +8,7 @@ import codecs import io +import itertools import sys import typing import zlib @@ -48,11 +49,42 @@ _zstandard_installed = False +MAX_DECODE_CHUNK_SIZE = 2**20 # 1 MiB + + +class Decompressor(typing.Protocol): + @property + def unconsumed_tail(self) -> bytes: ... + + def decompress(self, data: bytes, max_length: int) -> bytes: ... + + def flush(self) -> bytes: ... + + +class ZlibDecompressor: + """ + Drain a `zlib`/`gzip` decompressor in bounded pieces so a small compressed + input cannot inflate to an unbounded buffer in a single call. + """ + + def __init__(self, decompressor: Decompressor) -> None: + self.decompressor = decompressor + + def decompress(self, data: bytes) -> typing.Iterator[bytes]: + decompressed = self.decompressor.decompress(data, MAX_DECODE_CHUNK_SIZE) + while decompressed: + yield decompressed + decompressed = self.decompressor.decompress(self.decompressor.unconsumed_tail, MAX_DECODE_CHUNK_SIZE) + + def flush(self) -> bytes: + return self.decompressor.flush() + + class ContentDecoder: - def decode(self, data: bytes) -> bytes: + def decode(self, data: bytes) -> typing.Iterator[bytes]: raise NotImplementedError() # pragma: no cover - def flush(self) -> bytes: + def flush(self) -> typing.Iterator[bytes]: raise NotImplementedError() # pragma: no cover @@ -61,11 +93,11 @@ class IdentityDecoder(ContentDecoder): Handle unencoded data. """ - def decode(self, data: bytes) -> bytes: - return data + def decode(self, data: bytes) -> typing.Iterator[bytes]: + yield data - def flush(self) -> bytes: - return b"" + def flush(self) -> typing.Iterator[bytes]: + yield from () class DeflateDecoder(ContentDecoder): @@ -77,22 +109,23 @@ class DeflateDecoder(ContentDecoder): def __init__(self) -> None: self.first_attempt = True - self.decompressor = zlib.decompressobj() + self.decompressor = ZlibDecompressor(zlib.decompressobj()) - def decode(self, data: bytes) -> bytes: + def decode(self, data: bytes) -> typing.Iterator[bytes]: was_first_attempt = self.first_attempt self.first_attempt = False try: - return self.decompressor.decompress(data) + yield from self.decompressor.decompress(data) except zlib.error as exc: if was_first_attempt: - self.decompressor = zlib.decompressobj(-zlib.MAX_WBITS) - return self.decode(data) - raise DecodingError(str(exc)) from exc + self.decompressor = ZlibDecompressor(zlib.decompressobj(-zlib.MAX_WBITS)) + yield from self.decode(data) + else: + raise DecodingError(str(exc)) from exc - def flush(self) -> bytes: + def flush(self) -> typing.Iterator[bytes]: try: - return self.decompressor.flush() + yield self.decompressor.flush() except zlib.error as exc: # pragma: no cover raise DecodingError(str(exc)) from exc @@ -105,17 +138,17 @@ class GZipDecoder(ContentDecoder): """ def __init__(self) -> None: - self.decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16) + self.decompressor = ZlibDecompressor(zlib.decompressobj(zlib.MAX_WBITS | 16)) - def decode(self, data: bytes) -> bytes: + def decode(self, data: bytes) -> typing.Iterator[bytes]: try: - return self.decompressor.decompress(data) + yield from self.decompressor.decompress(data) except zlib.error as exc: raise DecodingError(str(exc)) from exc - def flush(self) -> bytes: + def flush(self) -> typing.Iterator[bytes]: try: - return self.decompressor.flush() + yield self.decompressor.flush() except zlib.error as exc: # pragma: no cover raise DecodingError(str(exc)) from exc @@ -140,26 +173,31 @@ def __init__(self) -> None: self.decompressor = brotli.Decompressor() self.seen_data = False - self._decompress: typing.Callable[[bytes], bytes] + self._decompress: typing.Callable[..., bytes] if hasattr(self.decompressor, "decompress"): # The 'brotlicffi' package. self._decompress = self.decompressor.decompress # pragma: no cover else: # The 'brotli' package. - self._decompress = self.decompressor.process # pragma: no cover + self._decompress = self.decompressor.process - def decode(self, data: bytes) -> bytes: + def decode(self, data: bytes) -> typing.Iterator[bytes]: if not data: - return b"" + return self.seen_data = True try: - return self._decompress(data) + # The C backend may allocate nearly twice the requested threshold. + output_buffer_limit = MAX_DECODE_CHUNK_SIZE // 2 + decompressed = self._decompress(data, output_buffer_limit=output_buffer_limit) + while decompressed: + yield decompressed + decompressed = self._decompress(b"", output_buffer_limit=output_buffer_limit) except brotli.error as exc: raise DecodingError(str(exc)) from exc - def flush(self) -> bytes: + def flush(self) -> typing.Iterator[bytes]: if not self.seen_data: - return b"" + return try: if hasattr(self.decompressor, "finish"): # Only available in the 'brotlicffi' package. @@ -168,9 +206,9 @@ def flush(self) -> bytes: # will never actually emit any data. However, it will potentially throw # errors if a truncated or damaged data stream has been used. self.decompressor.finish() # pragma: no cover - return b"" except brotli.error as exc: # pragma: no cover raise DecodingError(str(exc)) from exc + yield from () class ZStandardDecoder(ContentDecoder): @@ -189,30 +227,34 @@ def __init__(self) -> None: self.decompressor = ZstdDecompressor() self.seen_data = False - def decode(self, data: bytes) -> bytes: + def decode(self, data: bytes) -> typing.Iterator[bytes]: if not data: - return b"" + return self.seen_data = True - output = io.BytesIO() try: if self.decompressor.eof: data = self.decompressor.unused_data + data self.decompressor = ZstdDecompressor() - output.write(self.decompressor.decompress(data)) - while self.decompressor.eof and self.decompressor.unused_data: - unused_data = self.decompressor.unused_data + while True: + decompressed = self.decompressor.decompress(data, MAX_DECODE_CHUNK_SIZE) + while decompressed: + yield decompressed + if self.decompressor.needs_input or self.decompressor.eof: + break + decompressed = self.decompressor.decompress(b"", MAX_DECODE_CHUNK_SIZE) + if not (self.decompressor.eof and self.decompressor.unused_data): + break + data = self.decompressor.unused_data self.decompressor = ZstdDecompressor() - output.write(self.decompressor.decompress(unused_data)) except ZstdError as exc: raise DecodingError(str(exc)) from exc - return output.getvalue() - def flush(self) -> bytes: + def flush(self) -> typing.Iterator[bytes]: if not self.seen_data: - return b"" + return if not self.decompressor.eof: raise DecodingError("Zstandard data is incomplete") # pragma: no cover - return b"" + yield from () class MultiDecoder(ContentDecoder): @@ -233,16 +275,25 @@ def __init__(self, encodings: typing.Sequence[str]) -> None: # Note that we reverse the order for decoding. self.children: list[ContentDecoder] = [SUPPORTED_DECODERS[coding]() for coding in reversed(codings)] - def decode(self, data: bytes) -> bytes: + def decode(self, data: bytes) -> typing.Iterator[bytes]: + streams: typing.Iterator[bytes] = iter((data,)) for child in self.children: - data = child.decode(data) - return data + streams = self._pipe(child.decode, streams) + yield from streams - def flush(self) -> bytes: - data = b"" + def flush(self) -> typing.Iterator[bytes]: + streams: typing.Iterator[bytes] = iter(()) for child in self.children: - data = child.decode(data) + child.flush() - return data + streams = itertools.chain(self._pipe(child.decode, streams), child.flush()) + yield from streams + + @staticmethod + def _pipe( + decode: typing.Callable[[bytes], typing.Iterator[bytes]], + upstream: typing.Iterator[bytes], + ) -> typing.Iterator[bytes]: + for chunk in upstream: + yield from decode(chunk) class ByteChunker: diff --git a/src/httpx2/httpx2/_models.py b/src/httpx2/httpx2/_models.py index 9fb02337..0272f36d 100644 --- a/src/httpx2/httpx2/_models.py +++ b/src/httpx2/httpx2/_models.py @@ -1,13 +1,14 @@ from __future__ import annotations import codecs +import contextlib import datetime import email.message import json as jsonlib import re import typing import urllib.request -from collections.abc import Mapping +from collections.abc import AsyncGenerator, Mapping from http.cookiejar import Cookie, CookieJar from ._content import ByteStream, UnattachedStream, encode_request, encode_response @@ -889,12 +890,12 @@ def iter_bytes(self, chunk_size: int | None = None) -> typing.Iterator[bytes]: chunker = ByteChunker(chunk_size=chunk_size) with request_context(request=self._request): for raw_bytes in self.iter_raw(): - decoded = decoder.decode(raw_bytes) + for decoded in decoder.decode(raw_bytes): + for chunk in chunker.decode(decoded): + yield chunk + for decoded in decoder.flush(): for chunk in chunker.decode(decoded): - yield chunk - decoded = decoder.flush() - for chunk in chunker.decode(decoded): - yield chunk # pragma: no cover + yield chunk # pragma: no cover for chunk in chunker.flush(): yield chunk @@ -941,16 +942,17 @@ def iter_raw(self, chunk_size: int | None = None) -> typing.Iterator[bytes]: self._num_bytes_downloaded = 0 chunker = ByteChunker(chunk_size=chunk_size) - with request_context(request=self._request): - for raw_stream_bytes in self.stream: - self._num_bytes_downloaded += len(raw_stream_bytes) - for chunk in chunker.decode(raw_stream_bytes): - yield chunk - - for chunk in chunker.flush(): - yield chunk + try: + with request_context(request=self._request): + for raw_stream_bytes in self.stream: + self._num_bytes_downloaded += len(raw_stream_bytes) + for chunk in chunker.decode(raw_stream_bytes): + yield chunk - self.close() + for chunk in chunker.flush(): + yield chunk + finally: + self.close() def close(self) -> None: """ @@ -970,10 +972,11 @@ async def aread(self) -> bytes: Read and return the response content. """ if not hasattr(self, "_content"): - self._content = b"".join([part async for part in self.aiter_bytes()]) + async with contextlib.aclosing(self.aiter_bytes()) as parts: + self._content = b"".join([part async for part in parts]) return self._content - async def aiter_bytes(self, chunk_size: int | None = None) -> typing.AsyncIterator[bytes]: + async def aiter_bytes(self, chunk_size: int | None = None) -> typing.AsyncGenerator[bytes, None]: """ A byte-iterator over the decoded response content. This allows us to handle gzip, deflate, brotli, and zstd encoded responses. @@ -986,13 +989,14 @@ async def aiter_bytes(self, chunk_size: int | None = None) -> typing.AsyncIterat decoder = self._get_content_decoder() chunker = ByteChunker(chunk_size=chunk_size) with request_context(request=self._request): - async for raw_bytes in self.aiter_raw(): - decoded = decoder.decode(raw_bytes) + async with contextlib.aclosing(self.aiter_raw()) as raw_stream: + async for raw_bytes in raw_stream: + for decoded in decoder.decode(raw_bytes): + for chunk in chunker.decode(decoded): + yield chunk + for decoded in decoder.flush(): for chunk in chunker.decode(decoded): - yield chunk - decoded = decoder.flush() - for chunk in chunker.decode(decoded): - yield chunk # pragma: no cover + yield chunk # pragma: no cover for chunk in chunker.flush(): yield chunk @@ -1024,7 +1028,7 @@ async def aiter_lines(self) -> typing.AsyncIterator[str]: for line in decoder.flush(): yield line - async def aiter_raw(self, chunk_size: int | None = None) -> typing.AsyncIterator[bytes]: + async def aiter_raw(self, chunk_size: int | None = None) -> typing.AsyncGenerator[bytes, None]: """ A byte-iterator over the raw response content. """ @@ -1039,16 +1043,20 @@ async def aiter_raw(self, chunk_size: int | None = None) -> typing.AsyncIterator self._num_bytes_downloaded = 0 chunker = ByteChunker(chunk_size=chunk_size) - with request_context(request=self._request): - async for raw_stream_bytes in self.stream: - self._num_bytes_downloaded += len(raw_stream_bytes) - for chunk in chunker.decode(raw_stream_bytes): - yield chunk - - for chunk in chunker.flush(): - yield chunk + stream = self.stream.__aiter__() + try: + with request_context(request=self._request): + async for raw_stream_bytes in stream: + self._num_bytes_downloaded += len(raw_stream_bytes) + for chunk in chunker.decode(raw_stream_bytes): + yield chunk - await self.aclose() + for chunk in chunker.flush(): + yield chunk + finally: + if isinstance(stream, AsyncGenerator): + await stream.aclose() + await self.aclose() async def aclose(self) -> None: """ diff --git a/tests/httpx2/test_decoders.py b/tests/httpx2/test_decoders.py index 4bf5fb8f..0d463e95 100644 --- a/tests/httpx2/test_decoders.py +++ b/tests/httpx2/test_decoders.py @@ -2,9 +2,11 @@ import io import sys +import tracemalloc import typing import zlib +import brotli import chardet import pytest @@ -67,6 +69,33 @@ def test_gzip() -> None: assert response.content == body +@pytest.mark.parametrize("wbits", (zlib.MAX_WBITS | 16, zlib.MAX_WBITS)) +def test_decoding_bounds_peak_memory(wbits: int) -> None: + # A small compressed body that inflates to a large output must not be + # materialised in a single decode step. Even when the entire compressed + # body arrives as one raw chunk, peak memory stays far below the decoded + # size because decoding happens in bounded pieces. + decoded_size = 64 * 1024 * 1024 + compressor = zlib.compressobj(9, zlib.DEFLATED, wbits) + compressed = compressor.compress(b"\x00" * decoded_size) + compressor.flush() + encoding = b"gzip" if wbits & 16 else b"deflate" + + def raw() -> typing.Iterator[bytes]: + yield compressed # deliver the whole bomb as a single raw chunk + + response = httpx2.Response(200, headers=[(b"Content-Encoding", encoding)], content=raw()) + + total = 0 + tracemalloc.start() + for chunk in response.iter_bytes(): + total += len(chunk) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert total == decoded_size + assert peak < decoded_size // 8 + + def test_brotli() -> None: body = b"test 123" compressed_body = b"\x8b\x03\x80test 123\x03" @@ -80,6 +109,21 @@ def test_brotli() -> None: assert response.content == body +@pytest.mark.parametrize(("encoding", "compress"), [(b"br", brotli.compress), (b"zstd", zstd.compress)]) +def test_decoding_bounds_output_chunks(encoding: bytes, compress: typing.Callable[[bytes], bytes]) -> None: + body = b"\x00" * (2**20 + 1) + compressed_body = compress(body) + + def raw() -> typing.Iterator[bytes]: + yield compressed_body + + response = httpx2.Response(200, headers=[(b"Content-Encoding", encoding)], content=raw()) + chunks = list(response.iter_bytes()) + + assert b"".join(chunks) == body + assert max(map(len, chunks)) <= 2**20 + + def test_zstd() -> None: body = b"test 123" compressed_body = zstd.compress(body) @@ -226,6 +270,30 @@ def test_multi_brotli_zstd() -> None: assert response.content == body +def test_multi_zstd_gzip() -> None: + # A gzip layer ahead of zstd means gzip's end-of-stream flush feeds b"" into + # the zstd decoder after its frame is complete; it must not raise. + body = b"test 123" + inner = zstd.compress(body) + compressor = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS | 16) + compressed_body = compressor.compress(inner) + compressor.flush() + + headers = [(b"Content-Encoding", b"zstd, gzip")] + response = httpx2.Response(200, headers=headers, content=compressed_body) + assert response.content == body + + +def test_multi_brotli_gzip() -> None: + body = b"test 123" + inner = b"\x8b\x03\x80test 123\x03" + compressor = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS | 16) + compressed_body = compressor.compress(inner) + compressor.flush() + + headers = [(b"Content-Encoding", b"br, gzip")] + response = httpx2.Response(200, headers=headers, content=compressed_body) + assert response.content == body + + def test_multi_decode_links_limit() -> None: headers = [(b"Content-Encoding", b", ".join([b"gzip"] * 6))] with pytest.raises(httpx2.DecodingError, match="Cannot apply more than 5 content encodings"): @@ -405,3 +473,39 @@ def test_invalid_content_encoding_header() -> None: content=body, ) assert response.content == body + + +def test_streaming_decode_error_closes_response() -> None: + response: httpx2.Response | None = None + + def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal response + response = httpx2.Response(200, headers={"Content-Encoding": "gzip"}, content=iter([b"invalid"])) + return response + + with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: + with pytest.raises(httpx2.DecodingError): + client.get("https://example.org") + + assert response is not None + assert response.is_closed + + +@pytest.mark.anyio +async def test_async_streaming_decode_error_closes_response() -> None: + response: httpx2.Response | None = None + + async def content() -> typing.AsyncIterator[bytes]: + yield b"invalid" + + async def handler(request: httpx2.Request) -> httpx2.Response: + nonlocal response + response = httpx2.Response(200, headers={"Content-Encoding": "gzip"}, content=content()) + return response + + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as client: + with pytest.raises(httpx2.DecodingError): + await client.get("https://example.org") + + assert response is not None + assert response.is_closed diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 97fac974..89e64c07 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -84,7 +84,7 @@ def test_bench_queryparams_merge(benchmark: BenchmarkFixture) -> None: def test_bench_gzip_decode(benchmark: BenchmarkFixture) -> None: def decode() -> bytes: decoder = GZipDecoder() - return decoder.decode(GZIP_BODY) + decoder.flush() + return b"".join(decoder.decode(GZIP_BODY)) + b"".join(decoder.flush()) benchmark(decode)