diff --git a/CHANGES/13362.bugfix.rst b/CHANGES/13362.bugfix.rst new file mode 100644 index 00000000000..f5fc57dda3b --- /dev/null +++ b/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`. diff --git a/aiohttp/_websocket/reader_py.py b/aiohttp/_websocket/reader_py.py index b84dfac5298..9583d79c1df 100644 --- a/aiohttp/_websocket/reader_py.py +++ b/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, diff --git a/aiohttp/compression_utils.py b/aiohttp/compression_utils.py index 400df141470..5b6ffc7f1f9 100644 --- a/aiohttp/compression_utils.py +++ b/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 diff --git a/tests/test_compression_utils.py b/tests/test_compression_utils.py index 5deebc8470d..9e079b88c23 100644 --- a/tests/test_compression_utils.py +++ b/tests/test_compression_utils.py @@ -1,13 +1,21 @@ """Tests for compression utils.""" import gzip +import os import sys +import zlib +from typing import Any import pytest from aiohttp.compression_utils import ( + MAX_DECOMPRESS_MEMBERS, + MEMBER_WINDOW_MAX, + MEMBER_WINDOW_MIN, + TooManyMembersError, ZLibBackend, ZLibCompressor, + ZLibDecompressObjProtocol, ZLibDecompressor, ZSTDDecompressor, ) @@ -21,6 +29,21 @@ zstandard = None # type: ignore[assignment] +class _RecordingDecompressObj: + """decompressobj proxy recording the size of every buffer fed to it.""" + + def __init__(self, obj: ZLibDecompressObjProtocol, feeds: list[int]) -> None: + self._obj = obj + self._feeds = feeds + + def __getattr__(self, name: str) -> Any: + return getattr(self._obj, name) + + def decompress(self, data: Any, max_length: int = 0) -> bytes: + self._feeds.append(len(data)) + return self._obj.decompress(data, max_length) + + @pytest.mark.usefixtures("parametrize_zlib_backend") async def test_compression_round_trip_in_executor() -> None: """Ensure that compression and decompression work correctly in the executor.""" @@ -123,3 +146,146 @@ def test_zlib_gzip_multi_member_max_length_exhausted_preserves_unused_data() -> assert result1 == b"AAAA" result2 = d.decompress_sync(member3) assert result2 == b"BBBBCCCC" + + +def test_zlib_gzip_multi_member_walks_input_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Every member is decoded from a bounded window, not from the whole tail. + + A fresh decompressor copies everything past the member it decodes into + unused_data, so handing it the rest of the buffer at each boundary is + quadratic: this blob would push ~6GiB through unused_data. + """ + blob = gzip.compress(b"") * (MAX_DECOMPRESS_MEMBERS - 1) + d = ZLibDecompressor(encoding="gzip") + feeds: list[int] = [] + decompressobj = d._zlib_backend.decompressobj + + def recording(*, wbits: int) -> _RecordingDecompressObj: + return _RecordingDecompressObj(decompressobj(wbits=wbits), feeds) + + monkeypatch.setattr(d._zlib_backend, "decompressobj", recording) + assert d.decompress_sync(blob) == b"" + assert max(feeds) <= MEMBER_WINDOW_MAX + assert sum(feeds) < 8 * len(blob) + + +def test_zlib_deflate_multi_member_window_doubles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A member outgrowing the window is fed in doubling slices, then capped. + + The schedule has to adapt to member size, because the peer chooses it: a + fixed large window copies everything past a small member's end into + unused_data (a member of window+1 bytes costs a whole window per byte of + progress), while a fixed small one needs a call per 64 bytes of a large + member. + """ + co = zlib.compressobj(wbits=-15) + small = co.compress(b"x") + co.flush() + co = zlib.compressobj(wbits=-15) + # Incompressible, so the member's wire size exceeds the whole ramp below. + large = co.compress(os.urandom(200 * 1024)) + co.flush() + + d = ZLibDecompressor(encoding="deflate", suppress_deflate_header=True) + feeds: list[int] = [] + decompressobj = d._zlib_backend.decompressobj + + def recording(*, wbits: int) -> _RecordingDecompressObj: + return _RecordingDecompressObj(decompressobj(wbits=wbits), feeds) + + monkeypatch.setattr(d._zlib_backend, "decompressobj", recording) + out = d.decompress_sync(small + large) + + assert len(out) == 1 + 200 * 1024 + ramp = [MEMBER_WINDOW_MIN << i for i in range(11)] # 64 .. 65536 + assert ramp[-1] == MEMBER_WINDOW_MAX + assert feeds[: len(ramp)] == ramp + assert max(feeds) == MEMBER_WINDOW_MAX + + +def test_zlib_deflate_max_length_exhausted_mid_member() -> None: + """Output can run out partway through a member, splitting the input in two. + + The bytes the cap left unconsumed stay on the decompressor as + unconsumed_tail while the windows not yet fed become pending, so resuming + has to feed unconsumed_tail first or the member decodes as garbage. + """ + co = zlib.compressobj(wbits=-15) + first_member = co.compress(b"x") + co.flush() + body = b"A" * 100_000 + co = zlib.compressobj(wbits=-15) + # Compresses small enough to span several windows, expands far past the cap. + second_member = co.compress(body) + co.flush() + assert len(second_member) > MEMBER_WINDOW_MIN + + d = ZLibDecompressor(encoding="deflate", suppress_deflate_header=True) + out = d.decompress_sync(first_member + second_member, max_length=1000) + split_seen = bool(d._decompressor.unconsumed_tail) and ( + d._pending_unused_data is not None + ) + while d.data_available: + out += d.decompress_sync(b"", max_length=1000) + + assert split_seen, "walk never suspended mid-member" + assert out == b"x" + body + + +def test_zlib_deflate_member_flood_rejected() -> None: + """Zero-output members never decrement max_length, cap must stop them.""" + co = zlib.compressobj(wbits=-15) + empty_member = co.compress(b"") + co.flush() # 2-byte empty raw-deflate member + d = ZLibDecompressor(encoding="deflate", suppress_deflate_header=True) + with pytest.raises(TooManyMembersError): + d.decompress_sync(empty_member * 5000, max_length=262144) + + +def test_zlib_deflate_members_at_limit() -> None: + """The limit counts the member decoded before the walk started.""" + co = zlib.compressobj(wbits=-15) + empty_member = co.compress(b"") + co.flush() + d = ZLibDecompressor(encoding="deflate", suppress_deflate_header=True) + assert d.decompress_sync(empty_member * MAX_DECOMPRESS_MEMBERS) == b"" + + +def test_zlib_deflate_members_one_over_limit() -> None: + co = zlib.compressobj(wbits=-15) + empty_member = co.compress(b"") + co.flush() + d = ZLibDecompressor(encoding="deflate", suppress_deflate_header=True) + with pytest.raises(TooManyMembersError): + d.decompress_sync(empty_member * (MAX_DECOMPRESS_MEMBERS + 1)) + + +@pytest.mark.parametrize("max_length", (0, 262144), ids=("unlimited", "capped")) +def test_zlib_gzip_many_members(max_length: int) -> None: + """A call may decode up to the member limit, resuming across calls.""" + member = gzip.compress(b"A" * 64) + count = MAX_DECOMPRESS_MEMBERS + d = ZLibDecompressor(encoding="gzip") + out = d.decompress_sync(member * count, max_length=max_length) + while d.data_available: + out += d.decompress_sync(b"", max_length=max_length) + assert out == b"A" * 64 * count + + +def test_zlib_gzip_empty_members_interleaved_with_output() -> None: + blob = (gzip.compress(b"") + gzip.compress(b"DATA")) * 20 + d = ZLibDecompressor(encoding="gzip") + assert d.decompress_sync(blob) == b"DATA" * 20 + + +@pytest.mark.skipif(zstandard is None, reason="zstandard is not installed") +def test_zstd_frame_flood_rejected() -> None: + """Multi-frame zstd shares the walk, so it shares the limit.""" + d = ZSTDDecompressor() + with pytest.raises(TooManyMembersError): + d.decompress_sync(zstandard.compress(b"") * 5000) + + +@pytest.mark.skipif(zstandard is None, reason="zstandard is not installed") +def test_zstd_many_frames() -> None: + frame = zstandard.compress(b"B" * 64) + count = MAX_DECOMPRESS_MEMBERS + d = ZSTDDecompressor() + assert d.decompress_sync(frame * count) == b"B" * 64 * count diff --git a/tests/test_http_parser.py b/tests/test_http_parser.py index 197bf83cb48..259d88f7dfa 100644 --- a/tests/test_http_parser.py +++ b/tests/test_http_parser.py @@ -3389,6 +3389,26 @@ async def test_http_payload_gzip_many_small_members( assert b"".join(parts) == b"".join(out._buffer) assert out.is_eof() + async def test_http_payload_gzip_empty_member_flood( + self, protocol: BaseProtocol + ) -> None: + """A body of many empty members is rejected, not decoded. + + Empty members decode to nothing, so must fail from the member limit. + """ + payload = gzip.compress(b"") * 20000 + out = aiohttp.StreamReader( + protocol, DEFAULT_CHUNK_SIZE, loop=asyncio.get_running_loop() + ) + p = HttpPayloadParser( + out, + length=len(payload), + compression="gzip", + headers_parser=HeadersParser(), + ) + with pytest.raises(http_exceptions.ContentEncodingError): + p.feed_data(payload) + class TestDeflateBuffer: async def test_feed_data(self, protocol: BaseProtocol) -> None: diff --git a/tests/test_web_functional.py b/tests/test_web_functional.py index 7c5fde8af07..b8f61e1c2bc 100644 --- a/tests/test_web_functional.py +++ b/tests/test_web_functional.py @@ -2529,7 +2529,7 @@ async def test_bad_method_for_c_http_parser_not_hangs( aiohttp_client: AiohttpClient, ) -> None: app = web.Application() - timeout = aiohttp.ClientTimeout(sock_read=0.2) + timeout = aiohttp.ClientTimeout(sock_read=3) client = await aiohttp_client(app, timeout=timeout) resp = await client.request("GET1", "/") assert 400 == resp.status diff --git a/tests/test_websocket_parser.py b/tests/test_websocket_parser.py index 10cec0cc662..9e3d08a1649 100644 --- a/tests/test_websocket_parser.py +++ b/tests/test_websocket_parser.py @@ -2,6 +2,7 @@ import pickle import random import struct +import zlib from unittest import mock import pytest @@ -647,6 +648,78 @@ def test_compressed_frame_after_control_frame( assert out._buffer[1] == WSMessageText(data="hello", size=5, extra="") +# A complete raw-deflate member: BFINAL set, empty fixed-Huffman block, +# decoding to nothing. +EMPTY_DEFLATE_MEMBER = b"\x03\x00" + + +@pytest.mark.usefixtures("parametrize_zlib_backend") +def test_compressed_multi_block_message(out: WebSocketDataQueue) -> None: + """Blocks that leave BFINAL unset stay a single member.""" + reader = WebSocketReader(out, 4 * 1024 * 1024, compress=True, decode_text=True) + compressobj = zlib.compressobj(wbits=-15) + payload = b"".join( + compressobj.compress(b"block %d " % i) + + compressobj.flush(ZLibBackend.Z_SYNC_FLUSH) + for i in range(50) + ).removesuffix(WS_DEFLATE_TRAILING) + expected = b"".join(b"block %d " % i for i in range(50)) + + # FIN | RSV1 (compressed) | BINARY + error, _ = reader.feed_data( + PACK_LEN2(0x80 | 0x40 | WSMsgType.BINARY, 126, len(payload)) + payload + ) + + assert error is False + assert out._buffer[0] == WSMessageBinary( + data=expected, size=len(expected), extra="" + ) + + +@pytest.mark.usefixtures("parametrize_zlib_backend") +def test_compressed_multi_member_message(out: WebSocketDataQueue) -> None: + """Concatenated BFINAL members decode. + + A block following a BFINAL one starts a fresh deflate member, so this must + keep working well below the member limit. + """ + reader = WebSocketReader(out, 4 * 1024 * 1024, compress=True, decode_text=True) + payload = b"" + for i in range(50): + compressobj = zlib.compressobj(wbits=-15) + payload += compressobj.compress(b"part %d " % i) + compressobj.flush() + expected = b"".join(b"part %d " % i for i in range(50)) + + error, _ = reader.feed_data( + PACK_LEN2(0x80 | 0x40 | WSMsgType.BINARY, 126, len(payload)) + payload + ) + + assert error is False + assert out._buffer[0] == WSMessageBinary( + data=expected, size=len(expected), extra="" + ) + + +@pytest.mark.usefixtures("parametrize_zlib_backend") +def test_compressed_member_flood_rejected(out: WebSocketDataQueue) -> None: + """Past the member limit the message is rejected, not decoded. + + The reader hands the whole assembled message to one synchronous + decompress_sync() call. Empty members produce no output for max_msg_size + to bound, so must be limited by member count. + """ + max_msg_size = 262144 + reader = WebSocketReader(out, max_msg_size, compress=True, decode_text=True) + payload = EMPTY_DEFLATE_MEMBER * (2 * max_msg_size // 256) + + with pytest.raises(WebSocketError) as ctx: + reader._feed_data( + PACK_LEN2(0x80 | 0x40 | WSMsgType.BINARY, 126, len(payload)) + payload + ) + + assert ctx.value.code == WSCloseCode.MESSAGE_TOO_BIG + + @pytest.mark.parametrize("opcode", (WSMsgType.PING, WSMsgType.PONG, WSMsgType.CLOSE)) def test_control_frame_with_rsv1( parser: PatchableWebSocketReader, opcode: WSMsgType