gh-155305: Preserve user-provided sockets on transport errors

Author: RealBhupesh

Lib/asyncio/base_events.py

@@ -1083,6 +1083,7 @@ async def create_connection(
         connection in the background.  When successful, the coroutine
         returns a (transport, protocol) pair.
         """
+        sock_was_provided = sock is not None
         if server_hostname is not None and not ssl:
             raise ValueError('server_hostname is only meaningful with ssl')
 
@@ -1204,7 +1205,8 @@ async def create_connection(
         transport, protocol = await self._create_connection_transport(
             sock, protocol_factory, ssl, server_hostname,
             ssl_handshake_timeout=ssl_handshake_timeout,
-            ssl_shutdown_timeout=ssl_shutdown_timeout)
+            ssl_shutdown_timeout=ssl_shutdown_timeout,
+            sock_was_provided=sock_was_provided)
         if self._debug:
             # Get the socket from the transport because SSL transport closes
             # the old socket and creates a new SSL socket
@@ -1217,7 +1219,8 @@ async def _create_connection_transport(
             self, sock, protocol_factory, ssl,
             server_hostname, server_side=False,
             ssl_handshake_timeout=None,
-            ssl_shutdown_timeout=None, context=None):
+            ssl_shutdown_timeout=None, context=None,
+            sock_was_provided=False):
 
         try:
             sock.setblocking(False)
@@ -1236,8 +1239,10 @@ async def _create_connection_transport(
             else:
                 transport = self._make_socket_transport(sock, protocol, waiter, context=context)
         except:
-            # gh-153133: close the socket if the transport is never created.
-            sock.close()
+            # gh-153133: close internally created sockets if the transport is
+            # never created.
+            if not sock_was_provided:
+                sock.close()
             raise
 
         try:
@@ -1705,7 +1710,8 @@ async def connect_accepted_socket(
         transport, protocol = await self._create_connection_transport(
             sock, protocol_factory, ssl, '', server_side=True,
             ssl_handshake_timeout=ssl_handshake_timeout,
-            ssl_shutdown_timeout=ssl_shutdown_timeout)
+            ssl_shutdown_timeout=ssl_shutdown_timeout,
+            sock_was_provided=True)
         if self._debug:
             # Get the socket from the transport because SSL transport closes
             # the old socket and creates a new SSL socket

Lib/asyncio/unix_events.py

@@ -227,6 +227,7 @@ async def create_unix_connection(
             server_hostname=None,
             ssl_handshake_timeout=None,
             ssl_shutdown_timeout=None):
+        sock_was_provided = sock is not None
         assert server_hostname is None or isinstance(server_hostname, str)
         if ssl:
             if server_hostname is None:
@@ -268,7 +269,8 @@ async def create_unix_connection(
         transport, protocol = await self._create_connection_transport(
             sock, protocol_factory, ssl, server_hostname,
             ssl_handshake_timeout=ssl_handshake_timeout,
-            ssl_shutdown_timeout=ssl_shutdown_timeout)
+            ssl_shutdown_timeout=ssl_shutdown_timeout,
+            sock_was_provided=sock_was_provided)
         return transport, protocol
 
     async def create_unix_server(

Lib/test/test_asyncio/test_base_events.py

@@ -1341,9 +1341,9 @@ def getaddrinfo(*args, **kw):
                 self.loop.run_until_complete(coro)
             self.assertTrue(sock.close.called)
 
-    def test_create_connection_sock_transport_error_closes_sock(self):
-        # gh-153133: a user-provided socket is closed if the transport is
-        # never created.
+    def test_create_connection_sock_transport_error_does_not_close_sock(self):
+        # gh-155305: a user-provided socket remains owned by the caller when
+        # the transport is never created.
         sock = mock.Mock()
         sock.type = socket.SOCK_STREAM
 
@@ -1353,7 +1353,7 @@ def factory():
         coro = self.loop.create_connection(factory, sock=sock)
         with self.assertRaises(ZeroDivisionError):
             self.loop.run_until_complete(coro)
-        self.assertTrue(sock.close.called)
+        self.assertFalse(sock.close.called)
 
     @patch_socket
     def test_create_connection_transport_error_closes_sock(self, m_socket):

Lib/test/test_asyncio/test_events.py

@@ -914,6 +914,18 @@ def test_connect_accepted_socket_ssl_timeout_for_plain_socket(self):
                 'ssl_handshake_timeout is only meaningful with ssl'):
             self.loop.run_until_complete(coro)
 
+    def test_connect_accepted_socket_transport_error_does_not_close_sock(self):
+        sock = mock.Mock()
+        sock.type = socket.SOCK_STREAM
+
+        def factory():
+            raise ZeroDivisionError
+
+        coro = self.loop.connect_accepted_socket(factory, sock)
+        with self.assertRaises(ZeroDivisionError):
+            self.loop.run_until_complete(coro)
+        self.assertFalse(sock.close.called)
+
     @mock.patch('asyncio.base_events.socket')
     def create_server_multiple_hosts(self, family, hosts, mock_sock):
         async def getaddrinfo(host, port, *args, **kw):

Lib/test/test_asyncio/test_unix_events.py

@@ -393,6 +393,19 @@ def test_create_unix_connection_path_inetsock(self):
                                         'A UNIX Domain Stream.*was expected'):
                 self.loop.run_until_complete(coro)
 
+    def test_create_unix_connection_transport_error_does_not_close_sock(self):
+        sock = mock.Mock()
+        sock.family = socket.AF_UNIX
+        sock.type = socket.SOCK_STREAM
+
+        def factory():
+            raise ZeroDivisionError
+
+        coro = self.loop.create_unix_connection(factory, sock=sock)
+        with self.assertRaises(ZeroDivisionError):
+            self.loop.run_until_complete(coro)
+        self.assertFalse(sock.close.called)
+
     @mock.patch('asyncio.unix_events.socket')
     def test_create_unix_server_bind_error(self, m_socket):
         # Ensure that the socket is closed on any bind error