From 719c06a69e032673e26edfca991291db035af7d7 Mon Sep 17 00:00:00 2001
From: John Belmonte <john@neggie.net>
Date: Wed, 3 Feb 2021 21:24:55 +0900
Subject: [PATCH] server: close TCP connection after closing websocket (#150)

Fixes #115 and #147.
---
 Makefile                 |  2 +-
 tests/test_connection.py | 29 +++++++++++++++++++++++++++--
 trio_websocket/_impl.py  |  6 ++++++
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/Makefile b/Makefile
index 64dd996..fc78683 100644
--- a/Makefile
+++ b/Makefile
@@ -11,7 +11,7 @@ docs:
 	$(MAKE) -C docs html
 
 test:
-	$(PYTHON) -m pytest --cov=trio_websocket
+	$(PYTHON) -m pytest --cov=trio_websocket --no-cov-on-fail
 
 lint:
 	$(PYTHON) -m pylint trio_websocket/ tests/ autobahn/ examples/
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 493cc3a..8af07a9 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -39,10 +39,12 @@
 import trustme
 from async_generator import async_generator, yield_
 from trio.testing import memory_stream_pair
+from wsproto.events import CloseConnection
+
 try:
-    from trio.lowlevel import current_task
+    from trio.lowlevel import current_task  # pylint: disable=ungrouped-imports
 except ImportError:
-    from trio.hazmat import current_task
+    from trio.hazmat import current_task  # pylint: disable=ungrouped-imports
 
 from trio_websocket import (
     connect_websocket,
@@ -925,6 +927,29 @@ async def handler(request):
     await trio.sleep(.1)
 
 
+@fail_after(DEFAULT_TEST_MAX_DURATION)
+async def test_server_tcp_closed_on_close_connection_event(nursery):
+    """ensure server closes TCP immediately after receiving CloseConnection"""
+    server_stream_closed = trio.Event()
+
+    async def _close_stream_stub():
+        assert not server_stream_closed.is_set()
+        server_stream_closed.set()
+
+    async def handle_connection(request):
+        ws = await request.accept()
+        ws._close_stream = _close_stream_stub
+        await trio.sleep_forever()
+
+    server = await nursery.start(
+        partial(serve_websocket, handle_connection, HOST, 0, ssl_context=None))
+    client = await connect_websocket(nursery, HOST, server.port,
+                                     RESOURCE, use_ssl=False)
+    # send a CloseConnection event to server but leave client connected
+    await client._send(CloseConnection(code=1000))
+    await server_stream_closed.wait()
+
+
 async def test_finalization_dropped_exception(echo_server, autojump_clock):
     # Confirm that open_websocket finalization does not contribute to dropped
     # exceptions as described in https://github.com/python-trio/trio/issues/1559.
diff --git a/trio_websocket/_impl.py b/trio_websocket/_impl.py
index 36d6cbf..693c379 100644
--- a/trio_websocket/_impl.py
+++ b/trio_websocket/_impl.py
@@ -1094,6 +1094,12 @@ async def _handle_close_connection_event(self, event):
             await self._send(event.response())
         await self._close_web_socket(event.code, event.reason or None)
         self._close_handshake.set()
+        # RFC: "When a server is instructed to Close the WebSocket Connection
+        #   it SHOULD initiate a TCP Close immediately, and when a client is
+        #   instructed to do the same, it SHOULD wait for a TCP Close from the
+        #   server."
+        if self.is_server:
+            await self._close_stream()
 
     async def _handle_message_event(self, event):
         '''