Change send_message() behavior when close handshake is in progress

send_message() is changed to raise ConnectionClosed when a close
handshake is in progress.  Previously it would silently ignore
the call, which was an oversight, given that ConnectionClosed is
defined to cover connections "closed or in the process of closing".

Significantly, this fixes send_message() leaking a wsproto exception
(LocalProtocolError) with wsproto >= 1.22.0.

Fixes #175.

Author: belm0

tests/test_connection.py

@@ -40,7 +40,6 @@
 import wsproto
 from trio.testing import memory_stream_pair
 from wsproto.events import CloseConnection
-from wsproto.utilities import LocalProtocolError
 
 try:
     from trio.lowlevel import current_task  # pylint: disable=ungrouped-imports
@@ -928,14 +927,10 @@ async def handler(request: WebSocketRequest):
     await trio.sleep(.1)
 
 
-@pytest.mark.xfail(
-    reason='send_message() API oversight for closing-in-process case',
-    raises=None if WS_PROTO_VERSION < (1, 2, 0) else LocalProtocolError,
-    strict=True)
 async def test_remote_close_local_message_race(nursery, autojump_clock):
     """as remote initiates close, local attempts message (issue #175)
 
-    This exposes multiple problems in the trio-websocket API and implementation:
+    This exposed multiple problems in the trio-websocket API and implementation:
         * send_message() silently fails if a close is in progress.  This was
             likely an oversight in the API, since send_message() raises `ConnectionClosed`
             only in the already-closed case, yet `ConnectionClosed` is defined to cover
@@ -957,9 +952,7 @@ async def handler(request: WebSocketRequest):
     await client.send_message('foo')
     await client._for_testing_peer_closed_connection.wait()
     with pytest.raises(ConnectionClosed):
-        await client.send_message('bar')  # wsproto < 1.2.0: silently ignored
-                                          # wsproto >= 1.2.0: raises LocalProtocolError
-                                          # desired: raises ConnectionClosed
+        await client.send_message('bar')
 
 
 @fail_after(DEFAULT_TEST_MAX_DURATION)

trio_websocket/_impl.py

@@ -932,7 +932,7 @@ async def send_message(self, message):
 
         :param message: The message to send.
         :type message: str or bytes
-        :raises ConnectionClosed: if connection is already closed.
+        :raises ConnectionClosed: if connection is closed, or being closed
         '''
         if self._close_reason:
             raise ConnectionClosed(self._close_reason)
@@ -1096,10 +1096,13 @@ async def _handle_close_connection_event(self, event):
 
         :param wsproto.events.CloseConnection event:
         '''
-        if self._for_testing_peer_closed_connection:
-            self._for_testing_peer_closed_connection.set()
-            await trio.sleep(0)
         if self._wsproto.state == ConnectionState.REMOTE_CLOSING:
+            # Set _close_reason in advance, so that send_message() will raise
+            # ConnectionClosed during the close handshake.
+            self._close_reason = CloseReason(event.code, event.reason or None)
+            if self._for_testing_peer_closed_connection:
+                self._for_testing_peer_closed_connection.set()
+                await trio.sleep(0)
             await self._send(event.response())
         await self._close_web_socket(event.code, event.reason or None)
         self._close_handshake.set()