diff --git a/pymongo/pool_shared.py b/pymongo/pool_shared.py index c97b0eb217..fd1384a655 100644 --- a/pymongo/pool_shared.py +++ b/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: diff --git a/test/asynchronous/test_async_cancellation.py b/test/asynchronous/test_async_cancellation.py index a96e28d832..f435d5fc9e 100644 --- a/test/asynchronous/test_async_cancellation.py +++ b/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), + )