PYTHON-5867 - Close sockets on interruption or cancellation during async connection creation

pymongo/pool_shared.py

@@ -21,9 +21,12 @@
 import socket
 import ssl
 import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
 from typing import (
     TYPE_CHECKING,
     Any,
+    Callable,
     NoReturn,
     Optional,
     Union,
@@ -155,6 +158,20 @@ def _raise_connection_failure(
         raise AutoReconnect(msg) from error
 
 
+@contextmanager
+def _cleanup_on_error(cleanup: Callable[[], None]) -> Iterator[None]:
+    """Invoke `cleanup` if the context block exits with any exception.
+
+    Used on async connect paths to prevent raw socket/transport
+    leaks when the coroutine is canceled mid-handshake.
+    """
+    try:
+        yield
+    except BaseException:
+        cleanup()
+        raise
+
+
 class _CancellationContext:
     def __init__(self) -> None:
         self._cancelled = False
@@ -185,13 +202,10 @@ async def _async_create_connection(address: _Address, options: PoolOptions) -> s
         sock = socket.socket(socket.AF_UNIX)
         # SOCK_CLOEXEC not supported for Unix sockets.
         _set_non_inheritable_non_atomic(sock.fileno())
-        try:
+        with _cleanup_on_error(sock.close):
             sock.setblocking(False)
             await asyncio.get_running_loop().sock_connect(sock, host)
             return sock
-        except OSError:
-            sock.close()
-            raise
 
     # Don't try IPv6 if we don't support it. Also skip it if host
     # is 'localhost' (::1 is fine). Avoids slow connect issues
@@ -215,28 +229,27 @@ async def _async_create_connection(address: _Address, options: PoolOptions) -> s
         # Fallback when SOCK_CLOEXEC isn't available.
         _set_non_inheritable_non_atomic(sock.fileno())
         try:
-            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-            # CSOT: apply timeout to socket connect.
-            timeout = _csot.remaining()
-            if timeout is None:
-                timeout = options.connect_timeout
-            elif timeout <= 0:
-                raise socket.timeout("timed out")
-            sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True)
-            _set_keepalive_times(sock)
-            # Socket needs to be non-blocking during connection to not block the event loop
-            sock.setblocking(False)
-            await asyncio.wait_for(
-                asyncio.get_running_loop().sock_connect(sock, sa), timeout=timeout
-            )
-            sock.settimeout(timeout)
-            return sock
+            with _cleanup_on_error(sock.close):
+                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+                # CSOT: apply timeout to socket connect.
+                timeout = _csot.remaining()
+                if timeout is None:
+                    timeout = options.connect_timeout
+                elif timeout <= 0:
+                    raise socket.timeout("timed out")
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, True)
+                _set_keepalive_times(sock)
+                # Socket needs to be non-blocking during connection to not block the event loop
+                sock.setblocking(False)
+                await asyncio.wait_for(
+                    asyncio.get_running_loop().sock_connect(sock, sa), timeout=timeout
+                )
+                sock.settimeout(timeout)
+                return sock
         except asyncio.TimeoutError as e:
-            sock.close()
             err = socket.timeout("timed out")
             err.__cause__ = e
         except OSError as e:
-            sock.close()
             err = e  # type: ignore[assignment]
 
     if err is not None:
@@ -266,42 +279,37 @@ async def _async_configured_socket(
 
     host = address[0]
     try:
-        # We have to pass hostname / ip address to wrap_socket
-        # to use SSLContext.check_hostname.
-        if _has_sni(False):
-            loop = asyncio.get_running_loop()
-            ssl_sock = await loop.run_in_executor(
-                None,
-                functools.partial(ssl_context.wrap_socket, sock, server_hostname=host),  # type: ignore[assignment, misc, unused-ignore]
-            )
-        else:
-            loop = asyncio.get_running_loop()
-            ssl_sock = await loop.run_in_executor(None, ssl_context.wrap_socket, sock)  # type: ignore[assignment, misc, unused-ignore]
+        with _cleanup_on_error(sock.close):
+            # We have to pass hostname / ip address to wrap_socket
+            # to use SSLContext.check_hostname.
+            if _has_sni(False):
+                loop = asyncio.get_running_loop()
+                ssl_sock = await loop.run_in_executor(
+                    None,
+                    functools.partial(ssl_context.wrap_socket, sock, server_hostname=host),  # type: ignore[assignment, misc, unused-ignore]
+                )
+            else:
+                loop = asyncio.get_running_loop()
+                ssl_sock = await loop.run_in_executor(None, ssl_context.wrap_socket, sock)  # type: ignore[assignment, misc, unused-ignore]
     except _CertificateError:
-        sock.close()
-        # Raise _CertificateError directly like we do after match_hostname
-        # below.
+        # Raise _CertificateError directly like we do after match_hostname below.
         raise
     except (OSError, *SSLErrors) as exc:
-        sock.close()
         # We raise AutoReconnect for transient and permanent SSL handshake
         # failures alike. Permanent handshake failures, like protocol
         # mismatch, will be turned into ServerSelectionTimeoutErrors later.
         details = _get_timeout_details(options)
         _raise_connection_failure(address, exc, "SSL handshake failed: ", timeout_details=details)
-    if (
-        ssl_context.verify_mode
-        and not ssl_context.check_hostname
-        and not options.tls_allow_invalid_hostnames
-    ):
-        try:
+    with _cleanup_on_error(ssl_sock.close):
+        if (
+            ssl_context.verify_mode
+            and not ssl_context.check_hostname
+            and not options.tls_allow_invalid_hostnames
+        ):
             ssl.match_hostname(ssl_sock.getpeercert(), hostname=host)  # type:ignore[attr-defined, unused-ignore]
-        except _CertificateError:
-            ssl_sock.close()
-            raise
 
-    ssl_sock.settimeout(options.socket_timeout)
-    return ssl_sock
+        ssl_sock.settimeout(options.socket_timeout)
+        return ssl_sock
 
 
 async def _configured_protocol_interface(
@@ -320,11 +328,12 @@ async def _configured_protocol_interface(
     timeout = options.socket_timeout
 
     if ssl_context is None:
-        return AsyncNetworkingInterface(
-            await asyncio.get_running_loop().create_connection(
-                lambda: PyMongoProtocol(timeout=timeout), sock=sock
+        with _cleanup_on_error(sock.close):
+            return AsyncNetworkingInterface(
+                await asyncio.get_running_loop().create_connection(
+                    lambda: PyMongoProtocol(timeout=timeout), sock=sock
+                )
             )
-        )
 
     host = address[0]
     # asyncio does not support TLS session resumption natively (cpython#79152,
@@ -346,12 +355,13 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
     try:
         # We have to pass hostname / ip address to wrap_socket
         # to use SSLContext.check_hostname.
-        transport, protocol = await asyncio.get_running_loop().create_connection(  # type: ignore[call-overload]
-            lambda: PyMongoProtocol(timeout=timeout),
-            sock=sock,
-            server_hostname=host,
-            ssl=ssl_context,
-        )
+        with _cleanup_on_error(sock.close):
+            transport, protocol = await asyncio.get_running_loop().create_connection(  # type: ignore[call-overload]
+                lambda: PyMongoProtocol(timeout=timeout),
+                sock=sock,
+                server_hostname=host,
+                ssl=ssl_context,
+            )
     except _CertificateError:
         # Raise _CertificateError directly like we do after match_hostname
         # below.
@@ -362,26 +372,22 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # mismatch, will be turned into ServerSelectionTimeoutErrors later.
         details = _get_timeout_details(options)
         _raise_connection_failure(address, exc, "SSL handshake failed: ", timeout_details=details)
-
-    if (
-        ssl_context.verify_mode
-        and not ssl_context.check_hostname
-        and not options.tls_allow_invalid_hostnames
-    ):
-        try:
+    with _cleanup_on_error(transport.abort):
+        if (
+            ssl_context.verify_mode
+            and not ssl_context.check_hostname
+            and not options.tls_allow_invalid_hostnames
+        ):
             ssl.match_hostname(transport.get_extra_info("peercert"), hostname=host)  # type:ignore[attr-defined,unused-ignore]
-        except _CertificateError:
-            transport.abort()
-            raise
 
-    if ssl_session_cache is not None:
-        ssl_obj = transport.get_extra_info("ssl_object")
-        if ssl_obj is not None:
-            new_session = ssl_obj.session
-            if new_session is not None:
-                ssl_session_cache[0] = new_session
+        if ssl_session_cache is not None:
+            ssl_obj = transport.get_extra_info("ssl_object")
+            if ssl_obj is not None:
+                new_session = ssl_obj.session
+                if new_session is not None:
+                    ssl_session_cache[0] = new_session
 
-    return AsyncNetworkingInterface((transport, protocol))
+        return AsyncNetworkingInterface((transport, protocol))
 
 
 def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:

test/asynchronous/test_async_cancellation.py

@@ -17,16 +17,38 @@
 from __future__ import annotations
 
 import asyncio
+import functools
+import socket
+import ssl
 import sys
+from unittest.mock import patch
 
 from test.asynchronous.utils import async_get_pool
 from test.utils_shared import delay, one
 
 sys.path[0:0] = [""]
 
+from pymongo import pool_shared
+from pymongo.pool_options import PoolOptions
 from test.asynchronous import AsyncIntegrationTest, async_client_context, connected
 
 
+class _SocketLeakTracker:
+    """Track sockets created inside a target task to verify they're closed on cancellation."""
+
+    def __init__(self):
+        self.sockets = []
+        self.target_task = None
+        self.started = asyncio.Event()
+        self._socket_class = socket.socket
+
+    def track(self, *args, **kwargs):
+        s = self._socket_class(*args, **kwargs)
+        if asyncio.current_task() is self.target_task:
+            self.sockets.append(s)
+        return s
+
+
 class TestAsyncCancellation(AsyncIntegrationTest):
     async def test_async_cancellation_closes_connection(self):
         pool = await async_get_pool(self.client)
@@ -129,3 +151,96 @@ async def task():
             await task
 
         self.assertTrue(change_stream._closed)
+
+    async def _assert_cancel_closes_tracked_sockets(self, tracker, coro):
+        task = asyncio.create_task(coro)
+        tracker.target_task = task
+        await asyncio.wait_for(tracker.started.wait(), timeout=5)
+        task.cancel()
+        with self.assertRaises(asyncio.CancelledError):
+            await task
+        self.assertTrue(tracker.sockets, "expected at least one socket to be created")
+        for sock in tracker.sockets:
+            self.assertEqual(
+                sock.fileno(),
+                -1,
+                f"socket leaked across cancellation: {sock!r}",
+            )
+
+    async def test_cancellation_closes_socket_during_create_connection(self):
+        address = (await async_client_context.host, await async_client_context.port)
+        options = (await async_get_pool(self.client)).opts
+        tracker = _SocketLeakTracker()
+        loop = asyncio.get_running_loop()
+        real_sock_connect = loop.sock_connect
+        block_forever = asyncio.Event()
+
+        async def slow_sock_connect(sock, addr):
+            if sock in tracker.sockets:
+                tracker.started.set()
+                await block_forever.wait()
+                return None
+            return await real_sock_connect(sock, addr)
+
+        with (
+            patch.object(socket, "socket", tracker.track),
+            patch.object(loop, "sock_connect", slow_sock_connect),
+        ):
+            await self._assert_cancel_closes_tracked_sockets(
+                tracker,
+                pool_shared._async_create_connection(address, options),
+            )
+
+    async def test_cancellation_closes_socket_during_ssl_wrap_socket(self):
+        address = (await async_client_context.host, await async_client_context.port)
+        fake_ssl_context = ssl.create_default_context()
+        options = PoolOptions(ssl_context=fake_ssl_context)  # type: ignore[arg-type]
+        tracker = _SocketLeakTracker()
+        loop = asyncio.get_running_loop()
+        real_run_in_executor = loop.run_in_executor
+
+        def slow_run_in_executor(executor, func, *args):
+            # Need to unwrap the SNI branch here if present
+            inner = func.func if isinstance(func, functools.partial) else func
+            # Each `ctx.wrap_socket` access returns a fresh bound-method
+            # object, so we check the bound instance (__self__) instead
+            if (
+                getattr(inner, "__self__", None) is fake_ssl_context
+                and asyncio.current_task() is tracker.target_task
+            ):
+                tracker.started.set()
+                # Return a future that never completes for cancellation.
+                return asyncio.get_running_loop().create_future()
+            return real_run_in_executor(executor, func, *args)
+
+        with (
+            patch.object(socket, "socket", tracker.track),
+            patch.object(loop, "run_in_executor", slow_run_in_executor),
+        ):
+            await self._assert_cancel_closes_tracked_sockets(
+                tracker,
+                pool_shared._async_configured_socket(address, options),
+            )
+
+    async def test_cancellation_closes_socket_during_configured_protocol_interface(self):
+        address = (await async_client_context.host, await async_client_context.port)
+        options = PoolOptions()
+        tracker = _SocketLeakTracker()
+        loop = asyncio.get_running_loop()
+        real_create_connection = loop.create_connection
+        block_forever = asyncio.Event()
+
+        async def slow_create_connection(*args, **kwargs):
+            if asyncio.current_task() is tracker.target_task:
+                tracker.started.set()
+                await block_forever.wait()
+            return await real_create_connection(*args, **kwargs)
+
+        with (
+            patch.object(socket, "socket", tracker.track),
+            patch.object(loop, "create_connection", slow_create_connection),
+        ):
+            await self._assert_cancel_closes_tracked_sockets(
+                tracker,
+                pool_shared._configured_protocol_interface(address, options),
+            )