diff --git a/CHANGES/13393.breaking.rst b/CHANGES/13393.breaking.rst new file mode 100644 index 00000000000..dd34fe81535 --- /dev/null +++ b/CHANGES/13393.breaking.rst @@ -0,0 +1 @@ +The WebSocket receive queue now only holds a weak reference to the ``WebSocketReader`` while parsing is stalled; code constructing a reader directly and passing it to ``set_parser()`` must keep its own strong reference to it, or frames the reader stopped short of parsing are lost with it -- by :user:`Dreamsorcerer`. diff --git a/CHANGES/13393.bugfix.rst b/CHANGES/13393.bugfix.rst new file mode 100644 index 00000000000..64a8b54463a --- /dev/null +++ b/CHANGES/13393.bugfix.rst @@ -0,0 +1 @@ +Fixed excessive memory consumption with small WebSocket messages -- by :user:`Dreamsorcerer`. diff --git a/THREAT_MODEL.md b/THREAT_MODEL.md index da8c210b543..41d0ef51c7b 100644 --- a/THREAT_MODEL.md +++ b/THREAT_MODEL.md @@ -582,6 +582,21 @@ client-side, the writer adds masks to outgoing frames. fragment count (`max(1024, max_msg_size // 256)`) and pauses reading for backpressure once exceeded, mirroring the HTTP chunk-splits limit in `StreamReader` (PR #11894). +- **PR #13393** — queue accounting only counted payload bytes, so a flood + of empty/tiny frames could pin unbounded per-message object overhead in + `WebSocketDataQueue` before the `_limit` high-water mark fired. Each + queued message is now charged a 128-byte accounting overhead + (`MSG_SIZE_OVERHEAD`) and the reader stalls at a frame boundary once the + queue is over the mark: the unparsed remainder stays compressed in + `_tail`, the queue keeps only a weakref to the stalled reader, and + drains re-drive it at a `_limit // 2` low-water mark (each resume + re-slices the tail, so per-pop resumes would make draining a tiny-frame + burst quadratic CPU), and the transport stays paused until the stash is + exhausted (resuming earlier would admit a fresh socket read into `_tail` + per couple of drained messages, relocating the memory bound there). + Callers constructing a `WebSocketReader` for + `set_parser()` must hold a strong reference to it (both in-tree response + classes do, via `_parser`). --- diff --git a/aiohttp/_websocket/reader_c.pxd b/aiohttp/_websocket/reader_c.pxd index 918fea048ad..19c7b3150d9 100644 --- a/aiohttp/_websocket/reader_c.pxd +++ b/aiohttp/_websocket/reader_c.pxd @@ -44,11 +44,13 @@ cdef set MESSAGE_TYPES_WITH_CONTENT cdef tuple EMPTY_FRAME cdef tuple EMPTY_FRAME_ERROR +cdef unsigned int MSG_SIZE_OVERHEAD + cdef class WebSocketDataQueue: - cdef unsigned int _size + cdef readonly unsigned int _size cdef public object _protocol - cdef unsigned int _limit + cdef readonly unsigned int _limit cdef object _loop cdef bint _eof cdef object _waiter @@ -56,6 +58,7 @@ cdef class WebSocketDataQueue: cdef public object _buffer cdef object _get_buffer cdef object _put_buffer + cdef readonly object _stalled_reader cdef void _release_waiter(self) @@ -67,6 +70,11 @@ cdef class WebSocketDataQueue: cdef class WebSocketReader: + # cdef classes are not weak-referenceable without this; the queue parks a + # weakref here while parsing is stalled. + cdef object __weakref__ + cdef object _weak_self + cdef WebSocketDataQueue queue cdef unsigned int _max_msg_size cdef bint _decode_text diff --git a/aiohttp/_websocket/reader_py.py b/aiohttp/_websocket/reader_py.py index a295db0a4a7..c434303c710 100644 --- a/aiohttp/_websocket/reader_py.py +++ b/aiohttp/_websocket/reader_py.py @@ -3,11 +3,14 @@ import asyncio import builtins import sys +import weakref from collections import deque +from typing import Final from ..base_protocol import BaseProtocol from ..compression_utils import TooManyMembersError, ZLibDecompressor from ..helpers import _EXC_SENTINEL, set_exception +from ..log import ws_logger from ..streams import EofStream from .helpers import UNPACK_CLOSE_CODE, UNPACK_LEN3, websocket_mask from .models import ( @@ -60,6 +63,15 @@ TUPLE_NEW = tuple.__new__ +# Overhead added to each message to ensure that tiny messages can't use +# unreasonable amounts of memory. +MSG_SIZE_OVERHEAD: Final[int] = 128 + +STALLED_READER_COLLECTED: Final[str] = ( + "WebSocketReader was garbage collected while stalled; " + "callers of set_parser() must hold a strong reference" +) + cython_int = int # Typed to int in Python, but cython with use a signed int in the pxd @@ -82,6 +94,7 @@ def __init__( self._buffer: deque[WSMessage] = deque() self._get_buffer = self._buffer.popleft self._put_buffer = self._buffer.append + self._stalled_reader: "weakref.ref[WebSocketReader] | None" = None def is_eof(self) -> bool: return self._eof @@ -113,8 +126,10 @@ def feed_eof(self) -> None: self._exception = None # Break cyclic references def feed_data(self, data: "WSMessage") -> None: + # Unbox into the typed local before adding, so Cython keeps the sum in + # C instead of boxing MSG_SIZE_OVERHEAD for a Python-level add. size = data.size - self._size += size + self._size += size + MSG_SIZE_OVERHEAD self._put_buffer(data) self._release_waiter() if self._size > self._limit and not self._protocol._reading_paused: @@ -135,8 +150,33 @@ def _read_from_buffer(self) -> WSMessage: if self._buffer: data = self._get_buffer() size = data.size - self._size -= size - if self._size < self._limit and self._protocol._reading_paused: + self._size -= size + MSG_SIZE_OVERHEAD + if self._stalled_reader is not None and self._size <= self._limit // 2: + # Resume parsing once the queue drains to the low-water mark. + # Each resume re-slices the parser's unparsed tail, so waiting + # for headroom makes a drain cost one copy per batch of + # messages instead of one per message. + if (reader := self._stalled_reader()) is not None: + reader.feed_data(b"") + else: + # The stash died with the reader. Deliver what was already + # queued, then surface the contract violation on the next + # read instead of hanging. Log as well, since a caller that + # stops reading early never sees the deferred exception. A + # real failure that was already recorded stays the reported + # cause. + self._stalled_reader = None + ws_logger.warning(STALLED_READER_COLLECTED) + if self._exception is None: + self.set_exception(RuntimeError(STALLED_READER_COLLECTED)) + # Resuming the transport while a stash remains would admit + # another socket read into the tail for every couple of messages + # drained, moving the memory bound from the queue into the tail. + if ( + self._stalled_reader is None + and self._size < self._limit + and self._protocol._reading_paused + ): self._protocol.resume_reading() return data if self._exception is not None: @@ -155,6 +195,9 @@ def __init__( self.queue = queue self._max_msg_size = max_msg_size self._decode_text = decode_text + # Parked on the queue while parsing is stalled; created once so + # stalling does not allocate. + self._weak_self = weakref.ref(self) self._exc: Exception | None = None self._partial = bytearray() @@ -349,6 +392,7 @@ def _handle_frame( def _feed_data(self, data: bytes) -> None: """Return the next frame from the socket.""" + self.queue._stalled_reader = None if self._tail: data, self._tail = self._tail + data, b"" @@ -357,6 +401,14 @@ def _feed_data(self, data: bytes) -> None: data_cstr = data while True: + if start_pos < data_len and self.queue._size > self.queue._limit: + # Over the high-water mark with unparsed bytes left: stash the + # remainder and stall. Gating on unparsed bytes keeps a read + # that ended on a frame boundary from arming an empty stall, + # which would hold the transport paused with nothing to drain. + self.queue._stalled_reader = self._weak_self + break + # read header if self._state == READ_HEADER: if data_len - start_pos < 2: diff --git a/aiohttp/client.py b/aiohttp/client.py index 9739cc46019..f91b81e5712 100644 --- a/aiohttp/client.py +++ b/aiohttp/client.py @@ -1210,14 +1210,14 @@ async def _ws_connect( compress=compress, client_notakeover=notakeover, ) - parser = WebSocketReader( + ws_resp._parser = WebSocketReader( reader, max_msg_size, compress=bool(compress), decode_text=decode_text, ) cb = None if heartbeat is None else ws_resp._on_data_received - conn_proto.set_parser(parser, reader, data_received_cb=cb) + conn_proto.set_parser(ws_resp._parser, reader, data_received_cb=cb) return ws_resp def _prepare_headers(self, headers: LooseHeaders | None) -> "CIMultiDict[str]": diff --git a/aiohttp/client_ws.py b/aiohttp/client_ws.py index 40fbca676b7..1eb5491b8b1 100644 --- a/aiohttp/client_ws.py +++ b/aiohttp/client_ws.py @@ -7,7 +7,7 @@ from types import TracebackType from typing import Any, Final, Generic, Literal, overload -from ._websocket.reader import WebSocketDataQueue +from ._websocket.reader import WebSocketDataQueue, WebSocketReader from .client_exceptions import ClientError, ServerTimeoutError, WSMessageTypeError from .client_reqrep import ClientResponse from .helpers import calculate_timeout_when, frozen_dataclass_decorator, set_result @@ -79,6 +79,10 @@ def __init__( self._writer = writer self._reader = reader + # Set by ClientSession._ws_connect; owns the parser so a stalled + # reader parked on the queue by weakref stays alive while this + # response can still be drained. + self._parser: WebSocketReader | None = None self._protocol = protocol self._closed = False self._closing = False @@ -220,6 +224,8 @@ def _handle_ping_pong_exception(self, exc: BaseException) -> None: if self._closed: return self._set_closed() + # close() is never reached after this; release the parser here. + self._parser = None self._close_code = WSCloseCode.ABNORMAL_CLOSURE self._exception = exc self._response.close() @@ -337,25 +343,8 @@ async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bo self._set_closed() try: - await self._writer.close(code, message) - except asyncio.CancelledError: - self._close_code = WSCloseCode.ABNORMAL_CLOSURE - self._response.close() - raise - except Exception as exc: - self._close_code = WSCloseCode.ABNORMAL_CLOSURE - self._exception = exc - self._response.close() - return True - - if self._close_code: - self._response.close() - return True - - while True: try: - async with async_timeout.timeout(self._timeout.ws_close): - msg = await self._reader.read() + await self._writer.close(code, message) except asyncio.CancelledError: self._close_code = WSCloseCode.ABNORMAL_CLOSURE self._response.close() @@ -366,11 +355,33 @@ async def close(self, *, code: int = WSCloseCode.OK, message: bytes = b"") -> bo self._response.close() return True - if msg.type is WSMsgType.CLOSE: - self._close_code = msg.data + if self._close_code: self._response.close() return True + while True: + try: + async with async_timeout.timeout(self._timeout.ws_close): + msg = await self._reader.read() + except asyncio.CancelledError: + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + self._response.close() + raise + except Exception as exc: + self._close_code = WSCloseCode.ABNORMAL_CLOSURE + self._exception = exc + self._response.close() + return True + + if msg.type is WSMsgType.CLOSE: + self._close_code = msg.data + self._response.close() + return True + finally: + # Once closed the response can no longer be drained; release the + # parser and the stash it retains. + self._parser = None + @overload async def receive( self: "ClientWebSocketResponse[Literal[True]]", timeout: float | None = None @@ -425,7 +436,10 @@ async def receive( await self.close() return WS_CLOSED_MESSAGE except ClientError: - # Likely ServerDisconnectedError when connection is lost + # Likely ServerDisconnectedError when connection is lost. + # close() is not called on this path, so release the parser + # and the stash it retains here. + self._parser = None self._set_closed() self._close_code = WSCloseCode.ABNORMAL_CLOSURE return WS_CLOSED_MESSAGE diff --git a/aiohttp/web_ws.py b/aiohttp/web_ws.py index 0bda7ab4bdf..a71141fe940 100644 --- a/aiohttp/web_ws.py +++ b/aiohttp/web_ws.py @@ -82,6 +82,7 @@ class WebSocketResponse(StreamResponse, Generic[_DecodeText]): _ws_protocol: str | None = None _writer: WebSocketWriter | None = None _reader: WebSocketDataQueue | None = None + _parser: WebSocketReader | None = None _closed: bool = False _closing: bool = False _conn_lost: int = 0 @@ -247,6 +248,8 @@ def _handle_ping_pong_exception(self, exc: BaseException) -> None: if self._closed: return self._set_closed() + # close() is never reached after this; release the parser here. + self._parser = None self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) self._exception = exc if self._waiting and not self._closing and self._reader is not None: @@ -391,14 +394,16 @@ def _post_start( self._reader = WebSocketDataQueue( request._protocol, DEFAULT_CHUNK_SIZE, loop=loop ) - parser = WebSocketReader( + # Owns the parser so a stalled reader parked on the queue by weakref + # stays alive while this response can still be drained. + self._parser = WebSocketReader( self._reader, self._max_msg_size, compress=bool(self._compress), decode_text=self._decode_text, ) cb = None if self._heartbeat is None else self._on_data_received - request.protocol.set_parser(parser, data_received_cb=cb) + request.protocol.set_parser(self._parser, data_received_cb=cb) # disable HTTP keepalive for WebSocket request.protocol.keep_alive(False) @@ -524,48 +529,53 @@ async def close( self._set_closed() try: - await self._writer.close(code, message) - writer = self._payload_writer - assert writer is not None - if drain: - await writer.drain() - except (asyncio.CancelledError, asyncio.TimeoutError): - self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) - raise - except Exception as exc: - self._exception = exc - self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) - return True - - reader = self._reader - assert reader is not None - # we need to break `receive()` cycle before we can call - # `reader.read()` as `close()` may be called from different task - if self._waiting: - assert self._loop is not None - assert self._close_wait is None - self._close_wait = self._loop.create_future() - reader.feed_data(WS_CLOSING_MESSAGE) - await self._close_wait - - if self._closing: - self._close_transport() - return True + try: + await self._writer.close(code, message) + writer = self._payload_writer + assert writer is not None + if drain: + await writer.drain() + except (asyncio.CancelledError, asyncio.TimeoutError): + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + raise + except Exception as exc: + self._exception = exc + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + return True - try: - async with async_timeout.timeout(self._timeout): - while True: - msg = await reader.read() - if msg.type is WSMsgType.CLOSE: - self._set_code_close_transport(msg.data) - return True - except asyncio.CancelledError: - self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) - raise - except Exception as exc: - self._exception = exc - self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) - return True + reader = self._reader + assert reader is not None + # we need to break `receive()` cycle before we can call + # `reader.read()` as `close()` may be called from different task + if self._waiting: + assert self._loop is not None + assert self._close_wait is None + self._close_wait = self._loop.create_future() + reader.feed_data(WS_CLOSING_MESSAGE) + await self._close_wait + + if self._closing: + self._close_transport() + return True + + try: + async with async_timeout.timeout(self._timeout): + while True: + msg = await reader.read() + if msg.type is WSMsgType.CLOSE: + self._set_code_close_transport(msg.data) + return True + except asyncio.CancelledError: + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + raise + except Exception as exc: + self._exception = exc + self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE) + return True + finally: + # Once closed the response can no longer be drained; release the + # parser and the stash it retains. + self._parser = None def _set_closing(self, code: int) -> None: """Set the close code and mark the connection as closing.""" diff --git a/tests/test_client_ws_functional.py b/tests/test_client_ws_functional.py index 8eccc163150..628a55b49d2 100644 --- a/tests/test_client_ws_functional.py +++ b/tests/test_client_ws_functional.py @@ -1,5 +1,10 @@ import asyncio +import base64 +import contextlib +import gc +import hashlib import json +import socket import struct import sys import zlib @@ -21,7 +26,7 @@ from aiohttp._websocket.models import WS_DEFLATE_TRAILING, WSMessageBinary from aiohttp._websocket.reader import WebSocketDataQueue from aiohttp.client_ws import ClientWSTimeout -from aiohttp.http import WebSocketError, WSCloseCode +from aiohttp.http import WS_KEY, WebSocketError, WSCloseCode if sys.version_info >= (3, 11): import asyncio as async_timeout @@ -33,6 +38,90 @@ class PatchableWebSocketDataQueue(WebSocketDataQueue): """A WebSocketDataQueue that can be patched.""" +async def test_stashed_frames_survive_connection_loss( + unused_port_socket: socket.socket, +) -> None: + """Frames stashed by receive-queue backpressure outlive the connection. + + Mirror of the server-side test: a peer can pack more complete frames into + one read than the queue's high-water mark allows, so the parser stalls + with the rest in its tail. ``connection_lost()`` drops the protocol's + reference to the parser, so only ``ClientWebSocketResponse._parser`` keeps + it alive; without that the queue's weak link dies and the stash is lost. + """ + sent = 8000 + + writers: list[asyncio.StreamWriter] = [] + + async def raw_server( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + writers.append(writer) + request = await reader.readuntil(b"\r\n\r\n") + key = next( + line.split(b":", 1)[1].strip() + for line in request.split(b"\r\n") + if line.lower().startswith(b"sec-websocket-key") + ) + accept = base64.b64encode(hashlib.sha1(key + WS_KEY).digest()) + writer.write( + b"HTTP/1.1 101 Switching Protocols\r\n" + b"Upgrade: websocket\r\n" + b"Connection: Upgrade\r\n" + b"Sec-WebSocket-Accept: " + accept + b"\r\n\r\n" + # One oversized read: empty unmasked TEXT frames, 2 bytes each. + + b"\x81\x00" * sent + ) + await writer.drain() + + server = await asyncio.start_server(raw_server, sock=unused_port_socket) + port = unused_port_socket.getsockname()[1] + try: + async with aiohttp.ClientSession() as session: + ws = await session.ws_connect(f"http://127.0.0.1:{port}/") + queue = ws._reader + # Usually the whole burst rides along with the 101 response and + # is parsed before ws_connect() returns, so the wait is a + # fallback for split reads only. + for _ in range(1000): # pragma: no cover + if queue._stalled_reader is not None: + break + await asyncio.sleep(0.01) + assert queue._stalled_reader is not None, "parser never stalled" + stalled = len(queue._buffer) + + # Tear the connection down for real: closing the transport is what + # drives connection_lost(), which drops the protocol's reference to + # the parser. A paused transport never sees the peer's FIN, so this + # cannot be triggered by the peer going away. + connection = ws._conn + assert connection is not None + protocol = connection.protocol + assert protocol is not None + assert protocol.transport is not None + protocol.transport.close() + for _ in range(1000): # pragma: no branch + if protocol._payload_parser is None: + break + await asyncio.sleep(0) + assert protocol._payload_parser is None, "connection_lost never ran" + + count = 0 + while (msg := await ws.receive()).type is WSMsgType.TEXT: + count += 1 + # A clean drain ends with the queue's close, not an error. + assert msg.type is WSMsgType.CLOSED + finally: + for writer in writers: + writer.close() + with contextlib.suppress(Exception): + await writer.wait_closed() + server.close() + await server.wait_closed() + + assert stalled < count <= sent + + async def test_send_recv_text(aiohttp_client: AiohttpClient) -> None: async def handler(request: web.Request) -> web.WebSocketResponse: ws = web.WebSocketResponse() @@ -1678,3 +1767,75 @@ async def handler(request: web.Request) -> web.WebSocketResponse: assert msg.type is WSMsgType.ERROR, msg assert isinstance(msg.data, WebSocketError) assert msg.data.code == WSCloseCode.PROTOCOL_ERROR + + +async def test_stalled_parser_outlives_connection_lost( + aiohttp_client: AiohttpClient, +) -> None: + """Frames the parser stopped short of are delivered after the peer vanishes. + + Once the queue is over its high-water mark the parser stops mid-read and + parks the rest of the read on itself, reachable from the queue only through + a weak reference. Connection loss releases the protocol's reference to the + parser, so the response has to be the owner or those frames are collected + and the application silently sees a short stream. + """ + + async def handler(request: web.Request) -> web.WebSocketResponse: + ws = web.WebSocketResponse() + await ws.prepare(request) + # Hold the connection open. TestServer cancels the handler on + # connection loss; swallow it so the handler exits normally. + with contextlib.suppress(asyncio.CancelledError): + await ws.receive() + return ws + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + ws = await client.ws_connect("/") + + # Empty server->client TEXT frames: two bytes on the wire, but each one is + # charged MSG_SIZE_OVERHEAD in the queue, so a single read crosses the + # high-water mark and the parser stalls part way through it. + sent = 8000 + assert ws._conn is not None + protocol = ws._conn.protocol + assert protocol is not None + transport = protocol.transport + assert transport is not None + protocol.data_received(b"\x81\x00" * sent) + assert ws._reader._stalled_reader is not None, "parser never stalled" + + # Simulate the peer vanishing: the socket is gone and the protocol drops + # its reference to the parser, exactly as the event loop would do it. + transport.abort() + protocol.connection_lost(None) + for _ in range(3): # PyPy can need more than one pass + gc.collect() + + count = 0 + while (await ws.receive()).type is not WSMsgType.CLOSED: + count += 1 + assert count == sent + # Exhausting the queue releases the parser and the stash it retains. + assert ws._parser is None + + +async def test_close_releases_parser(aiohttp_client: AiohttpClient) -> None: + """A normal close handshake releases the parser and the state it holds.""" + + async def handler(request: web.Request) -> web.WebSocketResponse: + ws = web.WebSocketResponse() + await ws.prepare(request) + await ws.receive() # the peer's CLOSE + return ws + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + resp = await client.ws_connect("/") + + assert resp._parser is not None + await resp.close() + assert resp._parser is None diff --git a/tests/test_web_websocket_functional.py b/tests/test_web_websocket_functional.py index 6ec7ef8b6c0..353807d9364 100644 --- a/tests/test_web_websocket_functional.py +++ b/tests/test_web_websocket_functional.py @@ -2,7 +2,9 @@ import asyncio import contextlib +import gc import json +import socket import sys import weakref from typing import Literal, NoReturn @@ -80,6 +82,100 @@ async def second_handler(request: web.Request) -> web.Response: assert b"426" in data +async def test_stashed_frames_survive_connection_loss( + unused_port_socket: socket.socket, +) -> None: + """Frames stashed by receive-queue backpressure outlive the connection. + + A peer can pack far more complete frames into one read than the queue's + high-water mark allows, so the parser stops part-way and leaves the rest + in its tail. ``connection_lost()`` then drops the protocol's reference to + the parser, so only ``WebSocketResponse._parser`` keeps it alive; without + that the queue's weak link dies with the connection and every stashed + frame is silently lost. + """ + sent = 8000 + # (frames queued when the parser stalled, frames delivered in total) + received: asyncio.Future[tuple[int, int]] = ( + asyncio.get_running_loop().create_future() + ) + + async def handler(request: web.Request) -> web.WebSocketResponse: + ws = web.WebSocketResponse() + await ws.prepare(request) + try: + queue = ws._reader + assert queue is not None + + # Wait for the parser to stall with frames still stashed in its + # tail. The transport is paused from here, so the peer's FIN + # cannot be observed until reading resumes. + for _ in range(1000): # pragma: no branch + if queue._stalled_reader is not None: + break + await asyncio.sleep(0.01) + assert queue._stalled_reader is not None, "parser never stalled" + stalled = len(queue._buffer) + + # Tear the connection down for real: closing the transport is + # what drives connection_lost(), which drops the protocol's only + # reference to the parser, leaving ws._parser holding it. + protocol = request.protocol + transport = protocol.transport + assert transport is not None + transport.close() + for _ in range(1000): # pragma: no branch + if protocol._payload_parser is None: + break + await asyncio.sleep(0) + assert protocol._payload_parser is None, "connection_lost never ran" + + count = 0 + while (msg := await ws.receive()).type is WSMsgType.TEXT: + count += 1 + # A clean drain ends with the queue's close, not an error. + assert msg.type is WSMsgType.CLOSED + + received.set_result((stalled, count)) + except Exception as exc: # pragma: no cover + # Surface handler failures instead of an opaque timeout. + received.set_exception(exc) + raise + return ws + + # A plain AppRunner, not the aiohttp_server fixture: TestServer forces + # handler_cancellation=True, which kills the handler on connection loss. + # Production defaults to False, and that is the case where the stash + # still has to be drainable. + app = web.Application() + app.router.add_route("GET", "/ws", handler) + runner = web.AppRunner(app) + await runner.setup() + try: + await web.SockSite(runner, unused_port_socket).start() + port = unused_port_socket.getsockname()[1] + + reader, writer = await asyncio.open_connection("127.0.0.1", port) + try: + writer.write(_RAW_UPGRADE) + await writer.drain() + await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout=5) + + # Written as one call, but whether it arrives as one read is the + # kernel's choice. + writer.write(b"\x81\x80\x00\x00\x00\x00" * sent) + await writer.drain() + + stalled, count = await asyncio.wait_for(received, timeout=10) + assert stalled < count <= sent + finally: + writer.close() + with contextlib.suppress(ConnectionResetError): + await writer.wait_closed() + finally: + await runner.cleanup() + + async def test_partial_pipelined_request_after_failed_websocket_upgrade( aiohttp_server: AiohttpServer, ) -> None: @@ -1928,3 +2024,85 @@ async def handler(request: web.Request) -> web.Response: aiohttp.ServerDisconnectedError, aiohttp.ClientConnectionResetError ): await request_task + + +async def test_stalled_parser_outlives_connection_lost( + aiohttp_client: AiohttpClient, +) -> None: + """Frames the parser stopped short of are delivered after the peer vanishes. + + Once the queue is over its high-water mark the parser stops mid-read and + parks the rest of the read on itself, reachable from the queue only through + a weak reference. Connection loss releases the protocol's reference to the + parser, so the response has to be the owner or those frames are collected + and the handler silently sees a short stream. + """ + sent = 8000 + received: asyncio.Future[int] = asyncio.get_running_loop().create_future() + + async def handler(request: web.Request) -> web.WebSocketResponse: + ws = web.WebSocketResponse() + await ws.prepare(request) + try: + protocol = request.protocol + transport = request.transport + assert transport is not None + # Empty masked client->server TEXT frames: six bytes on the wire, + # but each one is charged MSG_SIZE_OVERHEAD in the queue, so a + # single read crosses the high-water mark and the parser stalls + # part way through. + protocol.data_received(b"\x81\x80\x00\x00\x00\x00" * sent) + queue = ws._reader + assert queue is not None + assert queue._stalled_reader is not None, "parser never stalled" + + # Simulate the peer vanishing: the socket is gone and the + # protocol drops its reference to the parser, as the event loop + # would do it. Detach the handler task first: TestServer forces + # handler_cancellation=True (production defaults to False), and + # the cancellation would land on this very task at its next yield. + protocol._task_handler = None + transport.abort() + protocol.connection_lost(None) + for _ in range(3): # PyPy can need more than one pass + gc.collect() + + count = 0 + while (await ws.receive()).type is not WSMsgType.CLOSED: + count += 1 + # Exhausting the queue releases the parser and its stash. + assert ws._parser is None + received.set_result(count) + except Exception as exc: # pragma: no cover + # Surface handler failures instead of an opaque timeout. + received.set_exception(exc) + raise + return ws + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + await client.ws_connect("/") + + assert await asyncio.wait_for(received, 10) == sent + + +async def test_close_releases_parser(aiohttp_client: AiohttpClient) -> None: + """A normal close handshake releases the parser and the state it holds.""" + released: asyncio.Future[bool] = asyncio.get_running_loop().create_future() + + async def handler(request: web.Request) -> web.WebSocketResponse: + ws = web.WebSocketResponse() + await ws.prepare(request) + assert ws._parser is not None + await ws.close() + released.set_result(ws._parser is None) + return ws + + app = web.Application() + app.router.add_get("/", handler) + client = await aiohttp_client(app) + ws = await client.ws_connect("/") + + await ws.receive() # the server's CLOSE + assert await asyncio.wait_for(released, 5) is True diff --git a/tests/test_websocket_parser.py b/tests/test_websocket_parser.py index 2cd383cdb77..4b3684ba864 100644 --- a/tests/test_websocket_parser.py +++ b/tests/test_websocket_parser.py @@ -1,8 +1,10 @@ import asyncio +import gc import pickle import random import struct import sys +import weakref import zlib from unittest import mock @@ -19,8 +21,10 @@ ) from aiohttp._websocket.models import WS_DEFLATE_TRAILING from aiohttp._websocket.reader import WebSocketDataQueue +from aiohttp._websocket.reader_py import MSG_SIZE_OVERHEAD from aiohttp.base_protocol import BaseProtocol from aiohttp.compression_utils import ZLibBackend, ZLibBackendWrapper +from aiohttp.helpers import DEFAULT_CHUNK_SIZE from aiohttp.http import HttpParser, WebSocketError, WSCloseCode, WSMsgType from aiohttp.http_websocket import ( WebSocketReader, @@ -30,6 +34,7 @@ WSMessagePong, WSMessageText, ) +from aiohttp.streams import EofStream class PatchableWebSocketReader(WebSocketReader): @@ -1009,3 +1014,356 @@ async def test_incomplete_frame_not_paused_for_normal_reads( for _ in range(payload_len // 4096): parser.feed_data(b"x" * 4096) assert protocol._reading_paused is False + + +def _compressed_burst(payload: bytes, count: int) -> bytes: + """`count` complete, independently deflated BINARY messages in one read.""" + return ( + build_frame(payload, WSMsgType.BINARY, ZLibBackend=ZLibBackend, mask=True) + * count + ) + + +async def test_empty_messages_apply_backpressure(protocol: BaseProtocol) -> None: + # Zero-length messages are not free: each one is a queued object. They used + # to add nothing to the queue's byte counter, so a peer could stream empty + # frames forever without the transport ever being paused. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 4 * 1024 * 1024, compress=False, decode_text=True) + + sent = 10000 + parser.feed_data(build_frame(b"", WSMsgType.TEXT, mask=True) * sent) + + assert protocol._reading_paused is True + assert len(out._buffer) < sent + # Overshoot is capped at the one message that crossed the mark. + assert out._size <= out._limit + MSG_SIZE_OVERHEAD + + +async def test_compressed_burst_stops_at_high_water(protocol: BaseProtocol) -> None: + # permessage-deflate reaches ~1000:1, so a single read can carry dozens of + # complete frames that each inflate to max_msg_size. pause_reading() only + # stops the transport from delivering more data; the parser has to stop + # too, or the whole read is inflated into the queue in one go. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=True, decode_text=False) + + payload = b"\0" * (512 * 1024) + burst = _compressed_burst(payload, 32) + assert len(burst) < 2**16 # the whole burst is one plausible socket read + + parser.feed_data(burst) + + assert protocol._reading_paused is True + # Overshoot is capped at the one message that crossed the mark, not the + # 16 MiB the peer asked us to inflate. + assert len(out._buffer) == 1 + assert out._buffer[0].size == len(payload) + + +async def test_read_arriving_over_high_water_inflates_nothing( + protocol: BaseProtocol, +) -> None: + # pause_reading() does not retract a read already in flight (the proactor + # transport completes its outstanding overlapped read, and a transport + # without flow control ignores the pause entirely). Such a read must not + # buy the peer one more inflated message before the parser stops again. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=True, decode_text=False) + + payload = b"\0" * (512 * 1024) + parser.feed_data(_compressed_burst(payload, 2)) + assert len(out._buffer) == 1 + + # Lands while the queue is still over the mark and nothing has drained. + parser.feed_data(_compressed_burst(payload, 2)) + assert len(out._buffer) == 1 + + # All four are still delivered once the application catches up. + for _ in range(4): + msg = await asyncio.wait_for(out.read(), 5) + assert msg.data == payload + + +async def test_backpressure_does_not_drop_stashed_frames( + protocol: BaseProtocol, +) -> None: + # The frames the parser stopped short of must still be delivered as the + # application drains the queue, with no further socket data to drive it. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=True, decode_text=False) + + payload = b"\0" * (512 * 1024) + parser.feed_data(_compressed_burst(payload, 32)) + + for _ in range(32): + msg = await asyncio.wait_for(out.read(), 5) + assert msg.data == payload + + assert not out._buffer + # feed_data() adds MSG_SIZE_OVERHEAD and _read_from_buffer() subtracts it; + # _size is an unsigned int under Cython, so if those ever drift the + # subtraction wraps to ~4G rather than going negative. _size < _limit is + # then permanently false and the connection wedges silently: the transport + # is never resumed and the stalled parser is never driven again. + assert out._size == 0 + assert protocol._reading_paused is False + + +async def test_drain_resumes_parsing_in_batches(protocol: BaseProtocol) -> None: + # The parser is re-driven only at the low-water mark. Each resume + # re-slices the whole unparsed tail, so re-driving on every pop would + # make draining a burst of tiny messages quadratic in its size. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=False, decode_text=True) + + sent = 8000 + parser.feed_data(build_frame(b"", WSMsgType.TEXT, mask=True) * sent) + stalled = len(out._buffer) + assert stalled < sent + + # Above the low-water mark a pop must not re-drive the parser. + await asyncio.wait_for(out.read(), 5) + assert len(out._buffer) == stalled - 1 + + # The pop that reaches the low-water mark refills the queue in one batch. + low_water_msgs = (out._limit // 2) // MSG_SIZE_OVERHEAD + for _ in range(stalled - 1 - low_water_msgs): + await asyncio.wait_for(out.read(), 5) + assert len(out._buffer) > low_water_msgs + assert out._size > out._limit + + +async def test_transport_stays_paused_while_stash_remains( + protocol: BaseProtocol, +) -> None: + # Resuming the transport before the stash is exhausted would admit a + # fresh socket read into the tail for every couple of messages drained, + # moving the memory bound from the queue into the tail. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=False, decode_text=True) + + sent = 8000 + parser.feed_data(build_frame(b"", WSMsgType.TEXT, mask=True) * sent) + assert protocol._reading_paused is True + + count = 0 + while out._stalled_reader is not None: + await asyncio.wait_for(out.read(), 5) + count += 1 + if out._stalled_reader is not None: + assert protocol._reading_paused is True, f"resumed after {count} pops" + assert count < sent + + # The pop that exhausted the stash resumed the transport. + assert protocol._reading_paused is False + + +async def test_read_ending_on_frame_boundary_does_not_stall( + protocol: BaseProtocol, +) -> None: + # A read that crosses the high-water mark but ends exactly on a frame + # boundary leaves nothing stashed, so the parser must not arm the stall; + # ordinary backpressure then resumes at the queue limit, not the low-water + # mark reserved for draining a stash. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 1024, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=False, decode_text=True) + + # 17 empty frames * MSG_SIZE_OVERHEAD crosses out._limit (2048); the whole + # burst is complete frames, so parsing ends with an empty tail. + parser.feed_data(build_frame(b"", WSMsgType.TEXT, mask=True) * 17) + assert out._size > out._limit + assert protocol._reading_paused is True + assert out._stalled_reader is None, "empty-tail read must not arm the stall" + + # Draining back under the limit resumes well above _limit // 2, so the + # low-water mark reserved for stashes is not in play. + await asyncio.wait_for(out.read(), 5) + await asyncio.wait_for(out.read(), 5) + assert out._limit // 2 < out._size < out._limit + assert protocol._reading_paused is False + + +async def test_backpressure_stash_survives_eof(protocol: BaseProtocol) -> None: + # A peer that bursts and then disconnects must not lose the tail: at EOF + # there is no resume to ride on, so draining has to keep driving the parser. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=True, decode_text=False) + + payload = b"\0" * (512 * 1024) + parser.feed_data(_compressed_burst(payload, 8)) + parser.feed_eof() + + for _ in range(8): + msg = await asyncio.wait_for(out.read(), 5) + assert msg.data == payload + + with pytest.raises(EofStream): + await asyncio.wait_for(out.read(), 5) + + +async def test_stalled_reader_reference_released_after_drain( + protocol: BaseProtocol, +) -> None: + # The queue holds the reader back only while parsing is stalled. Keeping it + # any longer would leave reader <-> queue as a cycle outliving the + # connection, reclaimable only by the collector rather than by refcounting. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=True, decode_text=False) + + payload = b"\0" * (512 * 1024) + parser.feed_data(_compressed_burst(payload, 4)) + stalled = out._stalled_reader + assert stalled is not None and stalled() is parser + + for _ in range(4): + await asyncio.wait_for(out.read(), 5) + + assert out._stalled_reader is None + + +async def test_set_exception_still_delivers_stashed_frames( + protocol: BaseProtocol, +) -> None: + # set_exception() is also the transport-died hook (WebSocketResponse._cancel). + # _read_from_buffer() only raises once the buffer is empty, so complete + # frames the parser stopped short of must still be delivered first, exactly + # as they were before the parser learned to stall. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, DEFAULT_CHUNK_SIZE, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=False, decode_text=True) + + sent = 8000 + parser.feed_data(build_frame(b"", WSMsgType.TEXT, mask=True) * sent) + stalled = len(out._buffer) + assert stalled < sent + + out.set_exception(ConnectionResetError()) + + for _ in range(sent): + await asyncio.wait_for(out.read(), 5) + with pytest.raises(ConnectionResetError): + await out.read() + + +async def test_parse_error_abandons_the_stash(protocol: BaseProtocol) -> None: + # A parse error poisons the reader, so re-driving it from a drain is a + # no-op and the error surfaces once the buffer is empty. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=True, decode_text=False) + + payload = b"\0" * (512 * 1024) + # The reserved opcode sits behind the stash, so it is only reached on a + # drain-driven resume. + parser.feed_data(_compressed_burst(payload, 3) + build_frame(b"", 0x3, mask=True)) + stalled = out._stalled_reader + assert stalled is not None and stalled() is parser + + for _ in range(3): + await asyncio.wait_for(out.read(), 5) + with pytest.raises(WebSocketError): + await asyncio.wait_for(out.read(), 5) + + +async def test_queue_does_not_keep_stalled_reader_alive( + protocol: BaseProtocol, +) -> None: + # The one case feed_eof() cannot clean up: the connection dies while + # parsing is stalled, then the application drops the response without + # draining. The queue's link back must be weak or reader <-> queue + # survives as a cycle that only the collector can reclaim. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=True, decode_text=False) + + parser.feed_data(_compressed_burst(b"\0" * (512 * 1024), 4)) + parser.feed_eof() + # `parser` deliberately never appears in an assert: pytest's assertion + # rewriting keeps temporaries for the failure message, and a lingering + # reference to it would mask the very thing under test. + stalled = out._stalled_reader + assert stalled is not None + + ref = weakref.ref(parser) + del parser, stalled + # `out` deliberately stays alive across the collection: a strong link back + # would keep the reader reachable from a live root, so this proves the link + # is weak rather than merely testing reclamation timing. Collect instead of + # disabling the gc -- PyPy has no refcounting to reclaim eagerly, and can + # need more than one pass to clear the weakref. + for _ in range(3): + gc.collect() + + assert ref() is None, "queue kept the stalled reader alive" + + +def _queue_with_collected_reader( + protocol: BaseProtocol, loop: asyncio.AbstractEventLoop, *, eof: bool +) -> WebSocketDataQueue: + """A queue whose reader stalled on a burst and was then collected.""" + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=True, decode_text=False) + parser.feed_data(_compressed_burst(b"\0" * (512 * 1024), 4)) + if eof: + parser.feed_eof() + del parser + for _ in range(3): + gc.collect() + return out + + +async def test_collected_stalled_reader_surfaces_contract_error( + protocol: BaseProtocol, caplog: pytest.LogCaptureFixture +) -> None: + # Every strong reference to the reader was dropped while it was stalled, + # so the stash is gone. Whatever was already queued is still delivered, + # then the contract violation surfaces instead of a silent short stream. + # It is also logged, for callers that stop reading before the exception. + loop = asyncio.get_running_loop() + out = _queue_with_collected_reader(protocol, loop, eof=True) + + for _ in range(len(out._buffer)): + await asyncio.wait_for(out.read(), 5) + assert "must hold a strong reference" in caplog.text + with pytest.raises(RuntimeError, match="must hold a strong reference"): + await out.read() + + +async def test_collected_stalled_reader_preserves_existing_exception( + protocol: BaseProtocol, +) -> None: + # set_exception() is also the transport-died hook; losing the reader on + # top of that must not replace the real cause with the contract error. + loop = asyncio.get_running_loop() + out = _queue_with_collected_reader(protocol, loop, eof=False) + out.set_exception(ConnectionResetError()) + + for _ in range(len(out._buffer)): + await asyncio.wait_for(out.read(), 5) + with pytest.raises(ConnectionResetError): + await out.read() + + +async def test_burst_under_high_water_is_parsed_in_one_read( + protocol: BaseProtocol, +) -> None: + # Ordinary pipelined traffic that fits under the mark must not be stalled + # or split across reads by the backpressure check. + loop = asyncio.get_running_loop() + out = WebSocketDataQueue(protocol, 2**16, loop=loop) + parser = WebSocketReader(out, 1024 * 1024, compress=False, decode_text=True) + + parser.feed_data(build_frame(b"hello", WSMsgType.TEXT, mask=True) * 50) + + assert len(out._buffer) == 50 + assert protocol._reading_paused is False