Abort TCP connection if it doesn't close fast enough.

Fix #112.

Author: aaugustin

websockets/protocol.py

@@ -66,7 +66,7 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
     The ``timeout`` parameter defines the maximum wait time in seconds for
     completing the closing handshake and, only on the client side, for
     terminating the TCP connection. :meth:`close()` will complete in at most
-    ``3 * timeout`` on the server side and ``4 * timeout`` on the client side.
+    ``4 * timeout`` on the server side and ``5 * timeout`` on the client side.
 
     The ``max_size`` parameter enforces the maximum size for incoming messages
     in bytes. The default value is 1MB. ``None`` disables the limit. If a
@@ -773,12 +773,23 @@ def close_connection(self, after_handshake=True):
             # Closing a transport is idempotent. If the transport was already
             # closed, for example from eof_received(), it's fine.
 
-            # Close the TCP connection.
+            # Close the TCP connection. Buffers are flushed asynchronously.
             logger.debug(
                 "%s x closing TCP connection", self.side)
             self.writer.close()
-            # There's little need to await self.wait_for_connection_lost()
-            # here. Closing the transport triggers self.connection_lost().
+
+            if (yield from self.wait_for_connection_lost()):
+                return
+            logger.debug(
+                "%s ! timed out waiting for TCP close", self.side)
+
+            # Abort the TCP connection. Buffers are discarded.
+            logger.debug(
+                "%s x aborting TCP connection", self.side)
+            self.writer.transport.abort()
+
+            # connection_lost() is called quickly after aborting.
+            yield from self.wait_for_connection_lost()
 
     @asyncio.coroutine
     def fail_connection(self, code=1011, reason=''):

websockets/test_protocol.py

@@ -70,6 +70,10 @@ def close(self):
             self.loop.call_soon(self.protocol.connection_lost, None)
         self._closing = True
 
+    def abort(self):
+        if self.protocol.state != CLOSED:
+            self.loop.call_soon(self.protocol.connection_lost, None)
+
 
 class CommonTests:
     """
@@ -789,7 +793,7 @@ def test_local_close_receive_close_frame_timeout(self):
             self.loop.run_until_complete(self.protocol.close(reason='close'))
         self.assertConnectionClosed(1006, '')
 
-    def test_local_close_connection_lost_timeout(self):
+    def test_local_close_connection_lost_timeout_after_write_eof(self):
         self.protocol.timeout = 10 * MS
         # If the client doesn't close its side of the TCP connection after we
         # half-close our side with write_eof(), time out in 10ms.
@@ -801,6 +805,21 @@ def test_local_close_connection_lost_timeout(self):
             self.loop.run_until_complete(self.protocol.close(reason='close'))
         self.assertConnectionClosed(1000, 'close')
 
+    def test_local_close_connection_lost_timeout_after_close(self):
+        self.protocol.timeout = 10 * MS
+        # If the client doesn't close its side of the TCP connection after we
+        # half-close our side with write_eof() and close it with close(), time
+        # out in 20ms.
+        # Check the timing within -1/+9ms for robustness.
+        with self.assertCompletesWithin(19 * MS, 29 * MS):
+            # HACK: disable write_eof => other end drops connection emulation.
+            self.transport._eof = True
+            # HACK: disable close => other end drops connection emulation.
+            self.transport._closing = True
+            self.receive_frame(self.close_frame)
+            self.loop.run_until_complete(self.protocol.close(reason='close'))
+        self.assertConnectionClosed(1000, 'close')
+
 
 class ClientTests(CommonTests, unittest.TestCase):
 
@@ -830,16 +849,34 @@ def test_local_close_receive_close_frame_timeout(self):
             self.loop.run_until_complete(self.protocol.close(reason='close'))
         self.assertConnectionClosed(1006, '')
 
-    def test_local_close_connection_lost_timeout(self):
+    def test_local_close_connection_lost_timeout_after_write_eof(self):
         self.protocol.timeout = 10 * MS
         # If the server doesn't half-close its side of the TCP connection
         # after we send a close frame, time out in 20ms:
         # - 10ms waiting for receiving a half-close
-        # - 10ms waiting for receiving a close
+        # - 10ms waiting for receiving a close after write_eof
         # Check the timing within -1/+9ms for robustness.
         with self.assertCompletesWithin(19 * MS, 29 * MS):
             # HACK: disable write_eof => other end drops connection emulation.
             self.transport._eof = True
             self.receive_frame(self.close_frame)
             self.loop.run_until_complete(self.protocol.close(reason='close'))
         self.assertConnectionClosed(1000, 'close')
+
+    def test_local_close_connection_lost_timeout_after_close(self):
+        self.protocol.timeout = 10 * MS
+        # If the client doesn't close its side of the TCP connection after we
+        # half-close our side with write_eof() and close it with close(), time
+        # out in 20ms.
+        # - 10ms waiting for receiving a half-close
+        # - 10ms waiting for receiving a close after write_eof
+        # - 10ms waiting for receiving a close after close
+        # Check the timing within -1/+9ms for robustness.
+        with self.assertCompletesWithin(29 * MS, 39 * MS):
+            # HACK: disable write_eof => other end drops connection emulation.
+            self.transport._eof = True
+            # HACK: disable close => other end drops connection emulation.
+            self.transport._closing = True
+            self.receive_frame(self.close_frame)
+            self.loop.run_until_complete(self.protocol.close(reason='close'))
+        self.assertConnectionClosed(1000, 'close')