Fix crash in transport.get_extra_info('socket') when tr is closed

Author: 1st1

tests/test_tcp.py

@@ -507,6 +507,44 @@ async def start_server():
 
         self.loop.run_until_complete(start_server())
 
+    def test_tcp_handle_exception_in_connection_made(self):
+        # Test that if connection_made raises an exception,
+        # 'create_connection' still returns.
+
+        # Silence error logging
+        self.loop.set_exception_handler(lambda *args: None)
+
+        fut = asyncio.Future(loop=self.loop)
+
+        async def server(reader, writer):
+            try:
+                await reader.read()
+            finally:
+                writer.close()
+
+        class Proto(asyncio.Protocol):
+            def connection_made(self, tr):
+                1 / 0
+
+        srv = self.loop.run_until_complete(asyncio.start_server(
+            server,
+            '127.0.0.1', 0,
+            family=socket.AF_INET,
+            loop=self.loop))
+
+        async def runner():
+            tr, pr = await asyncio.wait_for(
+                self.loop.create_connection(
+                    Proto, *srv.sockets[0].getsockname()),
+                timeout=1.0, loop=self.loop)
+            fut.set_result(None)
+            tr.close()
+
+        self.loop.run_until_complete(runner())
+        srv.close()
+        self.loop.run_until_complete(srv.wait_closed())
+        self.loop.run_until_complete(fut)
+
 
 class Test_UV_TCP(_TestTCP, tb.UVTestCase):
 
@@ -620,6 +658,40 @@ async def run():
         # we need to disable this check:
         self.skip_unclosed_handles_check()
 
+    def test_tcp_handle_abort_in_connection_made(self):
+        async def server(reader, writer):
+            try:
+                await reader.read()
+            finally:
+                writer.close()
+
+        class Proto(asyncio.Protocol):
+            def connection_made(self, tr):
+                tr.abort()
+
+        srv = self.loop.run_until_complete(asyncio.start_server(
+            server,
+            '127.0.0.1', 0,
+            family=socket.AF_INET,
+            loop=self.loop))
+
+        async def runner():
+            tr, pr = await asyncio.wait_for(
+                self.loop.create_connection(
+                    Proto, *srv.sockets[0].getsockname()),
+                timeout=1.0, loop=self.loop)
+
+            # Asyncio would return a closed socket, which we
+            # can't do: the transport was aborted, hence there
+            # is no FD to attach a socket to (to make
+            # get_extra_info() work).
+            self.assertIsNone(tr.get_extra_info('socket'))
+            tr.close()
+
+        self.loop.run_until_complete(runner())
+        srv.close()
+        self.loop.run_until_complete(srv.wait_closed())
+
 
 class Test_AIO_TCP(_TestTCP, tb.AIOTestCase):
     pass

uvloop/handles/basetransport.pxd

@@ -26,6 +26,7 @@ cdef class UVBaseTransport(UVSocketHandle):
     cdef inline _schedule_call_connection_made(self)
     cdef inline _schedule_call_connection_lost(self, exc)
 
+    cdef _wakeup_waiter(self)
     cdef _call_connection_made(self)
     cdef _call_connection_lost(self, exc)
 

uvloop/handles/basetransport.pyx

@@ -114,16 +114,34 @@ cdef class UVBaseTransport(UVSocketHandle):
                     'protocol': self._protocol,
                 })
 
+    cdef _wakeup_waiter(self):
+        if self._waiter is not None:
+            if not self._waiter.cancelled():
+                self._waiter.set_result(True)
+            self._waiter = None
+
     cdef _call_connection_made(self):
         cdef Py_ssize_t _loop_ready_len
         if self._protocol is None:
             raise RuntimeError(
                 'protocol is not set, cannot call connection_made()')
 
         _loop_ready_len = self._loop._ready_len
-        self._protocol.connection_made(self)
+
+        try:
+            self._protocol.connection_made(self)
+        except:
+            self._wakeup_waiter()
+            raise
+
         self._protocol_connected = 1
 
+        if self._closing:
+            # This might happen when "transport.abort()" is called
+            # from "Protocol.connection_made".
+            self._wakeup_waiter()
+            return
+
         if _loop_ready_len == self._loop._ready_len:
             # No new calls were scheduled by 'protocol.connection_made',
             # so it's safe to start reading right now.
@@ -142,10 +160,7 @@ cdef class UVBaseTransport(UVSocketHandle):
                                  <method_t>self._start_reading,
                                  self))
 
-        if self._waiter is not None:
-            if not self._waiter.cancelled():
-                self._waiter.set_result(True)
-            self._waiter = None
+        self._wakeup_waiter()
 
     cdef _call_connection_lost(self, exc):
         if self._waiter is not None:

uvloop/handles/handle.pyx

@@ -221,6 +221,7 @@ cdef class UVSocketHandle(UVHandle):
             int fd
             int err
 
+        self._ensure_alive()
         err = uv.uv_fileno(self._handle, <uv.uv_os_fd_t*>&fd)
         if err < 0:
             raise convert_error(err)
@@ -234,6 +235,9 @@ cdef class UVSocketHandle(UVHandle):
         if self.__cached_socket is not None:
             return self.__cached_socket
 
+        if not self._is_alive():
+            return None
+
         self.__cached_socket = self._new_socket()
         if self.__cached_socket.fileno() == self._fileno():
             raise RuntimeError('new socket shares fileno with the transport')