Set the close status code and reason more consistently.

Set them when the closing handshake is considered complete or aborted.

Author: aaugustin

docs/index.rst

@@ -305,6 +305,8 @@ Changelog
 * Cancelling :meth:`~websockets.protocol.WebSocketCommonProtocol.recv` no
   longer drops the next message.
 
+* Set the close status code and reason more consistently.
+
 * Improved tests.
 
 2.4

websockets/protocol.py

@@ -104,6 +104,7 @@ def __init__(self, *,
 
         self.subprotocol = None
 
+        # Code and reason must be set when the closing handshake completes.
         self.close_code = None
         self.close_reason = ''
 
@@ -161,7 +162,6 @@ def close(self, code=1000, reason=''):
         if self.state == OPEN:
             # 7.1.2. Start the WebSocket Closing Handshake
             # 7.1.3. The WebSocket Closing Handshake is Started
-            self.close_code, self.close_reason = code, reason
             frame_data = serialize_close(code, reason)
             yield from self.write_frame(OP_CLOSE, frame_data)
 
@@ -350,11 +350,13 @@ def read_data_frame(self, max_size):
             frame = yield from self.read_frame(max_size)
             # 5.5. Control Frames
             if frame.opcode == OP_CLOSE:
-                self.close_code, self.close_reason = parse_close(frame.data)
+                # Make sure the close frame is valid before echoing it.
+                code, reason = parse_close(frame.data)
                 if self.state == OPEN:
                     # 7.1.3. The WebSocket Closing Handshake is Started
                     yield from self.write_frame(OP_CLOSE, frame.data)
                 if not self.closing_handshake.done():
+                    self.close_code, self.close_reason = code, reason
                     self.closing_handshake.set_result(True)
                 return
             elif frame.opcode == OP_PING:
@@ -443,10 +445,6 @@ def close_connection(self, force=False):
 
     @asyncio.coroutine
     def fail_connection(self, code=1011, reason=''):
-        # Losing the connection usually results in a protocol error.
-        # Preserve the original error code in this case.
-        if self.close_code != 1006:
-            self.close_code, self.close_reason = code, reason
         # 7.1.7. Fail the WebSocket Connection
         logger.info("Failing the WebSocket connection: %d %s", code, reason)
         if self.state == OPEN:
@@ -458,6 +456,7 @@ def fail_connection(self, code=1011, reason=''):
                 frame_data = serialize_close(code, reason)
                 yield from self.write_frame(OP_CLOSE, frame_data)
         if not self.closing_handshake.done():
+            self.close_code, self.close_reason = code, reason
             self.closing_handshake.set_result(False)
         yield from self.close_connection()
 
@@ -472,8 +471,9 @@ def client_connected(self, reader, writer):
     def connection_lost(self, exc):
         # 7.1.4. The WebSocket Connection is Closed
         self.state = CLOSED
+        if not self.closing_handshake.done():
+            self.close_code, self.close_reason = 1006, ''
+            self.closing_handshake.set_result(False)
         if not self.connection_closed.done():
             self.connection_closed.set_result(None)
-        if self.close_code is None:
-            self.close_code = 1006
         super().connection_lost(exc)

websockets/test_protocol.py

@@ -391,7 +391,7 @@ def test_close_handshake_in_fragmented_text(self):
         self.receive_frame(Frame(True, OP_CLOSE, b''))
         self.receive_eof()
         self.assertIsNone(self.loop.run_until_complete(self.protocol.recv()))
-        self.assertConnectionClosed(1002, '')
+        self.assertConnectionClosed(1005, '')
 
     def test_connection_close_in_fragmented_text(self):
         self.receive_frame(Frame(False, OP_TEXT, 'ca'.encode('utf-8')))
@@ -458,20 +458,19 @@ def test_close_handshake_timeout(self):
             # response. The server will stop waiting for the close frame and
             # timeout.
             self.loop.run_until_complete(self.protocol.close(reason='close'))
-        self.assertConnectionClosed(1000, 'close')
+        self.assertConnectionClosed(1006, '')
 
     def test_client_close_race_with_failing_connection(self):
         self.make_drain_slow()
 
         # Fail the connection while answering a close frame from the client.
         self.loop.call_soon(self.receive_frame, self.client_close)
-        fail_connection = self.protocol.fail_connection(1000, 'server')
-        self.loop.call_later(MS, self.async, fail_connection)
+        self.loop.call_later(MS, self.async, self.protocol.fail_connection())
         next_message = self.loop.run_until_complete(self.protocol.recv())
 
         self.assertIsNone(next_message)
-        # The connection was closed before the close frame could be sent.
-        self.assertConnectionClosed(1006, '')
+        # The closing handshake was completed by fail_connection.
+        self.assertConnectionClosed(1011, '')
         self.assertOneFrameSent(*self.client_close)
 
     def test_close_protocol_error(self):
@@ -581,7 +580,7 @@ def test_close_handshake_timeout(self):
             # stop waiting for the close frame and timeout, then stop waiting
             # for the connection close and timeout again.
             self.loop.run_until_complete(self.protocol.close(reason='close'))
-        self.assertConnectionClosed(1000, 'close')
+        self.assertConnectionClosed(1006, '')
 
     def test_eof_received_timeout(self):
         # Timeout is expected in 10ms.
@@ -600,14 +599,13 @@ def test_server_close_race_with_failing_connection(self):
 
         # Fail the connection while answering a close frame from the server.
         self.loop.call_soon(self.receive_frame, self.server_close)
-        fail_connection = self.protocol.fail_connection(1000, 'client')
-        self.loop.call_later(MS, self.async, fail_connection)
+        self.loop.call_later(MS, self.async, self.protocol.fail_connection())
         self.loop.call_later(2 * MS, self.receive_eof)
         next_message = self.loop.run_until_complete(self.protocol.recv())
 
         self.assertIsNone(next_message)
-        # The connection was closed before the close frame could be sent.
-        self.assertConnectionClosed(1006, '')
+        # The closing handshake was completed by fail_connection.
+        self.assertConnectionClosed(1011, '')
         self.assertOneFrameSent(*self.server_close)
 
     def test_close_protocol_error(self):