Discard non-close messages in close timeout window

CHANGES/9506.bugfix.rst

@@ -0,0 +1,2 @@
+Fixed ``WebSocketResponse.close()`` to discard non-close messages within its timeout window after sending close.
+-- by :user:`lenard-mosys`

CONTRIBUTORS.txt

@@ -225,6 +225,7 @@ Lubomir Gelo
 Ludovic Gasc
 Luis Pedrosa
 Lukasz Marcin Dobrzanski
+Lénárd Szolnoki
 Makc Belousow
 Manuel Miranda
 Marat Sharafutdinov

aiohttp/web_ws.py

@@ -495,7 +495,11 @@ async def close(
 
         try:
             async with async_timeout.timeout(self._timeout):
-                msg = await reader.read()
+                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
@@ -504,14 +508,6 @@ async def close(
             self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
             return True
 
-        if msg.type is WSMsgType.CLOSE:
-            self._set_code_close_transport(msg.data)
-            return True
-
-        self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
-        self._exception = asyncio.TimeoutError()
-        return True
-
     def _set_closing(self, code: int) -> None:
         """Set the close code and mark the connection as closing."""
         self._closing = True

tests/test_web_websocket_functional.py

@@ -1243,7 +1243,11 @@ async def test_abnormal_closure_when_server_does_not_receive(
     """Test abnormal closure when the server closes and a message is pending."""
 
     async def handler(request: web.Request) -> web.WebSocketResponse:
-        ws = web.WebSocketResponse()
+        # Setting close timeout to 0, otherwise the server waits for a
+        # close response for 10 seconds by default.
+        # This would make the client's autoclose in resp.receive() to succeed,
+        # closing the connection cleanly from both sides.
+        ws = web.WebSocketResponse(timeout=0)
         await ws.prepare(request)
         await ws.close()
         return ws
@@ -1257,3 +1261,30 @@ async def handler(request: web.Request) -> web.WebSocketResponse:
     msg = await resp.receive()
     assert msg.type is aiohttp.WSMsgType.CLOSE
     assert resp.close_code == WSCloseCode.ABNORMAL_CLOSURE
+
+
+async def test_abnormal_closure_when_client_does_not_close(
+    aiohttp_client: AiohttpClient,
+) -> None:
+    """Test abnormal closure when the server closes and the client doesn't respond."""
+    close_code: Optional[WSCloseCode] = None
+
+    async def handler(request: web.Request) -> web.WebSocketResponse:
+        # Setting a short close timeout
+        ws = web.WebSocketResponse(timeout=0.1)
+        await ws.prepare(request)
+        await ws.close()
+
+        nonlocal close_code
+        if ws.close_code is not None:
+            close_code = WSCloseCode(ws.close_code)
+
+        return ws
+
+    app = web.Application()
+    app.router.add_route("GET", "/", handler)
+    client = await aiohttp_client(app)
+    async with client.ws_connect("/", autoclose=False):
+        await asyncio.sleep(0.2)
+    await client.server.close()
+    assert close_code == WSCloseCode.ABNORMAL_CLOSURE