Allow decompression to continue after exceeding max_length

aiohttp/_cparser.pxd

@@ -145,6 +145,7 @@ cdef extern from "llhttp.h":
 
     int llhttp_should_keep_alive(const llhttp_t* parser)
 
+    void llhttp_resume(llhttp_t* parser)
     void llhttp_resume_after_upgrade(llhttp_t* parser)
 
     llhttp_errno_t llhttp_get_errno(const llhttp_t* parser)

aiohttp/_http_parser.pyx

@@ -291,15 +291,19 @@ cdef class HttpParser:
         bint _response_with_body
         bint _read_until_eof
 
+        bytes   _tail
         bint    _started
         object  _url
         bytearray   _buf
         str     _path
         str     _reason
         list    _headers
+        bint    _last_had_more_data
         list    _raw_headers
         bint    _upgraded
         list    _messages
+        bint    _more_data_available
+        bint    _paused
         object  _payload
         bint    _payload_error
         object  _payload_exception
@@ -345,13 +349,17 @@ cdef class HttpParser:
         self._timer = timer
 
         self._buf = bytearray()
+        self._last_had_more_data = False
+        self._more_data_available = False
+        self._paused = False
         self._payload = None
         self._payload_error = 0
         self._payload_exception = payload_exception
         self._messages = []
 
         self._raw_name = EMPTY_BYTES
         self._raw_value = EMPTY_BYTES
+        self._tail = b""
         self._has_value = False
         self._header_name_size = 0
 
@@ -503,6 +511,10 @@ cdef class HttpParser:
 
     ### Public API ###
 
+    def pause_reading(self):
+        assert self._payload is not None
+        self._paused = True
+
     def feed_eof(self):
         cdef bytes desc
 
@@ -529,6 +541,23 @@ cdef class HttpParser:
             size_t nb
             cdef cparser.llhttp_errno_t errno
 
+        if self._tail:
+            data, self._tail = self._tail + data, EMPTY_BYTES
+
+        had_more_data = self._more_data_available
+        if self._more_data_available:
+            result = cb_on_body(self._cparser, EMPTY_BYTES, 0)
+            if result is cparser.HPE_PAUSED:
+                self._last_had_more_data = had_more_data
+                self._tail = data
+                return (), False, EMPTY_BYTES
+            # TODO: Do we need to handle error case (-1)?
+        # If the last pause had more data, then we probably paused at the
+        # end of the body. Therefore we need to continue with empty bytes.
+        if not data and not self._last_had_more_data:
+            return (), False, EMPTY_BYTES
+        self._last_had_more_data = False
+
         PyObject_GetBuffer(data, &self.py_buf, PyBUF_SIMPLE)
         data_len = <size_t>self.py_buf.len
 
@@ -539,12 +568,15 @@ cdef class HttpParser:
 
         if errno is cparser.HPE_PAUSED_UPGRADE:
             cparser.llhttp_resume_after_upgrade(self._cparser)
-
             nb = cparser.llhttp_get_error_pos(self._cparser) - <char*>self.py_buf.buf
+        elif errno is cparser.HPE_PAUSED:
+            cparser.llhttp_resume(self._cparser)
+            pos = cparser.llhttp_get_error_pos(self._cparser) - <char*>self.py_buf.buf
+            self._tail = data[pos:]
 
         PyBuffer_Release(&self.py_buf)
 
-        if errno not in (cparser.HPE_OK, cparser.HPE_PAUSED_UPGRADE):
+        if errno not in (cparser.HPE_OK, cparser.HPE_PAUSED, cparser.HPE_PAUSED_UPGRADE):
             if self._payload_error == 0:
                 if self._last_error is not None:
                     ex = self._last_error
@@ -569,7 +601,7 @@ cdef class HttpParser:
         if self._upgraded:
             return messages, True, data[nb:]
         else:
-            return messages, False, b""
+            return messages, False, EMPTY_BYTES
 
     def set_upgraded(self, val):
         self._upgraded = val
@@ -762,19 +794,26 @@ cdef int cb_on_body(cparser.llhttp_t* parser,
                     const char *at, size_t length) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
     cdef bytes body = at[:length]
-    try:
-        pyparser._payload.feed_data(body)
-    except BaseException as underlying_exc:
-        reraised_exc = underlying_exc
-        if pyparser._payload_exception is not None:
-            reraised_exc = pyparser._payload_exception(str(underlying_exc))
-
-        set_exception(pyparser._payload, reraised_exc, underlying_exc)
-
-        pyparser._payload_error = 1
-        return -1
-    else:
-        return 0
+    while body or pyparser._more_data_available:
+        try:
+            pyparser._more_data_available = pyparser._payload.feed_data(body)
+        except BaseException as underlying_exc:
+            reraised_exc = underlying_exc
+            if pyparser._payload_exception is not None:
+                reraised_exc = pyparser._payload_exception(str(underlying_exc))
+
+            set_exception(pyparser._payload, reraised_exc, underlying_exc)
+
+            pyparser._payload_error = 1
+            pyparser._paused = False
+            return -1
+        body = EMPTY_BYTES
+
+        if pyparser._paused:
+            pyparser._paused = False
+            return cparser.HPE_PAUSED
+    pyparser._paused = False
+    return 0
 
 
 cdef int cb_on_message_complete(cparser.llhttp_t* parser) except -1:

aiohttp/base_protocol.py

@@ -1,26 +1,35 @@
 import asyncio
-from typing import cast
+from typing import TYPE_CHECKING, Any, cast
 
 from .client_exceptions import ClientConnectionResetError
 from .helpers import set_exception
 from .tcp_helpers import tcp_nodelay
 
+if TYPE_CHECKING:
+    from .http_parser import HttpParser
+
 
 class BaseProtocol(asyncio.Protocol):
     __slots__ = (
         "_loop",
         "_paused",
+        "_parser",
         "_drain_waiter",
         "_connection_lost",
         "_reading_paused",
+        "_upgraded",
         "transport",
     )
 
-    def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
+    def __init__(
+        self, loop: asyncio.AbstractEventLoop, parser: "HttpParser[Any] | None" = None
+    ) -> None:
         self._loop: asyncio.AbstractEventLoop = loop
         self._paused = False
         self._drain_waiter: asyncio.Future[None] | None = None
         self._reading_paused = False
+        self._parser = parser
+        self._upgraded = False
 
         self.transport: asyncio.Transport | None = None
 
@@ -48,15 +57,27 @@ def resume_writing(self) -> None:
                 waiter.set_result(None)
 
     def pause_reading(self) -> None:
-        if not self._reading_paused and self.transport is not None:
+        self._reading_paused = True
+        # Parser shouldn't be paused on websockets.
+        if not self._upgraded:
+            assert self._parser is not None
+            self._parser.pause_reading()
+        if self.transport is not None:
             try:
                 self.transport.pause_reading()
             except (AttributeError, NotImplementedError, RuntimeError):
                 pass
-            self._reading_paused = True
 
-    def resume_reading(self) -> None:
-        if self._reading_paused and self.transport is not None:
+    def resume_reading(self, resume_parser: bool = True) -> None:
+        self._reading_paused = False
+
+        # This will resume parsing any unprocessed data from the last pause.
+        if not self._upgraded and resume_parser:
+            self.data_received(b"")
+
+        # Reading may have been paused again in the above call if there was a lot of
+        # compressed data still pending.
+        if not self._reading_paused and self.transport is not None:
             try:
                 self.transport.resume_reading()
             except (AttributeError, NotImplementedError, RuntimeError):

aiohttp/client_proto.py

@@ -26,7 +26,7 @@ class ResponseHandler(BaseProtocol, DataQueue[tuple[RawResponseMessage, StreamRe
     """Helper class to adapt between Protocol and StreamReader."""
 
     def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
-        BaseProtocol.__init__(self, loop=loop)
+        BaseProtocol.__init__(self, loop=loop, parser=None)
         DataQueue.__init__(self, loop)
 
         self._should_close = False
@@ -36,10 +36,7 @@ def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
         self._payload_parser: WebSocketReader | None = None
 
         self._timer = None
-
         self._tail = b""
-        self._upgraded = False
-        self._parser: HttpResponseParser | None = None
 
         self._read_timeout: float | None = None
         self._read_timeout_handle: asyncio.TimerHandle | None = None
@@ -190,8 +187,8 @@ def pause_reading(self) -> None:
         super().pause_reading()
         self._drop_timeout()
 
-    def resume_reading(self) -> None:
-        super().resume_reading()
+    def resume_reading(self, resume_parser: bool = True) -> None:
+        super().resume_reading(resume_parser)
         self._reschedule_timeout()
 
     def set_exception(
@@ -293,9 +290,6 @@ def _on_read_timeout(self) -> None:
     def data_received(self, data: bytes) -> None:
         self._reschedule_timeout()
 
-        if not data:
-            return
-
         # custom payload parser - currently always WebSocketReader
         if self._payload_parser is not None:
             eof, tail = self._payload_parser.feed_data(data)

aiohttp/compression_utils.py

@@ -34,7 +34,9 @@
 
 
 MAX_SYNC_CHUNK_SIZE = 4096
-DEFAULT_MAX_DECOMPRESS_SIZE = 2**25  # 32MiB
+# Matches the max size we receive from sockets:
+# https://github.com/python/cpython/blob/1857a40807daeae3a1bf5efb682de9c9ae6df845/Lib/asyncio/selector_events.py#L766
+DEFAULT_MAX_DECOMPRESS_SIZE = 256 * 1024
 
 # Unlimited decompression constants - different libraries use different conventions
 ZLIB_MAX_LENGTH_UNLIMITED = 0  # zlib uses 0 to mean unlimited
@@ -53,6 +55,9 @@ def flush(self, length: int = ..., /) -> bytes: ...
     @property
     def eof(self) -> bool: ...
 
+    @property
+    def unconsumed_tail(self) -> bytes: ...
+
 
 class ZLibBackendProtocol(Protocol):
     MAX_WBITS: int
@@ -179,6 +184,11 @@ async def decompress(
             )
         return self.decompress_sync(data, max_length)
 
+    @property
+    @abstractmethod
+    def data_available(self) -> bool:
+        """Return True if more output is available by passing b""."""
+
 
 class ZLibCompressor:
     def __init__(
@@ -271,7 +281,9 @@ def __init__(
     def decompress_sync(
         self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
     ) -> bytes:
-        return self._decompressor.decompress(data, max_length)
+        return self._decompressor.decompress(
+            self._decompressor.unconsumed_tail + data, max_length
+        )
 
     def flush(self, length: int = 0) -> bytes:
         return (
@@ -280,6 +292,10 @@ def flush(self, length: int = 0) -> bytes:
             else self._decompressor.flush()
         )
 
+    @property
+    def data_available(self) -> bool:
+        return bool(self._decompressor.unconsumed_tail)
+
     @property
     def eof(self) -> bool:
         return self._decompressor.eof
@@ -301,22 +317,31 @@ def __init__(
                 "Please install `Brotli` module"
             )
         self._obj = brotli.Decompressor()
+        self._last_empty = False
         super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size)
 
     def decompress_sync(
         self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
     ) -> bytes:
         """Decompress the given data."""
         if hasattr(self._obj, "decompress"):
-            return cast(bytes, self._obj.decompress(data, max_length))
-        return cast(bytes, self._obj.process(data, max_length))
+            result = cast(bytes, self._obj.decompress(data, max_length))
+        else:
+            result = cast(bytes, self._obj.process(data, max_length))
+        # Only way to know that brotli has no further data is checking we get no output
+        self._last_empty = result == b""
+        return result
 
     def flush(self) -> bytes:
         """Flush the decompressor."""
         if hasattr(self._obj, "flush"):
             return cast(bytes, self._obj.flush())
         return b""
 
+    @property
+    def data_available(self) -> bool:
+        return not self._obj.is_finished() and not self._last_empty
+
 
 class ZSTDDecompressor(DecompressionBaseHandler):
     def __init__(
@@ -346,3 +371,7 @@ def decompress_sync(
 
     def flush(self) -> bytes:
         return b""
+
+    @property
+    def data_available(self) -> bool:
+        return not self._obj.needs_input and not self._obj.eof

aiohttp/http_exceptions.py

@@ -73,10 +73,6 @@ class ContentLengthError(PayloadEncodingError):
     """Not enough data to satisfy content length header."""
 
 
-class DecompressSizeError(PayloadEncodingError):
-    """Decompressed size exceeds the configured limit."""
-
-
 class LineTooLong(BadHttpMessage):
     def __init__(self, line: bytes, limit: int) -> None:
         super().__init__(f"Got more than {limit} bytes when reading: {line!r}.")