Fix pipelining a rejected upgrade (aio-libs#13468)

Author: Dreamsorcerer

aiohttp/web_protocol.py

@@ -807,11 +807,24 @@ async def finish_response(
         prematurely.
         """
         request._finish()
-        if self._parser is not None:
+
+        # Handle feeding the message tail following an upgrade request that
+        # was declined.
+        # The upgrade request is the last request before the parser paused,
+        # so wait for self._messages to be empty.
+        # payload_parser is not None if the upgrade was accepted.
+        if (
+            self._upgraded
+            and not self._messages
+            and self._payload_parser is None
+            and self._parser is not None
+        ):
             self._parser.set_upgraded(False)
             self._upgraded = False
             if self._message_tail:
-                messages, _upgraded, tail = self._parser.feed_data(self._message_tail)
+                messages, upgraded, tail = self._parser.feed_data(self._message_tail)
+                # 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

tests/test_web_websocket_functional.py

@@ -14,6 +14,7 @@
 import aiohttp
 from aiohttp import WSServerHandshakeError, hdrs, web
 from aiohttp.http import WSCloseCode, WSMsgType
+from aiohttp.web_protocol import MAX_MSG_QUEUE_SIZE
 
 
 async def test_websocket_can_prepare(aiohttp_client: AiohttpClient) -> None:
@@ -127,6 +128,160 @@ async def second_handler(request: web.Request) -> web.Response:
         await writer.wait_closed()
 
 
+def _raw_get(path: str) -> bytes:
+    return f"GET {path} HTTP/1.1\r\nHost: localhost\r\n\r\n".encode("ascii")
+
+
+_RAW_UPGRADE = (
+    b"GET /ws HTTP/1.1\r\n"
+    b"Host: localhost\r\n"
+    b"Upgrade: websocket\r\n"
+    b"Connection: Upgrade\r\n"
+    b"Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
+    b"Sec-WebSocket-Version: 13\r\n"
+    b"\r\n"
+)
+
+
+def _masked_text_frame(payload: bytes) -> bytes:
+    """Build a client text frame; frames sent to a server must be masked."""
+    assert len(payload) < 126
+    mask = b"\x37\xfa\x21\x3d"
+    return (
+        b"\x81"
+        + bytes((0x80 | len(payload),))
+        + mask
+        + bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
+    )
+
+
+async def test_websocket_frames_pipelined_behind_request_burst(
+    aiohttp_server: AiohttpServer,
+) -> None:
+    """A websocket upgraded from within a request burst still reads frames.
+
+    The parser stops at the upgrade request and buffers everything after it,
+    which is websocket data once the handshake succeeds. Answering the requests
+    queued ahead of the upgrade must neither consume that buffer as HTTP nor
+    leave the transport paused by the pipeline queue after switching protocols.
+    """
+    # More than the queue holds, so reading pauses and the parser buffers.
+    pipelined_requests = MAX_MSG_QUEUE_SIZE + 8
+    handled: list[str] = []
+
+    async def handler(request: web.Request) -> web.Response:
+        handled.append(request.path)
+        return web.Response()
+
+    async def ws_handler(request: web.Request) -> web.WebSocketResponse:
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        await ws.send_str(await ws.receive_str())
+        return ws
+
+    app = web.Application()
+    app.router.add_get("/ws", ws_handler)
+    app.router.add_get("/{tail:.*}", handler)
+    server = await aiohttp_server(app)
+
+    reader, writer = await asyncio.open_connection(server.host, server.port)
+    try:
+        writer.write(
+            b"".join(_raw_get(f"/r{i}") for i in range(pipelined_requests))
+            + _RAW_UPGRADE
+            + _masked_text_frame(b"frame-ok")
+        )
+        await writer.drain()
+
+        # Without the fix the frame is eaten by the http parser and the paused
+        # transport is never resumed, so nothing is echoed back.
+        await asyncio.wait_for(reader.readuntil(b"\x81\x08frame-ok"), timeout=10)
+    finally:
+        writer.close()
+        with contextlib.suppress(ConnectionResetError, BrokenPipeError):
+            await writer.wait_closed()
+
+    assert handled == [f"/r{i}" for i in range(pipelined_requests)]
+
+
+async def test_pipelined_request_after_declined_upgrade_behind_burst(
+    aiohttp_server: AiohttpServer,
+) -> None:
+    """A declined upgrade replays its tail even when queued behind a request.
+
+    Only the upgrade request's own response settles whether the buffered bytes
+    are websocket data or pipelined HTTP, so an earlier request completing must
+    leave them alone and the declining response still has to replay them.
+    """
+    handled: list[str] = []
+
+    async def handler(request: web.Request) -> web.Response:
+        handled.append(request.path)
+        return web.Response(text=f"{request.path[1:]}-ok")
+
+    async def ws_handler(request: web.Request) -> NoReturn:
+        raise web.HTTPUpgradeRequired()
+
+    app = web.Application()
+    app.router.add_get("/ws", ws_handler)
+    app.router.add_get("/{tail:.*}", handler)
+    server = await aiohttp_server(app)
+
+    reader, writer = await asyncio.open_connection(server.host, server.port)
+    try:
+        writer.write(_raw_get("/first") + _RAW_UPGRADE + _raw_get("/second"))
+        await writer.drain()
+
+        # Without the replay the trailing request stalls until keep-alive expires.
+        data = await asyncio.wait_for(reader.readuntil(b"second-ok"), timeout=10)
+    finally:
+        writer.close()
+        with contextlib.suppress(ConnectionResetError, BrokenPipeError):
+            await writer.wait_closed()
+
+    assert handled == ["/first", "/second"]
+    assert data.count(b"HTTP/1.1 200 OK") == 2, data
+    assert b"426" in data, data
+
+
+async def test_pipelined_request_after_two_declined_upgrades(
+    aiohttp_server: AiohttpServer,
+) -> None:
+    """A second upgrade inside a replayed tail buffers its own remainder again.
+
+    Replaying a declined upgrade's tail can turn up another upgrade request,
+    which puts the parser back into upgraded mode. Losing that leaves the bytes
+    behind the second upgrade buffered with nothing left to replay them.
+    """
+    handled: list[str] = []
+
+    async def handler(request: web.Request) -> web.Response:
+        handled.append(request.path)
+        return web.Response(text="second-ok")
+
+    async def ws_handler(request: web.Request) -> NoReturn:
+        raise web.HTTPUpgradeRequired()
+
+    app = web.Application()
+    app.router.add_get("/ws", ws_handler)
+    app.router.add_get("/{tail:.*}", handler)
+    server = await aiohttp_server(app)
+
+    reader, writer = await asyncio.open_connection(server.host, server.port)
+    try:
+        writer.write(_RAW_UPGRADE + _RAW_UPGRADE + _raw_get("/second"))
+        await writer.drain()
+
+        data = await asyncio.wait_for(reader.readuntil(b"second-ok"), timeout=10)
+    finally:
+        writer.close()
+        with contextlib.suppress(ConnectionResetError, BrokenPipeError):
+            await writer.wait_closed()
+
+    assert handled == ["/second"]
+    assert data.count(b"HTTP/1.1 426 ") == 2, data
+
+
 async def test_handshake_connection_header_substring_not_a_token(
     aiohttp_client: AiohttpClient,
 ) -> None: