From c207cadb257e0959b192ffca0548eca8c71eb96a Mon Sep 17 00:00:00 2001 From: Sam Bull Date: Fri, 21 Aug 2026 03:27:31 +0100 Subject: [PATCH] Fix limit on message tail (#13501) --- CHANGES/13501.bugfix.rst | 1 + aiohttp/web_protocol.py | 46 +++++++- tests/test_web_functional.py | 218 +++++++++++++++++++++++++++++++++++ tests/test_web_protocol.py | 82 +++++++++++++ 4 files changed, 341 insertions(+), 6 deletions(-) create mode 100644 CHANGES/13501.bugfix.rst diff --git a/CHANGES/13501.bugfix.rst b/CHANGES/13501.bugfix.rst new file mode 100644 index 00000000000..4c45225e9c2 --- /dev/null +++ b/CHANGES/13501.bugfix.rst @@ -0,0 +1 @@ +Fixed a limit on message tail after an upgrade request -- by :user:`Dreamsorcerer`. diff --git a/aiohttp/web_protocol.py b/aiohttp/web_protocol.py index 33fcd7511ae..07da7326ddb 100644 --- a/aiohttp/web_protocol.py +++ b/aiohttp/web_protocol.py @@ -183,6 +183,7 @@ class RequestHandler(BaseProtocol, Generic[_Request]): "_msg_queue_resume_size", "_msg_queue_paused", "_message_tail", + "_read_bufsize", "_handler_waiter", "_waiter", "_task_handler", @@ -224,6 +225,7 @@ def __init__( # Low-water mark: resume reading once the queue drains to half the limit # so we refill in batches instead of churning pause/resume per request. self._msg_queue_resume_size = MAX_MSG_QUEUE_SIZE // 2 + self._read_bufsize = read_bufsize # Set before super().__init__ so _reading_paused_for_msg_queue() is safe # if BaseProtocol ever triggers a resume during init. self._msg_queue_paused = False @@ -454,6 +456,9 @@ def set_parser( self._payload_parser.feed_data(self._message_tail) self._message_tail = b"" + if self._msg_queue_paused: + self._resume_msg_queue_reading() + def eof_received(self) -> None: pass @@ -497,6 +502,11 @@ def data_received(self, data: bytes) -> None: # no parser, just store elif self._payload_parser is None and self._upgraded and data: self._message_tail += data + if ( + not self._msg_queue_paused + and len(self._message_tail) >= self._read_bufsize + ): + self._pause_msg_queue_reading() # feed payload elif data: @@ -520,6 +530,10 @@ def _pause_msg_queue_reading(self) -> None: pass def _resume_msg_queue_reading(self) -> None: + # Tested empty-first so a read_bufsize of 0 cannot wedge the connection. + if self._message_tail and len(self._message_tail) >= self._read_bufsize: + return + if not self._upgraded: # Reparse buffered pipelined requests while still marked paused so # a refill past the limit does not re-pause an already-paused @@ -822,19 +836,39 @@ async def finish_response( self._parser.set_upgraded(False) self._upgraded = False if self._message_tail: - messages, upgraded, tail = self._parser.feed_data(self._message_tail) + messages: Sequence[_MsgType] + try: + messages, upgraded, tail = self._parser.feed_data( + self._message_tail + ) + except HttpProcessingError as parse_exc: + # Garbage (or an oversized request line) buffered behind the + # upgrade: answer 400 instead of letting the error escape + # and lose this response, like data_received() does. + messages = [ + ( + _ErrInfo( + status=400, + exc=parse_exc, + message=parse_exc.message, + ), + EMPTY_PAYLOAD, + ) + ] + upgraded = False + tail = b"" # A further upgrade request in the tail buffers its own remainder. self._upgraded = upgraded self._message_tail = tail for msg, payload in messages: self._request_count += 1 self._messages.append((msg, payload)) - # Pause the transport, like in data_received(). - if ( - not self._msg_queue_paused - and len(self._messages) >= self._max_msg_queue_size - ): + if len(self._messages) >= self._max_msg_queue_size: + # Pause the transport, like in data_received(). self._pause_msg_queue_reading() + elif self._msg_queue_paused: + # Resume reading now the tail has been parsed. + self._resume_msg_queue_reading() # This shouldn't be possible. If a future refactor results in this # failing, then the code may need to be updated to set the waiter. assert self._waiter is None diff --git a/tests/test_web_functional.py b/tests/test_web_functional.py index 0f5f7d11109..c834c388a06 100644 --- a/tests/test_web_functional.py +++ b/tests/test_web_functional.py @@ -28,10 +28,13 @@ multipart, web, ) +from aiohttp._websocket.writer import WebSocketWriter from aiohttp.abc import AbstractResolver, ResolveResult +from aiohttp.base_protocol import BaseProtocol from aiohttp.compression_utils import ZLibBackend, ZLibCompressObjProtocol from aiohttp.hdrs import CONTENT_LENGTH, CONTENT_TYPE, TRANSFER_ENCODING from aiohttp.helpers import DEFAULT_CHUNK_SIZE, HeadersDictProxy +from aiohttp.http import WSMsgType from aiohttp.streams import StreamReader from aiohttp.typedefs import Handler, Middleware from aiohttp.web_protocol import MAX_MSG_QUEUE_SIZE, RequestHandler @@ -1901,6 +1904,221 @@ def raw_get(path: str) -> bytes: assert len(set(handled)) == len(handled) +async def test_upgrade_tail_is_byte_limited( + aiohttp_server: AiohttpServer, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Bytes buffered behind an in-flight upgrade must not grow unbounded. + + A request the parser flagged as an upgrade has no payload, so everything + arriving while its handler runs is buffered whole until the handler either + prepares a websocket or answers normally. Only ``read_bufsize`` of it may be + held before reading is paused, regardless of how much is sent. + """ + # Above the default so a ceiling that ignores read_bufsize fails the lower + # bound, and well below what is sent so no cap at all fails the upper one. + read_bufsize = 1024 * 1024 + handler_started = asyncio.Event() + release_handler = asyncio.Event() + reading_paused = asyncio.Event() + max_tail = 0 + data_received = RequestHandler.data_received + + def observe_data_received(self: RequestHandler[web.Request], data: bytes) -> None: + nonlocal max_tail + data_received(self, data) + if self._message_tail: + max_tail = max(max_tail, len(self._message_tail)) + if self._msg_queue_paused: + reading_paused.set() + + monkeypatch.setattr(RequestHandler, "data_received", observe_data_received) + + async def upgrade_handler(request: web.Request) -> web.Response: + handler_started.set() + await release_handler.wait() + return web.Response(text="declined") + + app = web.Application() + app.router.add_get("/upgrade", upgrade_handler) + server = await aiohttp_server(app, read_bufsize=read_bufsize) + + chunk = b"A" * (64 * 1024) + chunks = (4 * 1024 * 1024) // len(chunk) + + reader, writer = await asyncio.open_connection(server.host, server.port) + try: + writer.write( + b"GET /upgrade HTTP/1.1\r\nHost: localhost\r\n" + b"Connection: Upgrade\r\nUpgrade: websocket\r\n\r\n" + ) + await writer.drain() + await asyncio.wait_for(handler_started.wait(), 1) + + async def send_until_paused() -> None: + for _ in range(chunks): # pragma: no branch + if reading_paused.is_set(): + break + writer.write(chunk) + await writer.drain() + + sender = asyncio.create_task(send_until_paused()) + try: + # Only elapses if nothing caps the buffer, in which case the + # assertions below report what was actually buffered. + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(reading_paused.wait(), 5) + finally: + sender.cancel() + with suppress(asyncio.CancelledError): + await sender + finally: + release_handler.set() + writer.close() + with suppress(ConnectionResetError, BrokenPipeError): + await writer.wait_closed() + + assert reading_paused.is_set(), f"reading never paused, buffered {max_tail} bytes" + # pause_reading() only takes effect after the read in flight, and asyncio + # reads at most DEFAULT_CHUNK_SIZE per call, so one extra chunk may land. + assert read_bufsize <= max_tail < read_bufsize + DEFAULT_CHUNK_SIZE + + +async def test_upgrade_tail_resumes_reading_after_websocket_prepare( + aiohttp_server: AiohttpServer, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A websocket paused while its tail filled up must keep reading. + + A client may pipeline frames straight after the handshake request, filling + the tail buffer and pausing reading before the handler prepares. Once the + websocket owns the connection the tail is handed over, so reading has to + resume or whatever the client sent meanwhile is never read. + """ + handshake_seen = asyncio.Event() + reading_paused = asyncio.Event() + release_handler = asyncio.Event() + last_received = asyncio.Event() + received: list[str] = [] + data_received = RequestHandler.data_received + + def observe_data_received(self: RequestHandler[web.Request], data: bytes) -> None: + data_received(self, data) + if self._msg_queue_paused: + reading_paused.set() + + monkeypatch.setattr(RequestHandler, "data_received", observe_data_received) + + async def ws_handler(request: web.Request) -> web.WebSocketResponse: + handshake_seen.set() + await release_handler.wait() + ws = web.WebSocketResponse() + await ws.prepare(request) + async for msg in ws: # pragma: no branch + assert isinstance(msg.data, str) + received.append(msg.data) + if msg.data == "last": + last_received.set() + break + return ws + + app = web.Application() + app.router.add_get("/ws", ws_handler) + read_bufsize = 64 * 1024 + server = await aiohttp_server(app, read_bufsize=read_bufsize) + + # Enough pipelined frames to fill the tail, so reading must pause. + frame_payload = "B" * (32 * 1024) + frames = (read_bufsize // len(frame_payload)) + 2 + + reader, writer = await asyncio.open_connection(server.host, server.port) + try: + writer.write( + b"GET /ws HTTP/1.1\r\nHost: localhost\r\n" + b"Connection: Upgrade\r\nUpgrade: websocket\r\n" + b"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n" + b"Sec-WebSocket-Version: 13\r\n\r\n" + ) + await writer.drain() + await asyncio.wait_for(handshake_seen.wait(), 1) + + # Encode client frames off to the side, as tests/test_websocket_writer.py + # does, and put the bytes on the wire ourselves. + encoded = bytearray() + frame_transport = mock.create_autospec( + asyncio.Transport, spec_set=True, instance=True + ) + # Defaults to a truthy Mock, which WebSocketWriter reads as closing. + frame_transport.is_closing.return_value = False + frame_transport.write.side_effect = encoded.extend + ws_writer = WebSocketWriter( + mock.create_autospec(BaseProtocol, spec_set=True, instance=True), + frame_transport, + use_mask=True, + ) + + for _ in range(frames): + await ws_writer.send_frame(frame_payload.encode(), WSMsgType.TEXT) + writer.write(encoded) + await writer.drain() + await asyncio.wait_for(reading_paused.wait(), 5) + + # Sent while reading is paused, so this frame is only ever read if the + # handover to the websocket resumes the transport. + encoded.clear() + await ws_writer.send_frame(b"last", WSMsgType.TEXT) + writer.write(encoded) + await writer.drain() + release_handler.set() + + await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), 5) + await asyncio.wait_for(last_received.wait(), 5) + finally: + release_handler.set() + writer.close() + with suppress(ConnectionResetError, BrokenPipeError): + await writer.wait_closed() + + assert received == [frame_payload] * frames + ["last"] + + +async def test_bad_pipelined_data_behind_declined_upgrade_answers_400( + aiohttp_server: AiohttpServer, +) -> None: + """Junk buffered behind an upgrade must not swallow the upgrade response. + + The tail is only parsed once the handler answers, so a parse error there + has to be reported as a 400 like it would be from data_received(); letting + it escape loses the response that was about to be written. + """ + + async def upgrade_handler(request: web.Request) -> web.Response: + return web.Response(text="declined") + + app = web.Application() + app.router.add_get("/upgrade", upgrade_handler) + server = await aiohttp_server(app) + + reader, writer = await asyncio.open_connection(server.host, server.port) + try: + writer.write( + b"GET /upgrade HTTP/1.1\r\nHost: localhost\r\n" + b"Connection: Upgrade\r\nUpgrade: websocket\r\n\r\n" + b"\x00" * 64 + ) + await writer.drain() + response = await asyncio.wait_for(reader.read(), 5) + finally: + writer.close() + with suppress(ConnectionResetError, BrokenPipeError): + await writer.wait_closed() + + assert b"declined" in response, response + # A 400 only follows if the parse error was reported instead of escaping, + # and "declined" only arrives if it did not abort this response first. + assert b" 400 " in response, response + + async def test_declined_websocket_upgrade_reads_body( aiohttp_server: AiohttpServer, ) -> None: diff --git a/tests/test_web_protocol.py b/tests/test_web_protocol.py index 8968dea78b5..a88dedba3b9 100644 --- a/tests/test_web_protocol.py +++ b/tests/test_web_protocol.py @@ -99,6 +99,88 @@ def test_resume_msg_queue_reading_without_transport( assert handler._msg_queue_paused is False +def test_resume_msg_queue_reading_stays_paused_for_full_tail( + event_loop: asyncio.AbstractEventLoop, + dummy_manager: Server[BaseRequest], +) -> None: + """Resume is refused while an in-flight upgrade holds read_bufsize of tail. + + The tail is drained by set_parser()/finish_response(), so resuming before + then would let the buffer grow past its ceiling one read at a time. + """ + handler = RequestHandler(dummy_manager, loop=event_loop, read_bufsize=1024) + transport = mock.create_autospec(asyncio.Transport, spec_set=True, instance=True) + handler.transport = transport + handler._upgraded = True + handler._msg_queue_paused = True + handler._message_tail = b"x" * 1024 + + handler._resume_msg_queue_reading() + + assert handler._msg_queue_paused is True + transport.resume_reading.assert_not_called() + + +def test_resume_msg_queue_reading_with_room_left_in_tail( + event_loop: asyncio.AbstractEventLoop, + dummy_manager: Server[BaseRequest], +) -> None: + """A tail under read_bufsize does not hold the transport paused.""" + handler = RequestHandler(dummy_manager, loop=event_loop, read_bufsize=1024) + transport = mock.create_autospec(asyncio.Transport, spec_set=True, instance=True) + handler.transport = transport + handler._upgraded = True + handler._msg_queue_paused = True + handler._message_tail = b"x" * 1023 + + handler._resume_msg_queue_reading() + + assert handler._msg_queue_paused is False + transport.resume_reading.assert_called_once_with() + + +def test_resume_msg_queue_reading_with_zero_read_bufsize( + event_loop: asyncio.AbstractEventLoop, + dummy_manager: Server[BaseRequest], +) -> None: + """An empty tail resumes even when read_bufsize leaves it no room. + + Guards the degenerate ``read_bufsize=0`` case: comparing sizes alone would + match an empty tail and wedge the connection paused for good. + """ + handler = RequestHandler(dummy_manager, loop=event_loop, read_bufsize=0) + transport = mock.create_autospec(asyncio.Transport, spec_set=True, instance=True) + handler.transport = transport + handler._upgraded = True + handler._msg_queue_paused = True + + handler._resume_msg_queue_reading() + + assert handler._msg_queue_paused is False + transport.resume_reading.assert_called_once_with() + + +def test_set_parser_resumes_reading_paused_for_tail( + event_loop: asyncio.AbstractEventLoop, + dummy_manager: Server[BaseRequest], + dummy_reader: tuple[WebSocketReader, mock.Mock], +) -> None: + """Handing a full tail to the upgraded protocol resumes reading.""" + handler = RequestHandler(dummy_manager, loop=event_loop, read_bufsize=1024) + transport = mock.create_autospec(asyncio.Transport, spec_set=True, instance=True) + handler.transport = transport + handler._upgraded = True + handler._msg_queue_paused = True + handler._message_tail = b"x" * 1024 + + handler.set_parser(dummy_reader[0]) + + dummy_reader[1].feed_data.assert_called_once_with(b"x" * 1024) + assert handler._message_tail == b"" + assert handler._msg_queue_paused is False + transport.resume_reading.assert_called_once_with() + + def test_resume_reading_stays_paused_for_msg_queue( event_loop: asyncio.AbstractEventLoop, dummy_manager: Server[BaseRequest],