[pull] master from aio-libs:master

CHANGES/13362.bugfix.rst

@@ -0,0 +1 @@
+Reduced CPU consumption when encountering many concatenated members in a compressed payload and rejected large amounts of members -- by :user:`Dreamsorcerer`.

aiohttp/_websocket/reader_py.py

@@ -6,7 +6,7 @@
 from typing import Final
 
 from ..base_protocol import BaseProtocol
-from ..compression_utils import ZLibDecompressor
+from ..compression_utils import TooManyMembersError, ZLibDecompressor
 from ..helpers import _EXC_SENTINEL, set_exception
 from ..streams import EofStream
 from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask
@@ -248,14 +248,20 @@ def _handle_frame(
                 # but internally buffer more data such that the payload is
                 # >max_length, so we return one extra byte and if we're able
                 # to do that, then the message is too big.
-                payload_merged = self._decompressobj.decompress_sync(
-                    assembled_payload + WS_DEFLATE_TRAILING,
-                    (
-                        self._max_msg_size + 1
-                        if self._max_msg_size
-                        else self._max_msg_size
-                    ),
-                )
+                try:
+                    payload_merged = self._decompressobj.decompress_sync(
+                        assembled_payload + WS_DEFLATE_TRAILING,
+                        (
+                            self._max_msg_size + 1
+                            if self._max_msg_size
+                            else self._max_msg_size
+                        ),
+                    )
+                except TooManyMembersError as exc:
+                    raise WebSocketError(
+                        WSCloseCode.MESSAGE_TOO_BIG,
+                        "Compressed message has too many deflate members",
+                    ) from exc
                 if self._max_msg_size and len(payload_merged) > self._max_msg_size:
                     raise WebSocketError(
                         WSCloseCode.MESSAGE_TOO_BIG,

aiohttp/compression_utils.py

@@ -3,7 +3,7 @@
 import zlib
 from abc import ABC, abstractmethod
 from concurrent.futures import Executor
-from typing import Any, Final, Protocol, TypedDict, cast
+from typing import Any, Final, Generic, Protocol, TypedDict, TypeVar, cast
 
 if sys.version_info >= (3, 12):
     from collections.abc import Buffer
@@ -39,6 +39,21 @@
 ZLIB_MAX_LENGTH_UNLIMITED = 0  # zlib uses 0 to mean unlimited
 ZSTD_MAX_LENGTH_UNLIMITED = -1  # zstd uses -1 to mean unlimited
 
+# Concatenated members are decoded through a window that starts small and
+# doubles. A fresh decompressor copies everything past the member it decodes
+# into unused_data, so handing it the whole remaining buffer at every boundary
+# is quadratic over a stream of small members.
+MEMBER_WINDOW_MIN = 64
+MEMBER_WINDOW_MAX = 65536
+
+# Cap on concatenated members decoded in one call. Real payloads are unlikely
+# to have more than a few members.
+MAX_DECOMPRESS_MEMBERS = 1024
+
+
+class TooManyMembersError(ValueError):
+    """A stream concatenated more members than the caller allows."""
+
 
 class ZLibCompressObjProtocol(Protocol):
     def compress(self, data: Buffer) -> bytes: ...
@@ -155,6 +170,19 @@ def encoding_to_mode(
     return -ZLibBackend.MAX_WBITS if suppress_deflate_header else ZLibBackend.MAX_WBITS
 
 
+class MemberDecompressObjProtocol(Protocol):
+    def decompress(self, data: Buffer, max_length: int = ...) -> bytes: ...
+
+    @property
+    def eof(self) -> bool: ...
+
+    @property
+    def unused_data(self) -> bytes: ...
+
+
+_DecompressObjT = TypeVar("_DecompressObjT", bound=MemberDecompressObjProtocol)
+
+
 class DecompressionBaseHandler(ABC):
     def __init__(
         self,
@@ -190,6 +218,73 @@ def data_available(self) -> bool:
         """Return True if more output is available by passing b""."""
 
 
+class ConcatDecompressionHandler(DecompressionBaseHandler, Generic[_DecompressObjT]):
+    """Handler for a codec whose streams may concatenate independent members.
+
+    Concatenated gzip/deflate members and multi-frame zstd
+    (https://datatracker.ietf.org/doc/html/rfc8878#section-3.1.1) decode the
+    same way: a decompressor handles one member, then flags eof and leaves the
+    rest of the input in unused_data, so every member after it needs a fresh
+    one.
+    """
+
+    # Sentinel this codec's decompress() takes to mean "no output limit".
+    _unlimited: int
+    _decompressor: _DecompressObjT
+    # Input a max_length-capped walk stopped short of, fed back on the next call.
+    _pending_unused_data: bytes | None = None
+
+    @abstractmethod
+    def _new_decompressor(self) -> _DecompressObjT:
+        """Return a decompressor for the next member."""
+
+    def _decompress_members(self, first: bytes, max_length: int) -> bytes:
+        """Decode the members following the one ``first`` came from."""
+        remaining = memoryview(self._decompressor.unused_data)
+        parts = [first]
+        produced = len(first)
+        pos = 0
+        window = MEMBER_WINDOW_MIN
+        budget = max_length
+        members = 1
+
+        while pos < len(remaining):
+            if self._decompressor.eof:
+                members += 1
+                if members > MAX_DECOMPRESS_MEMBERS:
+                    raise TooManyMembersError(
+                        f"Compressed stream has more than "
+                        f"{MAX_DECOMPRESS_MEMBERS} members"
+                    )
+                # Replace the spent decompressor before the budget check below
+                # can break out of the loop: it still lists these bytes in its
+                # unused_data and would hand them back on the next call.
+                self._decompressor = self._new_decompressor()
+                window = MEMBER_WINDOW_MIN
+            if max_length != self._unlimited:
+                budget = max_length - produced
+                if budget <= 0:
+                    self._pending_unused_data = bytes(remaining[pos:])
+                    break
+
+            end = min(pos + window, len(remaining))
+            chunk = self._decompressor.decompress(remaining[pos:end], budget)
+            if chunk:
+                parts.append(chunk)
+                produced += len(chunk)
+
+            if self._decompressor.eof:
+                pos = end - len(self._decompressor.unused_data)
+            else:
+                pos = end
+                # Doubling the window on each iteration avoids too many calls
+                # when a large member is present, while protecting us from
+                # quadratic usage when members are of window+1 length.
+                window = min(window * 2, MEMBER_WINDOW_MAX)
+
+        return b"".join(parts)
+
+
 class ZLibCompressor:
     def __init__(
         self,
@@ -265,7 +360,9 @@ def flush(self, mode: int | None = None) -> bytes:
         )
 
 
-class ZLibDecompressor(DecompressionBaseHandler):
+class ZLibDecompressor(ConcatDecompressionHandler[ZLibDecompressObjProtocol]):
+    _unlimited = ZLIB_MAX_LENGTH_UNLIMITED
+
     def __init__(
         self,
         encoding: str | None = None,
@@ -276,9 +373,11 @@ def __init__(
         super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size)
         self._mode = encoding_to_mode(encoding, suppress_deflate_header)
         self._zlib_backend: Final = ZLibBackendWrapper(ZLibBackend._zlib_backend)
-        self._decompressor = self._zlib_backend.decompressobj(wbits=self._mode)
+        self._decompressor = self._new_decompressor()
         self._last_empty = False
-        self._pending_unused_data: bytes | None = None
+
+    def _new_decompressor(self) -> ZLibDecompressObjProtocol:
+        return self._zlib_backend.decompressobj(wbits=self._mode)
 
     def decompress_sync(
         self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
@@ -289,30 +388,20 @@ def decompress_sync(
         result = self._decompressor.decompress(
             self._decompressor.unconsumed_tail + data, max_length
         )
+
+        # Concatenated gzip/deflate stream: decode the members after this one.
+        if self._decompressor.eof and self._decompressor.unused_data:
+            result = self._decompress_members(result, max_length)
+
         # Only way to know that isal has no further data is checking we get no output
         self._last_empty = result == b""
 
-        # Handle concatenated gzip/deflate streams (multi-member).
-        # After a member ends, unused_data holds the start of the next member.
-        # Create a fresh decompressor for each subsequent member.
-        while self._decompressor.eof and self._decompressor.unused_data:
-            unused = self._decompressor.unused_data
-            self._decompressor = self._zlib_backend.decompressobj(wbits=self._mode)
-            if max_length != ZLIB_MAX_LENGTH_UNLIMITED:
-                max_length -= len(result)
-                if max_length <= 0:
-                    self._pending_unused_data = unused
-                    break
-            chunk = self._decompressor.decompress(unused, max_length)
-            self._last_empty = chunk == b""
-            result += chunk
-
         # Member ended exactly at chunk boundary — no unused_data, but the
         # next feed_data() call would fail on the spent decompressor.
         # Only reset for gzip; deflate's feed_eof() relies on eof=True to
         # confirm the stream is complete.
         if self._decompressor.eof and self._mode > self._zlib_backend.MAX_WBITS:
-            self._decompressor = self._zlib_backend.decompressobj(wbits=self._mode)
+            self._decompressor = self._new_decompressor()
 
         return result
 
@@ -384,7 +473,9 @@ def data_available(self) -> bool:
         return not self._obj.is_finished() and not self._last_empty
 
 
-class ZSTDDecompressor(DecompressionBaseHandler):
+class ZSTDDecompressor(ConcatDecompressionHandler["ZstdDecompressor"]):
+    _unlimited = ZSTD_MAX_LENGTH_UNLIMITED
+
     def __init__(
         self,
         executor: Executor | None = None,
@@ -395,9 +486,11 @@ def __init__(
                 "The zstd decompression is not available. "
                 "Please install `backports.zstd` module"
             )
-        self._obj = ZstdDecompressor()
-        self._pending_unused_data: bytes | None = None
         super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size)
+        self._decompressor = self._new_decompressor()
+
+    def _new_decompressor(self) -> "ZstdDecompressor":
+        return ZstdDecompressor()
 
     def decompress_sync(
         self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
@@ -412,28 +505,17 @@ def decompress_sync(
         if self._pending_unused_data is not None:
             data = self._pending_unused_data + data
             self._pending_unused_data = None
-        result = self._obj.decompress(data, zstd_max_length)
-
-        # Handle multi-frame zstd streams.
-        # https://datatracker.ietf.org/doc/html/rfc8878#section-3.1.1
-        # ZstdDecompressor handles one frame only. When a frame ends,
-        # eof becomes True and any trailing data goes to unused_data.
-        # We create a fresh decompressor to continue with the next frame.
-        while self._obj.eof and self._obj.unused_data:
-            unused_data = self._obj.unused_data
-            self._obj = ZstdDecompressor()
-            if zstd_max_length != ZSTD_MAX_LENGTH_UNLIMITED:
-                zstd_max_length -= len(result)
-                if zstd_max_length <= 0:
-                    self._pending_unused_data = unused_data
-                    break
-            result += self._obj.decompress(unused_data, zstd_max_length)
+        result = self._decompressor.decompress(data, zstd_max_length)
+
+        # Concatenated zstd stream: decode the frames after this one.
+        if self._decompressor.eof and self._decompressor.unused_data:
+            result = self._decompress_members(result, zstd_max_length)
 
         # Frame ended exactly at chunk boundary — no unused_data, but the
         # next feed_data() call would fail on the spent decompressor.
         # Prepare a fresh one for the next chunk.
-        if self._obj.eof:
-            self._obj = ZstdDecompressor()
+        if self._decompressor.eof:
+            self._decompressor = self._new_decompressor()
 
         return result
 
@@ -443,5 +525,5 @@ def flush(self) -> bytes:
     @property
     def data_available(self) -> bool:
         return (
-            not self._obj.needs_input and not self._obj.eof
+            not self._decompressor.needs_input and not self._decompressor.eof
         ) or self._pending_unused_data is not None