Discard non-close messages in close timeout window (#9508)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: J. Nick Koston <[email protected]>
(cherry picked from commit 274c54e)

CHANGES/9506.bugfix.rst

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

CONTRIBUTORS.txt

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

aiohttp/web_ws.py

@@ -473,7 +473,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
@@ -482,14 +486,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: WSCloseCode) -> None:
         """Set the close code and mark the connection as closing."""
         self._closing = True

tests/test_web_websocket_functional.py

@@ -1192,7 +1192,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
@@ -1206,3 +1210,70 @@ 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
+        assert 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
+
+
+async def test_normal_closure_while_client_sends_msg(
+    aiohttp_client: AiohttpClient,
+) -> None:
+    """Test abnormal closure when the server closes and the client doesn't respond."""
+    close_code: Optional[WSCloseCode] = None
+    got_close_code = asyncio.Event()
+
+    async def handler(request: web.Request) -> web.WebSocketResponse:
+        # Setting a short close timeout
+        ws = web.WebSocketResponse(timeout=0.2)
+        await ws.prepare(request)
+        await ws.close()
+
+        nonlocal close_code
+        assert ws.close_code is not None
+        close_code = WSCloseCode(ws.close_code)
+        got_close_code.set()
+
+        return ws
+
+    app = web.Application()
+    app.router.add_route("GET", "/", handler)
+    client = await aiohttp_client(app)
+    async with client.ws_connect("/", autoclose=False) as ws:
+        # send text and close message during server close timeout
+        await asyncio.sleep(0.1)
+        await ws.send_str("Hello")
+        await ws.close()
+    # wait for close code to be received by server
+    await asyncio.wait(
+        [
+            asyncio.create_task(asyncio.sleep(0.5)),
+            asyncio.create_task(got_close_code.wait()),
+        ],
+        return_when=asyncio.FIRST_COMPLETED,
+    )
+    await client.server.close()
+    assert close_code == WSCloseCode.OK