fix twisted websocket close handshake

rolo/serving/twisted.py

@@ -314,6 +314,7 @@ def __init__(self, request: Request):
         self.request = request
         self.wsproto = WSConnection(ConnectionType.SERVER)
         self.eventQueue = Queue()
+        self.upgraded = False
 
     @property
     def closed(self):
@@ -339,6 +340,9 @@ def dataReceived(self, data: bytes) -> None:
                 continue
             # TODO: filter other event types that are not expected by WebSocketAdapter
             if isinstance(event, events.CloseConnection):
+                # complete the closing handshake (RFC 6455 section 7): echo the close frame,
+                # then terminate the TCP connection, which is the server's job
+                self.wsSend(event.response())
                 self.close()
             self.eventQueue.put_nowait(event)
 
@@ -347,6 +351,8 @@ def wsSend(self, event: events.Event):
         if request.finished:
             return
         data = self.wsproto.send(event)
+        if isinstance(event, events.AcceptConnection):
+            self.upgraded = True
         request.transport.write(data)
 
     def wsReject(
@@ -361,6 +367,8 @@ def wsReject(
         # which is cleaner here, though perhaps inconsistent with the rest of the implementation.
         # TODO: set default twisted headers
         request = self.request
+        if request.finished:
+            return
 
         request.setResponseCode(statusCode)
         for k, v in extraHeaders.to_wsgi_list():
@@ -379,10 +387,16 @@ def wsClose(self, code: int = 1000, reason: t.Optional[str] = None):
             self.close()
 
     def close(self):
-        if not self.request.finished:
-            self.request.finish()
-            # special internal poison pill
-            self.eventQueue.put_nowait(events.CloseConnection(None))
+        if self.request.finished:
+            return
+        if self.upgraded:
+            # the 101 upgrade response was written raw to the transport, so ``Request.finish()``
+            # must not write its own (never started) HTTP response into the websocket stream
+            self.request.startedWriting = 1
+        self.request.finish()
+        self.request.transport.loseConnection()
+        # special internal poison pill
+        self.eventQueue.put_nowait(events.CloseConnection(None))
 
 
 class TwistedWebSocketAdapter(rolows.WebSocketAdapter):

tests/websocket/test_websockets.py

@@ -1,4 +1,5 @@
 import json
+import struct
 import threading
 from queue import Queue
 
@@ -82,6 +83,55 @@ def app(request: WebSocketRequest):
     assert received[1] == "bar"
 
 
+def test_close_handshake_client_initiated(serve_websocket_listener):
+    """When the client sends a close frame, the server has to echo the close frame back to
+    complete the closing handshake, and then terminate the TCP connection (RFC 6455 section 7)."""
+    disconnected = threading.Event()
+
+    @WebSocketRequest.listener
+    def app(request: WebSocketRequest):
+        with request.accept() as ws:
+            with pytest.raises(WebSocketDisconnectedError):
+                ws.receive()
+        disconnected.set()
+
+    server = serve_websocket_listener(app)
+
+    client = websocket.WebSocket()
+    client.connect(server.url.replace("http://", "ws://"))
+    client.send_close(websocket.STATUS_NORMAL)
+
+    frame = client.recv_frame()
+    assert frame.opcode == websocket.ABNF.OPCODE_CLOSE
+    assert struct.unpack("!H", frame.data[:2])[0] == websocket.STATUS_NORMAL
+
+    client.sock.settimeout(5)
+    assert client.sock.recv(1) == b"", "expected the server to terminate the TCP connection"
+    assert disconnected.wait(timeout=3)
+
+
+def test_close_handshake_server_initiated(serve_websocket_listener):
+    """When the server closes the websocket, the client has to receive a proper close frame,
+    followed by the termination of the TCP connection."""
+
+    @WebSocketRequest.listener
+    def app(request: WebSocketRequest):
+        with request.accept() as ws:
+            ws.send("hello")
+
+    server = serve_websocket_listener(app)
+
+    client = websocket.WebSocket()
+    client.connect(server.url.replace("http://", "ws://"))
+    assert client.recv() == "hello"
+
+    frame = client.recv_frame()
+    assert frame.opcode == websocket.ABNF.OPCODE_CLOSE
+
+    client.sock.settimeout(5)
+    assert client.sock.recv(1) == b"", "expected the server to terminate the TCP connection"
+
+
 def test_websocket_headers(serve_websocket_listener):
     @WebSocketRequest.listener
     def echo_headers(request: WebSocketRequest):