tcp: Cache peername & sockname in TCPTransport._call_connection_made

Author: 1st1

tests/test_tcp.py

@@ -362,13 +362,14 @@ async def test_client(addr):
             self.assertFalse(t._paused)
 
             sock = t.get_extra_info('socket')
-            self.assertTrue(isinstance(sock, socket.socket))
+            sockname = sock.getsockname()
+            peername = sock.getpeername()
 
+            self.assertTrue(isinstance(sock, socket.socket))
             self.assertEqual(t.get_extra_info('sockname'),
-                             sock.getsockname())
-
+                             sockname)
             self.assertEqual(t.get_extra_info('peername'),
-                             sock.getpeername())
+                             peername)
 
             t.write(b'OK')  # We want server to fail.
 
@@ -378,6 +379,13 @@ async def test_client(addr):
 
             await fut
 
+            # Test that peername and sockname are available after
+            # the transport is closed.
+            self.assertEqual(t.get_extra_info('peername'),
+                             peername)
+            self.assertEqual(t.get_extra_info('sockname'),
+                             sockname)
+
         async def start_server():
             srv = await asyncio.start_server(
                 handle_client,

tests/test_udp.py

@@ -112,7 +112,7 @@ def test_create_datagram_endpoint_sock(self):
             assert False, 'Can not create socket.'
 
         with sock:
-            f = self.loop.create_connection(
+            f = self.loop.create_datagram_endpoint(
                 lambda: MyDatagramProto(loop=self.loop), sock=sock)
             tr, pr = self.loop.run_until_complete(f)
             self.assertIsInstance(pr, MyDatagramProto)

uvloop/handles/tcp.pxd

@@ -8,6 +8,12 @@ cdef class TCPServer(UVStreamServer):
 
 
 cdef class TCPTransport(UVStream):
+    cdef:
+        bint __peername_set
+        bint __sockname_set
+        system.sockaddr_storage __peername
+        system.sockaddr_storage __sockname
+
     cdef bind(self, system.sockaddr* addr, unsigned int flags=*)
     cdef open(self, int sockfd)
     cdef connect(self, system.sockaddr* addr)

uvloop/handles/tcp.pyx

@@ -102,8 +102,53 @@ cdef class TCPTransport(UVStream):
         handle = TCPTransport.__new__(TCPTransport)
         handle._init(loop, protocol, server, waiter)
         __tcp_init_uv_handle(<UVStream>handle, loop)
+        handle.__peername_set = 0
+        handle.__sockname_set = 0
         return handle
 
+    cdef _call_connection_made(self):
+        # asyncio saves peername & sockname when transports are instantiated,
+        # so that they're accessible even after the transport is closed.
+        # We are doing the same thing here, except that we create Python
+        # objects lazily, on request in get_extra_info()
+
+        cdef:
+            int err
+            int buf_len
+
+        buf_len = sizeof(system.sockaddr_storage)
+        err = uv.uv_tcp_getsockname(<uv.uv_tcp_t*>self._handle,
+                                    <system.sockaddr*>&self.__sockname,
+                                    &buf_len)
+        if err >= 0:
+            # Ignore errors, this is an optional thing.
+            # If something serious is going on, the transport
+            # will crash later (in roughly the same way how
+            # an asyncio transport would.)
+            self.__sockname_set = 1
+
+        buf_len = sizeof(system.sockaddr_storage)
+        err = uv.uv_tcp_getpeername(<uv.uv_tcp_t*>self._handle,
+                                    <system.sockaddr*>&self.__peername,
+                                    &buf_len)
+        if err >= 0:
+            # Same as few lines above -- we don't really care
+            # about error case here.
+            self.__peername_set = 1
+
+        UVBaseTransport._call_connection_made(self)
+
+    def get_extra_info(self, name, default=None):
+        if name == 'sockname':
+            if self.__sockname_set:
+                return __convert_sockaddr_to_pyaddr(
+                    <system.sockaddr*>&self.__sockname)
+        elif name == 'peername':
+            if self.__peername_set:
+                return __convert_sockaddr_to_pyaddr(
+                    <system.sockaddr*>&self.__peername)
+        return super().get_extra_info(name, default)
+
     cdef _new_socket(self):
         return __tcp_get_socket(<UVSocketHandle>self)
 

uvloop/includes/uv.pxd

@@ -328,6 +328,8 @@ cdef extern from "../vendor/libuv/include/uv.h" nogil:
 
     int uv_tcp_getsockname(const uv_tcp_t* handle, system.sockaddr* name,
                            int* namelen)
+    int uv_tcp_getpeername(const uv_tcp_t* handle, system.sockaddr* name,
+                           int* namelen)
 
     int uv_tcp_connect(uv_connect_t* req, uv_tcp_t* handle,
                        const system.sockaddr* addr, uv_connect_cb cb)