asyncio, Tulip issue 205: Fix a race condition in BaseSelectorEventLoop.sock_connect()

There is a race condition in create_connection() used with wait_for() to have a
timeout. sock_connect() registers the file descriptor of the socket to be
notified of write event (if connect() raises BlockingIOError). When
create_connection() is cancelled with a TimeoutError, sock_connect() coroutine
gets the exception, but it doesn't unregister the file descriptor for write
event. create_connection() gets the TimeoutError and closes the socket.

If you call again create_connection(), the new socket will likely gets the same
file descriptor, which is still registered in the selector. When sock_connect()
calls add_writer(), it tries to modify the entry instead of creating a new one.

This issue was originally reported in the Trollius project, but the bug comes
from Tulip in fact (Trollius is based on Tulip):
https://bitbucket.org/enovance/trollius/issue/15/after-timeouterror-on-wait_for

This change fixes the race condition. It also makes sock_connect() more
reliable (and portable) is sock.connect() raises an InterruptedError.

1 files changed, 31 insertions, 13 deletions

diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
index 0434a70..33de92e 100644
--- a/Lib/asyncio/selector_events.py
+++ b/Lib/asyncio/selector_events.py
@@ -8,6 +8,7 @@ __all__ = ['BaseSelectorEventLoop']
 
 import collections
 import errno
+import functools
 import socket
 try:
     import ssl
@@ -345,26 +346,43 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
         except ValueError as err:
             fut.set_exception(err)
         else:
-            self._sock_connect(fut, False, sock, address)
+            self._sock_connect(fut, sock, address)
         return fut
 
-    def _sock_connect(self, fut, registered, sock, address):
+    def _sock_connect(self, fut, sock, address):
         fd = sock.fileno()
-        if registered:
-            self.remove_writer(fd)
+        try:
+            while True:
+                try:
+                    sock.connect(address)
+                except InterruptedError:
+                    continue
+                else:
+                    break
+        except BlockingIOError:
+            fut.add_done_callback(functools.partial(self._sock_connect_done,
+                                                    sock))
+            self.add_writer(fd, self._sock_connect_cb, fut, sock, address)
+        except Exception as exc:
+            fut.set_exception(exc)
+        else:
+            fut.set_result(None)
+
+    def _sock_connect_done(self, sock, fut):
+        self.remove_writer(sock.fileno())
+
+    def _sock_connect_cb(self, fut, sock, address):
         if fut.cancelled():
             return
+
         try:
-            if not registered:
-                # First time around.
-                sock.connect(address)
-            else:
-                err = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
-                if err != 0:
-                    # Jump to the except clause below.
-                    raise OSError(err, 'Connect call failed %s' % (address,))
+            err = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
+            if err != 0:
+                # Jump to any except clause below.
+                raise OSError(err, 'Connect call failed %s' % (address,))
         except (BlockingIOError, InterruptedError):
-            self.add_writer(fd, self._sock_connect, fut, True, sock, address)
+            # socket is still registered, the callback will be retried later
+            pass
         except Exception as exc:
             fut.set_exception(exc)
         else: