Repository: aaugustin/websockets (python)

Issue:
Basic example issue: Cannot write to a WebSocket in the CLOSING state
I'm running the basic example from here: https://websockets.readthedocs.io/en/stable/intro.html

When I run the client program, about 30% of the time I get an error (after successfully printing the greeting from the server):

```
> Brian
< Hello Brian!
Traceback (most recent call last):
  File "test_websockets_client.py", line 12, in <module>
    asyncio.get_event_loop().run_until_complete(hello())
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/asyncio/base_events.py", line 467, in run_until_complete
    return future.result()
  File "test_websockets_client.py", line 10, in hello
    print('< {}'.format(greeting))
  File "/Users/brianwong/evex-python2/websockets/client.py", line 390, in __aexit__
    yield from self.ws_client.close()
  File "/Users/brianwong/evex-python2/websockets/protocol.py", line 370, in close
    self.timeout, loop=self.loop)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/asyncio/tasks.py", line 358, in wait_for
    return fut.result()
  File "/Users/brianwong/evex-python2/websockets/protocol.py", line 642, in write_frame
    "in the {} state".format(self.state.name))
websockets.exceptions.InvalidState: Cannot write to a WebSocket in the CLOSING state
```

All help very much appreciated!

Pull request description:
Fix a race condition in the closing handshake.
Fix #339.

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index cf48a29..ab01df5 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -36,6 +36,9 @@ Also:
 
 * Fixed missing close code, which caused :exc:`TypeError` on connection close.
 
+* Fixed a race condition in the closing handshake that raised
+  :exc:`~exceptions.InvalidState`.
+
 * Stopped logging stack traces when the TCP connection dies prematurely.
 
 4.0
diff --git a/websockets/protocol.py b/websockets/protocol.py
index 563594b..bee5eaf 100644
--- a/websockets/protocol.py
+++ b/websockets/protocol.py
@@ -374,20 +374,16 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         ``code`` must be an :class:`int` and ``reason`` a :class:`str`.
 
         """
-        if self.state is State.OPEN:
-            # 7.1.2. Start the WebSocket Closing Handshake
-            # 7.1.3. The WebSocket Closing Handshake is Started
-            frame_data = serialize_close(code, reason)
-            try:
-                yield from asyncio.wait_for(
-                    self.write_frame(OP_CLOSE, frame_data),
-                    self.timeout, loop=self.loop)
-            except asyncio.TimeoutError:
-                # If the close frame cannot be sent because the send buffers
-                # are full, the closing handshake won't complete anyway.
-                # Cancel the data transfer task to shut down faster.
-                # Cancelling a task is idempotent.
-                self.transfer_data_task.cancel()
+        try:
+            yield from asyncio.wait_for(
+                self.write_close_frame(serialize_close(code, reason)),
+                self.timeout, loop=self.loop)
+        except asyncio.TimeoutError:
+            # If the close frame cannot be sent because the send buffers
+            # are full, the closing handshake won't complete anyway.
+            # Cancel the data transfer task to shut down faster.
+            # Cancelling a task is idempotent.
+            self.transfer_data_task.cancel()
 
         # If no close frame is received within the timeout, wait_for() cancels
         # the data transfer task and raises TimeoutError. Then transfer_data()
@@ -402,7 +398,8 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             # is cancelled before the timeout elapses (on Python ≥ 3.4.3).
             # This helps closing connections when shutting down a server.
             yield from asyncio.wait_for(
-                self.transfer_data_task, self.timeout, loop=self.loop)
+                self.transfer_data_task,
+                self.timeout, loop=self.loop)
         except asyncio.TimeoutError:
             pass
 
@@ -611,14 +608,13 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
 
             # 5.5. Control Frames
             if frame.opcode == OP_CLOSE:
-                # Make sure the close frame is valid before echoing it.
-                code, reason = parse_close(frame.data)
                 # 7.1.5.  The WebSocket Connection Close Code
                 # 7.1.6.  The WebSocket Connection Close Reason
-                self.close_code, self.close_reason = code, reason
-                if self.state is State.OPEN:
-                    # 7.1.3. The WebSocket Closing Handshake is Started
-                    yield from self.write_frame(OP_CLOSE, frame.data)
+                self.close_code, self.close_reason = parse_close(frame.data)
+                # Echo the original data instead of re-serializing it with
+                # serialize_close() because that fails when the close frame is
+                # empty and parse_close() synthetizes a 1005 close code.
+                yield from self.write_close_frame(frame.data)
                 return
 
             elif frame.opcode == OP_PING:
@@ -679,18 +675,12 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         return frame
 
     @asyncio.coroutine
-    def write_frame(self, opcode, data=b''):
+    def write_frame(self, opcode, data=b'', _expected_state=State.OPEN):
         # Defensive assertion for protocol compliance.
-        if self.state is not State.OPEN:                    # pragma: no cover
+        if self.state is not _expected_state:               # pragma: no cover
             raise InvalidState("Cannot write to a WebSocket "
                                "in the {} state".format(self.state.name))
 
-        # Make sure no other frame will be sent after a close frame. Do this
-        # before yielding control to avoid sending more than one close frame.
-        if opcode == OP_CLOSE:
-            self.state = State.CLOSING
-            logger.debug("%s - state = CLOSING", self.side)
-
         frame = Frame(True, opcode, data)
         logger.debug("%s > %s", self.side, frame)
         frame.write(
@@ -737,6 +727,25 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             except AttributeError:
                 return transport._closed
 
+    @asyncio.coroutine
+    def write_close_frame(self, data=b''):
+        """
+        Write a close frame if and only if the connection state is OPEN.
+
+        This dedicated coroutine must be used for writing close frames to
+        ensure that at most one close frame is sent on a given connection.
+
+        """
+        # Test and set the connection state before sending the close frame to
+        # avoid sending two frames in case of concurrent calls.
+        if self.state is State.OPEN:
+            # 7.1.3. The WebSocket Closing Handshake is Started
+            self.state = State.CLOSING
+            logger.debug("%s - state = CLOSING", self.side)
+
+            # 7.1.2. Start the WebSocket Closing Handshake
+            yield from self.write_frame(OP_CLOSE, data, State.CLOSING)
+
     @asyncio.coroutine
     def close_connection(self, after_handshake=True):
         """
@@ -845,9 +854,8 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             self.side, code, reason,
         )
         # Don't send a close frame if the connection is broken.
-        if self.state is State.OPEN and code != 1006:
-            frame_data = serialize_close(code, reason)
-            yield from self.write_frame(OP_CLOSE, frame_data)
+        if code != 1006:
+            yield from self.write_close_frame(serialize_close(code, reason))
 
     # asyncio.StreamReaderProtocol methods
 


Test patch:
diff --git a/websockets/test_protocol.py b/websockets/test_protocol.py
index 9048c64..5303cc9 100644
--- a/websockets/test_protocol.py
+++ b/websockets/test_protocol.py
@@ -796,16 +796,19 @@ class CommonTests:
         self.assertNoFrameSent()
 
     def test_simultaneous_close(self):
-        # Delay the incoming close frame until after we send the outgoing one.
-        self.loop.call_later(MS, self.receive_frame, self.remote_close)
-        self.loop.call_later(MS, self.receive_eof_if_client)
+        # Receive the incoming close frame right after self.protocol.close()
+        # starts executing. This reproduces the error described in:
+        # https://github.com/aaugustin/websockets/issues/339
+        self.loop.call_soon(self.receive_frame, self.remote_close)
+        self.loop.call_soon(self.receive_eof_if_client)
 
         self.loop.run_until_complete(self.protocol.close(reason='local'))
 
-        # The close code and reason are taken from the remote side because
-        # that's presumably more useful that the values from the local side.
         self.assertConnectionClosed(1000, 'remote')
-        self.assertOneFrameSent(*self.local_close)
+        # The current implementation sends a close frame in response to the
+        # close frame received from the remote end. It skips the close frame
+        # that should be sent as a result of calling close().
+        self.assertOneFrameSent(*self.remote_close)
 
     def test_close_preserves_incoming_frames(self):
         self.receive_frame(Frame(True, OP_TEXT, b'hello'))