Repository: aaugustin/websockets
Pull request #391 (2018-05-12 12:17:04)

Improve failing the WebSocket connection.
Mainly:

- Attach the reason why the connection was closed to ConnectionClosed exceptions
- Better mirror the specification in the implementation and improve compliance slightly in some edge cases


codecov[bot]: # [Codecov](https://codecov.io/gh/aaugustin/websockets/pull/391?src=pr&el=h1) Report
> Merging [#391](https://codecov.io/gh/aaugustin/websockets/pull/391?src=pr&el=desc) into [master](https://codecov.io/gh/aaugustin/websockets/commit/9eefbc25831837aac3fa7370ba5faff9b8c17bb1?src=pr&el=desc) will **decrease** coverage by `0.06%`.
> The diff coverage is `95.65%`.

[![Impacted file tree graph](https://codecov.io/gh/aaugustin/websockets/pull/391/graphs/tree.svg?src=pr&token=vqr9oExbGY&width=650&height=150)](https://codecov.io/gh/aaugustin/websockets/pull/391?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #391      +/-   ##
==========================================
- Coverage     100%   99.93%   -0.07%     
==========================================
  Files          29       29              
  Lines        3066     3082      +16     
  Branches      318      321       +3     
==========================================
+ Hits         3066     3080      +14     
- Misses          0        2       +2
```


| [Impacted Files](https://codecov.io/gh/aaugustin/websockets/pull/391?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [websockets/server.py](https://codecov.io/gh/aaugustin/websockets/pull/391/diff?src=pr&el=tree#diff-d2Vic29ja2V0cy9zZXJ2ZXIucHk=) | `100% <100%> (ø)` | :arrow_up: |
| [websockets/client.py](https://codecov.io/gh/aaugustin/websockets/pull/391/diff?src=pr&el=tree#diff-d2Vic29ja2V0cy9jbGllbnQucHk=) | `100% <100%> (ø)` | :arrow_up: |
| [websockets/protocol.py](https://codecov.io/gh/aaugustin/websockets/pull/391/diff?src=pr&el=tree#diff-d2Vic29ja2V0cy9wcm90b2NvbC5weQ==) | `99.39% <95%> (-0.61%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/aaugustin/websockets/pull/391?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/aaugustin/websockets/pull/391?src=pr&el=footer). Last update [9eefbc2...57d92e9](https://codecov.io/gh/aaugustin/websockets/pull/391?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

aaugustin: Unsure why test coverage isn't 100% on CircleCI. It is 100% locally.

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index ab01df5..c4fc289 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -32,6 +32,8 @@ Also:
 * If a :meth:`~protocol.WebSocketCommonProtocol.ping` doesn't receive a pong,
   it's cancelled when the connection is closed.
 
+* Reported the cause of :exc:`~exceptions.ConnectionClosed` exceptions.
+
 * Updated documentation with new features from Python 3.6.
 
 * Fixed missing close code, which caused :exc:`TypeError` on connection close.
@@ -41,6 +43,12 @@ Also:
 
 * Stopped logging stack traces when the TCP connection dies prematurely.
 
+* Prevented writing to a closing TCP connection during unclean shutdowns.
+
+* Made connection termination more robust to network congestion.
+
+* Prevented processing of incoming frames after failing the connection.
+
 4.0
 ...
 
diff --git a/docs/design.rst b/docs/design.rst
index f885b00..9974c1c 100644
--- a/docs/design.rst
+++ b/docs/design.rst
@@ -32,16 +32,18 @@ WebSocket connections go through a trivial state machine:
 Transitions happen in the following places:
 
 - ``CONNECTING -> OPEN``: in
-  :meth:`~protocol.WebSocketCommonProtocol.connection_open()`, which runs
-  when the :ref:`opening handshake <opening-handshake>` completes and the
-  WebSocket connection is established — not to be confused with
-  :meth:`~asyncio.Protocol.connection_made` which runs earlier, when the TCP
-  connection is established;
+  :meth:`~protocol.WebSocketCommonProtocol.connection_open()` which runs when
+  the :ref:`opening handshake <opening-handshake>` completes and the WebSocket
+  connection is established — not to be confused with
+  :meth:`~asyncio.Protocol.connection_made` which runs when the TCP connection
+  is established;
 - ``OPEN -> CLOSING``: in
   :meth:`~protocol.WebSocketCommonProtocol.write_frame()` immediately before
   sending a close frame; since receiving a close frame triggers sending a
   close frame, this does the right thing regardless of which side started the
-  :ref:`closing handshake <closing-handshake>`;
+  :ref:`closing handshake <closing-handshake>`; also in
+  :meth:`~protocol.WebSocketCommonProtocol.fail_connection()` which duplicates
+  a few lines of code from `write_close_frame()` and `write_frame()`;
 - ``* -> CLOSED``: in
   :meth:`~protocol.WebSocketCommonProtocol.connection_lost()` which is always
   called exactly once when the TCP connection is closed.
@@ -68,14 +70,36 @@ two tasks:
 - :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` runs
   :meth:`~protocol.WebSocketCommonProtocol.transfer_data()` which handles
   incoming data and lets :meth:`~protocol.WebSocketCommonProtocol.recv()`
-  consume it. It never exits with an exception but it may be cancelled.
-  See :ref:`data transfer <data-transfer>` below.
+  consume it. It may be cancelled to terminate the connection. It never exits
+  with an exception other than :exc:`~asyncio.CancelledError`. See :ref:`data
+  transfer <data-transfer>` below.
+
 - :attr:`~protocol.WebSocketCommonProtocol.close_connection_task` runs
   :meth:`~protocol.WebSocketCommonProtocol.close_connection()` which waits for
   the data transfer to terminate, then takes care of closing the TCP
-  connection. It never exits with an exception and it is never cancelled. See
+  connection. It must not be cancelled. It never exits with an exception. See
   :ref:`connection termination <connection-termination>` below.
 
+Besides, :meth:`~protocol.WebSocketCommonProtocol.fail_connection()` starts
+the same :attr:`~protocol.WebSocketCommonProtocol.close_connection_task` when
+the opening handshake fails, in order to close the TCP connection.
+
+Splitting the responsibilities between two tasks makes it easier to guarantee
+that ``websockets`` can terminate connections:
+
+- within a fixed timeout,
+- without leaking pending tasks,
+- without leaking open TCP connections,
+
+regardless of whether the connection terminates normally or abnormally.
+
+:attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` completes when no
+more data will be received on the connection. Under normal circumstances, it
+exits after exchanging close frames.
+
+:attr:`~protocol.WebSocketCommonProtocol.close_connection_task` completes when
+the TCP connection is closed.
+
 
 .. _opening-handshake:
 
@@ -121,9 +145,8 @@ lies in the negociation of extensions and, to a lesser extent, of the
 subprotocol. The server knows everything about both sides and decides what the
 parameters should be for the connection. The client merely applies them.
 
-If anything goes wrong during the opening handshake, ``websockets`` closes the
-TCP connection. This is the proper way to fail the WebSocket connection before
-it's established.
+If anything goes wrong during the opening handshake, ``websockets``
+:ref:`fails the connection <connection-failure>`.
 
 
 .. _data-transfer:
@@ -213,8 +236,8 @@ The right side of the diagram shows how ``websockets`` sends data.
 containing the message. Fragmentation isn't supported at this time.
 
 :meth:`~protocol.WebSocketCommonProtocol.ping()` writes a ping frame and
-returns a :class:`~asyncio.Future` which will be completed when a matching
-pong frame is received.
+yields a :class:`~asyncio.Future` which will be completed when a matching pong
+frame is received.
 
 :meth:`~protocol.WebSocketCommonProtocol.pong()` writes a pong frame.
 
@@ -243,9 +266,7 @@ state and sends a close frame. When the other side sends a close frame,
 :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` to terminate.
 
 If the other side doesn't send a close frame within the connection's timeout,
-:meth:`~protocol.WebSocketCommonProtocol.close()` cancels
-:attr:`~protocol.WebSocketCommonProtocol.transfer_data_task`, which has the
-same effect.
+``websockets`` :ref:`fails the connection <connection-failure>`.
 
 The closing handshake can take up to ``2 * timeout``: one ``timeout`` to write
 a close frame and one ``timeout`` to receive a close frame.
@@ -271,8 +292,8 @@ which may happen as a result of:
 - a timeout while waiting for the closing handshake to complete: this cancels
   :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task`;
 - a protocol error, including connection errors: depending on the exception,
-  :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` fails the
-  WebSocket connection with a suitable code and exits.
+  :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` :ref:`fails the
+  connection <connection-failure>`_ with a suitable code and exits.
 
 :attr:`~protocol.WebSocketCommonProtocol.close_connection_task` is separate
 from :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` to make it
@@ -290,6 +311,23 @@ aborting the connection. At this point the connection drops regardless of what
 happens on the network.
 
 
+.. _connection-failure:
+
+Connection failure
+------------------
+
+If the opening handshake doesn't complete successfully, ``websockets`` fails
+the connection by closing the TCP connection.
+
+Once the opening handshake has completed, ``websockets`` fails the connection
+by cancelling :attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` and
+sending a close frame if appropriate.
+
+:attr:`~protocol.WebSocketCommonProtocol.transfer_data_task` exits, unblocking
+:attr:`~protocol.WebSocketCommonProtocol.close_connection_task`, which closes
+the TCP connection.
+
+
 .. _cancellation:
 
 Cancellation
diff --git a/websockets/client.py b/websockets/client.py
index f8f7c77..92f29e9 100644
--- a/websockets/client.py
+++ b/websockets/client.py
@@ -404,7 +404,7 @@ class Connect:
                 extra_headers=protocol.extra_headers,
             )
         except Exception:
-            yield from protocol.close_connection(after_handshake=False)
+            yield from protocol.fail_connection()
             raise
 
         self.ws_client = protocol
diff --git a/websockets/protocol.py b/websockets/protocol.py
index bee5eaf..e6a0bef 100644
--- a/websockets/protocol.py
+++ b/websockets/protocol.py
@@ -198,6 +198,9 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         # Task running the data transfer.
         self.transfer_data_task = None
 
+        # Exception that occurred during data transfer, if any.
+        self.transfer_data_exc = None
+
         # Task closing the TCP connection.
         self.close_connection_task = None
 
@@ -271,15 +274,15 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         .. _EAFP: https://docs.python.org/3/glossary.html#term-eafp
 
         """
-        return self.state is State.OPEN
+        return self.state is State.OPEN and not self.transfer_data_task.done()
 
     @property
     def closed(self):
         """
         This property is ``True`` once the connection is closed.
 
-        Be aware that :attr:`open` and :attr`closed` are ``False`` when the
-        connection is in the OPENING or CLOSING state.
+        Be aware that both :attr:`open` and :attr`closed` are ``False`` during
+        the opening and closing sequences.
 
         """
         return self.state is State.CLOSED
@@ -317,6 +320,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         # Wait for a message until the connection is closed.
         next_message = asyncio_ensure_future(
             self.messages.get(), loop=self.loop)
+        # See https://bugs.python.org/issue23859 for cancellation handling.
         try:
             done, pending = yield from asyncio.wait(
                 [next_message, self.transfer_data_task],
@@ -381,9 +385,8 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         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()
+            # Fail the connection to shut down faster.
+            self.fail_connection()
 
         # If no close frame is received within the timeout, wait_for() cancels
         # the data transfer task and raises TimeoutError. Then transfer_data()
@@ -400,7 +403,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             yield from asyncio.wait_for(
                 self.transfer_data_task,
                 self.timeout, loop=self.loop)
-        except asyncio.TimeoutError:
+        except (asyncio.TimeoutError, asyncio.CancelledError):
             pass
 
         # Wait for the close connection task to close the TCP connection.
@@ -475,10 +478,18 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         """
         # Handle cases from most common to least common for performance.
         if self.state is State.OPEN:
-            return
+            # If self.transfer_data_task exited without a closing handshake,
+            # self.close_connection_task may be closing it, going straight
+            # from OPEN to CLOSED.
+            if self.transfer_data_task.done():
+                yield from asyncio.shield(self.close_connection_task)
+                raise ConnectionClosed(self.close_code, self.close_reason)
+            else:
+                return
 
         if self.state is State.CLOSED:
-            raise ConnectionClosed(self.close_code, self.close_reason)
+            raise ConnectionClosed(
+                self.close_code, self.close_reason) from self.transfer_data_exc
 
         if self.state is State.CLOSING:
             # If we started the closing handshake, wait for its completion to
@@ -488,7 +499,8 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             # that case self.close_connection_task will complete even faster.
             if self.close_code is None:
                 yield from asyncio.shield(self.close_connection_task)
-            raise ConnectionClosed(self.close_code, self.close_reason)
+            raise ConnectionClosed(
+                self.close_code, self.close_reason) from self.transfer_data_exc
 
         # Control may only reach this point in buggy third-party subclasses.
         assert self.state is State.CONNECTING
@@ -509,25 +521,42 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
                 if msg is None:
                     break
                 yield from self.messages.put(msg)
-        except asyncio.CancelledError:
-            # This happens if self.close() cancels self.transfer_data_task.
-            pass
-        except WebSocketProtocolError:
-            yield from self.fail_connection(1002)
-        except (ConnectionError, EOFError):
+
+        except asyncio.CancelledError as exc:
+            self.transfer_data_exc = exc
+            # If fail_connection() cancels this task, avoid logging the error
+            # twice and failing the connection again.
+            raise
+
+        except WebSocketProtocolError as exc:
+            self.transfer_data_exc = exc
+            self.fail_connection(1002)
+
+        except (ConnectionError, EOFError) as exc:
             # Reading data with self.reader.readexactly may raise:
             # - most subclasses of ConnectionError if the TCP connection
             #   breaks, is reset, or is aborted;
             # - IncompleteReadError, a subclass of EOFError, if fewer
             #   bytes are available than requested.
-            yield from self.fail_connection(1006)
-        except UnicodeDecodeError:
-            yield from self.fail_connection(1007)
-        except PayloadTooBig:
-            yield from self.fail_connection(1009)
-        except Exception:
-            logger.warning("Error in data transfer", exc_info=True)
-            yield from self.fail_connection(1011)
+            self.transfer_data_exc = exc
+            self.fail_connection(1006)
+
+        except UnicodeDecodeError as exc:
+            self.transfer_data_exc = exc
+            self.fail_connection(1007)
+
+        except PayloadTooBig as exc:
+            self.transfer_data_exc = exc
+            self.fail_connection(1009)
+
+        except Exception as exc:
+            # This shouldn't happen often because exceptions expected under
+            # regular circumstances are handled above. If it does, consider
+            # catching and handling more exceptions.
+            logger.error("Error in data transfer", exc_info=True)
+
+            self.transfer_data_exc = exc
+            self.fail_connection(1011)
 
     @asyncio.coroutine
     def read_message(self):
@@ -704,9 +733,10 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
                 yield from self.writer.drain()
         except ConnectionError:
             # Terminate the connection if the socket died.
-            yield from self.fail_connection(1006)
-            # And raise an exception, since the frame couldn't be sent.
-            raise ConnectionClosed(self.close_code, self.close_reason)
+            self.fail_connection()
+            # Wait until the connection is closed to raise ConnectionClosed
+            # with the correct code and reason.
+            yield from self.ensure_open()
 
     def writer_is_closing(self):
         """
@@ -747,7 +777,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
             yield from self.write_frame(OP_CLOSE, data, State.CLOSING)
 
     @asyncio.coroutine
-    def close_connection(self, after_handshake=True):
+    def close_connection(self):
         """
         7.1.1. Close the WebSocket Connection
 
@@ -755,14 +785,17 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         this coroutine in a task. It waits for the data transfer phase to
         complete then it closes the TCP connection cleanly.
 
-        When the opening handshake fails, the client or the server runs this
-        coroutine with ``after_handshake=False`` to close the TCP connection.
+        When the opening handshake fails, :meth:`fail_connection` does the
+        same. There's no data transfer phase in that case.
 
         """
         try:
             # Wait for the data transfer phase to complete.
-            if after_handshake:
-                yield from self.transfer_data_task
+            if self.transfer_data_task is not None:
+                try:
+                    yield from self.transfer_data_task
+                except asyncio.CancelledError:
+                    pass
 
             # Cancel all pending pings because they'll never receive a pong.
             for ping in self.pings.values():
@@ -777,8 +810,8 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
                     "%s - cancelled pending ping%s: %s",
                     self.side, plural, pings_hex)
 
-            # A client should wait for a TCP Close from the server.
-            if self.is_client and after_handshake:
+            # A client should wait for a TCP close from the server.
+            if self.is_client and self.transfer_data_task is not None:
                 if (yield from self.wait_for_connection_lost()):
                     return
                 logger.debug(
@@ -843,19 +876,72 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         # and the moment this coroutine resumes running.
         return self.connection_lost_waiter.done()
 
-    @asyncio.coroutine
-    def fail_connection(self, code=1011, reason=''):
+    def fail_connection(self, code=1006, reason=''):
         """
         7.1.7. Fail the WebSocket Connection
 
+        This requires:
+
+        1. Stopping all processing of incoming data, which means cancelling
+           :attr:`transfer_data_task`. The close code will be 1006 unless a
+           close frame was received earlier.
+
+        2. Sending a close frame with an appropriate code if the opening
+           handshake succeeded and the other side is likely to process it.
+
+        3. Closing the connection. :meth:`close_connection` takes care of
+           this once :attr:`transfer_data_task` exits after being cancelled.
+
+        (The specification describes these steps in the opposite order.)
+
+        Return a :class:`~asyncio.Task` that completes when the TCP connection
+        is closed.
+
         """
         logger.debug(
             "%s ! failing WebSocket connection: %d %s",
             self.side, code, reason,
         )
+
+        # Cancel transfer_data_task if the opening handshake succeeded.
+        # cancel() is idempotent and ignored if the task is done already.
+        if self.transfer_data_task is not None:
+            self.transfer_data_task.cancel()
+
+        # Send a close frame when the state is OPEN (a close frame was already
+        # sent if it's CLOSING), except when failing the connection because of
+        # an error reading from or writing to the network.
         # Don't send a close frame if the connection is broken.
-        if code != 1006:
-            yield from self.write_close_frame(serialize_close(code, reason))
+        if code != 1006 and self.state is State.OPEN:
+
+            frame_data = serialize_close(code, reason)
+
+            # Write the close frame without draining the write buffer.
+
+            # Keeping fail_connection() synchronous guarantees it can't
+            # get stuck and simplifies the implementation of the callers.
+            # Not drainig the write buffer is acceptable in this context.
+
+            # This duplicates a few lines of code from write_close_frame()
+            # and write_frame().
+
+            self.state = State.CLOSING
+            logger.debug("%s - state = CLOSING", self.side)
+
+            frame = Frame(True, OP_CLOSE, frame_data)
+            logger.debug("%s > %s", self.side, frame)
+            frame.write(
+                self.writer.write,
+                mask=self.is_client,
+                extensions=self.extensions,
+            )
+
+        # Start close_connection_task if the opening handshake didn't succeed.
+        if self.close_connection_task is None:
+            self.close_connection_task = asyncio_ensure_future(
+                self.close_connection(), loop=self.loop)
+
+        return self.close_connection_task
 
     # asyncio.StreamReaderProtocol methods
 
@@ -895,7 +981,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
            returned to have the same behavior on TLS and plain connections.
 
         3. The websockets protocol has its own closing handshake. Endpoints
-           close the TCP connection after sending a Close frame.
+           close the TCP connection after sending a close frame.
 
         As a consequence we revert to the previous, more useful behavior.
 
diff --git a/websockets/server.py b/websockets/server.py
index 85dea02..8db0482 100644
--- a/websockets/server.py
+++ b/websockets/server.py
@@ -134,7 +134,7 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
                     )
 
                 yield from self.write_http_response(*early_response)
-                yield from self.close_connection(after_handshake=False)
+                yield from self.fail_connection()
 
                 return
 
@@ -142,10 +142,12 @@ class WebSocketServerProtocol(WebSocketCommonProtocol):
                 yield from self.ws_handler(self, path)
             except Exception as exc:
                 if self._is_server_shutting_down(exc):
-                    yield from self.fail_connection(1001)
+                    if not self.closed:
+                        self.fail_connection(1001)
                 else:
                     logger.error("Error in connection handler", exc_info=True)
-                    yield from self.fail_connection(1011)
+                    if not self.closed:
+                        self.fail_connection(1011)
                 raise
 
             try:
@@ -585,8 +587,7 @@ class WebSocketServer:
             yield from asyncio.wait(
                 [websocket.handler_task for websocket in self.websockets] +
                 [websocket.close_connection_task
-                    for websocket in self.websockets
-                    if websocket.close_connection_task],
+                    for websocket in self.websockets],
                 loop=self.loop)
         yield from self.server.wait_closed()
 


Test patch:
diff --git a/websockets/test_protocol.py b/websockets/test_protocol.py
index 5303cc9..19ec7bc 100644
--- a/websockets/test_protocol.py
+++ b/websockets/test_protocol.py
@@ -728,13 +728,26 @@ class CommonTests:
 
         self.assertConnectionFailed(1006, '')
 
-    def test_ensure_connection_before_opening_handshake(self):
+    def test_ensure_open_before_opening_handshake(self):
         # Simulate a bug by forcibly reverting the protocol state.
         self.protocol.state = State.CONNECTING
 
         with self.assertRaises(InvalidState):
             self.loop.run_until_complete(self.protocol.ensure_open())
 
+    def test_ensure_open_during_unclean_close(self):
+        # Process connection_made in order to start transfer_data_task.
+        self.run_loop_once()
+
+        # Ensure the test terminates quickly.
+        self.loop.call_later(MS, self.receive_eof_if_client)
+
+        # Simulate the case when close() times out sending a close frame.
+        self.protocol.fail_connection()
+
+        with self.assertRaises(ConnectionClosed):
+            self.loop.run_until_complete(self.protocol.ensure_open())
+
     def test_legacy_recv(self):
         # By default legacy_recv in disabled.
         self.assertEqual(self.protocol.legacy_recv, False)