From 015a4f57bcd307b12ff0109cb0a010babdb8024c Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 25 Sep 2026 07:25:04 -0500 Subject: [PATCH 1/3] PYTHON-5805 Split TLS wrapping out of configured socket helpers _extract _async_wrap_socket_tls / _wrap_socket_tls so TLS can be applied to sockets obtained from a KMS connect callback, and shield the executor handshake so cancellation cannot orphan the wrapped socket. --- pymongo/pool_shared.py | 72 +++++++++++++++++++++++++------ test/asynchronous/test_pooling.py | 21 ++++++++- test/test_pooling.py | 21 ++++++++- 3 files changed, 98 insertions(+), 16 deletions(-) diff --git a/pymongo/pool_shared.py b/pymongo/pool_shared.py index 8cd546bda6..987f2f059a 100644 --- a/pymongo/pool_shared.py +++ b/pymongo/pool_shared.py @@ -304,16 +304,25 @@ async def _async_create_connection(address: _Address, options: PoolOptions) -> s raise OSError("getaddrinfo failed") -async def _async_configured_socket( - address: _Address, options: PoolOptions +def _close_late_socket(future: asyncio.Future[Any]) -> None: + """Close a socket produced after its awaiting task was cancelled.""" + if not future.cancelled() and future.exception() is None: + future.result().close() + + +async def _async_wrap_socket_tls( + sock: socket.socket, address: _Address, options: PoolOptions ) -> Union[socket.socket, _sslConn]: - """Given (host, port) and PoolOptions, return a raw configured socket. + """Given a connected socket, (host, port), and PoolOptions, apply TLS. + + The handshake, SNI, and certificate/hostname verification all target + ``address``, which may differ from the peer ``sock`` is connected to, e.g. + when ``sock`` tunnels through an HTTP proxy. Can raise socket.error, ConnectionFailure, or _CertificateError. - Sets socket's SSL and timeout options. + Sets the socket's SSL and timeout options. """ - sock = await _async_create_connection(address, options) ssl_context = options._ssl_context if ssl_context is None: @@ -326,13 +335,19 @@ async def _async_configured_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] - ) + wrap = 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] + wrap = functools.partial(ssl_context.wrap_socket, sock) # type: ignore[assignment, misc, unused-ignore] + # Shield the executor future: wrap_socket hands the fd to a new + # SSLSocket, so cancellation must not orphan the result it produces. + future = loop.run_in_executor(None, wrap) + try: + ssl_sock = await asyncio.shield(future) + except asyncio.CancelledError: + future.add_done_callback(_close_late_socket) + sock.close() + raise except _CertificateError: sock.close() # Raise _CertificateError directly like we do after match_hostname @@ -360,6 +375,19 @@ async def _async_configured_socket( return ssl_sock +async def _async_configured_socket( + address: _Address, options: PoolOptions +) -> Union[socket.socket, _sslConn]: + """Given (host, port) and PoolOptions, return a raw configured socket. + + Can raise socket.error, ConnectionFailure, or _CertificateError. + + Sets socket's SSL and timeout options. + """ + sock = await _async_create_connection(address, options) + return await _async_wrap_socket_tls(sock, address, options) + + async def _configured_protocol_interface( address: _Address, options: PoolOptions, @@ -510,14 +538,19 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket raise OSError("getaddrinfo failed") -def _configured_socket(address: _Address, options: PoolOptions) -> Union[socket.socket, _sslConn]: - """Given (host, port) and PoolOptions, return a raw configured socket. +def _wrap_socket_tls( + sock: socket.socket, address: _Address, options: PoolOptions +) -> Union[socket.socket, _sslConn]: + """Given a connected socket, (host, port), and PoolOptions, apply TLS. + + The handshake, SNI, and certificate/hostname verification all target + ``address``, which may differ from the peer ``sock`` is connected to, e.g. + when ``sock`` tunnels through an HTTP proxy. Can raise socket.error, ConnectionFailure, or _CertificateError. - Sets socket's SSL and timeout options. + Sets the socket's SSL and timeout options. """ - sock = _create_connection(address, options) ssl_context = options._ssl_context if ssl_context is None: @@ -559,6 +592,17 @@ def _configured_socket(address: _Address, options: PoolOptions) -> Union[socket. return ssl_sock +def _configured_socket(address: _Address, options: PoolOptions) -> Union[socket.socket, _sslConn]: + """Given (host, port) and PoolOptions, return a raw configured socket. + + Can raise socket.error, ConnectionFailure, or _CertificateError. + + Sets socket's SSL and timeout options. + """ + sock = _create_connection(address, options) + return _wrap_socket_tls(sock, address, options) + + def _configured_socket_interface( address: _Address, options: PoolOptions, diff --git a/test/asynchronous/test_pooling.py b/test/asynchronous/test_pooling.py index 661bd4e3d2..ed1e0558e0 100644 --- a/test/asynchronous/test_pooling.py +++ b/test/asynchronous/test_pooling.py @@ -34,13 +34,19 @@ from pymongo.hello import HelloCompat from pymongo.lock import _async_create_lock from pymongo.monitoring import _EventListeners +from pymongo.pool_shared import _async_wrap_socket_tls from test.asynchronous.utils import async_get_pool, async_joinall, flaky sys.path[0:0] = [""] from pymongo.asynchronous.pool import Pool, PoolOptions from pymongo.socket_checker import SocketChecker -from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest +from test.asynchronous import ( + AsyncIntegrationTest, + AsyncPyMongoTestCase, + async_client_context, + unittest, +) from test.asynchronous.helpers import ConcurrentRunner from test.utils_shared import CMAPListener, delay @@ -824,5 +830,18 @@ def test_certificate_error_is_not_labeled_overloaded(self): self.assertFalse(err.has_error_label("SystemOverloadedError")) +class TestWrapSocketTLS(AsyncPyMongoTestCase): + async def test_wrap_socket_tls_without_ssl_context_returns_same_socket(self): + options = PoolOptions(socket_timeout=7.5) + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + result = await _async_wrap_socket_tls(left, ("kms.example.com", 443), options) + + self.assertIs(result, left) + self.assertEqual(result.gettimeout(), 7.5) + + if __name__ == "__main__": unittest.main() diff --git a/test/test_pooling.py b/test/test_pooling.py index a3f0eaf589..f3df350bd1 100644 --- a/test/test_pooling.py +++ b/test/test_pooling.py @@ -34,13 +34,19 @@ from pymongo.hello import HelloCompat from pymongo.lock import _create_lock from pymongo.monitoring import _EventListeners +from pymongo.pool_shared import _wrap_socket_tls from test.utils import flaky, get_pool, joinall sys.path[0:0] = [""] from pymongo.socket_checker import SocketChecker from pymongo.synchronous.pool import Pool, PoolOptions -from test import IntegrationTest, client_context, unittest +from test import ( + IntegrationTest, + PyMongoTestCase, + client_context, + unittest, +) from test.helpers import ConcurrentRunner from test.utils_shared import CMAPListener, delay @@ -822,5 +828,18 @@ def test_certificate_error_is_not_labeled_overloaded(self): self.assertFalse(err.has_error_label("SystemOverloadedError")) +class TestWrapSocketTLS(PyMongoTestCase): + def test_wrap_socket_tls_without_ssl_context_returns_same_socket(self): + options = PoolOptions(socket_timeout=7.5) + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + result = _wrap_socket_tls(left, ("kms.example.com", 443), options) + + self.assertIs(result, left) + self.assertEqual(result.gettimeout(), 7.5) + + if __name__ == "__main__": unittest.main() From 676c0ef2289708146554bd071fe91e1e000fa8e7 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 25 Sep 2026 07:25:07 -0500 Subject: [PATCH 2/3] PYTHON-5805 Add kms_connect_callback and HTTP proxy connect helpers Add KMSConnectContext and the kms_connect_callback option to AutoEncryptionOpts and ClientEncryption so callers can route KMS connections through an HTTP proxy. The driver performs the KMS TLS handshake over the returned socket, so verification still targets the KMS host, and CSOT deadlines cover the callback. For ordinary proxies, callers can pass the new HTTPProxyKMSConnect or AsyncHTTPProxyKMSConnect helper instead of writing a callback. Enforce the CSOT deadline across the proxy tunnel, relay, and TLS handshake, and make the async callback contract strict: coroutine functions for the async API, plain callables rejected. --- doc/changelog.rst | 11 ++ pymongo/asynchronous/encryption.py | 150 ++++++++++++++-- pymongo/encryption_options.py | 277 ++++++++++++++++++++++++++++- pymongo/synchronous/encryption.py | 151 ++++++++++++++-- tools/synchro.py | 15 ++ 5 files changed, 579 insertions(+), 25 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 08d3908e62..de52cbfc00 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -16,6 +16,17 @@ PyMongo 4.19 brings a number of changes including: interpreter remain daemon threads and shutdown behavior is unchanged. Note that because these threads are non-daemon, a subinterpreter may block on teardown until any in-flight monitor work completes. +- Added support for routing Key Management Service (KMS) requests for + Client-Side Field Level Encryption and Queryable Encryption through an HTTP + proxy, using the new ``kms_connect_callback`` option on + :class:`~pymongo.encryption_options.AutoEncryptionOpts`, + :class:`~pymongo.encryption.ClientEncryption`, and + :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`. The callback + opens the connection and the driver performs the KMS TLS handshake over it, so + verification still targets the KMS host rather than the proxy. For an ordinary + HTTP proxy, pass :class:`~pymongo.encryption_options.HTTPProxyKMSConnect` or + :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect` instead of + writing a callback. Bug fixes ......... diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 9ba2758f78..bb0488d359 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -17,8 +17,11 @@ from __future__ import annotations import asyncio +import contextlib import functools +import inspect import socket +import ssl import time as time # noqa: PLC0414 # needed in sync version import uuid import weakref @@ -60,7 +63,9 @@ from pymongo.common import CONNECT_TIMEOUT from pymongo.daemon import _spawn_daemon from pymongo.encryption_options import ( + AsyncKMSConnectCallback, AutoEncryptionOpts, + KMSConnectContext, RangeOpts, StringOpts, # Re-exported for backwards compatibility: TextOpts is deprecated but must @@ -90,6 +95,8 @@ from pymongo.pool_options import PoolOptions from pymongo.pool_shared import ( _async_configured_socket, + _async_wrap_socket_tls, + _close_late_socket, _raise_connection_failure, ) from pymongo.read_concern import ReadConcern @@ -120,11 +127,108 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -async def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]: +def _close_rejected_kms_socket(obj: Any) -> None: + """Close a rejected kms_connect_callback return value, best effort. + + Nothing else will close it: _connect_kms raises before the result reaches + the caller's ``finally``. + """ + close = getattr(obj, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + +async def _connect_kms( + address: _Address, + opts: PoolOptions, + kms_connect_callback: Optional[AsyncKMSConnectCallback], + timeout: float, +) -> Union[socket.socket, _sslConn]: + """Connect to a KMS host and perform the TLS handshake over the socket. + + Uses ``kms_connect_callback`` when one is provided, otherwise connects + directly, and always verifies against ``address`` (the KMS host). + """ + if kms_connect_callback is None: + try: + return await _async_configured_socket(address, opts) + except Exception as exc: + _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) + + # TLS targets address, not the peer, so verification follows the KMS host. + # A plain callable would block the event loop before we could reject it, + # so check the callback first. + if not _IS_SYNC: + callback_any: Any = kms_connect_callback + is_coro = inspect.iscoroutinefunction(callback_any) + if not is_coro and callable(callback_any): + is_coro = inspect.iscoroutinefunction(callback_any.__call__) + if not is_coro: + raise ConfigurationError( + "kms_connect_callback must be a coroutine function for the async API." + ) + # Typed as Any so the generated synchronous flavor type-checks: the sync + # callback returns a plain socket, which is not awaitable. + result: Any = kms_connect_callback( + KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) + ) + remaining = _csot.remaining() + if remaining is None or _IS_SYNC: + # The synchronous API cannot interrupt a callback that has started + # running; honoring the deadline is the callback's contract there. + sock = await result + else: + # CSOT is cooperative: a callback that ignores the timeout could block + # past the deadline. Shield the task so stopping the wait does not + # cancel it mid-flight, and close any socket it yields later. + task = asyncio.ensure_future(result) + try: + sock = await asyncio.wait_for(asyncio.shield(task), remaining) + except asyncio.CancelledError: + task.add_done_callback(_close_late_socket) + raise + except asyncio.TimeoutError: + task.add_done_callback(_close_late_socket) + _raise_connection_failure( + address, + socket.timeout("timed out"), + timeout_details=_get_timeout_details(opts), + ) + if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return a connected, unwrapped " + f"socket.socket, not {type(sock)}; consider AsyncHTTPProxyKMSConnect." + ) + # wrap_socket refuses a non-blocking socket, so normalize the mode here. + try: + sock.getpeername() + except OSError: + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return an already connected socket." + ) from None + if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM: + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return a stream socket, not a datagram one." + ) + # The callback may have consumed much of the CSOT budget, and wrapping + # resets the socket timeout, so recompute the remaining time here and for + # the KMS request that follows. + sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) try: - return await _async_configured_socket(address, opts) + conn = await _async_wrap_socket_tls(sock, address, opts) + except asyncio.CancelledError: + # The executor may still be wrapping the socket; close it so a TLS + # proxy's relay threads wind down instead of leaking. + sock.close() + raise except Exception as exc: _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) + conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) + return conn class _EncryptionIO(AsyncMongoCryptCallback): # type: ignore[misc] @@ -179,13 +283,6 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: False, # disable_ocsp_endpoint_check _IS_SYNC, ) - # CSOT: set timeout for socket creation. - connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001) - opts = PoolOptions( - connect_timeout=connect_timeout, - socket_timeout=connect_timeout, - ssl_context=ctx, - ) address = parse_host(endpoint, _HTTPS_PORT) if address[0].endswith(".sock"): raise ConfigurationError(f"Invalid KMS endpoint {endpoint!r}") @@ -193,8 +290,21 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: if sleep_u: sleep_sec = float(sleep_u) / 1e6 await asyncio.sleep(sleep_sec) + # Set the connect timeout after the retry backoff so the budget + # reflects the sleep. + connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001) + opts = PoolOptions( + connect_timeout=connect_timeout, + socket_timeout=connect_timeout, + ssl_context=ctx, + ) try: - conn = await _connect_kms(address, opts) + conn = await _connect_kms( + address, + opts, + self.opts._kms_connect_callback, + connect_timeout, + ) try: await async_socket_sendall(conn, message) while kms_context.bytes_needed > 0: @@ -230,6 +340,8 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: conn.close() except MongoCryptError: raise # Propagate MongoCryptError errors directly. + except ConfigurationError: + raise # A callback contract violation is not transient. except Exception as exc: remaining = _csot.remaining() if isinstance(exc, NetworkTimeout) or (remaining is not None and remaining <= 0): @@ -526,6 +638,7 @@ def __init__( codec_options: CodecOptions[_DocumentTypeArg], kms_tls_options: Optional[Mapping[str, Any]] = None, key_expiration_ms: Optional[int] = None, + kms_connect_callback: Optional[AsyncKMSConnectCallback] = None, ) -> None: """Explicit client-side field level encryption. @@ -595,7 +708,21 @@ def __init__( :param key_expiration_ms: The cache expiration time for data encryption keys. Defaults to ``None`` which defers to libmongocrypt's default which is currently 60000. Set to 0 to disable key expiration. - + :param kms_connect_callback: A callable that opens the connection to a + KMS host, used to route KMS requests through an HTTP proxy. It + receives a :class:`~pymongo.encryption_options.KMSConnectContext` + and returns a connected, unwrapped :class:`socket.socket`, over + which the driver performs the KMS TLS handshake. The callback + must be a coroutine function for the asynchronous API; a plain callable is rejected before it can block the event loop. + When a CSOT timeout is active, the driver stops waiting at the + deadline and closes any socket the callback yields later. For an + ordinary HTTP proxy, pass + :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect`. + Defaults to ``None``, meaning the driver connects to KMS hosts + directly. + + .. versionchanged:: 4.19 + Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. .. versionchanged:: 4.0 @@ -639,6 +766,7 @@ def __init__( key_vault_namespace, kms_tls_options=kms_tls_options, key_expiration_ms=key_expiration_ms, + kms_connect_callback=kms_connect_callback, ) self._kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC) self._io_callbacks: Optional[_EncryptionIO] = _EncryptionIO( diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 0d90fe3b19..edb2690907 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -19,9 +19,16 @@ from __future__ import annotations +import asyncio +import functools +import socket +import ssl +import threading +import time import warnings -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Optional, TypedDict +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict from pymongo.uri_parser_shared import _parse_kms_tls_options @@ -55,6 +62,251 @@ def check_min_pymongocrypt() -> None: ) +@dataclass(frozen=True) +class KMSConnectContext: + """Information about a pending KMS connection. + + Passed to ``kms_connect_callback``, which must return a plain, unwrapped + :class:`socket.socket`. The driver performs the KMS TLS handshake over it, + verifying against ``host`` rather than the peer actually reached. + + Prefer :class:`HTTPProxyKMSConnect` or :class:`AsyncHTTPProxyKMSConnect` + over writing a callback. + + :param host: Hostname of the KMS server, and the TLS verification target. + :param port: Port of the KMS server. + :param timeout: Seconds left in the timeout budget, or the default KMS + connect timeout when no timeout is active. + + .. note:: ``timeoutMS`` does not constrain KMS requests for explicit + encryption, so ``timeout`` is always the default there. Automatic + encryption passes the remaining budget. This deviates from the Client + Side Operations Timeout specification; see PYTHON-6037. + + .. versionadded:: 4.19 + """ + + host: str + port: int + timeout: float + + +# A callback that opens a connection to a KMS host. +AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]] +KMSConnectCallback = Callable[[KMSConnectContext], socket.socket] + +# Cap the CONNECT response header so a silent proxy cannot grow the buffer without bound. +_MAX_CONNECT_HEADER = 8192 + + +def _close_completed_socket(future: asyncio.Future[socket.socket]) -> None: + """Close a socket produced after its awaiting task was cancelled.""" + if not future.cancelled() and future.exception() is None: + future.result().close() + + +def _remaining(deadline: float) -> float: + """Seconds left before ``deadline``.""" + left = deadline - time.monotonic() + if left <= 0: + raise socket.timeout("timed out connecting through the proxy") + return left + + +class HTTPProxyKMSConnect: + """Route KMS connections through an HTTP proxy, for the synchronous API. + + Pass an instance as ``kms_connect_callback`` to reach KMS hosts through a + forward proxy that speaks HTTP ``CONNECT``:: + + from pymongo.encryption_options import AutoEncryptionOpts, HTTPProxyKMSConnect + + opts = AutoEncryptionOpts( + kms_providers={"aws": aws_creds}, + key_vault_namespace="keyvault.datakeys", + kms_connect_callback=HTTPProxyKMSConnect("proxy.example.com", 8080), + ) + + To reach the proxy over TLS, pass an :class:`ssl.SSLContext`. It applies + only to the proxy connection; KMS TLS is still negotiated end to end:: + + import ssl + + proxy_tls = ssl.create_default_context(cafile="proxy-ca.pem") + callback = HTTPProxyKMSConnect("proxy.example.com", 8443, proxy_tls) + + Use :class:`AsyncHTTPProxyKMSConnect` with the asynchronous API. + + :param host: Hostname of the proxy. + :param port: Port of the proxy. + :param ssl_context: Optional :class:`ssl.SSLContext` for connecting to the + proxy over TLS. Defaults to ``None``, meaning a plain connection. + + .. versionadded:: 4.19 + """ + + def __init__(self, host: str, port: int, ssl_context: Optional[ssl.SSLContext] = None): + self.host = host + self.port = port + self.ssl_context = ssl_context + + def _tunnel(self, sock: socket.socket, context: KMSConnectContext, deadline: float) -> None: + # An IPv6 literal needs brackets to be a valid HTTP authority. + host = f"[{context.host}]" if ":" in context.host else context.host + target = f"{host}:{context.port}" + sock.sendall(f"CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n".encode()) + # Read a byte at a time: a bulk read could consume tunneled bytes from + # this same socket. Reapply the budget before each read so a trickling + # proxy cannot outlive the deadline. + response = bytearray() + while not response.endswith(b"\r\n\r\n"): + sock.settimeout(_remaining(deadline)) + chunk = sock.recv(1) + if not chunk: + raise OSError(f"proxy closed the connection while tunneling to {target}") + response += chunk + if len(response) > _MAX_CONNECT_HEADER: + raise OSError(f"proxy sent an oversized CONNECT response for {target}") + status = bytes(response).split(b"\r\n", 1)[0] + # A CONNECT is successful for any 2xx status, e.g. "HTTP/1.0 200" or + # "HTTP/1.1 201"; require a three-digit code and reject malformed lines. + parts = status.split(b" ", 2) + valid = ( + len(parts) >= 2 + and parts[0].startswith(b"HTTP/") + and len(parts[1]) == 3 + and parts[1].isdigit() + ) + if not valid or not 200 <= int(parts[1]) < 300: + raise OSError(f"proxy refused CONNECT to {target}: {status!r}") + + def _bridge(self, proxy: socket.socket) -> socket.socket: + """Relay a TLS proxy connection through a socketpair. + + Python cannot layer TLS over an :class:`ssl.SSLSocket`, so return the + plain end of a pair, using threads rather than tasks even in + :class:`AsyncHTTPProxyKMSConnect`: the event loop cannot read an + :class:`ssl.SSLSocket`. + """ + # Clear the CONNECT-phase timeout; the tunneled KMS request is governed + # by the driver's own timeout, not the elapsed connect budget. + proxy.settimeout(None) + driver_side, relay_side = socket.socketpair() + + def relay(src: socket.socket, dst: socket.socket) -> None: + # Daemon threads: any error, including the ValueError an + # SSLSocket.shutdown can raise in the teardown race, ends the relay. + try: + while True: + buf = src.recv(16384) + if not buf: + break + dst.sendall(buf) + except (OSError, ValueError): + pass + finally: + # Send EOF to the peer instead of closing a socket it may be reading. + try: + dst.shutdown(socket.SHUT_RDWR) + except (OSError, ValueError): + pass + src.close() + + try: + for pair in ((relay_side, proxy), (proxy, relay_side)): + threading.Thread(target=relay, args=pair, daemon=True).start() + except BaseException: + # Unblock any thread that did start, then drop every socket. + for sock in (proxy, relay_side, driver_side): + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + sock.close() + raise + return driver_side + + def __call__( + self, context: KMSConnectContext, deadline: Optional[float] = None + ) -> socket.socket: + # A configurable KMS host could inject CR/LF into the CONNECT request. + if "\r" in context.host or "\n" in context.host: + raise ConfigurationError( + f"KMS host must not contain control characters: {context.host!r}" + ) + # One deadline for all three phases; per-phase timeouts would multiply + # the caller's budget. + if deadline is None: + deadline = time.monotonic() + context.timeout + sock = self._connect_proxy(deadline) + try: + if self.ssl_context is not None: + sock.settimeout(_remaining(deadline)) + sock = self.ssl_context.wrap_socket(sock, server_hostname=self.host) + sock.settimeout(_remaining(deadline)) + self._tunnel(sock, context, deadline) + except BaseException: + sock.close() + raise + if self.ssl_context is None: + return sock + try: + return self._bridge(sock) + except BaseException: + sock.close() + raise + + def _connect_proxy(self, deadline: float) -> socket.socket: + # Recompute the budget per address, rather than socket.create_connection, + # which applies the timeout to every address. + last_error: Optional[OSError] = None + for family, socktype, proto, _, sockaddr in socket.getaddrinfo( + self.host, self.port, type=socket.SOCK_STREAM + ): + sock = socket.socket(family, socktype, proto) + try: + # Propagate the timeout from _remaining rather than report a + # connect error. + sock.settimeout(_remaining(deadline)) + except socket.timeout: + sock.close() + raise + try: + sock.connect(sockaddr) + except OSError as exc: + last_error = exc + sock.close() + continue + return sock + raise OSError( + f"could not connect to proxy {self.host}:{self.port}: {last_error}" + ) from last_error + + +class AsyncHTTPProxyKMSConnect(HTTPProxyKMSConnect): + """Route KMS connections through an HTTP proxy, for the asynchronous API. + + Behaves exactly like :class:`HTTPProxyKMSConnect`, but is a coroutine + callable and runs the blocking connect in a thread so the event loop stays + free. + + .. versionadded:: 4.19 + """ + + async def __call__(self, context: KMSConnectContext) -> socket.socket: # type: ignore[override] + # Capture the deadline before scheduling so time spent queued behind + # other executor work counts against the KMS budget. + deadline = time.monotonic() + context.timeout + connect = functools.partial(super().__call__, context, deadline) + future = asyncio.get_running_loop().run_in_executor(None, connect) + try: + return await asyncio.shield(future) + except asyncio.CancelledError: + # The thread runs on regardless, so close the socket it returns. + future.add_done_callback(_close_completed_socket) + raise + + class AutoEncryptionOpts: """Options to configure automatic client-side field level encryption.""" @@ -75,6 +327,7 @@ def __init__( bypass_query_analysis: bool = False, encrypted_fields_map: Optional[Mapping[str, Any]] = None, key_expiration_ms: Optional[int] = None, + kms_connect_callback: Optional[Callable[[KMSConnectContext], Any]] = None, ) -> None: """Options to configure automatic client-side field level encryption. @@ -212,7 +465,20 @@ def __init__( :param key_expiration_ms: The cache expiration time for data encryption keys. Defaults to ``None`` which defers to libmongocrypt's default which is currently 60000. Set to 0 to disable key expiration. - + :param kms_connect_callback: A callable that opens the connection to a + KMS host, used to route KMS requests through an HTTP proxy. It + receives a :class:`KMSConnectContext` and returns a connected, + unwrapped :class:`socket.socket`, over which the driver performs + the KMS TLS handshake. Must be a coroutine function for + :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient` and a + regular function for + :class:`~pymongo.synchronous.mongo_client.MongoClient`. For an + ordinary proxy, pass :class:`HTTPProxyKMSConnect` or + :class:`AsyncHTTPProxyKMSConnect`. Defaults to ``None``, meaning + the driver connects to KMS hosts directly. + + .. versionchanged:: 4.19 + Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. .. versionchanged:: 4.2 @@ -259,6 +525,11 @@ def __init__( self._async_kms_ssl_contexts: Optional[dict[str, SSLContext]] = None self._bypass_query_analysis = bypass_query_analysis self._key_expiration_ms = key_expiration_ms + if kms_connect_callback is not None and not callable(kms_connect_callback): + raise TypeError( + f"kms_connect_callback must be callable, not {type(kms_connect_callback)}" + ) + self._kms_connect_callback = kms_connect_callback def _kms_ssl_contexts(self, is_sync: bool) -> dict[str, SSLContext]: if is_sync: diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index e7d8a366ca..79a8566177 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -16,8 +16,12 @@ from __future__ import annotations +import asyncio +import contextlib import functools +import inspect import socket +import ssl import time as time # noqa: PLC0414 # needed in sync version import uuid import weakref @@ -56,6 +60,8 @@ from pymongo.daemon import _spawn_daemon from pymongo.encryption_options import ( AutoEncryptionOpts, + KMSConnectCallback, + KMSConnectContext, RangeOpts, StringOpts, # Re-exported for backwards compatibility: TextOpts is deprecated but must @@ -84,8 +90,10 @@ from pymongo.operations import UpdateOne from pymongo.pool_options import PoolOptions from pymongo.pool_shared import ( + _close_late_socket, _configured_socket, _raise_connection_failure, + _wrap_socket_tls, ) from pymongo.read_concern import ReadConcern from pymongo.results import DeleteResult @@ -119,11 +127,108 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]: +def _close_rejected_kms_socket(obj: Any) -> None: + """Close a rejected kms_connect_callback return value, best effort. + + Nothing else will close it: _connect_kms raises before the result reaches + the caller's ``finally``. + """ + close = getattr(obj, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + +def _connect_kms( + address: _Address, + opts: PoolOptions, + kms_connect_callback: Optional[KMSConnectCallback], + timeout: float, +) -> Union[socket.socket, _sslConn]: + """Connect to a KMS host and perform the TLS handshake over the socket. + + Uses ``kms_connect_callback`` when one is provided, otherwise connects + directly, and always verifies against ``address`` (the KMS host). + """ + if kms_connect_callback is None: + try: + return _configured_socket(address, opts) + except Exception as exc: + _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) + + # TLS targets address, not the peer, so verification follows the KMS host. + # A plain callable would block the event loop before we could reject it, + # so check the callback first. + if not _IS_SYNC: + callback_any: Any = kms_connect_callback + is_coro = inspect.iscoroutinefunction(callback_any) + if not is_coro and callable(callback_any): + is_coro = inspect.iscoroutinefunction(callback_any.__call__) + if not is_coro: + raise ConfigurationError( + "kms_connect_callback must be a coroutine function for the async API." + ) + # Typed as Any so the generated synchronous flavor type-checks: the sync + # callback returns a plain socket, which is not awaitable. + result: Any = kms_connect_callback( + KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) + ) + remaining = _csot.remaining() + if remaining is None or _IS_SYNC: + # The synchronous API cannot interrupt a callback that has started + # running; honoring the deadline is the callback's contract there. + sock = result + else: + # CSOT is cooperative: a callback that ignores the timeout could block + # past the deadline. Shield the task so stopping the wait does not + # cancel it mid-flight, and close any socket it yields later. + task = asyncio.ensure_future(result) + try: + sock = asyncio.wait_for(asyncio.shield(task), remaining) + except asyncio.CancelledError: + task.add_done_callback(_close_late_socket) + raise + except asyncio.TimeoutError: + task.add_done_callback(_close_late_socket) + _raise_connection_failure( + address, + socket.timeout("timed out"), + timeout_details=_get_timeout_details(opts), + ) + if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return a connected, unwrapped " + f"socket.socket, not {type(sock)}; consider HTTPProxyKMSConnect." + ) + # wrap_socket refuses a non-blocking socket, so normalize the mode here. + try: + sock.getpeername() + except OSError: + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return an already connected socket." + ) from None + if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM: + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return a stream socket, not a datagram one." + ) + # The callback may have consumed much of the CSOT budget, and wrapping + # resets the socket timeout, so recompute the remaining time here and for + # the KMS request that follows. + sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) try: - return _configured_socket(address, opts) + conn = _wrap_socket_tls(sock, address, opts) + except asyncio.CancelledError: + # The executor may still be wrapping the socket; close it so a TLS + # proxy's relay threads wind down instead of leaking. + sock.close() + raise except Exception as exc: _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) + conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) + return conn class _EncryptionIO(MongoCryptCallback): # type: ignore[misc] @@ -178,13 +283,6 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: False, # disable_ocsp_endpoint_check _IS_SYNC, ) - # CSOT: set timeout for socket creation. - connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001) - opts = PoolOptions( - connect_timeout=connect_timeout, - socket_timeout=connect_timeout, - ssl_context=ctx, - ) address = parse_host(endpoint, _HTTPS_PORT) if address[0].endswith(".sock"): raise ConfigurationError(f"Invalid KMS endpoint {endpoint!r}") @@ -192,8 +290,21 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: if sleep_u: sleep_sec = float(sleep_u) / 1e6 time.sleep(sleep_sec) + # Set the connect timeout after the retry backoff so the budget + # reflects the sleep. + connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001) + opts = PoolOptions( + connect_timeout=connect_timeout, + socket_timeout=connect_timeout, + ssl_context=ctx, + ) try: - conn = _connect_kms(address, opts) + conn = _connect_kms( + address, + opts, + self.opts._kms_connect_callback, + connect_timeout, + ) try: sendall(conn, message) while kms_context.bytes_needed > 0: @@ -229,6 +340,8 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: conn.close() except MongoCryptError: raise # Propagate MongoCryptError errors directly. + except ConfigurationError: + raise # A callback contract violation is not transient. except Exception as exc: remaining = _csot.remaining() if isinstance(exc, NetworkTimeout) or (remaining is not None and remaining <= 0): @@ -523,6 +636,7 @@ def __init__( codec_options: CodecOptions[_DocumentTypeArg], kms_tls_options: Optional[Mapping[str, Any]] = None, key_expiration_ms: Optional[int] = None, + kms_connect_callback: Optional[KMSConnectCallback] = None, ) -> None: """Explicit client-side field level encryption. @@ -592,7 +706,21 @@ def __init__( :param key_expiration_ms: The cache expiration time for data encryption keys. Defaults to ``None`` which defers to libmongocrypt's default which is currently 60000. Set to 0 to disable key expiration. - + :param kms_connect_callback: A callable that opens the connection to a + KMS host, used to route KMS requests through an HTTP proxy. It + receives a :class:`~pymongo.encryption_options.KMSConnectContext` + and returns a connected, unwrapped :class:`socket.socket`, over + which the driver performs the KMS TLS handshake. The callback + must be a regular function; the async API requires a coroutine function and rejects plain callables before they can block the event loop. + When a CSOT timeout is active, the driver stops waiting at the + deadline and closes any socket the callback yields later. For an + ordinary HTTP proxy, pass + :class:`~pymongo.encryption_options.HTTPProxyKMSConnect`. + Defaults to ``None``, meaning the driver connects to KMS hosts + directly. + + .. versionchanged:: 4.19 + Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. .. versionchanged:: 4.0 @@ -632,6 +760,7 @@ def __init__( key_vault_namespace, kms_tls_options=kms_tls_options, key_expiration_ms=key_expiration_ms, + kms_connect_callback=kms_connect_callback, ) self._kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC) self._io_callbacks: Optional[_EncryptionIO] = _EncryptionIO( diff --git a/tools/synchro.py b/tools/synchro.py index bebf92c005..fc99bb2747 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -72,6 +72,8 @@ "_a_grid_out_property": "_grid_out_property", "AsyncClientEncryption": "ClientEncryption", "AsyncMongoCryptCallback": "MongoCryptCallback", + "AsyncKMSConnectCallback": "KMSConnectCallback", + "AsyncHTTPProxyKMSConnect": "HTTPProxyKMSConnect", "AsyncExplicitEncrypter": "ExplicitEncrypter", "AsyncAutoEncrypter": "AutoEncrypter", "AsyncContextManager": "ContextManager", @@ -127,6 +129,7 @@ "AsyncNetworkingInterface": "NetworkingInterface", "_configured_protocol_interface": "_configured_socket_interface", "_async_configured_socket": "_configured_socket", + "_async_wrap_socket_tls": "_wrap_socket_tls", "SpecRunnerTask": "SpecRunnerThread", "AsyncMockConnection": "MockConnection", "AsyncMockPool": "MockPool", @@ -297,6 +300,18 @@ def translate_docstrings(lines: list[str]) -> list[str]: lines[i] = lines[i].replace("an asynchronous", "a") if "An asynchronous" in lines[i]: lines[i] = lines[i].replace("An asynchronous", "A") + # This sentence states the callback contract, whose meaning + # would invert under the async -> sync word replacements. + if ( + "must be a coroutine function for the asynchronous API; a plain callable is rejected before it can block the event loop" + in lines[i] + ): + lines[i] = lines[i].replace( + "must be a coroutine function for the asynchronous API; a plain callable is rejected before it can block the event loop", + "must be a regular function; the async API requires a " + "coroutine function and rejects plain callables before " + "they can block the event loop", + ) # This ensures docstring links are for `pymongo.X` instead of `pymongo.synchronous.X` if "pymongo.asynchronous" in lines[i] and "import" not in lines[i]: lines[i] = lines[i].replace("pymongo.asynchronous", "pymongo") From 47e250ce1bd7018345d4f68b2dcec049ecabf38d Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 25 Sep 2026 07:25:11 -0500 Subject: [PATCH 3/3] PYTHON-5805 Add tests for the KMS connect callback Cover the callback contract, proxy tunnel and relay, CONNECT status handling, CSOT deadline enforcement, and cancellation safety. --- test/asynchronous/test_encryption.py | 40 +- test/asynchronous/test_kms_connect.py | 891 ++++++++++++++++++++++++++ test/test_encryption.py | 39 +- test/test_kms_connect.py | 888 +++++++++++++++++++++++++ 4 files changed, 1852 insertions(+), 6 deletions(-) create mode 100644 test/asynchronous/test_kms_connect.py create mode 100644 test/test_kms_connect.py diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 128f26feb4..9888418f63 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -16,6 +16,7 @@ from __future__ import annotations +import asyncio import base64 import copy import http.client @@ -29,12 +30,16 @@ import ssl import sys import textwrap +import threading +import time import traceback import uuid import warnings +from asyncio.trsock import TransportSocket from collections.abc import Mapping from threading import Thread from typing import Any, Optional +from unittest import mock import pytest @@ -60,13 +65,22 @@ from bson.son import SON from pymongo import ReadPreference from pymongo.asynchronous import encryption -from pymongo.asynchronous.encryption import Algorithm, AsyncClientEncryption, QueryType +from pymongo.asynchronous.encryption import ( + Algorithm, + AsyncClientEncryption, + QueryType, + _connect_kms, + _EncryptionIO, + _wrap_encryption_errors, +) from pymongo.asynchronous.helpers import anext from pymongo.asynchronous.mongo_client import AsyncMongoClient from pymongo.cursor_shared import CursorType from pymongo.encryption_options import ( _HAVE_PYMONGOCRYPT, + AsyncHTTPProxyKMSConnect, AutoEncryptionOpts, + HTTPProxyKMSConnect, RangeOpts, StringOpts, TextOpts, @@ -85,6 +99,8 @@ WriteError, ) from pymongo.operations import InsertOne, ReplaceOne, UpdateOne +from pymongo.pool_options import PoolOptions +from pymongo.ssl_support import get_ssl_context from pymongo.write_concern import WriteConcern from test import ( unittest, @@ -222,6 +238,9 @@ async def test_init_kms_tls_options(self): self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED) +# KMS connect callback unit and prose tests live in test_kms_connect.py. + + class TestClientOptions(AsyncPyMongoTestCase): async def test_default(self): client = self.simple_client(connect=False) @@ -316,9 +335,15 @@ def create_client_encryption( key_vault_client: AsyncMongoClient, codec_options: CodecOptions, kms_tls_options: Optional[Mapping[str, Any]] = None, + kms_connect_callback: Optional[Any] = None, ): client_encryption = AsyncClientEncryption( - kms_providers, key_vault_namespace, key_vault_client, codec_options, kms_tls_options + kms_providers, + key_vault_namespace, + key_vault_client, + codec_options, + kms_tls_options, + kms_connect_callback=kms_connect_callback, ) self.addAsyncCleanup(client_encryption.close) return client_encryption @@ -331,9 +356,15 @@ def unmanaged_create_client_encryption( key_vault_client: AsyncMongoClient, codec_options: CodecOptions, kms_tls_options: Optional[Mapping[str, Any]] = None, + kms_connect_callback: Optional[Any] = None, ): client_encryption = AsyncClientEncryption( - kms_providers, key_vault_namespace, key_vault_client, codec_options, kms_tls_options + kms_providers, + key_vault_namespace, + key_vault_client, + codec_options, + kms_tls_options, + kms_connect_callback=kms_connect_callback, ) return client_encryption @@ -1988,6 +2019,9 @@ async def test_invalid_hostname_in_kms_certificate(self): await self.client_encrypted.create_data_key("aws", master_key=key) +# KMS connect callback unit and prose tests live in test_kms_connect.py. + + # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-tls-options-tests class TestKmsTLSOptions(AsyncEncryptionIntegrationTest): @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py new file mode 100644 index 0000000000..23eb091a00 --- /dev/null +++ b/test/asynchronous/test_kms_connect.py @@ -0,0 +1,891 @@ +"""Tests for the KMS connect callback and HTTP proxy support.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import http.client +import socket +import ssl +import threading +import time +import unittest +from asyncio.trsock import TransportSocket +from typing import Any +from unittest import mock + +import pytest + +import pymongo +from bson.binary import Binary +from pymongo.asynchronous.encryption import ( + AsyncClientEncryption, + _connect_kms, + _EncryptionIO, + _wrap_encryption_errors, +) +from pymongo.encryption_options import ( + AsyncHTTPProxyKMSConnect, + AutoEncryptionOpts, + HTTPProxyKMSConnect, + KMSConnectContext, +) +from pymongo.errors import ConfigurationError, EncryptionError, NetworkTimeout +from pymongo.pool_options import PoolOptions +from pymongo.ssl_support import get_ssl_context +from test.asynchronous import AsyncPyMongoTestCase +from test.asynchronous.test_encryption import OPTS, AsyncEncryptionIntegrationTest +from test.helpers_shared import AWS_CREDS, CA_PEM, CLIENT_PEM + +_IS_SYNC = False + +pytestmark = pytest.mark.encryption + +KMS_PROXY_HOST = "127.0.0.1" +KMS_PROXY_PORT = 9004 +KMS_TLS_PROXY_PORT = 9005 + +AWS_MASTER_KEY = { + "region": "us-east-1", + "key": "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0", +} + + +class TestKmsConnectCallbackUnit(AsyncPyMongoTestCase): + """Contract checks for kms_connect_callback that need no KMS server.""" + + @staticmethod + def _pool_options(): + return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=None) + + async def test_init_kms_connect_callback(self): + opts = AutoEncryptionOpts({}, "k.d") + self.assertIsNone(opts._kms_connect_callback) + + async def callback(context): + raise AssertionError("not called") + + opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) + self.assertIs(opts._kms_connect_callback, callback) + + for bad in [1, "not-callable", object()]: + with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): + AutoEncryptionOpts({}, "k.d", kms_connect_callback=bad) # type: ignore[arg-type] + + context = KMSConnectContext(host="kms.example.com", port=443, timeout=9.5) + self.assertEqual(context.host, "kms.example.com") + self.assertEqual(context.port, 443) + self.assertEqual(context.timeout, 9.5) + with self.assertRaises(dataclasses.FrozenInstanceError): + context.host = "evil.example.com" # type: ignore[misc] + + async def test_non_socket_return_raises_configuration_error(self): + async def callback(context): + return "not-a-socket" + + with self.assertRaisesRegex(ConfigurationError, "must return a connected"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + async def test_already_wrapped_socket_is_rejected(self): + # ssl.SSLSocket passes isinstance but cannot be TLS-wrapped again. + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + left, right = socket.socketpair() + self.addCleanup(right.close) + # No peer needed to produce a genuine ssl.SSLSocket. + wrapped = ctx.wrap_socket(left, do_handshake_on_connect=False, server_hostname="x") + self.addCleanup(wrapped.close) + + async def callback(context): + return wrapped + + with self.assertRaisesRegex(ConfigurationError, "unwrapped"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + async def test_context_receives_host_port_and_timeout(self): + received = [] + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + async def callback(context): + received.append(context) + return left + + # ssl_context=None returns the socket unchanged, so a plain socket is accepted. + conn = await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 12.5) + self.assertIs(conn, left) + + self.assertEqual(len(received), 1) + self.assertEqual(received[0].host, "kms.example.com") + self.assertEqual(received[0].port, 443) + self.assertEqual(received[0].timeout, 12.5) + + async def test_non_blocking_socket_from_callback_is_accepted(self): + # Without the driver normalizing the mode, this raises ValueError. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def serve(): + try: + conn, _ = listener.accept() + server_ctx.wrap_socket(conn, server_side=True).close() + except OSError: + pass + + threading.Thread(target=serve, daemon=True).start() + + # Built as the driver does, for the flavor-correct type; the local cert won't verify. + client_ctx = get_ssl_context(None, None, None, None, True, True, False, _IS_SYNC) + options = PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=client_ctx) + + def connect(): + sock = socket.create_connection(listener.getsockname(), timeout=10) + sock.setblocking(False) + return sock + + async def callback(context): + if _IS_SYNC: + return connect() + return await asyncio.get_running_loop().run_in_executor(None, connect) + + conn = await _connect_kms(listener.getsockname(), options, callback, 10.0) + self.addCleanup(conn.close) + self.assertIsNotNone(conn.gettimeout()) + + async def test_asyncio_transport_socket_is_rejected(self): + # get_extra_info("socket") is a TransportSocket, not a socket.socket. + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + async def callback(context): + return TransportSocket(left) + + with self.assertRaisesRegex(ConfigurationError, "TransportSocket"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + async def test_http_proxy_helper_tunnels_and_reports_refusal(self): + # Covers the CONNECT handshake without KMS credentials. + accepted = [] + + def stub(listener, reply): + try: + conn, _ = listener.accept() + except OSError: + return + accepted.append(conn.recv(4096)) + conn.sendall(reply) + conn.close() + + def run_stub(reply): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + threading.Thread(target=stub, args=(listener, reply), daemon=True).start() + return listener.getsockname() + + host, port = run_stub(b"HTTP/1.1 200 Connection Established\r\n\r\n") + callback = AsyncHTTPProxyKMSConnect(host, port) + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = await callback(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + self.assertEqual(accepted[0].split(b"\r\n")[0], b"CONNECT kms.example.com:443 HTTP/1.1") + + host, port = run_stub(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n") + with self.assertRaisesRegex(OSError, "refused CONNECT"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + + # Any 2xx status is a successful tunnel, not just HTTP/1.1 200. + host, port = run_stub(b"HTTP/1.0 200 Connection Established\r\n\r\n") + sock = await AsyncHTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + + # A status code must be exactly three digits, with no zero padding. + for reply in (b"HTTP/1.1 2000 Evil\r\n\r\n", b"HTTP/1.1 00200 Evil\r\n\r\n"): + host, port = run_stub(reply) + with self.assertRaisesRegex(OSError, "refused CONNECT"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + + async def test_control_characters_in_kms_host_are_rejected(self): + # Reject CR/LF in the configurable host before it reaches CONNECT. + callback = AsyncHTTPProxyKMSConnect("proxy.example.com", 8080) + context = KMSConnectContext(host="kms.example.com\r\nX-Injected: 1", port=443, timeout=10) + with self.assertRaisesRegex(ConfigurationError, "control characters"): + await callback(context) + + async def test_cancelled_tls_wrap_closes_late_socket(self): + # A cancelled wrap can leave the executor producing an SSLSocket; the + # done callback must close it. + if _IS_SYNC: + raise unittest.SkipTest("the cancel-safe wrap is an async path") + from pymongo.pool_shared import _close_late_socket + + left, right = socket.socketpair() + future = asyncio.get_running_loop().create_future() + future.set_result(left) + self.assertNotEqual(left.fileno(), -1) + _close_late_socket(future) + self.assertEqual(left.fileno(), -1) + self.addCleanup(right.close) + + async def test_tls_proxy_helper_bridges_the_tunnel(self): + # Covers the TLS-proxy path and the socketpair relay without KMS creds. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + request = b"" + while b"\r\n\r\n" not in request: + chunk = tls.recv(4096) + if not chunk: + return + request += chunk + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + # The tunneled peer speaks only after the client does, as a TLS server would. + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + client_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + + sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + else: + await asyncio.get_running_loop().run_in_executor(None, sock.sendall, b"ping") + data = await asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"echo:ping") + + async def test_bridge_does_not_inherit_the_connect_deadline(self): + # The relay must outlast the much shorter CONNECT deadline. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + request = b"" + while b"\r\n\r\n" not in request: + chunk = tls.recv(4096) + if not chunk: + return + request += chunk + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + # Outlast the connect-phase deadline before the KMS side speaks. + time.sleep(2.0) + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + client_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=1.0) + + sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + else: + await asyncio.get_running_loop().run_in_executor(None, sock.sendall, b"ping") + data = await asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"echo:ping") + + async def test_non_coroutine_callback_is_rejected(self): + # A plain def must be rejected before it blocks the event loop. + if _IS_SYNC: + raise unittest.SkipTest("a regular function is correct for the sync API") + + entered = [] + + def callback(context): + entered.append(context) + return None + + with self.assertRaisesRegex(ConfigurationError, "coroutine function"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + self.assertEqual(entered, [], "invalid callback must not be entered") + + async def test_proxy_closing_before_connect_reply_raises(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + # Read the CONNECT request, then hang up without replying. + conn.recv(4096) + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "proxy closed the connection"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + + async def test_tunnel_keeps_bytes_sent_with_the_connect_reply(self): + # A proxy may coalesce its 200 with tunneled bytes; reading past the header would drop them. + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + conn.recv(4096) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\nearly-bytes") + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = await AsyncHTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + data = sock.recv(64) + else: + data = await asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"early-bytes") + + async def test_unconnected_socket_from_callback_is_rejected(self): + # An unconnected socket would fail later as a transient error and be retried. + bare = socket.socket() + self.addCleanup(bare.close) + + async def callback(context): + return bare + + with self.assertRaisesRegex(ConfigurationError, "already connected"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + async def test_ipv6_host_is_bracketed_in_connect(self): + accepted = [] + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + accepted.append(conn.recv(4096)) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="::1", port=443, timeout=10) + sock = await AsyncHTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + self.assertEqual(accepted[0].split(b"\r\n")[0], b"CONNECT [::1]:443 HTTP/1.1") + + async def test_oversized_connect_response_is_rejected(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + conn.recv(4096) + # Never sends the terminator. + while True: + conn.sendall(b"x" * 1024) + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "oversized CONNECT response"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + + async def test_remaining_raises_once_the_deadline_passes(self): + from pymongo.encryption_options import _remaining + + self.assertGreater(_remaining(time.monotonic() + 5), 0) + with self.assertRaises(socket.timeout): + _remaining(time.monotonic() - 1) + + async def test_datagram_socket_from_callback_is_rejected(self): + # TLS on a connected UDP socket raises NotImplementedError, which would be retried. + left = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + right = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.addCleanup(left.close) + self.addCleanup(right.close) + right.bind(("127.0.0.1", 0)) + left.connect(right.getsockname()) + + async def callback(context): + return left + + with self.assertRaisesRegex(ConfigurationError, "stream socket"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + async def test_kms_request_does_not_retry_a_contract_violation(self): + # _connect_kms has no retry loop; the no-retry guarantee is in + # kms_request, so exercise that instead. + calls = [] + + async def callback(context): + calls.append(context) + return "not-a-socket" + + opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) + io = _EncryptionIO(None, mock.MagicMock(), None, opts) + + class StubKmsContext: + endpoint = "kms.example.com:443" + message = b"request" + kms_provider = "aws" + usleep = 0 + bytes_needed = 1 + + def feed(self, data): + raise AssertionError("should not reach the socket") + + def fail(self): + raise AssertionError("a contract violation must not be retried") + + with self.assertRaises(ConfigurationError): + await io.kms_request(StubKmsContext()) + self.assertEqual(len(calls), 1) + + async def test_contract_violation_surfaces_as_encryption_error(self): + # Callers see EncryptionError with ConfigurationError as its cause. + with self.assertRaises(EncryptionError) as caught: + with _wrap_encryption_errors(): + raise ConfigurationError("kms_connect_callback must return ...") + self.assertIsInstance(caught.exception.__cause__, ConfigurationError) + + async def test_bridge_failure_closes_the_proxy_socket(self): + # A failure inside _bridge must not strand the connected proxy socket. + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + tls.recv(4096) + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + captured = [] + + def failing_bridge(self, proxy): + captured.append(proxy) + raise OSError("no file descriptors") + + host, port = listener.getsockname() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + + with mock.patch.object(HTTPProxyKMSConnect, "_bridge", failing_bridge): + with self.assertRaisesRegex(OSError, "no file descriptors"): + await AsyncHTTPProxyKMSConnect(host, port, ctx)(context) + + self.assertEqual(captured[0].fileno(), -1, "proxy socket was left open") + + async def test_network_error_from_callback_propagates(self): + async def callback(context): + raise OSError("proxy unreachable") + + # Not a ConfigurationError, so kms_request retries it. + with self.assertRaises(OSError): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + async def test_csot_deadline_stops_a_hung_callback(self): + # A callback that ignores the timeout cannot block past the CSOT + # deadline, and a socket it yields later must be closed. + if _IS_SYNC: + raise unittest.SkipTest("the sync API cannot interrupt a callback") + + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + async def hung_callback(context): + await asyncio.sleep(0.5) + return left + + with self.assertRaises(NetworkTimeout): + with pymongo.timeout(0.1): + await _connect_kms( + ("kms.example.com", 443), self._pool_options(), hung_callback, 10.0 + ) + self.assertNotEqual(left.fileno(), -1) + # Let the shielded callback finish; the driver closes the late result. + await asyncio.sleep(0.75) + self.assertEqual(left.fileno(), -1) + + async def test_cancelling_kms_connect_closes_the_callback_socket(self): + # Cancelling during the TLS handshake must close the callback's socket, + # so a TLS proxy's relay threads wind down. + if _IS_SYNC: + raise unittest.SkipTest("cancellation is an async-only behavior") + + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + gate = threading.Event() + eof = threading.Event() + + def stub_server(): + conn = None + try: + conn, _ = listener.accept() + # The cancel may land before or after the executor starts the + # handshake. Peek for the ClientHello without consuming it, or + # for EOF if the driver closed it, before wrap_socket detaches conn. + while True: + data = conn.recv(4096, socket.MSG_PEEK) + if not data: + eof.set() + return + if data[:1] == b"\x16": # TLS handshake record + break + # Hold the handshake open until the test has cancelled. + if not gate.wait(5): + return + tls = server_ctx.wrap_socket(conn, server_side=True, do_handshake_on_connect=False) + try: + tls.do_handshake() + # A discarded connection may end in a reset rather than a + # clean EOF; either proves the driver closed it. + while tls.recv(4096): + pass + except OSError: + pass + eof.set() + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_server, daemon=True).start() + + client_ctx = get_ssl_context(None, None, None, None, True, True, False, _IS_SYNC) + options = PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=client_ctx) + socks = [] + + async def callback(context): + sock = await asyncio.get_running_loop().run_in_executor( + None, + lambda: socket.create_connection(listener.getsockname(), timeout=10), + ) + socks.append(sock) + return sock + + # The sync flavor returns a socket instead of a coroutine, so both + # error codes are needed depending on the flavor being checked. + connect = _connect_kms(listener.getsockname(), options, callback, 10.0) + task = asyncio.ensure_future(connect) # type: ignore[type-var,arg-type] + for _ in range(100): + if socks: + break + await asyncio.sleep(0.01) + self.assertTrue(socks, "callback was never invoked") + # Bias the cancel to land mid-handshake; the stub handles the earlier + # window too. + await asyncio.sleep(0.1) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + # The late SSLSocket (or raw socket) must be closed; the stub sees EOF. + gate.set() + for _ in range(50): + if eof.is_set(): + break + await asyncio.sleep(0.1) + self.assertTrue(eof.is_set(), "driver never closed the callback socket") + + async def test_client_encryption_accepts_callback(self): + async def callback(context): + raise AssertionError("not called") + + client = self.simple_client() + encryption = AsyncClientEncryption( + {"local": {"key": b"\x00" * 96}}, + "keyvault.datakeys", + client, + OPTS, + kms_connect_callback=callback, + ) + self.addAsyncCleanup(encryption.close) + self.assertIs(encryption._io_callbacks.opts._kms_connect_callback, callback) + + async def test_client_encryption_rejects_non_callable(self): + client = self.simple_client() + with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): + AsyncClientEncryption( + {"local": {"key": b"\x00" * 96}}, + "keyvault.datakeys", + client, + OPTS, + kms_connect_callback="not-callable", # type: ignore[arg-type] + ) + + +class TestKmsConnectCallbackProse(AsyncEncryptionIntegrationTest): + @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") + async def asyncSetUp(self): + await super().asyncSetUp() + self.callback_calls: list[Any] = [] + + async def plain_callback(self, context): + self.callback_calls.append(context) + return await AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) + + async def tls_callback(self, context): + self.callback_calls.append(context) + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + # PYTHON-5040 tracks re-enabling verification: the evergreen-tools CA + # lacks an Authority Key Identifier newer OpenSSL requires. + ctx.verify_mode = ssl.CERT_NONE + callback = AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) + return await callback(context) + + async def proxy_request(self, method, path, tls=False): + """Call the proxy's control endpoints and return the body.""" + if _IS_SYNC: + return self._proxy_request(method, path, tls) + return await asyncio.get_running_loop().run_in_executor( + None, self._proxy_request, method, path, tls + ) + + def _proxy_request(self, method, path, tls=False): + if tls: + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + # PYTHON-5040 tracks re-enabling verification once the test CA cert + # is fixed; the evergreen-tools CA lacks an Authority Key Identifier + # that newer OpenSSL requires, so verification fails on Windows 3.14. + ctx.verify_mode = ssl.CERT_NONE + conn = http.client.HTTPSConnection( + f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=ctx + ) + else: + conn = http.client.HTTPConnection(f"{KMS_PROXY_HOST}:{KMS_PROXY_PORT}") + try: + conn.request(method, path) + return conn.getresponse().read().decode() + finally: + conn.close() + + async def connect_count(self, tls=False): + body = await self.proxy_request("GET", "/metrics", tls=tls) + # One "key value" per line; the server also emits connect_target. + for line in body.splitlines(): + key, _, value = line.partition(" ") + if key == "connect_count": + return int(value) + raise AssertionError(f"no connect_count in metrics body: {body!r}") + + async def test_01_plain_http_proxy(self): + await self.proxy_request("POST", "/reset") + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(await self.connect_count(), 1) + + async def test_02_https_proxy(self): + await self.proxy_request("POST", "/reset", tls=True) + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.tls_callback, + ) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(await self.connect_count(tls=True), 1) + + async def test_03_auto_encryption_through_proxy(self): + await self.client.keyvault.datakeys.drop() + await self.client.db.coll.drop() + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + data_key_id = await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + schema = { + "bsonType": "object", + "properties": { + "encrypted_string": { + "encrypt": { + "keyId": [data_key_id], + "bsonType": "string", + "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic", + } + } + }, + } + + await self.proxy_request("POST", "/reset") + opts = AutoEncryptionOpts( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + schema_map={"db.coll": schema}, + kms_connect_callback=self.plain_callback, + ) + client_encrypted = await self.async_rs_or_single_client(auto_encryption_opts=opts) + + await client_encrypted.db.coll.insert_one({"_id": 1, "encrypted_string": "hello"}) + decrypted = await client_encrypted.db.coll.find_one({"_id": 1}) + self.assertEqual(decrypted["encrypted_string"], "hello") + + raw = await self.client.db.coll.find_one({"_id": 1}) + self.assertIsInstance(raw["encrypted_string"], Binary) + + # The decrypt reuses the cached key, so exactly one KMS request follows + # the reset. + self.assertEqual(await self.connect_count(), 1) + + async def test_04_callback_error(self): + async def failing_callback(context): + raise OSError("proxy is on fire") + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=failing_callback, + ) + with self.assertRaisesRegex(EncryptionError, "proxy is on fire"): + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + + @unittest.skip( + "PYTHON-6037 ClientEncryption does not support timeoutMS, so the " + "callback always receives the default KMS connect timeout" + ) + async def test_05_callback_receives_timeout(self): + key_vault_client = await self.async_rs_or_single_client(timeoutMS=1000) + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + key_vault_client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + + self.assertTrue(self.callback_calls, "callback was never invoked") + for context in self.callback_calls: + # Checks only the spec's non-zero requirement, which cannot fail. + self.assertIsNotNone(context.timeout) + self.assertGreater(context.timeout, 0) + + async def test_06_retry_after_network_error(self): + state = {"calls": 0} + + async def flaky_callback(context): + state["calls"] += 1 + if state["calls"] == 1: + raise OSError("first attempt fails") + return await AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=flaky_callback, + ) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(state["calls"], 2) diff --git a/test/test_encryption.py b/test/test_encryption.py index adb6005ea1..12553b6091 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -16,6 +16,7 @@ from __future__ import annotations +import asyncio import base64 import copy import http.client @@ -29,12 +30,16 @@ import ssl import sys import textwrap +import threading +import time import traceback import uuid import warnings +from asyncio.trsock import TransportSocket from collections.abc import Mapping from threading import Thread from typing import Any, Optional +from unittest import mock import pytest @@ -63,6 +68,7 @@ from pymongo.encryption_options import ( _HAVE_PYMONGOCRYPT, AutoEncryptionOpts, + HTTPProxyKMSConnect, RangeOpts, StringOpts, TextOpts, @@ -81,8 +87,17 @@ WriteError, ) from pymongo.operations import InsertOne, ReplaceOne, UpdateOne +from pymongo.pool_options import PoolOptions +from pymongo.ssl_support import get_ssl_context from pymongo.synchronous import encryption -from pymongo.synchronous.encryption import Algorithm, ClientEncryption, QueryType +from pymongo.synchronous.encryption import ( + Algorithm, + ClientEncryption, + QueryType, + _connect_kms, + _EncryptionIO, + _wrap_encryption_errors, +) from pymongo.synchronous.helpers import next from pymongo.synchronous.mongo_client import MongoClient from pymongo.write_concern import WriteConcern @@ -222,6 +237,9 @@ def test_init_kms_tls_options(self): self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED) +# KMS connect callback unit and prose tests live in test_kms_connect.py. + + class TestClientOptions(PyMongoTestCase): def test_default(self): client = self.simple_client(connect=False) @@ -316,9 +334,15 @@ def create_client_encryption( key_vault_client: MongoClient, codec_options: CodecOptions, kms_tls_options: Optional[Mapping[str, Any]] = None, + kms_connect_callback: Optional[Any] = None, ): client_encryption = ClientEncryption( - kms_providers, key_vault_namespace, key_vault_client, codec_options, kms_tls_options + kms_providers, + key_vault_namespace, + key_vault_client, + codec_options, + kms_tls_options, + kms_connect_callback=kms_connect_callback, ) self.addCleanup(client_encryption.close) return client_encryption @@ -331,9 +355,15 @@ def unmanaged_create_client_encryption( key_vault_client: MongoClient, codec_options: CodecOptions, kms_tls_options: Optional[Mapping[str, Any]] = None, + kms_connect_callback: Optional[Any] = None, ): client_encryption = ClientEncryption( - kms_providers, key_vault_namespace, key_vault_client, codec_options, kms_tls_options + kms_providers, + key_vault_namespace, + key_vault_client, + codec_options, + kms_tls_options, + kms_connect_callback=kms_connect_callback, ) return client_encryption @@ -1980,6 +2010,9 @@ def test_invalid_hostname_in_kms_certificate(self): self.client_encrypted.create_data_key("aws", master_key=key) +# KMS connect callback unit and prose tests live in test_kms_connect.py. + + # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-tls-options-tests class TestKmsTLSOptions(EncryptionIntegrationTest): @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py new file mode 100644 index 0000000000..04338a495c --- /dev/null +++ b/test/test_kms_connect.py @@ -0,0 +1,888 @@ +"""Tests for the KMS connect callback and HTTP proxy support.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import http.client +import socket +import ssl +import threading +import time +import unittest +from asyncio.trsock import TransportSocket +from typing import Any +from unittest import mock + +import pytest + +import pymongo +from bson.binary import Binary +from pymongo.encryption_options import ( + AutoEncryptionOpts, + HTTPProxyKMSConnect, + KMSConnectContext, +) +from pymongo.errors import ConfigurationError, EncryptionError, NetworkTimeout +from pymongo.pool_options import PoolOptions +from pymongo.ssl_support import get_ssl_context +from pymongo.synchronous.encryption import ( + ClientEncryption, + _connect_kms, + _EncryptionIO, + _wrap_encryption_errors, +) +from test import PyMongoTestCase +from test.helpers_shared import AWS_CREDS, CA_PEM, CLIENT_PEM +from test.test_encryption import OPTS, EncryptionIntegrationTest + +_IS_SYNC = True + +pytestmark = pytest.mark.encryption + +KMS_PROXY_HOST = "127.0.0.1" +KMS_PROXY_PORT = 9004 +KMS_TLS_PROXY_PORT = 9005 + +AWS_MASTER_KEY = { + "region": "us-east-1", + "key": "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0", +} + + +class TestKmsConnectCallbackUnit(PyMongoTestCase): + """Contract checks for kms_connect_callback that need no KMS server.""" + + @staticmethod + def _pool_options(): + return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=None) + + def test_init_kms_connect_callback(self): + opts = AutoEncryptionOpts({}, "k.d") + self.assertIsNone(opts._kms_connect_callback) + + def callback(context): + raise AssertionError("not called") + + opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) + self.assertIs(opts._kms_connect_callback, callback) + + for bad in [1, "not-callable", object()]: + with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): + AutoEncryptionOpts({}, "k.d", kms_connect_callback=bad) # type: ignore[arg-type] + + context = KMSConnectContext(host="kms.example.com", port=443, timeout=9.5) + self.assertEqual(context.host, "kms.example.com") + self.assertEqual(context.port, 443) + self.assertEqual(context.timeout, 9.5) + with self.assertRaises(dataclasses.FrozenInstanceError): + context.host = "evil.example.com" # type: ignore[misc] + + def test_non_socket_return_raises_configuration_error(self): + def callback(context): + return "not-a-socket" + + with self.assertRaisesRegex(ConfigurationError, "must return a connected"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + def test_already_wrapped_socket_is_rejected(self): + # ssl.SSLSocket passes isinstance but cannot be TLS-wrapped again. + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + left, right = socket.socketpair() + self.addCleanup(right.close) + # No peer needed to produce a genuine ssl.SSLSocket. + wrapped = ctx.wrap_socket(left, do_handshake_on_connect=False, server_hostname="x") + self.addCleanup(wrapped.close) + + def callback(context): + return wrapped + + with self.assertRaisesRegex(ConfigurationError, "unwrapped"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + def test_context_receives_host_port_and_timeout(self): + received = [] + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + def callback(context): + received.append(context) + return left + + # ssl_context=None returns the socket unchanged, so a plain socket is accepted. + conn = _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 12.5) + self.assertIs(conn, left) + + self.assertEqual(len(received), 1) + self.assertEqual(received[0].host, "kms.example.com") + self.assertEqual(received[0].port, 443) + self.assertEqual(received[0].timeout, 12.5) + + def test_non_blocking_socket_from_callback_is_accepted(self): + # Without the driver normalizing the mode, this raises ValueError. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def serve(): + try: + conn, _ = listener.accept() + server_ctx.wrap_socket(conn, server_side=True).close() + except OSError: + pass + + threading.Thread(target=serve, daemon=True).start() + + # Built as the driver does, for the flavor-correct type; the local cert won't verify. + client_ctx = get_ssl_context(None, None, None, None, True, True, False, _IS_SYNC) + options = PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=client_ctx) + + def connect(): + sock = socket.create_connection(listener.getsockname(), timeout=10) + sock.setblocking(False) + return sock + + def callback(context): + if _IS_SYNC: + return connect() + return asyncio.get_running_loop().run_in_executor(None, connect) + + conn = _connect_kms(listener.getsockname(), options, callback, 10.0) + self.addCleanup(conn.close) + self.assertIsNotNone(conn.gettimeout()) + + def test_asyncio_transport_socket_is_rejected(self): + # get_extra_info("socket") is a TransportSocket, not a socket.socket. + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + def callback(context): + return TransportSocket(left) + + with self.assertRaisesRegex(ConfigurationError, "TransportSocket"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + def test_http_proxy_helper_tunnels_and_reports_refusal(self): + # Covers the CONNECT handshake without KMS credentials. + accepted = [] + + def stub(listener, reply): + try: + conn, _ = listener.accept() + except OSError: + return + accepted.append(conn.recv(4096)) + conn.sendall(reply) + conn.close() + + def run_stub(reply): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + threading.Thread(target=stub, args=(listener, reply), daemon=True).start() + return listener.getsockname() + + host, port = run_stub(b"HTTP/1.1 200 Connection Established\r\n\r\n") + callback = HTTPProxyKMSConnect(host, port) + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = callback(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + self.assertEqual(accepted[0].split(b"\r\n")[0], b"CONNECT kms.example.com:443 HTTP/1.1") + + host, port = run_stub(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n") + with self.assertRaisesRegex(OSError, "refused CONNECT"): + HTTPProxyKMSConnect(host, port)(context) + + # Any 2xx status is a successful tunnel, not just HTTP/1.1 200. + host, port = run_stub(b"HTTP/1.0 200 Connection Established\r\n\r\n") + sock = HTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + + # A status code must be exactly three digits, with no zero padding. + for reply in (b"HTTP/1.1 2000 Evil\r\n\r\n", b"HTTP/1.1 00200 Evil\r\n\r\n"): + host, port = run_stub(reply) + with self.assertRaisesRegex(OSError, "refused CONNECT"): + HTTPProxyKMSConnect(host, port)(context) + + def test_control_characters_in_kms_host_are_rejected(self): + # Reject CR/LF in the configurable host before it reaches CONNECT. + callback = HTTPProxyKMSConnect("proxy.example.com", 8080) + context = KMSConnectContext(host="kms.example.com\r\nX-Injected: 1", port=443, timeout=10) + with self.assertRaisesRegex(ConfigurationError, "control characters"): + callback(context) + + def test_cancelled_tls_wrap_closes_late_socket(self): + # A cancelled wrap can leave the executor producing an SSLSocket; the + # done callback must close it. + if _IS_SYNC: + raise unittest.SkipTest("the cancel-safe wrap is an async path") + from pymongo.pool_shared import _close_late_socket + + left, right = socket.socketpair() + future = asyncio.get_running_loop().create_future() + future.set_result(left) + self.assertNotEqual(left.fileno(), -1) + _close_late_socket(future) + self.assertEqual(left.fileno(), -1) + self.addCleanup(right.close) + + def test_tls_proxy_helper_bridges_the_tunnel(self): + # Covers the TLS-proxy path and the socketpair relay without KMS creds. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + request = b"" + while b"\r\n\r\n" not in request: + chunk = tls.recv(4096) + if not chunk: + return + request += chunk + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + # The tunneled peer speaks only after the client does, as a TLS server would. + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + client_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + + sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + else: + asyncio.get_running_loop().run_in_executor(None, sock.sendall, b"ping") + data = asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"echo:ping") + + def test_bridge_does_not_inherit_the_connect_deadline(self): + # The relay must outlast the much shorter CONNECT deadline. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + request = b"" + while b"\r\n\r\n" not in request: + chunk = tls.recv(4096) + if not chunk: + return + request += chunk + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + # Outlast the connect-phase deadline before the KMS side speaks. + time.sleep(2.0) + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + client_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=1.0) + + sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + else: + asyncio.get_running_loop().run_in_executor(None, sock.sendall, b"ping") + data = asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"echo:ping") + + def test_non_coroutine_callback_is_rejected(self): + # A plain def must be rejected before it blocks the event loop. + if _IS_SYNC: + raise unittest.SkipTest("a regular function is correct for the sync API") + + entered = [] + + def callback(context): + entered.append(context) + return None + + with self.assertRaisesRegex(ConfigurationError, "coroutine function"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + self.assertEqual(entered, [], "invalid callback must not be entered") + + def test_proxy_closing_before_connect_reply_raises(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + # Read the CONNECT request, then hang up without replying. + conn.recv(4096) + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "proxy closed the connection"): + HTTPProxyKMSConnect(host, port)(context) + + def test_tunnel_keeps_bytes_sent_with_the_connect_reply(self): + # A proxy may coalesce its 200 with tunneled bytes; reading past the header would drop them. + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + conn.recv(4096) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\nearly-bytes") + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = HTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + data = sock.recv(64) + else: + data = asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"early-bytes") + + def test_unconnected_socket_from_callback_is_rejected(self): + # An unconnected socket would fail later as a transient error and be retried. + bare = socket.socket() + self.addCleanup(bare.close) + + def callback(context): + return bare + + with self.assertRaisesRegex(ConfigurationError, "already connected"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + def test_ipv6_host_is_bracketed_in_connect(self): + accepted = [] + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + accepted.append(conn.recv(4096)) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="::1", port=443, timeout=10) + sock = HTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + self.assertEqual(accepted[0].split(b"\r\n")[0], b"CONNECT [::1]:443 HTTP/1.1") + + def test_oversized_connect_response_is_rejected(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + conn.recv(4096) + # Never sends the terminator. + while True: + conn.sendall(b"x" * 1024) + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "oversized CONNECT response"): + HTTPProxyKMSConnect(host, port)(context) + + def test_remaining_raises_once_the_deadline_passes(self): + from pymongo.encryption_options import _remaining + + self.assertGreater(_remaining(time.monotonic() + 5), 0) + with self.assertRaises(socket.timeout): + _remaining(time.monotonic() - 1) + + def test_datagram_socket_from_callback_is_rejected(self): + # TLS on a connected UDP socket raises NotImplementedError, which would be retried. + left = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + right = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.addCleanup(left.close) + self.addCleanup(right.close) + right.bind(("127.0.0.1", 0)) + left.connect(right.getsockname()) + + def callback(context): + return left + + with self.assertRaisesRegex(ConfigurationError, "stream socket"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + def test_kms_request_does_not_retry_a_contract_violation(self): + # _connect_kms has no retry loop; the no-retry guarantee is in + # kms_request, so exercise that instead. + calls = [] + + def callback(context): + calls.append(context) + return "not-a-socket" + + opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) + io = _EncryptionIO(None, mock.MagicMock(), None, opts) + + class StubKmsContext: + endpoint = "kms.example.com:443" + message = b"request" + kms_provider = "aws" + usleep = 0 + bytes_needed = 1 + + def feed(self, data): + raise AssertionError("should not reach the socket") + + def fail(self): + raise AssertionError("a contract violation must not be retried") + + with self.assertRaises(ConfigurationError): + io.kms_request(StubKmsContext()) + self.assertEqual(len(calls), 1) + + def test_contract_violation_surfaces_as_encryption_error(self): + # Callers see EncryptionError with ConfigurationError as its cause. + with self.assertRaises(EncryptionError) as caught: + with _wrap_encryption_errors(): + raise ConfigurationError("kms_connect_callback must return ...") + self.assertIsInstance(caught.exception.__cause__, ConfigurationError) + + def test_bridge_failure_closes_the_proxy_socket(self): + # A failure inside _bridge must not strand the connected proxy socket. + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + tls.recv(4096) + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + captured = [] + + def failing_bridge(self, proxy): + captured.append(proxy) + raise OSError("no file descriptors") + + host, port = listener.getsockname() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + + with mock.patch.object(HTTPProxyKMSConnect, "_bridge", failing_bridge): + with self.assertRaisesRegex(OSError, "no file descriptors"): + HTTPProxyKMSConnect(host, port, ctx)(context) + + self.assertEqual(captured[0].fileno(), -1, "proxy socket was left open") + + def test_network_error_from_callback_propagates(self): + def callback(context): + raise OSError("proxy unreachable") + + # Not a ConfigurationError, so kms_request retries it. + with self.assertRaises(OSError): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + def test_csot_deadline_stops_a_hung_callback(self): + # A callback that ignores the timeout cannot block past the CSOT + # deadline, and a socket it yields later must be closed. + if _IS_SYNC: + raise unittest.SkipTest("the sync API cannot interrupt a callback") + + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + def hung_callback(context): + time.sleep(0.5) + return left + + with self.assertRaises(NetworkTimeout): + with pymongo.timeout(0.1): + _connect_kms(("kms.example.com", 443), self._pool_options(), hung_callback, 10.0) + self.assertNotEqual(left.fileno(), -1) + # Let the shielded callback finish; the driver closes the late result. + time.sleep(0.75) + self.assertEqual(left.fileno(), -1) + + def test_cancelling_kms_connect_closes_the_callback_socket(self): + # Cancelling during the TLS handshake must close the callback's socket, + # so a TLS proxy's relay threads wind down. + if _IS_SYNC: + raise unittest.SkipTest("cancellation is an async-only behavior") + + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + gate = threading.Event() + eof = threading.Event() + + def stub_server(): + conn = None + try: + conn, _ = listener.accept() + # The cancel may land before or after the executor starts the + # handshake. Peek for the ClientHello without consuming it, or + # for EOF if the driver closed it, before wrap_socket detaches conn. + while True: + data = conn.recv(4096, socket.MSG_PEEK) + if not data: + eof.set() + return + if data[:1] == b"\x16": # TLS handshake record + break + # Hold the handshake open until the test has cancelled. + if not gate.wait(5): + return + tls = server_ctx.wrap_socket(conn, server_side=True, do_handshake_on_connect=False) + try: + tls.do_handshake() + # A discarded connection may end in a reset rather than a + # clean EOF; either proves the driver closed it. + while tls.recv(4096): + pass + except OSError: + pass + eof.set() + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_server, daemon=True).start() + + client_ctx = get_ssl_context(None, None, None, None, True, True, False, _IS_SYNC) + options = PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=client_ctx) + socks = [] + + def callback(context): + sock = asyncio.get_running_loop().run_in_executor( + None, + lambda: socket.create_connection(listener.getsockname(), timeout=10), + ) + socks.append(sock) + return sock + + # The sync flavor returns a socket instead of a coroutine, so both + # error codes are needed depending on the flavor being checked. + connect = _connect_kms(listener.getsockname(), options, callback, 10.0) + task = asyncio.ensure_future(connect) # type: ignore[type-var,arg-type] + for _ in range(100): + if socks: + break + time.sleep(0.01) + self.assertTrue(socks, "callback was never invoked") + # Bias the cancel to land mid-handshake; the stub handles the earlier + # window too. + time.sleep(0.1) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + task + # The late SSLSocket (or raw socket) must be closed; the stub sees EOF. + gate.set() + for _ in range(50): + if eof.is_set(): + break + time.sleep(0.1) + self.assertTrue(eof.is_set(), "driver never closed the callback socket") + + def test_client_encryption_accepts_callback(self): + def callback(context): + raise AssertionError("not called") + + client = self.simple_client() + encryption = ClientEncryption( + {"local": {"key": b"\x00" * 96}}, + "keyvault.datakeys", + client, + OPTS, + kms_connect_callback=callback, + ) + self.addCleanup(encryption.close) + self.assertIs(encryption._io_callbacks.opts._kms_connect_callback, callback) + + def test_client_encryption_rejects_non_callable(self): + client = self.simple_client() + with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): + ClientEncryption( + {"local": {"key": b"\x00" * 96}}, + "keyvault.datakeys", + client, + OPTS, + kms_connect_callback="not-callable", # type: ignore[arg-type] + ) + + +class TestKmsConnectCallbackProse(EncryptionIntegrationTest): + @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") + def setUp(self): + super().setUp() + self.callback_calls: list[Any] = [] + + def plain_callback(self, context): + self.callback_calls.append(context) + return HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) + + def tls_callback(self, context): + self.callback_calls.append(context) + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + # PYTHON-5040 tracks re-enabling verification: the evergreen-tools CA + # lacks an Authority Key Identifier newer OpenSSL requires. + ctx.verify_mode = ssl.CERT_NONE + callback = HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) + return callback(context) + + def proxy_request(self, method, path, tls=False): + """Call the proxy's control endpoints and return the body.""" + if _IS_SYNC: + return self._proxy_request(method, path, tls) + return asyncio.get_running_loop().run_in_executor( + None, self._proxy_request, method, path, tls + ) + + def _proxy_request(self, method, path, tls=False): + if tls: + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + # PYTHON-5040 tracks re-enabling verification once the test CA cert + # is fixed; the evergreen-tools CA lacks an Authority Key Identifier + # that newer OpenSSL requires, so verification fails on Windows 3.14. + ctx.verify_mode = ssl.CERT_NONE + conn = http.client.HTTPSConnection( + f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=ctx + ) + else: + conn = http.client.HTTPConnection(f"{KMS_PROXY_HOST}:{KMS_PROXY_PORT}") + try: + conn.request(method, path) + return conn.getresponse().read().decode() + finally: + conn.close() + + def connect_count(self, tls=False): + body = self.proxy_request("GET", "/metrics", tls=tls) + # One "key value" per line; the server also emits connect_target. + for line in body.splitlines(): + key, _, value = line.partition(" ") + if key == "connect_count": + return int(value) + raise AssertionError(f"no connect_count in metrics body: {body!r}") + + def test_01_plain_http_proxy(self): + self.proxy_request("POST", "/reset") + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(self.connect_count(), 1) + + def test_02_https_proxy(self): + self.proxy_request("POST", "/reset", tls=True) + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.tls_callback, + ) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(self.connect_count(tls=True), 1) + + def test_03_auto_encryption_through_proxy(self): + self.client.keyvault.datakeys.drop() + self.client.db.coll.drop() + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + data_key_id = encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + schema = { + "bsonType": "object", + "properties": { + "encrypted_string": { + "encrypt": { + "keyId": [data_key_id], + "bsonType": "string", + "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic", + } + } + }, + } + + self.proxy_request("POST", "/reset") + opts = AutoEncryptionOpts( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + schema_map={"db.coll": schema}, + kms_connect_callback=self.plain_callback, + ) + client_encrypted = self.rs_or_single_client(auto_encryption_opts=opts) + + client_encrypted.db.coll.insert_one({"_id": 1, "encrypted_string": "hello"}) + decrypted = client_encrypted.db.coll.find_one({"_id": 1}) + self.assertEqual(decrypted["encrypted_string"], "hello") + + raw = self.client.db.coll.find_one({"_id": 1}) + self.assertIsInstance(raw["encrypted_string"], Binary) + + # The decrypt reuses the cached key, so exactly one KMS request follows + # the reset. + self.assertEqual(self.connect_count(), 1) + + def test_04_callback_error(self): + def failing_callback(context): + raise OSError("proxy is on fire") + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=failing_callback, + ) + with self.assertRaisesRegex(EncryptionError, "proxy is on fire"): + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + + @unittest.skip( + "PYTHON-6037 ClientEncryption does not support timeoutMS, so the " + "callback always receives the default KMS connect timeout" + ) + def test_05_callback_receives_timeout(self): + key_vault_client = self.rs_or_single_client(timeoutMS=1000) + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + key_vault_client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + + self.assertTrue(self.callback_calls, "callback was never invoked") + for context in self.callback_calls: + # Checks only the spec's non-zero requirement, which cannot fail. + self.assertIsNotNone(context.timeout) + self.assertGreater(context.timeout, 0) + + def test_06_retry_after_network_error(self): + state = {"calls": 0} + + def flaky_callback(context): + state["calls"] += 1 + if state["calls"] == 1: + raise OSError("first attempt fails") + return HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=flaky_callback, + ) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(state["calls"], 2)