fix reason not being set after local close (#159)

During a local-initiated close handshake, the following incorrect
behavior was observed:
  * `closed` attribute would be `None`
  * `send_message()` would be silently ignored (wsproto < 0.2.0) or leak a
`LocalProtocolException` (wsproto >= 0.2.0)

Upon local-initiated close, `closed` will now have the reason, and
`send_message()` will raise `ConnectionClosed`.

Fixes #158.

Author: belm0

tests/test_connection.py

@@ -791,7 +791,7 @@ async def handler(request):
     server = await nursery.start(
         partial(serve_websocket, handler, HOST, 0, ssl_context=None))
 
-    # connection should close when server handler exists
+    # connection should close when server handler exits
     with trio.fail_after(2):
         async with open_websocket(
                 HOST, server.port, '/', use_ssl=False) as connection:
@@ -955,6 +955,35 @@ async def handler(request: WebSocketRequest):
         await client.send_message('bar')
 
 
+async def test_message_after_local_close_race(nursery):
+    """test message send during local-initiated close handshake (issue #158)"""
+
+    async def handler(request: WebSocketRequest):
+        await request.accept()
+        await trio.sleep_forever()
+
+    server = await nursery.start(
+        partial(serve_websocket, handler, HOST, 0, ssl_context=None))
+
+    client = await connect_websocket(nursery, HOST, server.port,
+                                     RESOURCE, use_ssl=False)
+    orig_send = client._send
+    close_sent = trio.Event()
+
+    async def _send_wrapper(event):
+        if isinstance(event, CloseConnection):
+            close_sent.set()
+        return await orig_send(event)
+
+    client._send = _send_wrapper
+    assert not client.closed
+    nursery.start_soon(client.aclose)
+    await close_sent.wait()
+    assert client.closed
+    with pytest.raises(ConnectionClosed):
+        await client.send_message('hello')
+
+
 @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"""

trio_websocket/_impl.py

@@ -493,7 +493,7 @@ def __init__(self, code, reason):
         Constructor.
 
         :param int code:
-        :param str reason:
+        :param Optional[str] reason:
         '''
         self._code = code
         try:
@@ -838,17 +838,22 @@ async def aclose(self, code=1000, reason=None):  # pylint: disable=arguments-dif
         with _preserve_current_exception():
             await self._aclose(code, reason)
 
-    async def _aclose(self, code=1000, reason=None):
+    async def _aclose(self, code, reason):
         if self._close_reason:
             # Per AsyncResource interface, calling aclose() on a closed resource
             # should succeed.
             return
         try:
             if self._wsproto.state == ConnectionState.OPEN:
+                # Our side is initiating the close, so send a close connection
+                # event to peer, while setting the local close reason to normal.
+                self._close_reason = CloseReason(1000, None)
                 await self._send(CloseConnection(code=code, reason=reason))
             elif self._wsproto.state in (ConnectionState.CONNECTING,
                     ConnectionState.REJECTING):
                 self._close_handshake.set()
+            # TODO: shouldn't the receive channel be closed earlier, so that
+            #  get_message() during send of the CloseConneciton event fails?
             await self._recv_channel.aclose()
             await self._close_handshake.wait()
         except ConnectionClosed: