Fix limit on message tail (aio-libs#13501)

Author: Dreamsorcerer

CHANGES/13501.bugfix.rst

@@ -0,0 +1 @@
+Fixed a limit on message tail after an upgrade request -- by :user:`Dreamsorcerer`.

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

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: