diff --git a/rolo/asgi.py b/rolo/asgi.py index af15c3f..827da58 100644 --- a/rolo/asgi.py +++ b/rolo/asgi.py @@ -345,6 +345,10 @@ class ASGIWebSocketAdapter(rolows.WebSocketAdapter): _receive: "ASGIReceiveCallable" _send: "ASGISendCallable" + close_timeout: float = 5 + """Seconds to wait for the client to echo the server's close frame, before the ASGI app returns and leaves + the TCP connection to the server.""" + def __init__( self, scope: "WebsocketScope", @@ -357,12 +361,35 @@ def __init__( self._send = send self._loop = loop self._last_received_at = time.monotonic() + self._close_sent = False + self._disconnected = False async def asgi_send_async(self, event: "_WebsocketResponse"): await self._send(event) async def asgi_receive_async(self) -> "_WebsocketRequest": - return await self._receive() + event = await self._receive() + if event["type"] == "websocket.disconnect": + self._disconnected = True + return event + + async def wait_closed(self) -> None: + """ + Waits for the client to echo the close frame the server sent, so the close handshake completes before the + ASGI app returns. An ASGI server can terminate the TCP connection as soon as the app returns, and the + client's close frame then arriving on a closed socket makes the kernel reset the connection (RFC 6455 + section 5.5.1). Frames the client sends in the meantime are discarded. + """ + if not self._close_sent or self._disconnected: + return + try: + await asyncio.wait_for(self._receive_until_disconnect(), self.close_timeout) + except asyncio.TimeoutError: + pass + + async def _receive_until_disconnect(self) -> None: + while not self._disconnected: + await self.asgi_receive_async() def asgi_send(self, event: "_WebsocketResponse", timeout: float = None) -> None: """ @@ -496,6 +523,7 @@ def accept( ) def close(self, code: int = 1001, reason: str = None, timeout: float = None): + self._close_sent = True self.asgi_send( { "type": "websocket.close", @@ -660,6 +688,8 @@ async def handle_websocket( ) raise + await environ["rolo.websocket"].wait_closed() + async def handle_lifespan( self, scope: "HTTPScope", receive: "ASGIReceiveCallable", send: "ASGISendCallable" ): diff --git a/tests/test_asgi.py b/tests/test_asgi.py index 52ef44d..ab254d7 100644 --- a/tests/test_asgi.py +++ b/tests/test_asgi.py @@ -175,10 +175,14 @@ class IterableResponse: def __init__(self, data: list[bytes]): self.data = data self.closed = False + # holds the server after the last packet, so it can't close the iterable before + # the client has checked that it's still open + self.release = threading.Event() def __iter__(self): for packet in self.data: yield packet + assert self.release.wait(timeout=5) def close(self): # should be called through the werkzeug layers @@ -198,6 +202,7 @@ def app(request: Request) -> Response: assert next(gen) == b"foo" assert next(gen) == b"bar" assert not iterable.closed + iterable.release.set() with pytest.raises(StopIteration): next(gen) diff --git a/tests/websocket/test_websockets.py b/tests/websocket/test_websockets.py index d7a94a5..6ddf4a8 100644 --- a/tests/websocket/test_websockets.py +++ b/tests/websocket/test_websockets.py @@ -221,7 +221,11 @@ def app(request: WebSocketRequest): frame = client.recv_frame() assert frame.opcode == websocket.ABNF.OPCODE_CLOSE + # a server that terminates the TCP connection without waiting for the echo makes the kernel reset the + # connection once the echo arrives, and the reset then fails the read below + time.sleep(0.1) client.send_close(status=struct.unpack("!H", frame.data[:2])[0]) + time.sleep(0.1) client.sock.settimeout(5) assert client.sock.recv(1) == b"", "expected the server to terminate the TCP connection" @@ -273,6 +277,7 @@ def echo_headers(request: WebSocketRequest): assert headers["Connection"] == "Upgrade" assert headers["Authorization"] == "Basic let-me-in" assert headers["CasedHeader"] == "hello" + client.close() def test_websocket_reject(serve_websocket_listener): @@ -318,6 +323,7 @@ def echo_headers(request: WebSocketRequest): assert received.get(timeout=5) == "textbar" assert received.get(timeout=5) == b"bar" + client.close() def test_last_received_at_message(serve_websocket_listener): @@ -436,6 +442,7 @@ def echo_headers(request: WebSocketRequest): e = match.get(timeout=5) assert e.match("Cannot send data type over websocket") + client.close() def test_router_integration(serve_websocket_listener): @@ -457,6 +464,7 @@ def _handler(request: WebSocketRequest, request_args: dict): assert client.recv() == "foo" assert client.recv() == "id=bar" assert "CasedHeader" in json.loads(client.recv()) + client.close() def test_send_many_messages(serve_websocket_listener):