From c37efd53eb7e89aa36d67431598f3fd1be1fd7fd Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 12:42:43 -0500 Subject: [PATCH 01/50] PYTHON-5805 Split TLS wrapping out of configured socket helpers --- pymongo/pool_shared.py | 50 +++++++++++++++++++++++++------ test/asynchronous/test_pooling.py | 21 ++++++++++++- test/test_pooling.py | 21 ++++++++++++- tools/synchro.py | 1 + 4 files changed, 82 insertions(+), 11 deletions(-) diff --git a/pymongo/pool_shared.py b/pymongo/pool_shared.py index 410ffd8189..ec529e810c 100644 --- a/pymongo/pool_shared.py +++ b/pymongo/pool_shared.py @@ -259,16 +259,19 @@ async def _async_create_connection(address: _Address, options: PoolOptions) -> s raise OSError("getaddrinfo failed") -async def _async_configured_socket( - address: _Address, options: PoolOptions +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, for + example 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: @@ -315,6 +318,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, @@ -465,14 +481,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, for + example 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: @@ -514,6 +535,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 063f5f06ec..5b97fb108f 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 @@ -759,5 +765,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 64146a0e13..55f450d57a 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 @@ -757,5 +763,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() diff --git a/tools/synchro.py b/tools/synchro.py index bebf92c005..0dcf46bcc7 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -127,6 +127,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", From 967f17a8a365fa09919ed6acc3bcc93ff6ebb60d Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 15:19:33 -0500 Subject: [PATCH 02/50] PYTHON-5805 Add KMSConnectContext and AutoEncryptionOpts.kms_connect_callback --- pymongo/encryption_options.py | 145 ++++++++++++++++++++++++++- test/asynchronous/test_encryption.py | 29 ++++++ test/test_encryption.py | 29 ++++++ 3 files changed, 200 insertions(+), 3 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index f2fcd47c65..0f91bf9739 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -19,8 +19,10 @@ from __future__ import annotations -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Optional, TypedDict +import socket +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 @@ -54,6 +56,124 @@ def check_min_pymongocrypt() -> None: ) +@dataclass(frozen=True) +class KMSConnectContext: + """Information about a pending KMS connection. + + An instance is passed to the ``kms_connect_callback`` configured on + :class:`AutoEncryptionOpts`, :class:`~pymongo.encryption.ClientEncryption`, + or :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`. + + The callback opens a connection to ``host``:``port`` and returns it. The + driver then performs the KMS TLS handshake over that socket, so the + callback must return a plain, unwrapped :class:`socket.socket`. Certificate + and hostname verification target ``host``, not whatever peer the socket is + actually connected to, which is what makes proxying safe. + + To reach a KMS host through an HTTP proxy, connect to the proxy and issue + an HTTP ``CONNECT`` request:: + + import socket + + def connect_through_proxy(context): + sock = socket.create_connection( + ("proxy.example.com", 8080), timeout=context.timeout + ) + target = f"{context.host}:{context.port}" + sock.sendall( + f"CONNECT {target} HTTP/1.1\\r\\nHost: {target}\\r\\n\\r\\n".encode() + ) + response = b"" + while b"\\r\\n\\r\\n" not in response: + chunk = sock.recv(4096) + if not chunk: + sock.close() + raise OSError("proxy closed the connection") + response += chunk + if not response.startswith(b"HTTP/1.1 200"): + sock.close() + raise OSError(f"proxy CONNECT failed: {response.splitlines()[0]!r}") + return sock + + opts = AutoEncryptionOpts( + kms_providers={"aws": aws_creds}, + key_vault_namespace="keyvault.datakeys", + kms_connect_callback=connect_through_proxy, + ) + + For :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption` and + :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient`, the callback + must be a coroutine function. It must not block the event loop, so drive + the connection with :mod:`asyncio` or hand the blocking work to a thread + with :func:`asyncio.to_thread`. + + Reaching the proxy itself over TLS takes one extra step. Python cannot + layer a second TLS session over an :class:`ssl.SSLSocket`, so the callback + cannot return its TLS connection to the proxy directly. Bridge it to a + :func:`socket.socketpair` and return the plain end instead:: + + import socket, ssl, threading + + def connect_through_tls_proxy(context): + proxy_ctx = ssl.create_default_context(cafile="proxy-ca.pem") + raw = socket.create_connection( + ("proxy.example.com", 8443), timeout=context.timeout + ) + proxy = proxy_ctx.wrap_socket(raw, server_hostname="proxy.example.com") + target = f"{context.host}:{context.port}" + proxy.sendall( + f"CONNECT {target} HTTP/1.1\\r\\nHost: {target}\\r\\n\\r\\n".encode() + ) + response = b"" + while b"\\r\\n\\r\\n" not in response: + chunk = proxy.recv(4096) + if not chunk: + proxy.close() + raise OSError("proxy closed the connection") + response += chunk + if not response.startswith(b"HTTP/1.1 200"): + proxy.close() + raise OSError(f"proxy CONNECT failed: {response.splitlines()[0]!r}") + + driver_side, relay_side = socket.socketpair() + + def relay(src, dst): + try: + while True: + buf = src.recv(16384) + if not buf: + break + dst.sendall(buf) + except OSError: + pass + finally: + src.close() + dst.close() + + threading.Thread(target=relay, args=(relay_side, proxy), daemon=True).start() + threading.Thread(target=relay, args=(proxy, relay_side), daemon=True).start() + return driver_side + + :param host: Hostname of the KMS server, and the target of TLS certificate + and hostname verification. + :param port: Port of the KMS server. + :param timeout: Seconds remaining before the operation's timeout expires, + or ``None`` when no timeout applies. + + .. versionadded:: 4.18 + """ + + host: str + port: int + timeout: Optional[float] + + +# A callback that opens a connection to a KMS host. The async driver requires a +# coroutine function; the synchronous driver requires a regular function. +AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]] +KMSConnectCallback = Callable[[KMSConnectContext], socket.socket] + + class AutoEncryptionOpts: """Options to configure automatic client-side field level encryption.""" @@ -74,6 +194,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. @@ -211,7 +332,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`; the driver then performs the KMS + TLS handshake over it. Must be a coroutine function for + :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient` and a + regular function for + :class:`~pymongo.synchronous.mongo_client.MongoClient`. See + :class:`KMSConnectContext` for a worked HTTP ``CONNECT`` example. + Defaults to ``None``, meaning the driver connects to KMS hosts + directly. + + .. versionchanged:: 4.18 + Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. .. versionchanged:: 4.2 @@ -258,6 +392,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/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index daeb18607a..aa00f843eb 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -18,6 +18,7 @@ import base64 import copy +import dataclasses import http.client import json import os @@ -213,6 +214,34 @@ async def test_init_kms_tls_options(self): self.assertEqual(ctx.check_hostname, True) self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED) + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + async def test_init_kms_connect_callback(self): + from pymongo.encryption_options import KMSConnectContext + + # Default is None. + opts = AutoEncryptionOpts({}, "k.d") + self.assertIsNone(opts._kms_connect_callback) + + # A callable is accepted and stored unchanged. + async def callback(context): + raise AssertionError("not called") + + opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) + self.assertIs(opts._kms_connect_callback, callback) + + # Non-callables are rejected eagerly. + 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] + + # The context is frozen and carries host, port, and timeout. + 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] + class TestClientOptions(AsyncPyMongoTestCase): async def test_default(self): diff --git a/test/test_encryption.py b/test/test_encryption.py index 744db01b1b..0fb24afe76 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -18,6 +18,7 @@ import base64 import copy +import dataclasses import http.client import json import os @@ -213,6 +214,34 @@ def test_init_kms_tls_options(self): self.assertEqual(ctx.check_hostname, True) self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED) + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + def test_init_kms_connect_callback(self): + from pymongo.encryption_options import KMSConnectContext + + # Default is None. + opts = AutoEncryptionOpts({}, "k.d") + self.assertIsNone(opts._kms_connect_callback) + + # A callable is accepted and stored unchanged. + def callback(context): + raise AssertionError("not called") + + opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) + self.assertIs(opts._kms_connect_callback, callback) + + # Non-callables are rejected eagerly. + 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] + + # The context is frozen and carries host, port, and timeout. + 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] + class TestClientOptions(PyMongoTestCase): def test_default(self): From 38ecefc3f3924a683fcadae835b6fe4472d07fce Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 17:25:08 -0500 Subject: [PATCH 03/50] PYTHON-5805 Route KMS connections through kms_connect_callback --- pymongo/asynchronous/encryption.py | 49 +++++++++++++++++-- pymongo/synchronous/encryption.py | 49 +++++++++++++++++-- test/asynchronous/test_encryption.py | 73 +++++++++++++++++++++++++++- test/test_encryption.py | 73 +++++++++++++++++++++++++++- tools/synchro.py | 1 + 5 files changed, 237 insertions(+), 8 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 524ae45c11..813044f57a 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -20,6 +20,7 @@ import contextlib import enum import socket +import ssl import time as time # noqa: PLC0414 # needed in sync version import uuid import weakref @@ -63,7 +64,9 @@ from pymongo.common import CONNECT_TIMEOUT from pymongo.daemon import _spawn_daemon from pymongo.encryption_options import ( + AsyncKMSConnectCallback, AutoEncryptionOpts, + KMSConnectContext, RangeOpts, TextOpts, check_min_pymongocrypt, @@ -82,6 +85,7 @@ from pymongo.pool_options import PoolOptions from pymongo.pool_shared import ( _async_configured_socket, + _async_wrap_socket_tls, _raise_connection_failure, ) from pymongo.read_concern import ReadConcern @@ -112,9 +116,41 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -async def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]: +class _KMSCallbackContractError(Exception): + """Raised when a kms_connect_callback violates its contract. + + Unlike a network error from the callback, this is a programming error, so + it is never retried. + """ + + +async def _connect_kms( + address: _Address, + opts: PoolOptions, + kms_connect_callback: Optional[AsyncKMSConnectCallback] = None, + timeout: Optional[float] = None, +) -> Union[socket.socket, _sslConn]: + 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)) + + # Let the caller open the connection, then apply TLS ourselves so that SNI + # and certificate verification still target the KMS host even when the + # socket actually terminates at a proxy. + sock = await kms_connect_callback( + KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) + ) + if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): + raise _KMSCallbackContractError( + "kms_connect_callback must return a connected, unwrapped " + f"socket.socket, not {type(sock)}. TLS cannot be layered over an " + "already-wrapped socket; to reach the proxy over TLS, relay through " + "a socket.socketpair and return the plain end." + ) try: - return await _async_configured_socket(address, opts) + return await _async_wrap_socket_tls(sock, address, opts) except Exception as exc: _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) @@ -197,7 +233,12 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: sleep_sec = float(sleep_u) / 1e6 await asyncio.sleep(sleep_sec) 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: @@ -233,6 +274,8 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: conn.close() except MongoCryptError: raise # Propagate MongoCryptError errors directly. + except _KMSCallbackContractError: + 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): diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 014d162e2b..365de2b1be 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -19,6 +19,7 @@ import contextlib import enum import socket +import ssl import time as time # noqa: PLC0414 # needed in sync version import uuid import weakref @@ -59,6 +60,8 @@ from pymongo.daemon import _spawn_daemon from pymongo.encryption_options import ( AutoEncryptionOpts, + KMSConnectCallback, + KMSConnectContext, RangeOpts, TextOpts, check_min_pymongocrypt, @@ -78,6 +81,7 @@ from pymongo.pool_shared import ( _configured_socket, _raise_connection_failure, + _wrap_socket_tls, ) from pymongo.read_concern import ReadConcern from pymongo.results import BulkWriteResult, DeleteResult @@ -111,9 +115,41 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]: +class _KMSCallbackContractError(Exception): + """Raised when a kms_connect_callback violates its contract. + + Unlike a network error from the callback, this is a programming error, so + it is never retried. + """ + + +def _connect_kms( + address: _Address, + opts: PoolOptions, + kms_connect_callback: Optional[KMSConnectCallback] = None, + timeout: Optional[float] = None, +) -> Union[socket.socket, _sslConn]: + 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)) + + # Let the caller open the connection, then apply TLS ourselves so that SNI + # and certificate verification still target the KMS host even when the + # socket actually terminates at a proxy. + sock = kms_connect_callback( + KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) + ) + if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): + raise _KMSCallbackContractError( + "kms_connect_callback must return a connected, unwrapped " + f"socket.socket, not {type(sock)}. TLS cannot be layered over an " + "already-wrapped socket; to reach the proxy over TLS, relay through " + "a socket.socketpair and return the plain end." + ) try: - return _configured_socket(address, opts) + return _wrap_socket_tls(sock, address, opts) except Exception as exc: _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) @@ -196,7 +232,12 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: sleep_sec = float(sleep_u) / 1e6 time.sleep(sleep_sec) 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: @@ -232,6 +273,8 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: conn.close() except MongoCryptError: raise # Propagate MongoCryptError errors directly. + except _KMSCallbackContractError: + 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): diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index aa00f843eb..88fae34533 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -60,7 +60,13 @@ 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, + _KMSCallbackContractError, +) from pymongo.asynchronous.helpers import anext from pymongo.asynchronous.mongo_client import AsyncMongoClient from pymongo.cursor_shared import CursorType @@ -79,6 +85,7 @@ WriteError, ) from pymongo.operations import InsertOne, ReplaceOne, UpdateOne +from pymongo.pool_options import PoolOptions from pymongo.write_concern import WriteConcern from test import ( unittest, @@ -243,6 +250,70 @@ async def callback(context): context.host = "evil.example.com" # type: ignore[misc] +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_non_socket_return_is_not_retried(self): + async def callback(context): + return "not-a-socket" + + with self.assertRaisesRegex(_KMSCallbackContractError, "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 subclasses socket.socket, but wrap_socket cannot layer + # TLS over it, so it must be rejected with an actionable message rather + # than failing later with an opaque handshake error. + 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) + # do_handshake_on_connect=False means no peer is 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(_KMSCallbackContractError, "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 + + # With ssl_context=None the socket is returned unchanged, which also + # confirms 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_network_error_from_callback_propagates(self): + async def callback(context): + raise OSError("proxy unreachable") + + # Not wrapped in _KMSCallbackContractError, so kms_request's broad + # handler can treat it as transient and let libmongocrypt retry. + with self.assertRaises(OSError): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + class TestClientOptions(AsyncPyMongoTestCase): async def test_default(self): client = self.simple_client(connect=False) diff --git a/test/test_encryption.py b/test/test_encryption.py index 0fb24afe76..de68668922 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -75,8 +75,15 @@ WriteError, ) from pymongo.operations import InsertOne, ReplaceOne, UpdateOne +from pymongo.pool_options import PoolOptions from pymongo.synchronous import encryption -from pymongo.synchronous.encryption import Algorithm, ClientEncryption, QueryType +from pymongo.synchronous.encryption import ( + Algorithm, + ClientEncryption, + QueryType, + _connect_kms, + _KMSCallbackContractError, +) from pymongo.synchronous.helpers import next from pymongo.synchronous.mongo_client import MongoClient from pymongo.write_concern import WriteConcern @@ -243,6 +250,70 @@ def callback(context): context.host = "evil.example.com" # type: ignore[misc] +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_non_socket_return_is_not_retried(self): + def callback(context): + return "not-a-socket" + + with self.assertRaisesRegex(_KMSCallbackContractError, "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 subclasses socket.socket, but wrap_socket cannot layer + # TLS over it, so it must be rejected with an actionable message rather + # than failing later with an opaque handshake error. + 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) + # do_handshake_on_connect=False means no peer is 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(_KMSCallbackContractError, "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 + + # With ssl_context=None the socket is returned unchanged, which also + # confirms 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_network_error_from_callback_propagates(self): + def callback(context): + raise OSError("proxy unreachable") + + # Not wrapped in _KMSCallbackContractError, so kms_request's broad + # handler can treat it as transient and let libmongocrypt retry. + with self.assertRaises(OSError): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + class TestClientOptions(PyMongoTestCase): def test_default(self): client = self.simple_client(connect=False) diff --git a/tools/synchro.py b/tools/synchro.py index 0dcf46bcc7..bc81c1e65e 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -72,6 +72,7 @@ "_a_grid_out_property": "_grid_out_property", "AsyncClientEncryption": "ClientEncryption", "AsyncMongoCryptCallback": "MongoCryptCallback", + "AsyncKMSConnectCallback": "KMSConnectCallback", "AsyncExplicitEncrypter": "ExplicitEncrypter", "AsyncAutoEncrypter": "AutoEncrypter", "AsyncContextManager": "ContextManager", From 8c24872ced80f57426164ef44b8c1b67e9f01678 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 18:57:26 -0500 Subject: [PATCH 04/50] PYTHON-5805 Add kms_connect_callback to ClientEncryption --- pymongo/asynchronous/encryption.py | 15 +++++++++- pymongo/synchronous/encryption.py | 15 +++++++++- test/asynchronous/test_encryption.py | 44 ++++++++++++++++++++++++++-- test/test_encryption.py | 44 ++++++++++++++++++++++++++-- 4 files changed, 112 insertions(+), 6 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 813044f57a..3693929858 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -639,6 +639,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. @@ -708,7 +709,18 @@ 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`; the + driver then performs the KMS TLS handshake over it. Must be a + coroutine function. See + :class:`~pymongo.encryption_options.KMSConnectContext` for a worked + HTTP ``CONNECT`` example. Defaults to ``None``, meaning the driver + connects to KMS hosts directly. + + .. versionchanged:: 4.18 + Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. .. versionchanged:: 4.0 @@ -752,6 +764,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/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 365de2b1be..611772aaa0 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -636,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. @@ -705,7 +706,18 @@ 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`; the + driver then performs the KMS TLS handshake over it. Must be a + coroutine function. See + :class:`~pymongo.encryption_options.KMSConnectContext` for a worked + HTTP ``CONNECT`` example. Defaults to ``None``, meaning the driver + connects to KMS hosts directly. + + .. versionchanged:: 4.18 + Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. .. versionchanged:: 4.0 @@ -745,6 +757,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/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 88fae34533..2b27e33f1b 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -313,6 +313,34 @@ async def callback(context): with self.assertRaises(OSError): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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) + + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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 TestClientOptions(AsyncPyMongoTestCase): async def test_default(self): @@ -352,9 +380,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 @@ -367,9 +401,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 diff --git a/test/test_encryption.py b/test/test_encryption.py index de68668922..022635ad4e 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -313,6 +313,34 @@ def callback(context): with self.assertRaises(OSError): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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) + + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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 TestClientOptions(PyMongoTestCase): def test_default(self): @@ -352,9 +380,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 @@ -367,9 +401,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 From 84c76a20cc644a3c21a7329f00d84eef49ee0ac6 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 20:48:08 -0500 Subject: [PATCH 05/50] PYTHON-5805 Fix sync ClientEncryption docstring to say 'regular function' --- pymongo/asynchronous/encryption.py | 5 ++--- pymongo/synchronous/encryption.py | 5 ++--- tools/synchro.py | 1 + 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 3693929858..7d2a877533 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -713,9 +713,8 @@ def __init__( 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`; the - driver then performs the KMS TLS handshake over it. Must be a - coroutine function. See - :class:`~pymongo.encryption_options.KMSConnectContext` for a worked + driver then performs the KMS TLS handshake over it. Must be a coroutine function. + See :class:`~pymongo.encryption_options.KMSConnectContext` for a worked HTTP ``CONNECT`` example. Defaults to ``None``, meaning the driver connects to KMS hosts directly. diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 611772aaa0..48bd4f098a 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -710,9 +710,8 @@ def __init__( 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`; the - driver then performs the KMS TLS handshake over it. Must be a - coroutine function. See - :class:`~pymongo.encryption_options.KMSConnectContext` for a worked + driver then performs the KMS TLS handshake over it. Must be a regular function. + See :class:`~pymongo.encryption_options.KMSConnectContext` for a worked HTTP ``CONNECT`` example. Defaults to ``None``, meaning the driver connects to KMS hosts directly. diff --git a/tools/synchro.py b/tools/synchro.py index bc81c1e65e..d0cf8f902f 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -143,6 +143,7 @@ "dns.asyncresolver.resolve": "dns.resolver.resolve", "__aenter__": "__enter__", "__aexit__": "__exit__", + "Must be a coroutine function.": "Must be a regular function.", } docstring_replacements: dict[tuple[str, str], str] = { From 0b1f785a0bd7ddfae68e0c2f1f23f338a9f94da2 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 20:56:48 -0500 Subject: [PATCH 06/50] PYTHON-5805 Add synchro tripwire for kms_connect_callback docstring mapping --- tools/synchro.py | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/tools/synchro.py b/tools/synchro.py index d0cf8f902f..8844b20e55 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -143,6 +143,10 @@ "dns.asyncresolver.resolve": "dns.resolver.resolve", "__aenter__": "__enter__", "__aexit__": "__exit__", + # Prose substitution, not an identifier/token like the rest of this dict. Matching is + # line-based (see translate_docstrings), so this sentence must stay on a single line in + # pymongo/asynchronous/encryption.py or the replacement silently stops firing. Guarded by + # check_kms_connect_callback_docstring(), called from main() below. "Must be a coroutine function.": "Must be a regular function.", } @@ -336,6 +340,37 @@ def process_ignores(lines: list[str]) -> list[str]: return lines +def check_kms_connect_callback_docstring() -> None: + """Guard the "Must be a coroutine function." -> "Must be a regular function." mapping. + + That replacement in `replacements` above only fires if the sentence sits on a single + line in the async source (see translate_docstrings). A future edit or re-wrap could + silently break the match, leaving the generated synchronous docs telling users to write + an `async def` callback. Fail loudly instead of leaving that undetected. + """ + sync_encryption = Path(_pymongo_dest_base) / "encryption.py" + if not sync_encryption.is_file(): + # Nothing to check yet, e.g. a partial/filtered run that didn't touch this file. + return + content = sync_encryption.read_text() + if "Must be a coroutine function." in content: + raise RuntimeError( + f"{sync_encryption} still says 'Must be a coroutine function.' after synchro. " + "The 'Must be a coroutine function.' -> 'Must be a regular function.' entry in " + "tools/synchro.py's `replacements` dict didn't fire, most likely because the " + "sentence in pymongo/asynchronous/encryption.py got wrapped across multiple " + "lines. Keep it on one line, or fix the replacement mapping." + ) + if "kms_connect_callback" in content and "Must be a regular function." not in content: + raise RuntimeError( + f"{sync_encryption} defines kms_connect_callback but its docstring no longer " + "says 'Must be a regular function.'. Check that the " + "'Must be a coroutine function.' -> 'Must be a regular function.' entry is still " + "present in tools/synchro.py's `replacements` dict and that the docstring wording " + "in pymongo/asynchronous/encryption.py hasn't changed out from under it." + ) + + def unasync_directory(files: list[str], src: str, dest: str, replacements: dict[str, str]) -> None: unasync_files( files, @@ -404,6 +439,8 @@ def main() -> None: generated_tests, ) + check_kms_connect_callback_docstring() + generated_files = generated_pymongo + generated_gridfs + generated_tests if is_ci and generated_files: From 6f2b0e18bec6553fb758bc3ad9c80179c3dff212 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 18 Aug 2026 05:22:31 -0500 Subject: [PATCH 07/50] PYTHON-5805 Add prose tests for KMS connect callback Implements all six cases of spec section 28 "KMS Connect Callback": plain and TLS proxy tunneling via kms_connect_callback, auto encryption through a proxy, callback error propagation, timeout visibility on KMSConnectContext, and retry after a callback network error. --- test/asynchronous/test_encryption.py | 232 +++++++++++++++++++++++++++ test/test_encryption.py | 232 +++++++++++++++++++++++++++ 2 files changed, 464 insertions(+) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 2b27e33f1b..55c794314b 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 dataclasses @@ -29,6 +30,7 @@ import ssl import sys import textwrap +import threading import traceback import uuid import warnings @@ -2056,6 +2058,236 @@ async def test_invalid_hostname_in_kms_certificate(self): await self.client_encrypted.create_data_key("aws", master_key=key) +_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", +} + + +def _http_connect(context, proxy_port, proxy_ssl_context=None): + """Open an HTTP CONNECT tunnel to context.host:context.port. + + Returns a plain socket. When proxy_ssl_context is given, the connection to + the proxy is TLS and the tunnel is bridged to a socketpair, because TLS + cannot be layered over an ssl.SSLSocket. + """ + sock: socket.socket = socket.create_connection( + (_KMS_PROXY_HOST, proxy_port), timeout=context.timeout + ) + try: + if proxy_ssl_context is not None: + sock = proxy_ssl_context.wrap_socket(sock, server_hostname=_KMS_PROXY_HOST) + target = f"{context.host}:{context.port}" + sock.sendall(f"CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n".encode()) + response = b"" + while b"\r\n\r\n" not in response: + chunk = sock.recv(4096) + if not chunk: + raise OSError("proxy closed the connection before responding") + response += chunk + status = response.split(b"\r\n", 1)[0] + if not status.startswith(b"HTTP/1.1 200"): + raise OSError(f"proxy CONNECT failed: {status!r}") + except Exception: + sock.close() + raise + + if proxy_ssl_context is None: + return sock + + driver_side, relay_side = socket.socketpair() + + def relay(src, dst): + try: + while True: + buf = src.recv(16384) + if not buf: + break + dst.sendall(buf) + except OSError: + pass + finally: + for s in (src, dst): + try: + s.close() + except OSError: + pass + + threading.Thread(target=relay, args=(relay_side, sock), daemon=True).start() + threading.Thread(target=relay, args=(sock, relay_side), daemon=True).start() + return driver_side + + +# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-connect-callback +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) + if not _IS_SYNC: + return await asyncio.to_thread(_http_connect, context, _KMS_PROXY_PORT) + return _http_connect(context, _KMS_PROXY_PORT) + + async def tls_callback(self, context): + self.callback_calls.append(context) + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + if not _IS_SYNC: + return await asyncio.to_thread(_http_connect, context, _KMS_TLS_PROXY_PORT, ctx) + return _http_connect(context, _KMS_TLS_PROXY_PORT, ctx) + + async def proxy_request(self, method, path, tls=False): + """Call the proxy's control endpoints and return the body.""" + if tls: + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + 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) + # The body is one "key value" pair per line. Only connect_count is + # required by the spec; the server also emits connect_target lines. + 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) + + self.assertGreaterEqual(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.assertRaises(EncryptionError): + await encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + + 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: + 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") + if not _IS_SYNC: + return await asyncio.to_thread(_http_connect, context, _KMS_PROXY_PORT) + return _http_connect(context, _KMS_PROXY_PORT) + + 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) + + # 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/test_encryption.py b/test/test_encryption.py index 022635ad4e..9c5fe13f42 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 dataclasses @@ -29,6 +30,7 @@ import ssl import sys import textwrap +import threading import traceback import uuid import warnings @@ -2048,6 +2050,236 @@ def test_invalid_hostname_in_kms_certificate(self): self.client_encrypted.create_data_key("aws", master_key=key) +_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", +} + + +def _http_connect(context, proxy_port, proxy_ssl_context=None): + """Open an HTTP CONNECT tunnel to context.host:context.port. + + Returns a plain socket. When proxy_ssl_context is given, the connection to + the proxy is TLS and the tunnel is bridged to a socketpair, because TLS + cannot be layered over an ssl.SSLSocket. + """ + sock: socket.socket = socket.create_connection( + (_KMS_PROXY_HOST, proxy_port), timeout=context.timeout + ) + try: + if proxy_ssl_context is not None: + sock = proxy_ssl_context.wrap_socket(sock, server_hostname=_KMS_PROXY_HOST) + target = f"{context.host}:{context.port}" + sock.sendall(f"CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n".encode()) + response = b"" + while b"\r\n\r\n" not in response: + chunk = sock.recv(4096) + if not chunk: + raise OSError("proxy closed the connection before responding") + response += chunk + status = response.split(b"\r\n", 1)[0] + if not status.startswith(b"HTTP/1.1 200"): + raise OSError(f"proxy CONNECT failed: {status!r}") + except Exception: + sock.close() + raise + + if proxy_ssl_context is None: + return sock + + driver_side, relay_side = socket.socketpair() + + def relay(src, dst): + try: + while True: + buf = src.recv(16384) + if not buf: + break + dst.sendall(buf) + except OSError: + pass + finally: + for s in (src, dst): + try: + s.close() + except OSError: + pass + + threading.Thread(target=relay, args=(relay_side, sock), daemon=True).start() + threading.Thread(target=relay, args=(sock, relay_side), daemon=True).start() + return driver_side + + +# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-connect-callback +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) + if not _IS_SYNC: + return asyncio.to_thread(_http_connect, context, _KMS_PROXY_PORT) + return _http_connect(context, _KMS_PROXY_PORT) + + def tls_callback(self, context): + self.callback_calls.append(context) + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + if not _IS_SYNC: + return asyncio.to_thread(_http_connect, context, _KMS_TLS_PROXY_PORT, ctx) + return _http_connect(context, _KMS_TLS_PROXY_PORT, ctx) + + def proxy_request(self, method, path, tls=False): + """Call the proxy's control endpoints and return the body.""" + if tls: + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + 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) + # The body is one "key value" pair per line. Only connect_count is + # required by the spec; the server also emits connect_target lines. + 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) + + self.assertGreaterEqual(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.assertRaises(EncryptionError): + encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + + 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: + 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") + if not _IS_SYNC: + return asyncio.to_thread(_http_connect, context, _KMS_PROXY_PORT) + return _http_connect(context, _KMS_PROXY_PORT) + + 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) + + # 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") From 516b2675fb2ff12db89a1b500b86da08d185d75f Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 18 Aug 2026 05:42:57 -0500 Subject: [PATCH 08/50] PYTHON-5805 Fix KMSConnectContext.timeout docstring and document case 5 gap The docstring promised the driver could pass timeout=None, but the only producer (max(_csot.clamp_remaining(...), 0.001)) is always a positive float. Reworded to describe actual behavior without narrowing the Optional[float] type. Also added a comment on the case 5 timeout assertion in TestKmsConnectCallbackProse recording that explicit ClientEncryption operations set no CSOT deadline, so timeoutMS on the key-vault client does not currently tighten the value asserted. --- pymongo/encryption_options.py | 5 +++-- test/asynchronous/test_encryption.py | 5 +++++ test/test_encryption.py | 5 +++++ 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 0f91bf9739..6089147e23 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -157,8 +157,9 @@ def relay(src, dst): :param host: Hostname of the KMS server, and the target of TLS certificate and hostname verification. :param port: Port of the KMS server. - :param timeout: Seconds remaining before the operation's timeout expires, - or ``None`` when no timeout applies. + :param timeout: Seconds remaining in the operation's timeout budget when + one is active, otherwise the driver's default KMS connect timeout. + Always a positive number; the driver never passes ``None``. .. versionadded:: 4.18 """ diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 55c794314b..279c6a896c 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -2263,6 +2263,11 @@ async def test_05_callback_receives_timeout(self): self.assertTrue(self.callback_calls, "callback was never invoked") for context in self.callback_calls: + # This only checks the spec's literal requirement (a non-zero + # timeout). timeoutMS=1000 above does not currently tighten this + # value: explicit ClientEncryption operations establish no CSOT + # deadline, so this always falls back to the driver's default KMS + # connect timeout regardless of key_vault_client's timeoutMS. self.assertIsNotNone(context.timeout) self.assertGreater(context.timeout, 0) diff --git a/test/test_encryption.py b/test/test_encryption.py index 9c5fe13f42..f361229426 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -2255,6 +2255,11 @@ def test_05_callback_receives_timeout(self): self.assertTrue(self.callback_calls, "callback was never invoked") for context in self.callback_calls: + # This only checks the spec's literal requirement (a non-zero + # timeout). timeoutMS=1000 above does not currently tighten this + # value: explicit ClientEncryption operations establish no CSOT + # deadline, so this always falls back to the driver's default KMS + # connect timeout regardless of key_vault_client's timeoutMS. self.assertIsNotNone(context.timeout) self.assertGreater(context.timeout, 0) From 640b991bdf42099c12dd166831785685d23e12e3 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 18 Aug 2026 05:49:13 -0500 Subject: [PATCH 09/50] PYTHON-5805 Add changelog entry for kms_connect_callback --- doc/changelog.rst | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/doc/changelog.rst b/doc/changelog.rst index fb7d300b2e..8175670a89 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -39,6 +39,17 @@ PyMongo 4.18 brings a number of changes including: - Fixed a bug on Windows, and on macOS when using PyOpenSSL, where ``SSL_CERT_FILE``/``SSL_CERT_DIR`` were merged with, rather than replacing, the OS/certifi certificate store. +- 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 to the KMS host and the driver performs the KMS TLS + handshake over it, so certificate and hostname verification continue to + target the KMS host rather than the proxy. See + :class:`~pymongo.encryption_options.KMSConnectContext` for an HTTP + ``CONNECT`` example. Changes in Version 4.17.0 (2026/04/20) -------------------------------------- From 9c1c5449e7d0cc84795ad441c82e3d43392c7b37 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 18 Aug 2026 07:27:20 -0500 Subject: [PATCH 10/50] PYTHON-5805 Guard async callback contract and fix proxy example docs --- pymongo/asynchronous/encryption.py | 31 +++++++++++++++++++++++++++- pymongo/encryption_options.py | 28 +++++++++++++++++++++---- pymongo/synchronous/encryption.py | 31 +++++++++++++++++++++++++++- test/asynchronous/test_encryption.py | 2 +- test/test_encryption.py | 2 +- 5 files changed, 86 insertions(+), 8 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 7d2a877533..ca6e8de650 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -19,6 +19,7 @@ import asyncio import contextlib import enum +import inspect import socket import ssl import time as time # noqa: PLC0414 # needed in sync version @@ -124,6 +125,19 @@ class _KMSCallbackContractError(Exception): """ +def _close_rejected_kms_socket(obj: Any) -> None: + """Close a rejected kms_connect_callback return value, if it can be closed. + + The caller may have handed us a live socket, and nothing else will close it: + _connect_kms raises before its result reaches the caller's ``finally``. The + value can be anything a user returned, so closing is strictly best effort. + """ + close = getattr(obj, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + async def _connect_kms( address: _Address, opts: PoolOptions, @@ -139,10 +153,25 @@ async def _connect_kms( # Let the caller open the connection, then apply TLS ourselves so that SNI # and certificate verification still target the KMS host even when the # socket actually terminates at a proxy. - sock = await kms_connect_callback( + result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) + # _IS_SYNC is True in the generated synchronous flavor, where a regular + # function is the correct thing to pass and nothing is awaited. + if not _IS_SYNC and not inspect.isawaitable(result): + _close_rejected_kms_socket(result) + # "async" and "def" are deliberately split across the next two source + # lines: tools/synchro.py deletes "async " from any line that spells + # those two words together, which would garble this message in the + # generated synchronous file even though the branch never runs there. + raise _KMSCallbackContractError( + "kms_connect_callback must be a coroutine function (an 'async" + " def') for the async driver, but calling it returned " + f"{type(result)}, which is not awaitable." + ) + sock = await result if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): + _close_rejected_kms_socket(sock) raise _KMSCallbackContractError( "kms_connect_callback must return a connected, unwrapped " f"socket.socket, not {type(sock)}. TLS cannot be layered over an " diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 6089147e23..2a218d7b57 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -101,11 +101,26 @@ def connect_through_proxy(context): kms_connect_callback=connect_through_proxy, ) + The socket must be in blocking or timeout mode. + :meth:`ssl.SSLContext.wrap_socket` rejects a non-blocking socket, which + rules out the transports returned by :func:`asyncio.open_connection` and + :meth:`asyncio.loop.create_connection`. + For :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption` and :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient`, the callback - must be a coroutine function. It must not block the event loop, so drive - the connection with :mod:`asyncio` or hand the blocking work to a thread - with :func:`asyncio.to_thread`. + must be a coroutine function. Keep the event loop free by running the + blocking connect in a thread:: + + import asyncio + + async def async_connect_through_proxy(context): + return await asyncio.to_thread(connect_through_proxy, context) + + opts = AutoEncryptionOpts( + kms_providers={"aws": aws_creds}, + key_vault_namespace="keyvault.datakeys", + kms_connect_callback=async_connect_through_proxy, + ) Reaching the proxy itself over TLS takes one extra step. Python cannot layer a second TLS session over an :class:`ssl.SSLSocket`, so the callback @@ -147,8 +162,13 @@ def relay(src, dst): except OSError: pass finally: + # Unblock the sibling thread with EOF instead of closing a + # socket it may be reading, then close only this side. + try: + dst.shutdown(socket.SHUT_RDWR) + except OSError: + pass src.close() - dst.close() threading.Thread(target=relay, args=(relay_side, proxy), daemon=True).start() threading.Thread(target=relay, args=(proxy, relay_side), daemon=True).start() diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 48bd4f098a..7beed0c954 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -18,6 +18,7 @@ import contextlib import enum +import inspect import socket import ssl import time as time # noqa: PLC0414 # needed in sync version @@ -123,6 +124,19 @@ class _KMSCallbackContractError(Exception): """ +def _close_rejected_kms_socket(obj: Any) -> None: + """Close a rejected kms_connect_callback return value, if it can be closed. + + The caller may have handed us a live socket, and nothing else will close it: + _connect_kms raises before its result reaches the caller's ``finally``. The + value can be anything a user returned, so closing is strictly best effort. + """ + close = getattr(obj, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + def _connect_kms( address: _Address, opts: PoolOptions, @@ -138,10 +152,25 @@ def _connect_kms( # Let the caller open the connection, then apply TLS ourselves so that SNI # and certificate verification still target the KMS host even when the # socket actually terminates at a proxy. - sock = kms_connect_callback( + result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) + # _IS_SYNC is True in the generated synchronous flavor, where a regular + # function is the correct thing to pass and nothing is awaited. + if not _IS_SYNC and not inspect.isawaitable(result): + _close_rejected_kms_socket(result) + # "async" and "def" are deliberately split across the next two source + # lines: tools/synchro.py deletes "async " from any line that spells + # those two words together, which would garble this message in the + # generated synchronous file even though the branch never runs there. + raise _KMSCallbackContractError( + "kms_connect_callback must be a coroutine function (an 'async" + " def') for the async driver, but calling it returned " + f"{type(result)}, which is not awaitable." + ) + sock = result if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): + _close_rejected_kms_socket(sock) raise _KMSCallbackContractError( "kms_connect_callback must return a connected, unwrapped " f"socket.socket, not {type(sock)}. TLS cannot be layered over an " diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 279c6a896c..da56b8f3e7 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -2247,7 +2247,7 @@ async def failing_callback(context): OPTS, kms_connect_callback=failing_callback, ) - with self.assertRaises(EncryptionError): + with self.assertRaisesRegex(EncryptionError, "proxy is on fire"): await encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) async def test_05_callback_receives_timeout(self): diff --git a/test/test_encryption.py b/test/test_encryption.py index f361229426..8f81f41958 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -2239,7 +2239,7 @@ def failing_callback(context): OPTS, kms_connect_callback=failing_callback, ) - with self.assertRaises(EncryptionError): + with self.assertRaisesRegex(EncryptionError, "proxy is on fire"): encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) def test_05_callback_receives_timeout(self): From 93c144512990669c906a51e246b74745402009dc Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 18 Aug 2026 09:35:46 -0500 Subject: [PATCH 11/50] PYTHON-5805 Skip prose case 5 as a known CSOT spec discrepancy Case 5 asserts the KMS connect callback receives a non-zero timeout. That cannot fail in PyMongo: ClientEncryption does not support timeoutMS and explicit encryption operations establish no CSOT deadline, so the callback always receives the default KMS connect timeout. Skip the case rather than leave it passing vacuously, and record the deviation on KMSConnectContext, which the CSOT specification requires for any blocking section timeoutMS does not cover. Tracked in PYTHON-6037. --- pymongo/encryption_options.py | 8 ++++++++ test/asynchronous/test_encryption.py | 15 ++++++++++----- test/test_encryption.py | 15 ++++++++++----- 3 files changed, 28 insertions(+), 10 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 2a218d7b57..1f2e55167f 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -181,6 +181,14 @@ def relay(src, dst): one is active, otherwise the driver's default KMS connect timeout. Always a positive number; the driver never passes ``None``. + .. note:: ``timeoutMS`` configured on a + :class:`~pymongo.encryption.ClientEncryption` or on its key vault client + does not currently constrain KMS requests, so for explicit encryption + ``timeout`` is always the default KMS connect timeout. Automatic + encryption is unaffected and passes the remaining budget. This is a + known deviation from the Client Side Operations Timeout specification, + tracked in PYTHON-6037. + .. versionadded:: 4.18 """ diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index da56b8f3e7..b6d1fc1aa7 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -2250,6 +2250,10 @@ async def failing_callback(context): 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( @@ -2263,11 +2267,12 @@ async def test_05_callback_receives_timeout(self): self.assertTrue(self.callback_calls, "callback was never invoked") for context in self.callback_calls: - # This only checks the spec's literal requirement (a non-zero - # timeout). timeoutMS=1000 above does not currently tighten this - # value: explicit ClientEncryption operations establish no CSOT - # deadline, so this always falls back to the driver's default KMS - # connect timeout regardless of key_vault_client's timeoutMS. + # Skipped: this would only check the spec's literal requirement + # (a non-zero timeout), which cannot fail here. timeoutMS=1000 + # above does not tighten this value, because explicit + # ClientEncryption operations establish no CSOT deadline, so it + # always falls back to the driver's default KMS connect timeout. + # Un-skip once PYTHON-6037 adds timeoutMS to ClientEncryption. self.assertIsNotNone(context.timeout) self.assertGreater(context.timeout, 0) diff --git a/test/test_encryption.py b/test/test_encryption.py index 8f81f41958..0f72859e84 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -2242,6 +2242,10 @@ def failing_callback(context): 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( @@ -2255,11 +2259,12 @@ def test_05_callback_receives_timeout(self): self.assertTrue(self.callback_calls, "callback was never invoked") for context in self.callback_calls: - # This only checks the spec's literal requirement (a non-zero - # timeout). timeoutMS=1000 above does not currently tighten this - # value: explicit ClientEncryption operations establish no CSOT - # deadline, so this always falls back to the driver's default KMS - # connect timeout regardless of key_vault_client's timeoutMS. + # Skipped: this would only check the spec's literal requirement + # (a non-zero timeout), which cannot fail here. timeoutMS=1000 + # above does not tighten this value, because explicit + # ClientEncryption operations establish no CSOT deadline, so it + # always falls back to the driver's default KMS connect timeout. + # Un-skip once PYTHON-6037 adds timeoutMS to ClientEncryption. self.assertIsNotNone(context.timeout) self.assertGreater(context.timeout, 0) From ac823c8afeaa1e40578cbd18b83bf6150627c997 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 19 Aug 2026 05:24:00 -0500 Subject: [PATCH 12/50] PYTHON-5805 Tighten prose in docs, changelog, and test comments Factor the duplicated HTTP CONNECT handshake out of the two KMSConnectContext examples so the TLS-proxy example shows only what is different about it, the socketpair relay. Trim the CSOT deviation note, the changelog entry, and the case 5 comment, which restated the reason already carried by the skip decorator. No behavior change. --- doc/changelog.rst | 8 +-- pymongo/encryption_options.py | 103 +++++++++++---------------- test/asynchronous/test_encryption.py | 9 +-- test/test_encryption.py | 9 +-- 4 files changed, 51 insertions(+), 78 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 8175670a89..d366d6661a 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -45,11 +45,9 @@ PyMongo 4.18 brings a number of changes including: :class:`~pymongo.encryption_options.AutoEncryptionOpts`, :class:`~pymongo.encryption.ClientEncryption`, and :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`. The callback - opens the connection to the KMS host and the driver performs the KMS TLS - handshake over it, so certificate and hostname verification continue to - target the KMS host rather than the proxy. See - :class:`~pymongo.encryption_options.KMSConnectContext` for an HTTP - ``CONNECT`` example. + opens the connection and the driver performs the KMS TLS handshake over it, so + verification still targets the KMS host rather than the proxy. See + :class:`~pymongo.encryption_options.KMSConnectContext` for an example. Changes in Version 4.17.0 (2026/04/20) -------------------------------------- diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 1f2e55167f..6a15fbd5eb 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -64,21 +64,21 @@ class KMSConnectContext: :class:`AutoEncryptionOpts`, :class:`~pymongo.encryption.ClientEncryption`, or :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`. - The callback opens a connection to ``host``:``port`` and returns it. The - driver then performs the KMS TLS handshake over that socket, so the - callback must return a plain, unwrapped :class:`socket.socket`. Certificate - and hostname verification target ``host``, not whatever peer the socket is - actually connected to, which is what makes proxying safe. + The callback connects to ``host``:``port`` and returns a plain, unwrapped + :class:`socket.socket`. The driver then performs the KMS TLS handshake over + it, verifying the certificate and hostname against ``host`` rather than the + peer the socket actually reached. That is what makes proxying safe. - To reach a KMS host through an HTTP proxy, connect to the proxy and issue - an HTTP ``CONNECT`` request:: + The socket must be in blocking or timeout mode. + :meth:`ssl.SSLContext.wrap_socket` rejects non-blocking sockets, which rules + out :func:`asyncio.open_connection` and + :meth:`asyncio.loop.create_connection`. + + To reach a KMS host through an HTTP proxy, tunnel with ``CONNECT``:: import socket - def connect_through_proxy(context): - sock = socket.create_connection( - ("proxy.example.com", 8080), timeout=context.timeout - ) + def open_tunnel(sock, context): target = f"{context.host}:{context.port}" sock.sendall( f"CONNECT {target} HTTP/1.1\\r\\nHost: {target}\\r\\n\\r\\n".encode() @@ -87,12 +87,20 @@ def connect_through_proxy(context): while b"\\r\\n\\r\\n" not in response: chunk = sock.recv(4096) if not chunk: - sock.close() raise OSError("proxy closed the connection") response += chunk if not response.startswith(b"HTTP/1.1 200"): + raise OSError(f"CONNECT failed: {response.splitlines()[0]!r}") + + def connect_through_proxy(context): + sock = socket.create_connection( + ("proxy.example.com", 8080), timeout=context.timeout + ) + try: + open_tunnel(sock, context) + except OSError: sock.close() - raise OSError(f"proxy CONNECT failed: {response.splitlines()[0]!r}") + raise return sock opts = AutoEncryptionOpts( @@ -101,54 +109,29 @@ def connect_through_proxy(context): kms_connect_callback=connect_through_proxy, ) - The socket must be in blocking or timeout mode. - :meth:`ssl.SSLContext.wrap_socket` rejects a non-blocking socket, which - rules out the transports returned by :func:`asyncio.open_connection` and - :meth:`asyncio.loop.create_connection`. - - For :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption` and - :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient`, the callback - must be a coroutine function. Keep the event loop free by running the - blocking connect in a thread:: - - import asyncio + The async API requires a coroutine function. Run the blocking connect in a + thread to keep the event loop free:: async def async_connect_through_proxy(context): return await asyncio.to_thread(connect_through_proxy, context) - opts = AutoEncryptionOpts( - kms_providers={"aws": aws_creds}, - key_vault_namespace="keyvault.datakeys", - kms_connect_callback=async_connect_through_proxy, - ) - - Reaching the proxy itself over TLS takes one extra step. Python cannot - layer a second TLS session over an :class:`ssl.SSLSocket`, so the callback - cannot return its TLS connection to the proxy directly. Bridge it to a - :func:`socket.socketpair` and return the plain end instead:: + Reaching the proxy itself over TLS needs one extra step, because Python + cannot layer TLS over an :class:`ssl.SSLSocket`. Relay the proxy connection + through a :func:`socket.socketpair` and return the plain end:: - import socket, ssl, threading + import ssl, threading def connect_through_tls_proxy(context): - proxy_ctx = ssl.create_default_context(cafile="proxy-ca.pem") + ctx = ssl.create_default_context(cafile="proxy-ca.pem") raw = socket.create_connection( ("proxy.example.com", 8443), timeout=context.timeout ) - proxy = proxy_ctx.wrap_socket(raw, server_hostname="proxy.example.com") - target = f"{context.host}:{context.port}" - proxy.sendall( - f"CONNECT {target} HTTP/1.1\\r\\nHost: {target}\\r\\n\\r\\n".encode() - ) - response = b"" - while b"\\r\\n\\r\\n" not in response: - chunk = proxy.recv(4096) - if not chunk: - proxy.close() - raise OSError("proxy closed the connection") - response += chunk - if not response.startswith(b"HTTP/1.1 200"): + proxy = ctx.wrap_socket(raw, server_hostname="proxy.example.com") + try: + open_tunnel(proxy, context) + except OSError: proxy.close() - raise OSError(f"proxy CONNECT failed: {response.splitlines()[0]!r}") + raise driver_side, relay_side = socket.socketpair() @@ -162,16 +145,15 @@ def relay(src, dst): except OSError: pass finally: - # Unblock the sibling thread with EOF instead of closing a - # socket it may be reading, then close only this side. + # EOF the peer instead of closing a socket it may be reading. try: dst.shutdown(socket.SHUT_RDWR) except OSError: pass src.close() - threading.Thread(target=relay, args=(relay_side, proxy), daemon=True).start() - threading.Thread(target=relay, args=(proxy, relay_side), daemon=True).start() + for pair in ((relay_side, proxy), (proxy, relay_side)): + threading.Thread(target=relay, args=pair, daemon=True).start() return driver_side :param host: Hostname of the KMS server, and the target of TLS certificate @@ -181,13 +163,12 @@ def relay(src, dst): one is active, otherwise the driver's default KMS connect timeout. Always a positive number; the driver never passes ``None``. - .. note:: ``timeoutMS`` configured on a - :class:`~pymongo.encryption.ClientEncryption` or on its key vault client - does not currently constrain KMS requests, so for explicit encryption - ``timeout`` is always the default KMS connect timeout. Automatic - encryption is unaffected and passes the remaining budget. This is a - known deviation from the Client Side Operations Timeout specification, - tracked in PYTHON-6037. + .. note:: ``timeoutMS`` on a + :class:`~pymongo.encryption.ClientEncryption` or its key vault client + does not constrain KMS requests, so for explicit encryption ``timeout`` + is always the default. Automatic encryption passes the remaining budget. + This deviates from the Client Side Operations Timeout specification and + is tracked in PYTHON-6037. .. versionadded:: 4.18 """ diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index b6d1fc1aa7..2998b8a24f 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -2267,12 +2267,9 @@ async def test_05_callback_receives_timeout(self): self.assertTrue(self.callback_calls, "callback was never invoked") for context in self.callback_calls: - # Skipped: this would only check the spec's literal requirement - # (a non-zero timeout), which cannot fail here. timeoutMS=1000 - # above does not tighten this value, because explicit - # ClientEncryption operations establish no CSOT deadline, so it - # always falls back to the driver's default KMS connect timeout. - # Un-skip once PYTHON-6037 adds timeoutMS to ClientEncryption. + # Only checks the spec's literal non-zero requirement, which + # cannot fail: timeoutMS does not tighten this value because + # explicit ClientEncryption operations set no CSOT deadline. self.assertIsNotNone(context.timeout) self.assertGreater(context.timeout, 0) diff --git a/test/test_encryption.py b/test/test_encryption.py index 0f72859e84..a1a5119fa0 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -2259,12 +2259,9 @@ def test_05_callback_receives_timeout(self): self.assertTrue(self.callback_calls, "callback was never invoked") for context in self.callback_calls: - # Skipped: this would only check the spec's literal requirement - # (a non-zero timeout), which cannot fail here. timeoutMS=1000 - # above does not tighten this value, because explicit - # ClientEncryption operations establish no CSOT deadline, so it - # always falls back to the driver's default KMS connect timeout. - # Un-skip once PYTHON-6037 adds timeoutMS to ClientEncryption. + # Only checks the spec's literal non-zero requirement, which + # cannot fail: timeoutMS does not tighten this value because + # explicit ClientEncryption operations set no CSOT deadline. self.assertIsNotNone(context.timeout) self.assertGreater(context.timeout, 0) From b3d0c299104171f17df4e7dc35cc808547f85501 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 19 Aug 2026 20:42:11 -0500 Subject: [PATCH 13/50] PYTHON-5805 Use ConfigurationError and drop the synchro docstring tripwire Raise ConfigurationError for kms_connect_callback contract violations instead of a private sentinel exception. It is the right public type for a misconfigured callback, and the no-retry clause in kms_request now keys off it. Remove the flavor-specific 'Must be a coroutine function.' sentence from the ClientEncryption parameter docs. KMSConnectContext already documents both flavors, so the synchro replacement entry and the tripwire that guarded it are no longer needed. Rewording the awaitable-guard message also removes the split string literals that dodged synchro's rewriting. --- pymongo/asynchronous/encryption.py | 25 +++++------------- pymongo/synchronous/encryption.py | 25 +++++------------- test/asynchronous/test_encryption.py | 7 +++-- test/test_encryption.py | 7 +++-- tools/synchro.py | 38 ---------------------------- 5 files changed, 18 insertions(+), 84 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index ca6e8de650..504f144381 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -117,14 +117,6 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -class _KMSCallbackContractError(Exception): - """Raised when a kms_connect_callback violates its contract. - - Unlike a network error from the callback, this is a programming error, so - it is never retried. - """ - - def _close_rejected_kms_socket(obj: Any) -> None: """Close a rejected kms_connect_callback return value, if it can be closed. @@ -160,19 +152,14 @@ async def _connect_kms( # function is the correct thing to pass and nothing is awaited. if not _IS_SYNC and not inspect.isawaitable(result): _close_rejected_kms_socket(result) - # "async" and "def" are deliberately split across the next two source - # lines: tools/synchro.py deletes "async " from any line that spells - # those two words together, which would garble this message in the - # generated synchronous file even though the branch never runs there. - raise _KMSCallbackContractError( - "kms_connect_callback must be a coroutine function (an 'async" - " def') for the async driver, but calling it returned " - f"{type(result)}, which is not awaitable." + raise ConfigurationError( + "kms_connect_callback must be a coroutine function for the async " + f"API, but calling it returned {type(result)}, which is not awaitable." ) sock = await result if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): _close_rejected_kms_socket(sock) - raise _KMSCallbackContractError( + raise ConfigurationError( "kms_connect_callback must return a connected, unwrapped " f"socket.socket, not {type(sock)}. TLS cannot be layered over an " "already-wrapped socket; to reach the proxy over TLS, relay through " @@ -303,7 +290,7 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: conn.close() except MongoCryptError: raise # Propagate MongoCryptError errors directly. - except _KMSCallbackContractError: + except ConfigurationError: raise # A callback contract violation is not transient. except Exception as exc: remaining = _csot.remaining() @@ -742,7 +729,7 @@ def __init__( 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`; the - driver then performs the KMS TLS handshake over it. Must be a coroutine function. + driver then performs the KMS TLS handshake over it. See :class:`~pymongo.encryption_options.KMSConnectContext` for a worked HTTP ``CONNECT`` example. Defaults to ``None``, meaning the driver connects to KMS hosts directly. diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 7beed0c954..e7df3688cf 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -116,14 +116,6 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -class _KMSCallbackContractError(Exception): - """Raised when a kms_connect_callback violates its contract. - - Unlike a network error from the callback, this is a programming error, so - it is never retried. - """ - - def _close_rejected_kms_socket(obj: Any) -> None: """Close a rejected kms_connect_callback return value, if it can be closed. @@ -159,19 +151,14 @@ def _connect_kms( # function is the correct thing to pass and nothing is awaited. if not _IS_SYNC and not inspect.isawaitable(result): _close_rejected_kms_socket(result) - # "async" and "def" are deliberately split across the next two source - # lines: tools/synchro.py deletes "async " from any line that spells - # those two words together, which would garble this message in the - # generated synchronous file even though the branch never runs there. - raise _KMSCallbackContractError( - "kms_connect_callback must be a coroutine function (an 'async" - " def') for the async driver, but calling it returned " - f"{type(result)}, which is not awaitable." + raise ConfigurationError( + "kms_connect_callback must be a coroutine function for the async " + f"API, but calling it returned {type(result)}, which is not awaitable." ) sock = result if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): _close_rejected_kms_socket(sock) - raise _KMSCallbackContractError( + raise ConfigurationError( "kms_connect_callback must return a connected, unwrapped " f"socket.socket, not {type(sock)}. TLS cannot be layered over an " "already-wrapped socket; to reach the proxy over TLS, relay through " @@ -302,7 +289,7 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: conn.close() except MongoCryptError: raise # Propagate MongoCryptError errors directly. - except _KMSCallbackContractError: + except ConfigurationError: raise # A callback contract violation is not transient. except Exception as exc: remaining = _csot.remaining() @@ -739,7 +726,7 @@ def __init__( 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`; the - driver then performs the KMS TLS handshake over it. Must be a regular function. + driver then performs the KMS TLS handshake over it. See :class:`~pymongo.encryption_options.KMSConnectContext` for a worked HTTP ``CONNECT`` example. Defaults to ``None``, meaning the driver connects to KMS hosts directly. diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 2998b8a24f..786669f319 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -67,7 +67,6 @@ AsyncClientEncryption, QueryType, _connect_kms, - _KMSCallbackContractError, ) from pymongo.asynchronous.helpers import anext from pymongo.asynchronous.mongo_client import AsyncMongoClient @@ -263,7 +262,7 @@ async def test_non_socket_return_is_not_retried(self): async def callback(context): return "not-a-socket" - with self.assertRaisesRegex(_KMSCallbackContractError, "must return a connected"): + 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): @@ -283,7 +282,7 @@ async def test_already_wrapped_socket_is_rejected(self): async def callback(context): return wrapped - with self.assertRaisesRegex(_KMSCallbackContractError, "unwrapped"): + 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): @@ -310,7 +309,7 @@ async def test_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") - # Not wrapped in _KMSCallbackContractError, so kms_request's broad + # Not a ConfigurationError, so kms_request's broad # handler can treat it as transient and let libmongocrypt retry. with self.assertRaises(OSError): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) diff --git a/test/test_encryption.py b/test/test_encryption.py index a1a5119fa0..e3e0d561b3 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -84,7 +84,6 @@ ClientEncryption, QueryType, _connect_kms, - _KMSCallbackContractError, ) from pymongo.synchronous.helpers import next from pymongo.synchronous.mongo_client import MongoClient @@ -263,7 +262,7 @@ def test_non_socket_return_is_not_retried(self): def callback(context): return "not-a-socket" - with self.assertRaisesRegex(_KMSCallbackContractError, "must return a connected"): + 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): @@ -283,7 +282,7 @@ def test_already_wrapped_socket_is_rejected(self): def callback(context): return wrapped - with self.assertRaisesRegex(_KMSCallbackContractError, "unwrapped"): + 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): @@ -310,7 +309,7 @@ def test_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") - # Not wrapped in _KMSCallbackContractError, so kms_request's broad + # Not a ConfigurationError, so kms_request's broad # handler can treat it as transient and let libmongocrypt retry. with self.assertRaises(OSError): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) diff --git a/tools/synchro.py b/tools/synchro.py index 8844b20e55..bc81c1e65e 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -143,11 +143,6 @@ "dns.asyncresolver.resolve": "dns.resolver.resolve", "__aenter__": "__enter__", "__aexit__": "__exit__", - # Prose substitution, not an identifier/token like the rest of this dict. Matching is - # line-based (see translate_docstrings), so this sentence must stay on a single line in - # pymongo/asynchronous/encryption.py or the replacement silently stops firing. Guarded by - # check_kms_connect_callback_docstring(), called from main() below. - "Must be a coroutine function.": "Must be a regular function.", } docstring_replacements: dict[tuple[str, str], str] = { @@ -340,37 +335,6 @@ def process_ignores(lines: list[str]) -> list[str]: return lines -def check_kms_connect_callback_docstring() -> None: - """Guard the "Must be a coroutine function." -> "Must be a regular function." mapping. - - That replacement in `replacements` above only fires if the sentence sits on a single - line in the async source (see translate_docstrings). A future edit or re-wrap could - silently break the match, leaving the generated synchronous docs telling users to write - an `async def` callback. Fail loudly instead of leaving that undetected. - """ - sync_encryption = Path(_pymongo_dest_base) / "encryption.py" - if not sync_encryption.is_file(): - # Nothing to check yet, e.g. a partial/filtered run that didn't touch this file. - return - content = sync_encryption.read_text() - if "Must be a coroutine function." in content: - raise RuntimeError( - f"{sync_encryption} still says 'Must be a coroutine function.' after synchro. " - "The 'Must be a coroutine function.' -> 'Must be a regular function.' entry in " - "tools/synchro.py's `replacements` dict didn't fire, most likely because the " - "sentence in pymongo/asynchronous/encryption.py got wrapped across multiple " - "lines. Keep it on one line, or fix the replacement mapping." - ) - if "kms_connect_callback" in content and "Must be a regular function." not in content: - raise RuntimeError( - f"{sync_encryption} defines kms_connect_callback but its docstring no longer " - "says 'Must be a regular function.'. Check that the " - "'Must be a coroutine function.' -> 'Must be a regular function.' entry is still " - "present in tools/synchro.py's `replacements` dict and that the docstring wording " - "in pymongo/asynchronous/encryption.py hasn't changed out from under it." - ) - - def unasync_directory(files: list[str], src: str, dest: str, replacements: dict[str, str]) -> None: unasync_files( files, @@ -439,8 +403,6 @@ def main() -> None: generated_tests, ) - check_kms_connect_callback_docstring() - generated_files = generated_pymongo + generated_gridfs + generated_tests if is_ci and generated_files: From defb5c5b05c4d0211836d7878cd1140f38cca76d Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 06:34:09 -0500 Subject: [PATCH 14/50] PYTHON-5805 Normalize the callback socket's blocking mode ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback has no reason to care which mode it leaves the socket in. Set the timeout in _connect_kms rather than pushing the requirement onto the caller. Do it there and not in _async_wrap_socket_tls, which is shared with every MongoDB connection and currently handshakes under the connect-derived timeout; forcing socket_timeout for all callers would change behavior on that path. Narrow the docstring accordingly. The real constraint is a real socket the event loop is not managing, not the blocking mode: asyncio streams and transports are not sockets, and the socket under one stays registered with the loop. --- pymongo/asynchronous/encryption.py | 4 ++++ pymongo/encryption_options.py | 10 ++++---- pymongo/synchronous/encryption.py | 4 ++++ test/asynchronous/test_encryption.py | 36 ++++++++++++++++++++++++++++ test/test_encryption.py | 36 ++++++++++++++++++++++++++++ 5 files changed, 86 insertions(+), 4 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 504f144381..45c5e7d736 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -165,6 +165,10 @@ async def _connect_kms( "already-wrapped socket; to reach the proxy over TLS, relay through " "a socket.socketpair and return the plain end." ) + # ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback + # has no reason to care which mode it left the socket in, so normalize it + # here rather than pushing the requirement onto the caller. + sock.settimeout(opts.socket_timeout) try: return await _async_wrap_socket_tls(sock, address, opts) except Exception as exc: diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 6a15fbd5eb..b5003300d5 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -69,10 +69,12 @@ class KMSConnectContext: it, verifying the certificate and hostname against ``host`` rather than the peer the socket actually reached. That is what makes proxying safe. - The socket must be in blocking or timeout mode. - :meth:`ssl.SSLContext.wrap_socket` rejects non-blocking sockets, which rules - out :func:`asyncio.open_connection` and - :meth:`asyncio.loop.create_connection`. + The callback must return a real socket that the event loop is not managing. + :func:`asyncio.open_connection` and :meth:`asyncio.loop.create_connection` + return a stream or transport rather than a socket, and the socket + underneath one stays registered with the running loop, so neither can be + used here. The driver sets the socket's timeout itself, so the mode the + callback leaves it in does not matter. To reach a KMS host through an HTTP proxy, tunnel with ``CONNECT``:: diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index e7df3688cf..314047b0cc 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -164,6 +164,10 @@ def _connect_kms( "already-wrapped socket; to reach the proxy over TLS, relay through " "a socket.socketpair and return the plain end." ) + # ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback + # has no reason to care which mode it left the socket in, so normalize it + # here rather than pushing the requirement onto the caller. + sock.settimeout(opts.socket_timeout) try: return _wrap_socket_tls(sock, address, opts) except Exception as exc: diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 786669f319..5da6a83a52 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -87,6 +87,7 @@ ) 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, @@ -305,6 +306,41 @@ async def callback(context): self.assertEqual(received[0].port, 443) self.assertEqual(received[0].timeout, 12.5) + async def test_non_blocking_socket_from_callback_is_accepted(self): + # ssl.SSLContext.wrap_socket refuses a non-blocking socket. The driver + # normalizes the mode so a callback need not care which mode it leaves + # the socket in. Without that, this handshake 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() + + # Build the context the way the driver does, so PoolOptions gets the + # flavor-correct type. Verification is off because the local test + # server's certificate is not what the driver would expect. + 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) + + async def callback(context): + sock = socket.create_connection(listener.getsockname(), timeout=10) + sock.setblocking(False) + return sock + + conn = await _connect_kms(listener.getsockname(), options, callback, 10.0) + self.addCleanup(conn.close) + self.assertIsNotNone(conn.gettimeout()) + async def test_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") diff --git a/test/test_encryption.py b/test/test_encryption.py index e3e0d561b3..c3a2a820e2 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -78,6 +78,7 @@ ) 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, @@ -305,6 +306,41 @@ def callback(context): self.assertEqual(received[0].port, 443) self.assertEqual(received[0].timeout, 12.5) + def test_non_blocking_socket_from_callback_is_accepted(self): + # ssl.SSLContext.wrap_socket refuses a non-blocking socket. The driver + # normalizes the mode so a callback need not care which mode it leaves + # the socket in. Without that, this handshake 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() + + # Build the context the way the driver does, so PoolOptions gets the + # flavor-correct type. Verification is off because the local test + # server's certificate is not what the driver would expect. + 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 callback(context): + sock = socket.create_connection(listener.getsockname(), timeout=10) + sock.setblocking(False) + return sock + + conn = _connect_kms(listener.getsockname(), options, callback, 10.0) + self.addCleanup(conn.close) + self.assertIsNotNone(conn.gettimeout()) + def test_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") From fba651bf957d69bfbaf60e35e153c873fc4aab35 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 06:52:37 -0500 Subject: [PATCH 15/50] PYTHON-5805 Name the asyncio transport case in the callback error transport.get_extra_info('socket') returns an asyncio TransportSocket, which is not a socket.socket, so the existing contract check already rejects it. Say so in the message and point at loop.sock_connect, rather than leaving the caller to work out why their socket was refused. Also correct the docstring: loop.sock_connect leaves an ordinary socket behind once it completes, so it is usable here. Only streams, transports, and the transport socket underneath them are not. --- pymongo/asynchronous/encryption.py | 9 ++++++--- pymongo/encryption_options.py | 10 ++++++---- pymongo/synchronous/encryption.py | 9 ++++++--- test/asynchronous/test_encryption.py | 16 ++++++++++++++++ test/test_encryption.py | 16 ++++++++++++++++ 5 files changed, 50 insertions(+), 10 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 45c5e7d736..425b7d3434 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -161,9 +161,12 @@ async def _connect_kms( _close_rejected_kms_socket(sock) raise ConfigurationError( "kms_connect_callback must return a connected, unwrapped " - f"socket.socket, not {type(sock)}. TLS cannot be layered over an " - "already-wrapped socket; to reach the proxy over TLS, relay through " - "a socket.socketpair and return the plain end." + f"socket.socket, not {type(sock)}. An asyncio stream or transport, " + "including the object from transport.get_extra_info('socket'), " + "cannot be used; connect with loop.sock_connect instead. TLS cannot " + "be layered over an already-wrapped socket either; to reach the " + "proxy over TLS, relay through a socket.socketpair and return the " + "plain end." ) # ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback # has no reason to care which mode it left the socket in, so normalize it diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index b5003300d5..ffda7cd216 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -71,10 +71,12 @@ class KMSConnectContext: The callback must return a real socket that the event loop is not managing. :func:`asyncio.open_connection` and :meth:`asyncio.loop.create_connection` - return a stream or transport rather than a socket, and the socket - underneath one stays registered with the running loop, so neither can be - used here. The driver sets the socket's timeout itself, so the mode the - callback leaves it in does not matter. + return a stream or transport rather than a socket, and the object from + ``transport.get_extra_info("socket")`` is a transport socket the loop still + owns, so none of them can be used here. :meth:`asyncio.loop.sock_connect` + is fine: it leaves an ordinary socket behind once it completes. The driver + sets the socket's timeout itself, so the mode the callback leaves it in + does not matter. To reach a KMS host through an HTTP proxy, tunnel with ``CONNECT``:: diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 314047b0cc..2df2c8df9c 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -160,9 +160,12 @@ def _connect_kms( _close_rejected_kms_socket(sock) raise ConfigurationError( "kms_connect_callback must return a connected, unwrapped " - f"socket.socket, not {type(sock)}. TLS cannot be layered over an " - "already-wrapped socket; to reach the proxy over TLS, relay through " - "a socket.socketpair and return the plain end." + f"socket.socket, not {type(sock)}. An asyncio stream or transport, " + "including the object from transport.get_extra_info('socket'), " + "cannot be used; connect with loop.sock_connect instead. TLS cannot " + "be layered over an already-wrapped socket either; to reach the " + "proxy over TLS, relay through a socket.socketpair and return the " + "plain end." ) # ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback # has no reason to care which mode it left the socket in, so normalize it diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 5da6a83a52..f3f77fc9dc 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -341,6 +341,22 @@ async def callback(context): self.addCleanup(conn.close) self.assertIsNotNone(conn.gettimeout()) + async def test_asyncio_transport_socket_is_rejected(self): + # transport.get_extra_info("socket") returns an asyncio TransportSocket, + # not a socket.socket, and the loop still owns it. Reject it with a + # message that names the mistake rather than failing later. + from asyncio.trsock import TransportSocket + + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + async def callback(context): + return TransportSocket(left) + + with self.assertRaisesRegex(ConfigurationError, "asyncio stream or transport"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + async def test_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") diff --git a/test/test_encryption.py b/test/test_encryption.py index c3a2a820e2..5d86756339 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -341,6 +341,22 @@ def callback(context): self.addCleanup(conn.close) self.assertIsNotNone(conn.gettimeout()) + def test_asyncio_transport_socket_is_rejected(self): + # transport.get_extra_info("socket") returns an asyncio TransportSocket, + # not a socket.socket, and the loop still owns it. Reject it with a + # message that names the mistake rather than failing later. + from asyncio.trsock import TransportSocket + + left, right = socket.socketpair() + self.addCleanup(left.close) + self.addCleanup(right.close) + + def callback(context): + return TransportSocket(left) + + with self.assertRaisesRegex(ConfigurationError, "asyncio stream or transport"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + def test_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") From 4c20867606787140b3523a6f482873e137a0be82 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 07:27:07 -0500 Subject: [PATCH 16/50] PYTHON-5805 Tighten prose further in docstrings, comments, and errors Merge the two KMSConnectContext paragraphs that both described what the callback returns, and shorten the asyncio guidance to the three facts a caller needs. Shorten the contract error: an exception should point at the mistake, not restate the docstring. Drop four test comments that the assertion on the next line already states, and condense the ones that explain something non-obvious. --- pymongo/asynchronous/encryption.py | 9 +++---- pymongo/encryption_options.py | 38 +++++++++++++--------------- pymongo/synchronous/encryption.py | 9 +++---- test/asynchronous/test_encryption.py | 24 ++++++------------ test/test_encryption.py | 24 ++++++------------ 5 files changed, 39 insertions(+), 65 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 425b7d3434..84f7e278de 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -161,12 +161,9 @@ async def _connect_kms( _close_rejected_kms_socket(sock) raise ConfigurationError( "kms_connect_callback must return a connected, unwrapped " - f"socket.socket, not {type(sock)}. An asyncio stream or transport, " - "including the object from transport.get_extra_info('socket'), " - "cannot be used; connect with loop.sock_connect instead. TLS cannot " - "be layered over an already-wrapped socket either; to reach the " - "proxy over TLS, relay through a socket.socketpair and return the " - "plain end." + f"socket.socket, not {type(sock)}. Streams, transports and " + "transport sockets cannot be used; try loop.sock_connect. For a " + "TLS proxy, relay through a socket.socketpair and return the plain end." ) # ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback # has no reason to care which mode it left the socket in, so normalize it diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index ffda7cd216..9af375d71a 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -65,18 +65,15 @@ class KMSConnectContext: or :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`. The callback connects to ``host``:``port`` and returns a plain, unwrapped - :class:`socket.socket`. The driver then performs the KMS TLS handshake over - it, verifying the certificate and hostname against ``host`` rather than the - peer the socket actually reached. That is what makes proxying safe. - - The callback must return a real socket that the event loop is not managing. - :func:`asyncio.open_connection` and :meth:`asyncio.loop.create_connection` - return a stream or transport rather than a socket, and the object from - ``transport.get_extra_info("socket")`` is a transport socket the loop still - owns, so none of them can be used here. :meth:`asyncio.loop.sock_connect` - is fine: it leaves an ordinary socket behind once it completes. The driver - sets the socket's timeout itself, so the mode the callback leaves it in - does not matter. + :class:`socket.socket`. The driver performs the KMS TLS handshake over it, + verifying against ``host`` rather than the peer actually reached, which is + what makes proxying safe. The driver also sets the socket's timeout, so the + mode the callback leaves it in does not matter. + + Streams and transports do not work: :func:`asyncio.open_connection` and + :meth:`asyncio.loop.create_connection` do not return sockets, and + ``transport.get_extra_info("socket")`` returns one the loop still owns. + :meth:`asyncio.loop.sock_connect` is fine. To reach a KMS host through an HTTP proxy, tunnel with ``CONNECT``:: @@ -119,9 +116,9 @@ def connect_through_proxy(context): async def async_connect_through_proxy(context): return await asyncio.to_thread(connect_through_proxy, context) - Reaching the proxy itself over TLS needs one extra step, because Python - cannot layer TLS over an :class:`ssl.SSLSocket`. Relay the proxy connection - through a :func:`socket.socketpair` and return the plain end:: + Reaching the proxy over TLS needs one extra step, because Python cannot + layer TLS over an :class:`ssl.SSLSocket`. Relay through a + :func:`socket.socketpair` and return the plain end:: import ssl, threading @@ -167,12 +164,11 @@ def relay(src, dst): one is active, otherwise the driver's default KMS connect timeout. Always a positive number; the driver never passes ``None``. - .. note:: ``timeoutMS`` on a - :class:`~pymongo.encryption.ClientEncryption` or its key vault client - does not constrain KMS requests, so for explicit encryption ``timeout`` - is always the default. Automatic encryption passes the remaining budget. - This deviates from the Client Side Operations Timeout specification and - is tracked in PYTHON-6037. + .. note:: ``timeoutMS`` on a :class:`~pymongo.encryption.ClientEncryption` + or its key vault client does not constrain KMS requests, so ``timeout`` + is always the default for explicit encryption. Automatic encryption + passes the remaining budget. A known deviation from the Client Side + Operations Timeout specification, tracked in PYTHON-6037. .. versionadded:: 4.18 """ diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 2df2c8df9c..124159a686 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -160,12 +160,9 @@ def _connect_kms( _close_rejected_kms_socket(sock) raise ConfigurationError( "kms_connect_callback must return a connected, unwrapped " - f"socket.socket, not {type(sock)}. An asyncio stream or transport, " - "including the object from transport.get_extra_info('socket'), " - "cannot be used; connect with loop.sock_connect instead. TLS cannot " - "be layered over an already-wrapped socket either; to reach the " - "proxy over TLS, relay through a socket.socketpair and return the " - "plain end." + f"socket.socket, not {type(sock)}. Streams, transports and " + "transport sockets cannot be used; try loop.sock_connect. For a " + "TLS proxy, relay through a socket.socketpair and return the plain end." ) # ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback # has no reason to care which mode it left the socket in, so normalize it diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index f3f77fc9dc..20e51a68f9 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -227,23 +227,19 @@ async def test_init_kms_tls_options(self): async def test_init_kms_connect_callback(self): from pymongo.encryption_options import KMSConnectContext - # Default is None. opts = AutoEncryptionOpts({}, "k.d") self.assertIsNone(opts._kms_connect_callback) - # A callable is accepted and stored unchanged. async def callback(context): raise AssertionError("not called") opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) self.assertIs(opts._kms_connect_callback, callback) - # Non-callables are rejected eagerly. 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] - # The context is frozen and carries host, port, and timeout. context = KMSConnectContext(host="kms.example.com", port=443, timeout=9.5) self.assertEqual(context.host, "kms.example.com") self.assertEqual(context.port, 443) @@ -268,8 +264,7 @@ async def callback(context): async def test_already_wrapped_socket_is_rejected(self): # ssl.SSLSocket subclasses socket.socket, but wrap_socket cannot layer - # TLS over it, so it must be rejected with an actionable message rather - # than failing later with an opaque handshake error. + # TLS over it, so it needs its own rejection. ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) ctx.check_hostname = False ctx.verify_mode = ssl.CERT_NONE @@ -307,9 +302,8 @@ async def callback(context): self.assertEqual(received[0].timeout, 12.5) async def test_non_blocking_socket_from_callback_is_accepted(self): - # ssl.SSLContext.wrap_socket refuses a non-blocking socket. The driver - # normalizes the mode so a callback need not care which mode it leaves - # the socket in. Without that, this handshake raises ValueError. + # wrap_socket refuses a non-blocking socket; the driver normalizes the + # mode. Without that, this handshake raises ValueError. server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) server_ctx.load_cert_chain(CLIENT_PEM) listener = socket.socket() @@ -326,9 +320,8 @@ def serve(): threading.Thread(target=serve, daemon=True).start() - # Build the context the way the driver does, so PoolOptions gets the - # flavor-correct type. Verification is off because the local test - # server's certificate is not what the driver would expect. + # Built the way the driver does, so PoolOptions gets the flavor-correct + # type. Verification is off: the local cert is not what it expects. 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) @@ -342,9 +335,8 @@ async def callback(context): self.assertIsNotNone(conn.gettimeout()) async def test_asyncio_transport_socket_is_rejected(self): - # transport.get_extra_info("socket") returns an asyncio TransportSocket, - # not a socket.socket, and the loop still owns it. Reject it with a - # message that names the mistake rather than failing later. + # transport.get_extra_info("socket") is an asyncio TransportSocket, not + # a socket.socket, and the loop still owns it. from asyncio.trsock import TransportSocket left, right = socket.socketpair() @@ -354,7 +346,7 @@ async def test_asyncio_transport_socket_is_rejected(self): async def callback(context): return TransportSocket(left) - with self.assertRaisesRegex(ConfigurationError, "asyncio stream or transport"): + with self.assertRaisesRegex(ConfigurationError, "Streams, transports"): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) async def test_network_error_from_callback_propagates(self): diff --git a/test/test_encryption.py b/test/test_encryption.py index 5d86756339..6b97628971 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -227,23 +227,19 @@ def test_init_kms_tls_options(self): def test_init_kms_connect_callback(self): from pymongo.encryption_options import KMSConnectContext - # Default is None. opts = AutoEncryptionOpts({}, "k.d") self.assertIsNone(opts._kms_connect_callback) - # A callable is accepted and stored unchanged. def callback(context): raise AssertionError("not called") opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) self.assertIs(opts._kms_connect_callback, callback) - # Non-callables are rejected eagerly. 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] - # The context is frozen and carries host, port, and timeout. context = KMSConnectContext(host="kms.example.com", port=443, timeout=9.5) self.assertEqual(context.host, "kms.example.com") self.assertEqual(context.port, 443) @@ -268,8 +264,7 @@ def callback(context): def test_already_wrapped_socket_is_rejected(self): # ssl.SSLSocket subclasses socket.socket, but wrap_socket cannot layer - # TLS over it, so it must be rejected with an actionable message rather - # than failing later with an opaque handshake error. + # TLS over it, so it needs its own rejection. ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) ctx.check_hostname = False ctx.verify_mode = ssl.CERT_NONE @@ -307,9 +302,8 @@ def callback(context): self.assertEqual(received[0].timeout, 12.5) def test_non_blocking_socket_from_callback_is_accepted(self): - # ssl.SSLContext.wrap_socket refuses a non-blocking socket. The driver - # normalizes the mode so a callback need not care which mode it leaves - # the socket in. Without that, this handshake raises ValueError. + # wrap_socket refuses a non-blocking socket; the driver normalizes the + # mode. Without that, this handshake raises ValueError. server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) server_ctx.load_cert_chain(CLIENT_PEM) listener = socket.socket() @@ -326,9 +320,8 @@ def serve(): threading.Thread(target=serve, daemon=True).start() - # Build the context the way the driver does, so PoolOptions gets the - # flavor-correct type. Verification is off because the local test - # server's certificate is not what the driver would expect. + # Built the way the driver does, so PoolOptions gets the flavor-correct + # type. Verification is off: the local cert is not what it expects. 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) @@ -342,9 +335,8 @@ def callback(context): self.assertIsNotNone(conn.gettimeout()) def test_asyncio_transport_socket_is_rejected(self): - # transport.get_extra_info("socket") returns an asyncio TransportSocket, - # not a socket.socket, and the loop still owns it. Reject it with a - # message that names the mistake rather than failing later. + # transport.get_extra_info("socket") is an asyncio TransportSocket, not + # a socket.socket, and the loop still owns it. from asyncio.trsock import TransportSocket left, right = socket.socketpair() @@ -354,7 +346,7 @@ def test_asyncio_transport_socket_is_rejected(self): def callback(context): return TransportSocket(left) - with self.assertRaisesRegex(ConfigurationError, "asyncio stream or transport"): + with self.assertRaisesRegex(ConfigurationError, "Streams, transports"): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) def test_network_error_from_callback_propagates(self): From b9e06666c4114e2f078b1c7eaa109ea95ab7a368 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 08:24:52 -0500 Subject: [PATCH 17/50] PYTHON-5805 Say module rather than flavor in the _IS_SYNC comment --- pymongo/asynchronous/encryption.py | 2 +- pymongo/synchronous/encryption.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 84f7e278de..93c4259871 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -148,7 +148,7 @@ async def _connect_kms( result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) - # _IS_SYNC is True in the generated synchronous flavor, where a regular + # _IS_SYNC is True in the generated synchronous module, where a regular # function is the correct thing to pass and nothing is awaited. if not _IS_SYNC and not inspect.isawaitable(result): _close_rejected_kms_socket(result) diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 124159a686..3e1710301c 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -147,7 +147,7 @@ def _connect_kms( result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) - # _IS_SYNC is True in the generated synchronous flavor, where a regular + # _IS_SYNC is True in the generated synchronous module, where a regular # function is the correct thing to pass and nothing is awaited. if not _IS_SYNC and not inspect.isawaitable(result): _close_rejected_kms_socket(result) From b1beab6390e203083f29f6182970c0669c64b9d8 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 08:50:34 -0500 Subject: [PATCH 18/50] PYTHON-5805 Add HTTPProxyKMSConnect so callers need not write CONNECT Writing a callback by hand meant implementing the CONNECT handshake and, for a TLS proxy, a socketpair bridge with two relay threads, because Python cannot layer TLS over an ssl.SSLSocket. That is the most delicate code in the feature and every user would have copied it from a docstring. Ship it instead. HTTPProxyKMSConnect and its async variant handle plain and TLS proxies and are usable directly as kms_connect_callback. The callback option is unchanged and remains the escape hatch for cases the helper does not cover, such as proxy authentication. The prose tests now drive the shipped class rather than a private copy, so the spec tests cover the public API. KMSConnectContext's docstring drops from 117 lines to 44, since it no longer carries two worked examples. --- pymongo/encryption_options.py | 199 ++++++++++++++++----------- test/asynchronous/test_encryption.py | 110 +++++++-------- test/test_encryption.py | 110 +++++++-------- tools/synchro.py | 1 + 4 files changed, 212 insertions(+), 208 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 9af375d71a..3454ffd949 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -19,7 +19,10 @@ from __future__ import annotations +import asyncio import socket +import ssl +import threading from collections.abc import Awaitable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict @@ -75,87 +78,14 @@ class KMSConnectContext: ``transport.get_extra_info("socket")`` returns one the loop still owns. :meth:`asyncio.loop.sock_connect` is fine. - To reach a KMS host through an HTTP proxy, tunnel with ``CONNECT``:: - - import socket - - def open_tunnel(sock, context): - target = f"{context.host}:{context.port}" - sock.sendall( - f"CONNECT {target} HTTP/1.1\\r\\nHost: {target}\\r\\n\\r\\n".encode() - ) - response = b"" - while b"\\r\\n\\r\\n" not in response: - chunk = sock.recv(4096) - if not chunk: - raise OSError("proxy closed the connection") - response += chunk - if not response.startswith(b"HTTP/1.1 200"): - raise OSError(f"CONNECT failed: {response.splitlines()[0]!r}") - - def connect_through_proxy(context): - sock = socket.create_connection( - ("proxy.example.com", 8080), timeout=context.timeout - ) - try: - open_tunnel(sock, context) - except OSError: - sock.close() - raise - return sock + For an ordinary HTTP proxy, use :class:`HTTPProxyKMSConnect` or + :class:`AsyncHTTPProxyKMSConnect` rather than writing this yourself. Supply + a callback directly only when you need something they do not cover, such as + proxy authentication. - opts = AutoEncryptionOpts( - kms_providers={"aws": aws_creds}, - key_vault_namespace="keyvault.datakeys", - kms_connect_callback=connect_through_proxy, - ) - - The async API requires a coroutine function. Run the blocking connect in a - thread to keep the event loop free:: - - async def async_connect_through_proxy(context): - return await asyncio.to_thread(connect_through_proxy, context) - - Reaching the proxy over TLS needs one extra step, because Python cannot - layer TLS over an :class:`ssl.SSLSocket`. Relay through a - :func:`socket.socketpair` and return the plain end:: - - import ssl, threading - - def connect_through_tls_proxy(context): - ctx = ssl.create_default_context(cafile="proxy-ca.pem") - raw = socket.create_connection( - ("proxy.example.com", 8443), timeout=context.timeout - ) - proxy = ctx.wrap_socket(raw, server_hostname="proxy.example.com") - try: - open_tunnel(proxy, context) - except OSError: - proxy.close() - raise - - driver_side, relay_side = socket.socketpair() - - def relay(src, dst): - try: - while True: - buf = src.recv(16384) - if not buf: - break - dst.sendall(buf) - except OSError: - pass - finally: - # EOF the peer instead of closing a socket it may be reading. - try: - dst.shutdown(socket.SHUT_RDWR) - except OSError: - pass - src.close() - - for pair in ((relay_side, proxy), (proxy, relay_side)): - threading.Thread(target=relay, args=pair, daemon=True).start() - return driver_side + The asynchronous API requires a coroutine function. Run any blocking + connect in a thread with :func:`asyncio.to_thread` so the event loop stays + free. :param host: Hostname of the KMS server, and the target of TLS certificate and hostname verification. @@ -184,6 +114,115 @@ def relay(src, dst): KMSConnectCallback = Callable[[KMSConnectContext], socket.socket] +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 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 itself over TLS, pass an :class:`ssl.SSLContext`. It is + used only for the connection to the proxy; the driver still negotiates KMS + TLS end to end through the tunnel:: + + 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.18 + """ + + 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) -> None: + target = f"{context.host}:{context.port}" + sock.sendall(f"CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n".encode()) + response = b"" + while b"\r\n\r\n" not in response: + chunk = sock.recv(4096) + if not chunk: + raise OSError(f"proxy closed the connection while tunneling to {target}") + response += chunk + status = response.split(b"\r\n", 1)[0] + if not status.startswith(b"HTTP/1.1 200"): + 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. + + The driver wraps what we return in KMS TLS, and Python cannot layer TLS + over an :class:`ssl.SSLSocket`, so hand back the plain end of a pair and + pump bytes between it and the proxy connection. + """ + driver_side, relay_side = socket.socketpair() + + def relay(src: socket.socket, dst: socket.socket) -> None: + try: + while True: + buf = src.recv(16384) + if not buf: + break + dst.sendall(buf) + except OSError: + pass + finally: + # EOF the peer instead of closing a socket it may be reading. + try: + dst.shutdown(socket.SHUT_RDWR) + except OSError: + pass + src.close() + + for pair in ((relay_side, proxy), (proxy, relay_side)): + threading.Thread(target=relay, args=pair, daemon=True).start() + return driver_side + + def __call__(self, context: KMSConnectContext) -> socket.socket: + sock = socket.create_connection((self.host, self.port), timeout=context.timeout) + try: + if self.ssl_context is not None: + sock = self.ssl_context.wrap_socket(sock, server_hostname=self.host) + self._tunnel(sock, context) + except BaseException: + sock.close() + raise + if self.ssl_context is None: + return sock + return self._bridge(sock) + + +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.18 + """ + + async def __call__(self, context: KMSConnectContext) -> socket.socket: # type: ignore[override] + return await asyncio.to_thread(super().__call__, context) + + class AutoEncryptionOpts: """Options to configure automatic client-side field level encryption.""" diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 20e51a68f9..d2da3f6a01 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -71,7 +71,14 @@ 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, AutoEncryptionOpts, RangeOpts, TextOpts +from pymongo.encryption_options import ( + _HAVE_PYMONGOCRYPT, + AsyncHTTPProxyKMSConnect, + AutoEncryptionOpts, + KMSConnectContext, + RangeOpts, + TextOpts, +) from pymongo.errors import ( AutoReconnect, BulkWriteError, @@ -349,6 +356,40 @@ async def callback(context): with self.assertRaisesRegex(ConfigurationError, "Streams, transports"): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + async def test_http_proxy_helper_tunnels_and_reports_refusal(self): + # Exercise the shipped helper against a stub proxy, so the CONNECT + # handshake is covered 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) + async def test_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") @@ -2111,60 +2152,6 @@ async def test_invalid_hostname_in_kms_certificate(self): } -def _http_connect(context, proxy_port, proxy_ssl_context=None): - """Open an HTTP CONNECT tunnel to context.host:context.port. - - Returns a plain socket. When proxy_ssl_context is given, the connection to - the proxy is TLS and the tunnel is bridged to a socketpair, because TLS - cannot be layered over an ssl.SSLSocket. - """ - sock: socket.socket = socket.create_connection( - (_KMS_PROXY_HOST, proxy_port), timeout=context.timeout - ) - try: - if proxy_ssl_context is not None: - sock = proxy_ssl_context.wrap_socket(sock, server_hostname=_KMS_PROXY_HOST) - target = f"{context.host}:{context.port}" - sock.sendall(f"CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n".encode()) - response = b"" - while b"\r\n\r\n" not in response: - chunk = sock.recv(4096) - if not chunk: - raise OSError("proxy closed the connection before responding") - response += chunk - status = response.split(b"\r\n", 1)[0] - if not status.startswith(b"HTTP/1.1 200"): - raise OSError(f"proxy CONNECT failed: {status!r}") - except Exception: - sock.close() - raise - - if proxy_ssl_context is None: - return sock - - driver_side, relay_side = socket.socketpair() - - def relay(src, dst): - try: - while True: - buf = src.recv(16384) - if not buf: - break - dst.sendall(buf) - except OSError: - pass - finally: - for s in (src, dst): - try: - s.close() - except OSError: - pass - - threading.Thread(target=relay, args=(relay_side, sock), daemon=True).start() - threading.Thread(target=relay, args=(sock, relay_side), daemon=True).start() - return driver_side - - # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-connect-callback class TestKmsConnectCallbackProse(AsyncEncryptionIntegrationTest): @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") @@ -2174,17 +2161,14 @@ async def asyncSetUp(self): async def plain_callback(self, context): self.callback_calls.append(context) - if not _IS_SYNC: - return await asyncio.to_thread(_http_connect, context, _KMS_PROXY_PORT) - return _http_connect(context, _KMS_PROXY_PORT) + 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 - if not _IS_SYNC: - return await asyncio.to_thread(_http_connect, context, _KMS_TLS_PROXY_PORT, ctx) - return _http_connect(context, _KMS_TLS_PROXY_PORT, ctx) + 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.""" @@ -2323,9 +2307,7 @@ async def flaky_callback(context): state["calls"] += 1 if state["calls"] == 1: raise OSError("first attempt fails") - if not _IS_SYNC: - return await asyncio.to_thread(_http_connect, context, _KMS_PROXY_PORT) - return _http_connect(context, _KMS_PROXY_PORT) + return await AsyncHTTPProxyKMSConnect(_KMS_PROXY_HOST, _KMS_PROXY_PORT)(context) encryption = self.create_client_encryption( {"aws": AWS_CREDS}, diff --git a/test/test_encryption.py b/test/test_encryption.py index 6b97628971..0fc62f071b 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -62,7 +62,14 @@ from bson.son import SON from pymongo import ReadPreference from pymongo.cursor_shared import CursorType -from pymongo.encryption_options import _HAVE_PYMONGOCRYPT, AutoEncryptionOpts, RangeOpts, TextOpts +from pymongo.encryption_options import ( + _HAVE_PYMONGOCRYPT, + AutoEncryptionOpts, + HTTPProxyKMSConnect, + KMSConnectContext, + RangeOpts, + TextOpts, +) from pymongo.errors import ( AutoReconnect, BulkWriteError, @@ -349,6 +356,40 @@ def callback(context): with self.assertRaisesRegex(ConfigurationError, "Streams, transports"): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + def test_http_proxy_helper_tunnels_and_reports_refusal(self): + # Exercise the shipped helper against a stub proxy, so the CONNECT + # handshake is covered 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) + def test_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") @@ -2103,60 +2144,6 @@ def test_invalid_hostname_in_kms_certificate(self): } -def _http_connect(context, proxy_port, proxy_ssl_context=None): - """Open an HTTP CONNECT tunnel to context.host:context.port. - - Returns a plain socket. When proxy_ssl_context is given, the connection to - the proxy is TLS and the tunnel is bridged to a socketpair, because TLS - cannot be layered over an ssl.SSLSocket. - """ - sock: socket.socket = socket.create_connection( - (_KMS_PROXY_HOST, proxy_port), timeout=context.timeout - ) - try: - if proxy_ssl_context is not None: - sock = proxy_ssl_context.wrap_socket(sock, server_hostname=_KMS_PROXY_HOST) - target = f"{context.host}:{context.port}" - sock.sendall(f"CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n".encode()) - response = b"" - while b"\r\n\r\n" not in response: - chunk = sock.recv(4096) - if not chunk: - raise OSError("proxy closed the connection before responding") - response += chunk - status = response.split(b"\r\n", 1)[0] - if not status.startswith(b"HTTP/1.1 200"): - raise OSError(f"proxy CONNECT failed: {status!r}") - except Exception: - sock.close() - raise - - if proxy_ssl_context is None: - return sock - - driver_side, relay_side = socket.socketpair() - - def relay(src, dst): - try: - while True: - buf = src.recv(16384) - if not buf: - break - dst.sendall(buf) - except OSError: - pass - finally: - for s in (src, dst): - try: - s.close() - except OSError: - pass - - threading.Thread(target=relay, args=(relay_side, sock), daemon=True).start() - threading.Thread(target=relay, args=(sock, relay_side), daemon=True).start() - return driver_side - - # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-connect-callback class TestKmsConnectCallbackProse(EncryptionIntegrationTest): @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") @@ -2166,17 +2153,14 @@ def setUp(self): def plain_callback(self, context): self.callback_calls.append(context) - if not _IS_SYNC: - return asyncio.to_thread(_http_connect, context, _KMS_PROXY_PORT) - return _http_connect(context, _KMS_PROXY_PORT) + 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 - if not _IS_SYNC: - return asyncio.to_thread(_http_connect, context, _KMS_TLS_PROXY_PORT, ctx) - return _http_connect(context, _KMS_TLS_PROXY_PORT, ctx) + 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.""" @@ -2315,9 +2299,7 @@ def flaky_callback(context): state["calls"] += 1 if state["calls"] == 1: raise OSError("first attempt fails") - if not _IS_SYNC: - return asyncio.to_thread(_http_connect, context, _KMS_PROXY_PORT) - return _http_connect(context, _KMS_PROXY_PORT) + return HTTPProxyKMSConnect(_KMS_PROXY_HOST, _KMS_PROXY_PORT)(context) encryption = self.create_client_encryption( {"aws": AWS_CREDS}, diff --git a/tools/synchro.py b/tools/synchro.py index bc81c1e65e..04a433763d 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -73,6 +73,7 @@ "AsyncClientEncryption": "ClientEncryption", "AsyncMongoCryptCallback": "MongoCryptCallback", "AsyncKMSConnectCallback": "KMSConnectCallback", + "AsyncHTTPProxyKMSConnect": "HTTPProxyKMSConnect", "AsyncExplicitEncrypter": "ExplicitEncrypter", "AsyncAutoEncrypter": "AutoEncrypter", "AsyncContextManager": "ContextManager", From 6d4d6d3fa25a52d898d37f8fa8fbf85c1dbe91e5 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 08:51:42 -0500 Subject: [PATCH 19/50] PYTHON-5805 Point the changelog at the proxy helper --- doc/changelog.rst | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index d366d6661a..70e483e247 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -46,8 +46,10 @@ PyMongo 4.18 brings a number of changes including: :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. See - :class:`~pymongo.encryption_options.KMSConnectContext` for an example. + 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. Changes in Version 4.17.0 (2026/04/20) -------------------------------------- From 8dc628d45da827f7d42fe36ee9b52a275d931640 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 08:59:33 -0500 Subject: [PATCH 20/50] PYTHON-5805 Use run_in_executor to match the library's existing idiom asyncio.to_thread is run_in_executor(None, ...) plus a contextvars copy. The propagation buys nothing here, since the helper takes its timeout as an argument and never reads the CSOT contextvar in the thread, so the only effect was introducing a second idiom for a job auth_oidc.py already does one way when it runs a user-supplied callback off the loop. --- pymongo/encryption_options.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 3454ffd949..7d7435bc9c 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -20,6 +20,7 @@ from __future__ import annotations import asyncio +import functools import socket import ssl import threading @@ -84,8 +85,8 @@ class KMSConnectContext: proxy authentication. The asynchronous API requires a coroutine function. Run any blocking - connect in a thread with :func:`asyncio.to_thread` so the event loop stays - free. + connect in a thread, with :meth:`asyncio.loop.run_in_executor` or + :func:`asyncio.to_thread`, so the event loop stays free. :param host: Hostname of the KMS server, and the target of TLS certificate and hostname verification. @@ -220,7 +221,10 @@ class AsyncHTTPProxyKMSConnect(HTTPProxyKMSConnect): """ async def __call__(self, context: KMSConnectContext) -> socket.socket: # type: ignore[override] - return await asyncio.to_thread(super().__call__, context) + # run_in_executor rather than asyncio.to_thread, matching how + # auth_oidc.py runs a user-supplied callback off the event loop. + connect = functools.partial(super().__call__, context) + return await asyncio.get_running_loop().run_in_executor(None, connect) class AutoEncryptionOpts: From 2ce429871d208433c82fe20d219549eaf34d24ff Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 10:48:40 -0500 Subject: [PATCH 21/50] PYTHON-5805 Make the CSOT deviation note a full sentence --- pymongo/encryption_options.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 7d7435bc9c..96d152d8f4 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -98,8 +98,8 @@ class KMSConnectContext: .. note:: ``timeoutMS`` on a :class:`~pymongo.encryption.ClientEncryption` or its key vault client does not constrain KMS requests, so ``timeout`` is always the default for explicit encryption. Automatic encryption - passes the remaining budget. A known deviation from the Client Side - Operations Timeout specification, tracked in PYTHON-6037. + passes the remaining budget. This is a known deviation from the Client + Side Operations Timeout specification, tracked in PYTHON-6037. .. versionadded:: 4.18 """ From d9d5a7119b82b8b1c2939a4f88a82bc3482209ab Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 12:36:39 -0500 Subject: [PATCH 22/50] PYTHON-5805 Record why _bridge uses threads rather than tasks --- pymongo/encryption_options.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 96d152d8f4..fea14ef7c8 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -172,6 +172,12 @@ def _bridge(self, proxy: socket.socket) -> socket.socket: The driver wraps what we return in KMS TLS, and Python cannot layer TLS over an :class:`ssl.SSLSocket`, so hand back the plain end of a pair and pump bytes between it and the proxy connection. + + Threads rather than asyncio tasks: ``proxy`` may be an + :class:`ssl.SSLSocket`, which the event loop refuses to read, so tasks + would mean reimplementing the connect and CONNECT handshake on streams. + A KMS connection happens once per data key and is then cached, so the + threads are short-lived and rare. """ driver_side, relay_side = socket.socketpair() From d38bd007b722bd508827621939cba065e2feb860 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 12:39:42 -0500 Subject: [PATCH 23/50] PYTHON-5805 Name the async class in the threads-over-tasks note --- pymongo/encryption_options.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index fea14ef7c8..ac1c246a83 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -173,9 +173,10 @@ def _bridge(self, proxy: socket.socket) -> socket.socket: over an :class:`ssl.SSLSocket`, so hand back the plain end of a pair and pump bytes between it and the proxy connection. - Threads rather than asyncio tasks: ``proxy`` may be an - :class:`ssl.SSLSocket`, which the event loop refuses to read, so tasks - would mean reimplementing the connect and CONNECT handshake on streams. + Threads rather than asyncio tasks in :class:`AsyncHTTPProxyKMSConnect`: + ``proxy`` may be an :class:`ssl.SSLSocket`, which the event loop refuses + to read, so tasks would mean reimplementing the connect and CONNECT + handshake on streams. A KMS connection happens once per data key and is then cached, so the threads are short-lived and rare. """ From 5446bfce80fd575f0f6f38d35ce289b06da09a67 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 12:51:34 -0500 Subject: [PATCH 24/50] PYTHON-5805 Cut prose to the house budgets --- pymongo/asynchronous/encryption.py | 18 +++---- pymongo/encryption_options.py | 71 +++++++++------------------- pymongo/synchronous/encryption.py | 18 +++---- test/asynchronous/test_encryption.py | 33 +++++-------- test/test_encryption.py | 33 +++++-------- 5 files changed, 58 insertions(+), 115 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 93c4259871..898b8cc1b7 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -142,32 +142,26 @@ async def _connect_kms( except Exception as exc: _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) - # Let the caller open the connection, then apply TLS ourselves so that SNI - # and certificate verification still target the KMS host even when the - # socket actually terminates at a proxy. + # TLS is applied here, against address, so verification targets the KMS + # host even when the socket terminates at a proxy. result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) - # _IS_SYNC is True in the generated synchronous module, where a regular - # function is the correct thing to pass and nothing is awaited. + # The synchronous module takes a regular function and awaits nothing. if not _IS_SYNC and not inspect.isawaitable(result): _close_rejected_kms_socket(result) raise ConfigurationError( "kms_connect_callback must be a coroutine function for the async " - f"API, but calling it returned {type(result)}, which is not awaitable." + f"API, but returned {type(result)}." ) sock = await result 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)}. Streams, transports and " - "transport sockets cannot be used; try loop.sock_connect. For a " - "TLS proxy, relay through a socket.socketpair and return the plain end." + f"socket.socket, not {type(sock)}; consider HTTPProxyKMSConnect." ) - # ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback - # has no reason to care which mode it left the socket in, so normalize it - # here rather than pushing the requirement onto the caller. + # wrap_socket refuses a non-blocking socket, so normalize the mode here. sock.settimeout(opts.socket_timeout) try: return await _async_wrap_socket_tls(sock, address, opts) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index ac1c246a83..e9e384ede6 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -64,42 +64,23 @@ def check_min_pymongocrypt() -> None: class KMSConnectContext: """Information about a pending KMS connection. - An instance is passed to the ``kms_connect_callback`` configured on - :class:`AutoEncryptionOpts`, :class:`~pymongo.encryption.ClientEncryption`, - or :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`. - - The callback connects to ``host``:``port`` and returns a plain, unwrapped - :class:`socket.socket`. The driver performs the KMS TLS handshake over it, - verifying against ``host`` rather than the peer actually reached, which is - what makes proxying safe. The driver also sets the socket's timeout, so the - mode the callback leaves it in does not matter. - - Streams and transports do not work: :func:`asyncio.open_connection` and - :meth:`asyncio.loop.create_connection` do not return sockets, and - ``transport.get_extra_info("socket")`` returns one the loop still owns. - :meth:`asyncio.loop.sock_connect` is fine. - - For an ordinary HTTP proxy, use :class:`HTTPProxyKMSConnect` or - :class:`AsyncHTTPProxyKMSConnect` rather than writing this yourself. Supply - a callback directly only when you need something they do not cover, such as - proxy authentication. - - The asynchronous API requires a coroutine function. Run any blocking - connect in a thread, with :meth:`asyncio.loop.run_in_executor` or - :func:`asyncio.to_thread`, so the event loop stays free. - - :param host: Hostname of the KMS server, and the target of TLS certificate - and hostname verification. + Passed to ``kms_connect_callback``, which must return a plain, unwrapped + :class:`socket.socket`. The driver performs the KMS TLS handshake over it + against ``host``, not the peer actually reached, which is what makes + proxying safe. + + 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 remaining in the operation's timeout budget when - one is active, otherwise the driver's default KMS connect timeout. - Always a positive number; the driver never passes ``None``. + :param timeout: Seconds left in the timeout budget, else the default KMS + connect timeout. Never ``None``. - .. note:: ``timeoutMS`` on a :class:`~pymongo.encryption.ClientEncryption` - or its key vault client does not constrain KMS requests, so ``timeout`` - is always the default for explicit encryption. Automatic encryption - passes the remaining budget. This is a known deviation from the Client - Side Operations Timeout specification, tracked in PYTHON-6037. + .. 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.18 """ @@ -129,9 +110,8 @@ class HTTPProxyKMSConnect: kms_connect_callback=HTTPProxyKMSConnect("proxy.example.com", 8080), ) - To reach the proxy itself over TLS, pass an :class:`ssl.SSLContext`. It is - used only for the connection to the proxy; the driver still negotiates KMS - TLS end to end through the tunnel:: + 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 @@ -169,16 +149,10 @@ def _tunnel(self, sock: socket.socket, context: KMSConnectContext) -> None: def _bridge(self, proxy: socket.socket) -> socket.socket: """Relay a TLS proxy connection through a socketpair. - The driver wraps what we return in KMS TLS, and Python cannot layer TLS - over an :class:`ssl.SSLSocket`, so hand back the plain end of a pair and - pump bytes between it and the proxy connection. - - Threads rather than asyncio tasks in :class:`AsyncHTTPProxyKMSConnect`: - ``proxy`` may be an :class:`ssl.SSLSocket`, which the event loop refuses - to read, so tasks would mean reimplementing the connect and CONNECT - handshake on streams. - A KMS connection happens once per data key and is then cached, so the - threads are short-lived and rare. + Python cannot layer TLS over an :class:`ssl.SSLSocket`, so return the + plain end of a pair. Threads rather than tasks, even in + :class:`AsyncHTTPProxyKMSConnect`, because the event loop cannot read + an :class:`ssl.SSLSocket`. """ driver_side, relay_side = socket.socketpair() @@ -228,8 +202,7 @@ class AsyncHTTPProxyKMSConnect(HTTPProxyKMSConnect): """ async def __call__(self, context: KMSConnectContext) -> socket.socket: # type: ignore[override] - # run_in_executor rather than asyncio.to_thread, matching how - # auth_oidc.py runs a user-supplied callback off the event loop. + # run_in_executor, as auth_oidc.py does for user callbacks. connect = functools.partial(super().__call__, context) return await asyncio.get_running_loop().run_in_executor(None, connect) diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 3e1710301c..16d7eba081 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -141,32 +141,26 @@ def _connect_kms( except Exception as exc: _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) - # Let the caller open the connection, then apply TLS ourselves so that SNI - # and certificate verification still target the KMS host even when the - # socket actually terminates at a proxy. + # TLS is applied here, against address, so verification targets the KMS + # host even when the socket terminates at a proxy. result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) - # _IS_SYNC is True in the generated synchronous module, where a regular - # function is the correct thing to pass and nothing is awaited. + # The synchronous module takes a regular function and awaits nothing. if not _IS_SYNC and not inspect.isawaitable(result): _close_rejected_kms_socket(result) raise ConfigurationError( "kms_connect_callback must be a coroutine function for the async " - f"API, but calling it returned {type(result)}, which is not awaitable." + f"API, but returned {type(result)}." ) sock = result 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)}. Streams, transports and " - "transport sockets cannot be used; try loop.sock_connect. For a " - "TLS proxy, relay through a socket.socketpair and return the plain end." + f"socket.socket, not {type(sock)}; consider HTTPProxyKMSConnect." ) - # ssl.SSLContext.wrap_socket refuses a non-blocking socket, and a callback - # has no reason to care which mode it left the socket in, so normalize it - # here rather than pushing the requirement onto the caller. + # wrap_socket refuses a non-blocking socket, so normalize the mode here. sock.settimeout(opts.socket_timeout) try: return _wrap_socket_tls(sock, address, opts) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index d2da3f6a01..eba0540f78 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -270,15 +270,13 @@ async def callback(context): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) async def test_already_wrapped_socket_is_rejected(self): - # ssl.SSLSocket subclasses socket.socket, but wrap_socket cannot layer - # TLS over it, so it needs its own rejection. + # 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) - # do_handshake_on_connect=False means no peer is needed to produce a - # genuine ssl.SSLSocket. + # 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) @@ -298,7 +296,7 @@ async def callback(context): received.append(context) return left - # With ssl_context=None the socket is returned unchanged, which also + # ssl_context=None returns the socket unchanged, so this also # confirms a plain socket is accepted. conn = await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 12.5) self.assertIs(conn, left) @@ -309,8 +307,7 @@ async def callback(context): self.assertEqual(received[0].timeout, 12.5) async def test_non_blocking_socket_from_callback_is_accepted(self): - # wrap_socket refuses a non-blocking socket; the driver normalizes the - # mode. Without that, this handshake raises ValueError. + # 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() @@ -327,8 +324,8 @@ def serve(): threading.Thread(target=serve, daemon=True).start() - # Built the way the driver does, so PoolOptions gets the flavor-correct - # type. Verification is off: the local cert is not what it expects. + # Built as the driver does, for the flavor-correct type; the local + # cert would not 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) @@ -342,8 +339,7 @@ async def callback(context): self.assertIsNotNone(conn.gettimeout()) async def test_asyncio_transport_socket_is_rejected(self): - # transport.get_extra_info("socket") is an asyncio TransportSocket, not - # a socket.socket, and the loop still owns it. + # get_extra_info("socket") is a TransportSocket, not a socket.socket. from asyncio.trsock import TransportSocket left, right = socket.socketpair() @@ -353,12 +349,11 @@ async def test_asyncio_transport_socket_is_rejected(self): async def callback(context): return TransportSocket(left) - with self.assertRaisesRegex(ConfigurationError, "Streams, transports"): + 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): - # Exercise the shipped helper against a stub proxy, so the CONNECT - # handshake is covered without KMS credentials. + # Covers the CONNECT handshake without KMS credentials. accepted = [] def stub(listener, reply): @@ -394,8 +389,7 @@ async def test_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") - # Not a ConfigurationError, so kms_request's broad - # handler can treat it as transient and let libmongocrypt retry. + # 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) @@ -2188,8 +2182,7 @@ async def proxy_request(self, method, path, tls=False): async def connect_count(self, tls=False): body = await self.proxy_request("GET", "/metrics", tls=tls) - # The body is one "key value" pair per line. Only connect_count is - # required by the spec; the server also emits connect_target lines. + # One "key value" per line; the server also emits connect_target. for line in body.splitlines(): key, _, value = line.partition(" ") if key == "connect_count": @@ -2294,9 +2287,7 @@ async def test_05_callback_receives_timeout(self): self.assertTrue(self.callback_calls, "callback was never invoked") for context in self.callback_calls: - # Only checks the spec's literal non-zero requirement, which - # cannot fail: timeoutMS does not tighten this value because - # explicit ClientEncryption operations set no CSOT deadline. + # Checks only the spec's non-zero requirement, which cannot fail. self.assertIsNotNone(context.timeout) self.assertGreater(context.timeout, 0) diff --git a/test/test_encryption.py b/test/test_encryption.py index 0fc62f071b..c1327c5c5e 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -270,15 +270,13 @@ def callback(context): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) def test_already_wrapped_socket_is_rejected(self): - # ssl.SSLSocket subclasses socket.socket, but wrap_socket cannot layer - # TLS over it, so it needs its own rejection. + # 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) - # do_handshake_on_connect=False means no peer is needed to produce a - # genuine ssl.SSLSocket. + # 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) @@ -298,7 +296,7 @@ def callback(context): received.append(context) return left - # With ssl_context=None the socket is returned unchanged, which also + # ssl_context=None returns the socket unchanged, so this also # confirms a plain socket is accepted. conn = _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 12.5) self.assertIs(conn, left) @@ -309,8 +307,7 @@ def callback(context): self.assertEqual(received[0].timeout, 12.5) def test_non_blocking_socket_from_callback_is_accepted(self): - # wrap_socket refuses a non-blocking socket; the driver normalizes the - # mode. Without that, this handshake raises ValueError. + # 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() @@ -327,8 +324,8 @@ def serve(): threading.Thread(target=serve, daemon=True).start() - # Built the way the driver does, so PoolOptions gets the flavor-correct - # type. Verification is off: the local cert is not what it expects. + # Built as the driver does, for the flavor-correct type; the local + # cert would not 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) @@ -342,8 +339,7 @@ def callback(context): self.assertIsNotNone(conn.gettimeout()) def test_asyncio_transport_socket_is_rejected(self): - # transport.get_extra_info("socket") is an asyncio TransportSocket, not - # a socket.socket, and the loop still owns it. + # get_extra_info("socket") is a TransportSocket, not a socket.socket. from asyncio.trsock import TransportSocket left, right = socket.socketpair() @@ -353,12 +349,11 @@ def test_asyncio_transport_socket_is_rejected(self): def callback(context): return TransportSocket(left) - with self.assertRaisesRegex(ConfigurationError, "Streams, transports"): + 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): - # Exercise the shipped helper against a stub proxy, so the CONNECT - # handshake is covered without KMS credentials. + # Covers the CONNECT handshake without KMS credentials. accepted = [] def stub(listener, reply): @@ -394,8 +389,7 @@ def test_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") - # Not a ConfigurationError, so kms_request's broad - # handler can treat it as transient and let libmongocrypt retry. + # Not a ConfigurationError, so kms_request retries it. with self.assertRaises(OSError): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) @@ -2180,8 +2174,7 @@ def proxy_request(self, method, path, tls=False): def connect_count(self, tls=False): body = self.proxy_request("GET", "/metrics", tls=tls) - # The body is one "key value" pair per line. Only connect_count is - # required by the spec; the server also emits connect_target lines. + # One "key value" per line; the server also emits connect_target. for line in body.splitlines(): key, _, value = line.partition(" ") if key == "connect_count": @@ -2286,9 +2279,7 @@ def test_05_callback_receives_timeout(self): self.assertTrue(self.callback_calls, "callback was never invoked") for context in self.callback_calls: - # Only checks the spec's literal non-zero requirement, which - # cannot fail: timeoutMS does not tighten this value because - # explicit ClientEncryption operations set no CSOT deadline. + # Checks only the spec's non-zero requirement, which cannot fail. self.assertIsNotNone(context.timeout) self.assertGreater(context.timeout, 0) From 8d32dbc4513aa59eb4a1cbf6588581b3d2434529 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 13:00:46 -0500 Subject: [PATCH 25/50] PYTHON-5805 Point the callback docs at the proxy helper The worked CONNECT example moved to HTTPProxyKMSConnect when the helper landed, so the cross-reference to KMSConnectContext was stale. --- pymongo/asynchronous/encryption.py | 9 +++++---- pymongo/encryption_options.py | 8 ++++---- pymongo/synchronous/encryption.py | 9 +++++---- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 898b8cc1b7..26fac7cdf2 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -727,10 +727,11 @@ def __init__( 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`; the - driver then performs the KMS TLS handshake over it. - See :class:`~pymongo.encryption_options.KMSConnectContext` for a worked - HTTP ``CONNECT`` example. Defaults to ``None``, meaning the driver - connects to KMS hosts directly. + driver then performs the KMS TLS handshake over it. For an ordinary + HTTP proxy, pass + :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect`. + Defaults to ``None``, meaning the driver connects to KMS hosts + directly. .. versionchanged:: 4.18 Added the `kms_connect_callback` parameter. diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index e9e384ede6..109699017c 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -372,10 +372,10 @@ def __init__( TLS handshake over it. Must be a coroutine function for :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient` and a regular function for - :class:`~pymongo.synchronous.mongo_client.MongoClient`. See - :class:`KMSConnectContext` for a worked HTTP ``CONNECT`` example. - Defaults to ``None``, meaning the driver connects to KMS hosts - directly. + :class:`~pymongo.synchronous.mongo_client.MongoClient`. For an + ordinary HTTP proxy, pass :class:`HTTPProxyKMSConnect` or + :class:`AsyncHTTPProxyKMSConnect`. Defaults to ``None``, meaning + the driver connects to KMS hosts directly. .. versionchanged:: 4.18 Added the `kms_connect_callback` parameter. diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 16d7eba081..ec04730eb3 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -724,10 +724,11 @@ def __init__( 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`; the - driver then performs the KMS TLS handshake over it. - See :class:`~pymongo.encryption_options.KMSConnectContext` for a worked - HTTP ``CONNECT`` example. Defaults to ``None``, meaning the driver - connects to KMS hosts directly. + driver then performs the KMS TLS handshake over it. For an ordinary + HTTP proxy, pass + :class:`~pymongo.encryption_options.HTTPProxyKMSConnect`. + Defaults to ``None``, meaning the driver connects to KMS hosts + directly. .. versionchanged:: 4.18 Added the `kms_connect_callback` parameter. From a709ce1f48314efbe2828df3a0b6c577cb5af90b Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 17:30:01 -0500 Subject: [PATCH 26/50] PYTHON-5805 Hoist test imports to module scope --- test/asynchronous/test_encryption.py | 5 +---- test/test_encryption.py | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index eba0540f78..b2af2ec71f 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -34,6 +34,7 @@ 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 @@ -232,8 +233,6 @@ async def test_init_kms_tls_options(self): @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") async def test_init_kms_connect_callback(self): - from pymongo.encryption_options import KMSConnectContext - opts = AutoEncryptionOpts({}, "k.d") self.assertIsNone(opts._kms_connect_callback) @@ -340,8 +339,6 @@ async def callback(context): async def test_asyncio_transport_socket_is_rejected(self): # get_extra_info("socket") is a TransportSocket, not a socket.socket. - from asyncio.trsock import TransportSocket - left, right = socket.socketpair() self.addCleanup(left.close) self.addCleanup(right.close) diff --git a/test/test_encryption.py b/test/test_encryption.py index c1327c5c5e..39b3868328 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -34,6 +34,7 @@ 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 @@ -232,8 +233,6 @@ def test_init_kms_tls_options(self): @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") def test_init_kms_connect_callback(self): - from pymongo.encryption_options import KMSConnectContext - opts = AutoEncryptionOpts({}, "k.d") self.assertIsNone(opts._kms_connect_callback) @@ -340,8 +339,6 @@ def callback(context): def test_asyncio_transport_socket_is_rejected(self): # get_extra_info("socket") is a TransportSocket, not a socket.socket. - from asyncio.trsock import TransportSocket - left, right = socket.socketpair() self.addCleanup(left.close) self.addCleanup(right.close) From 082d5eda2bdfc326f6dffbc16c11b1e027ef62b0 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 20 Aug 2026 17:50:13 -0500 Subject: [PATCH 27/50] PYTHON-5805 Drop underscores from the proxy test constants --- test/asynchronous/test_encryption.py | 30 ++++++++++++++-------------- test/test_encryption.py | 30 ++++++++++++++-------------- 2 files changed, 30 insertions(+), 30 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index b2af2ec71f..bb18e8789f 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -2133,11 +2133,11 @@ async def test_invalid_hostname_in_kms_certificate(self): await self.client_encrypted.create_data_key("aws", master_key=key) -_KMS_PROXY_HOST = "127.0.0.1" -_KMS_PROXY_PORT = 9004 -_KMS_TLS_PROXY_PORT = 9005 +KMS_PROXY_HOST = "127.0.0.1" +KMS_PROXY_PORT = 9004 +KMS_TLS_PROXY_PORT = 9005 -_AWS_MASTER_KEY = { +AWS_MASTER_KEY = { "region": "us-east-1", "key": "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0", } @@ -2152,13 +2152,13 @@ async def asyncSetUp(self): async def plain_callback(self, context): self.callback_calls.append(context) - return await AsyncHTTPProxyKMSConnect(_KMS_PROXY_HOST, _KMS_PROXY_PORT)(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 - callback = AsyncHTTPProxyKMSConnect(_KMS_PROXY_HOST, _KMS_TLS_PROXY_PORT, ctx) + callback = AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) return await callback(context) async def proxy_request(self, method, path, tls=False): @@ -2167,10 +2167,10 @@ async def proxy_request(self, method, path, tls=False): ctx = ssl.create_default_context(cafile=CA_PEM) ctx.check_hostname = False conn = http.client.HTTPSConnection( - f"{_KMS_PROXY_HOST}:{_KMS_TLS_PROXY_PORT}", context=ctx + f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=ctx ) else: - conn = http.client.HTTPConnection(f"{_KMS_PROXY_HOST}:{_KMS_PROXY_PORT}") + conn = http.client.HTTPConnection(f"{KMS_PROXY_HOST}:{KMS_PROXY_PORT}") try: conn.request(method, path) return conn.getresponse().read().decode() @@ -2195,7 +2195,7 @@ async def test_01_plain_http_proxy(self): OPTS, kms_connect_callback=self.plain_callback, ) - await encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + 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): @@ -2207,7 +2207,7 @@ async def test_02_https_proxy(self): OPTS, kms_connect_callback=self.tls_callback, ) - await encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + 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): @@ -2221,7 +2221,7 @@ async def test_03_auto_encryption_through_proxy(self): OPTS, kms_connect_callback=self.plain_callback, ) - data_key_id = await encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + data_key_id = await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) schema = { "bsonType": "object", "properties": { @@ -2265,7 +2265,7 @@ async def failing_callback(context): kms_connect_callback=failing_callback, ) with self.assertRaisesRegex(EncryptionError, "proxy is on fire"): - await encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) @unittest.skip( "PYTHON-6037 ClientEncryption does not support timeoutMS, so the " @@ -2280,7 +2280,7 @@ async def test_05_callback_receives_timeout(self): OPTS, kms_connect_callback=self.plain_callback, ) - await encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + 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: @@ -2295,7 +2295,7 @@ 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) + return await AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) encryption = self.create_client_encryption( {"aws": AWS_CREDS}, @@ -2304,7 +2304,7 @@ async def flaky_callback(context): OPTS, kms_connect_callback=flaky_callback, ) - await encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + 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 39b3868328..ae4a6a8c4e 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -2125,11 +2125,11 @@ def test_invalid_hostname_in_kms_certificate(self): self.client_encrypted.create_data_key("aws", master_key=key) -_KMS_PROXY_HOST = "127.0.0.1" -_KMS_PROXY_PORT = 9004 -_KMS_TLS_PROXY_PORT = 9005 +KMS_PROXY_HOST = "127.0.0.1" +KMS_PROXY_PORT = 9004 +KMS_TLS_PROXY_PORT = 9005 -_AWS_MASTER_KEY = { +AWS_MASTER_KEY = { "region": "us-east-1", "key": "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0", } @@ -2144,13 +2144,13 @@ def setUp(self): def plain_callback(self, context): self.callback_calls.append(context) - return HTTPProxyKMSConnect(_KMS_PROXY_HOST, _KMS_PROXY_PORT)(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 - callback = HTTPProxyKMSConnect(_KMS_PROXY_HOST, _KMS_TLS_PROXY_PORT, ctx) + callback = HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) return callback(context) def proxy_request(self, method, path, tls=False): @@ -2159,10 +2159,10 @@ def proxy_request(self, method, path, tls=False): ctx = ssl.create_default_context(cafile=CA_PEM) ctx.check_hostname = False conn = http.client.HTTPSConnection( - f"{_KMS_PROXY_HOST}:{_KMS_TLS_PROXY_PORT}", context=ctx + f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=ctx ) else: - conn = http.client.HTTPConnection(f"{_KMS_PROXY_HOST}:{_KMS_PROXY_PORT}") + conn = http.client.HTTPConnection(f"{KMS_PROXY_HOST}:{KMS_PROXY_PORT}") try: conn.request(method, path) return conn.getresponse().read().decode() @@ -2187,7 +2187,7 @@ def test_01_plain_http_proxy(self): OPTS, kms_connect_callback=self.plain_callback, ) - encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) self.assertGreaterEqual(self.connect_count(), 1) def test_02_https_proxy(self): @@ -2199,7 +2199,7 @@ def test_02_https_proxy(self): OPTS, kms_connect_callback=self.tls_callback, ) - encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + 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): @@ -2213,7 +2213,7 @@ def test_03_auto_encryption_through_proxy(self): OPTS, kms_connect_callback=self.plain_callback, ) - data_key_id = encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + data_key_id = encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) schema = { "bsonType": "object", "properties": { @@ -2257,7 +2257,7 @@ def failing_callback(context): kms_connect_callback=failing_callback, ) with self.assertRaisesRegex(EncryptionError, "proxy is on fire"): - encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) @unittest.skip( "PYTHON-6037 ClientEncryption does not support timeoutMS, so the " @@ -2272,7 +2272,7 @@ def test_05_callback_receives_timeout(self): OPTS, kms_connect_callback=self.plain_callback, ) - encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + 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: @@ -2287,7 +2287,7 @@ 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) + return HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) encryption = self.create_client_encryption( {"aws": AWS_CREDS}, @@ -2296,7 +2296,7 @@ def flaky_callback(context): OPTS, kms_connect_callback=flaky_callback, ) - encryption.create_data_key("aws", master_key=_AWS_MASTER_KEY) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) self.assertGreaterEqual(state["calls"], 2) From 1dce3057918188b9210f2afe5e4abe105c253aa0 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 21 Aug 2026 05:43:27 -0500 Subject: [PATCH 28/50] PYTHON-5805 Collapse the remaining comments to one line each --- pymongo/asynchronous/encryption.py | 3 +-- pymongo/encryption_options.py | 3 +-- pymongo/synchronous/encryption.py | 3 +-- test/asynchronous/test_encryption.py | 6 ++---- test/test_encryption.py | 6 ++---- 5 files changed, 7 insertions(+), 14 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 26fac7cdf2..d9803d0663 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -142,8 +142,7 @@ async def _connect_kms( except Exception as exc: _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) - # TLS is applied here, against address, so verification targets the KMS - # host even when the socket terminates at a proxy. + # TLS targets address, not the peer, so verification follows the KMS host. result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 109699017c..460cd3de5e 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -90,8 +90,7 @@ class KMSConnectContext: timeout: Optional[float] -# A callback that opens a connection to a KMS host. The async driver requires a -# coroutine function; the synchronous driver requires a regular function. +# A callback that opens a connection to a KMS host. AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]] KMSConnectCallback = Callable[[KMSConnectContext], socket.socket] diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index ec04730eb3..dfe668c7b1 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -141,8 +141,7 @@ def _connect_kms( except Exception as exc: _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) - # TLS is applied here, against address, so verification targets the KMS - # host even when the socket terminates at a proxy. + # TLS targets address, not the peer, so verification follows the KMS host. result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index bb18e8789f..3712b1310b 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -295,8 +295,7 @@ async def callback(context): received.append(context) return left - # ssl_context=None returns the socket unchanged, so this also - # confirms a plain socket is accepted. + # 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) @@ -323,8 +322,7 @@ def serve(): threading.Thread(target=serve, daemon=True).start() - # Built as the driver does, for the flavor-correct type; the local - # cert would not verify. + # 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) diff --git a/test/test_encryption.py b/test/test_encryption.py index ae4a6a8c4e..54b3313733 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -295,8 +295,7 @@ def callback(context): received.append(context) return left - # ssl_context=None returns the socket unchanged, so this also - # confirms a plain socket is accepted. + # 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) @@ -323,8 +322,7 @@ def serve(): threading.Thread(target=serve, daemon=True).start() - # Built as the driver does, for the flavor-correct type; the local - # cert would not verify. + # 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) From 869e91fb0dd127388d9683a7d646f514c8406182 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 21 Aug 2026 06:56:02 -0500 Subject: [PATCH 29/50] PYTHON-5805 Cover the proxy helper's tunnel and relay Adds unit tests for the TLS-proxy bridge, a proxy that hangs up before replying, and a plain def passed to the async API. The last closes a gap a reviewer flagged: the coroutine guard had no permanent test. --- test/asynchronous/test_encryption.py | 78 ++++++++++++++++++++++++++++ test/test_encryption.py | 78 ++++++++++++++++++++++++++++ 2 files changed, 156 insertions(+) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 3712b1310b..1e156b965d 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -380,6 +380,84 @@ def run_stub(reply): with self.assertRaisesRegex(OSError, "refused CONNECT"): await AsyncHTTPProxyKMSConnect(host, port)(context) + 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(): + 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 tunnelled peer speaks only after the client does, as a + # TLS server would. + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + + 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) + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + + async def test_non_coroutine_callback_is_rejected(self): + # The async API needs a coroutine function; a plain def must not be + # awaited and retried. + if _IS_SYNC: + raise unittest.SkipTest("a regular function is correct for the sync API") + + left, right = socket.socketpair() + self.addCleanup(right.close) + + def callback(context): + return left + + with self.assertRaisesRegex(ConfigurationError, "coroutine function"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + 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_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") diff --git a/test/test_encryption.py b/test/test_encryption.py index 54b3313733..0c806b7887 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -380,6 +380,84 @@ def run_stub(reply): with self.assertRaisesRegex(OSError, "refused CONNECT"): HTTPProxyKMSConnect(host, port)(context) + 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(): + 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 tunnelled peer speaks only after the client does, as a + # TLS server would. + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + + 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) + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + + def test_non_coroutine_callback_is_rejected(self): + # The async API needs a coroutine function; a plain def must not be + # awaited and retried. + if _IS_SYNC: + raise unittest.SkipTest("a regular function is correct for the sync API") + + left, right = socket.socketpair() + self.addCleanup(right.close) + + def callback(context): + return left + + with self.assertRaisesRegex(ConfigurationError, "coroutine function"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + 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_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") From 50fca5cec52faf59f23a7abd8d3dcc03c29f1dac Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 21 Aug 2026 07:22:54 -0500 Subject: [PATCH 30/50] PYTHON-5805 Stop _tunnel consuming tunnelled bytes A bulk read of the CONNECT response could swallow bytes a proxy sent in the same segment, and the driver reads those from the same socket. Read to the header boundary instead. Also close the stub proxies' sockets, so the tests raise no ResourceWarning. --- pymongo/encryption_options.py | 7 +++++-- test/asynchronous/test_encryption.py | 27 +++++++++++++++++++++++++++ test/test_encryption.py | 27 +++++++++++++++++++++++++++ 3 files changed, 59 insertions(+), 2 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 460cd3de5e..ace0e6f8f7 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -135,9 +135,12 @@ def __init__(self, host: str, port: int, ssl_context: Optional[ssl.SSLContext] = def _tunnel(self, sock: socket.socket, context: KMSConnectContext) -> None: target = f"{context.host}:{context.port}" sock.sendall(f"CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n".encode()) + # A byte at a time: a bulk read could consume tunnelled bytes sent in + # the same segment as the response, and the driver reads those from + # this same socket. response = b"" - while b"\r\n\r\n" not in response: - chunk = sock.recv(4096) + while not response.endswith(b"\r\n\r\n"): + chunk = sock.recv(1) if not chunk: raise OSError(f"proxy closed the connection while tunneling to {target}") response += chunk diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 1e156b965d..0f6210e948 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -404,6 +404,7 @@ def stub_proxy(): # TLS server would. tls.sendall(b"echo:" + tls.recv(64)) tls.close() + conn.close() except OSError: pass @@ -458,6 +459,32 @@ def stub_proxy(): 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 tunnelled bytes; reading past the + # header would silently 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) + self.assertEqual(sock.recv(64), b"early-bytes") + async def test_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") diff --git a/test/test_encryption.py b/test/test_encryption.py index 0c806b7887..c8be355423 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -404,6 +404,7 @@ def stub_proxy(): # TLS server would. tls.sendall(b"echo:" + tls.recv(64)) tls.close() + conn.close() except OSError: pass @@ -458,6 +459,32 @@ def stub_proxy(): 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 tunnelled bytes; reading past the + # header would silently 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) + self.assertEqual(sock.recv(64), b"early-bytes") + def test_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") From 7d9158a1947d2b327467b347f35825d6ca0f5537 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 21 Aug 2026 08:09:57 -0500 Subject: [PATCH 31/50] PYTHON-5805 Address review findings on the proxy helper Bracket IPv6 hosts in CONNECT, cap the response header, and share one deadline across connect, proxy TLS and the tunnel rather than giving each phase the full budget. Reject an unconnected socket, which TLS accepts and then fails as a retryable error. Compute the KMS timeout after libmongocrypt's retry backoff, not before. --- pymongo/asynchronous/encryption.py | 20 ++++--- pymongo/encryption_options.py | 32 +++++++++-- pymongo/synchronous/encryption.py | 20 ++++--- test/asynchronous/test_encryption.py | 79 +++++++++++++++++++++++++++- test/test_encryption.py | 79 +++++++++++++++++++++++++++- 5 files changed, 212 insertions(+), 18 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index d9803d0663..d78a05d5fe 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -161,6 +161,13 @@ async def _connect_kms( 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 sock.settimeout(opts.socket_timeout) try: return await _async_wrap_socket_tls(sock, address, opts) @@ -233,18 +240,19 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: False, # disable_ocsp_endpoint_check _IS_SYNC, ) - # CSOT: set timeout for socket creation. + address = parse_host(endpoint, _HTTPS_PORT) + sleep_u = kms_context.usleep + if sleep_u: + sleep_sec = float(sleep_u) / 1e6 + await asyncio.sleep(sleep_sec) + # CSOT: set timeout for socket creation. After the retry backoff above, + # so the budget reflects what the sleep consumed. 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) - sleep_u = kms_context.usleep - if sleep_u: - sleep_sec = float(sleep_u) / 1e6 - await asyncio.sleep(sleep_sec) try: conn = await _connect_kms( address, diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index ace0e6f8f7..4eac109c47 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -24,6 +24,7 @@ import socket import ssl import threading +import time from collections.abc import Awaitable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict @@ -94,6 +95,20 @@ class KMSConnectContext: AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]] KMSConnectCallback = Callable[[KMSConnectContext], socket.socket] +# Largest CONNECT response header accepted, so a proxy that never sends the +# terminator cannot grow the buffer without bound. +_MAX_CONNECT_HEADER = 8192 + + +def _remaining(deadline: Optional[float]) -> Optional[float]: + """Seconds left before ``deadline``, or None when there is no deadline.""" + if deadline is None: + return None + 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. @@ -133,18 +148,22 @@ def __init__(self, host: str, port: int, ssl_context: Optional[ssl.SSLContext] = self.ssl_context = ssl_context def _tunnel(self, sock: socket.socket, context: KMSConnectContext) -> None: - target = f"{context.host}:{context.port}" + # 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()) # A byte at a time: a bulk read could consume tunnelled bytes sent in # the same segment as the response, and the driver reads those from # this same socket. - response = b"" + response = bytearray() while not response.endswith(b"\r\n\r\n"): chunk = sock.recv(1) if not chunk: raise OSError(f"proxy closed the connection while tunneling to {target}") response += chunk - status = response.split(b"\r\n", 1)[0] + 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] if not status.startswith(b"HTTP/1.1 200"): raise OSError(f"proxy refused CONNECT to {target}: {status!r}") @@ -180,10 +199,15 @@ def relay(src: socket.socket, dst: socket.socket) -> None: return driver_side def __call__(self, context: KMSConnectContext) -> socket.socket: - sock = socket.create_connection((self.host, self.port), timeout=context.timeout) + # One deadline for all three phases; a timeout per phase would let the + # total run to several times the caller's budget. + deadline = None if context.timeout is None else time.monotonic() + context.timeout + sock = socket.create_connection((self.host, self.port), timeout=_remaining(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) except BaseException: sock.close() diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index dfe668c7b1..6b7c89bb46 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -160,6 +160,13 @@ def _connect_kms( 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 sock.settimeout(opts.socket_timeout) try: return _wrap_socket_tls(sock, address, opts) @@ -232,18 +239,19 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: False, # disable_ocsp_endpoint_check _IS_SYNC, ) - # CSOT: set timeout for socket creation. + address = parse_host(endpoint, _HTTPS_PORT) + sleep_u = kms_context.usleep + if sleep_u: + sleep_sec = float(sleep_u) / 1e6 + time.sleep(sleep_sec) + # CSOT: set timeout for socket creation. After the retry backoff above, + # so the budget reflects what the sleep consumed. 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) - sleep_u = kms_context.usleep - if sleep_u: - sleep_sec = float(sleep_u) / 1e6 - time.sleep(sleep_sec) try: conn = _connect_kms( address, diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 0f6210e948..d56f88fafc 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -31,6 +31,7 @@ import sys import textwrap import threading +import time import traceback import uuid import warnings @@ -390,6 +391,7 @@ async def test_tls_proxy_helper_bridges_the_tunnel(self): self.addCleanup(listener.close) def stub_proxy(): + conn = None try: conn, _ = listener.accept() tls = server_ctx.wrap_socket(conn, server_side=True) @@ -404,9 +406,11 @@ def stub_proxy(): # TLS server would. tls.sendall(b"echo:" + tls.recv(64)) tls.close() - conn.close() except OSError: pass + finally: + if conn is not None: + conn.close() threading.Thread(target=stub_proxy, daemon=True).start() @@ -485,6 +489,79 @@ def stub_proxy(): sock.settimeout(10) self.assertEqual(sock.recv(64), b"early-bytes") + async def test_unconnected_socket_from_callback_is_rejected(self): + # The contract says connected; an unconnected socket would otherwise + # 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.assertIsNone(_remaining(None)) + left = _remaining(time.monotonic() + 5) + assert left is not None + self.assertGreater(left, 0) + with self.assertRaises(socket.timeout): + _remaining(time.monotonic() - 1) + async def test_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") diff --git a/test/test_encryption.py b/test/test_encryption.py index c8be355423..96f7555b3e 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -31,6 +31,7 @@ import sys import textwrap import threading +import time import traceback import uuid import warnings @@ -390,6 +391,7 @@ def test_tls_proxy_helper_bridges_the_tunnel(self): self.addCleanup(listener.close) def stub_proxy(): + conn = None try: conn, _ = listener.accept() tls = server_ctx.wrap_socket(conn, server_side=True) @@ -404,9 +406,11 @@ def stub_proxy(): # TLS server would. tls.sendall(b"echo:" + tls.recv(64)) tls.close() - conn.close() except OSError: pass + finally: + if conn is not None: + conn.close() threading.Thread(target=stub_proxy, daemon=True).start() @@ -485,6 +489,79 @@ def stub_proxy(): sock.settimeout(10) self.assertEqual(sock.recv(64), b"early-bytes") + def test_unconnected_socket_from_callback_is_rejected(self): + # The contract says connected; an unconnected socket would otherwise + # 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.assertIsNone(_remaining(None)) + left = _remaining(time.monotonic() + 5) + assert left is not None + self.assertGreater(left, 0) + with self.assertRaises(socket.timeout): + _remaining(time.monotonic() - 1) + def test_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") From 24f1fd88d5fbb178e860805996e9f4b0dfbb173d Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 21 Aug 2026 12:00:23 -0500 Subject: [PATCH 32/50] PYTHON-5805 Address the second review pass Reject datagram sockets, which pass the connected check and then fail TLS as a retryable error. Make the callback timeout a plain float, matching its documentation. Clean up in _bridge when thread startup fails, and keep blocking work in the test helpers off the event loop. Replace the non-retry test: it drove _connect_kms, which has no retry loop, so it could not have caught a regression. It now drives kms_request. --- pymongo/asynchronous/encryption.py | 9 +++- pymongo/encryption_options.py | 29 ++++++++---- pymongo/synchronous/encryption.py | 9 +++- test/asynchronous/test_encryption.py | 68 +++++++++++++++++++++++++--- test/test_encryption.py | 68 +++++++++++++++++++++++++--- 5 files changed, 158 insertions(+), 25 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index d78a05d5fe..cd0fbaf888 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -133,8 +133,8 @@ def _close_rejected_kms_socket(obj: Any) -> None: async def _connect_kms( address: _Address, opts: PoolOptions, - kms_connect_callback: Optional[AsyncKMSConnectCallback] = None, - timeout: Optional[float] = None, + kms_connect_callback: Optional[AsyncKMSConnectCallback], + timeout: float, ) -> Union[socket.socket, _sslConn]: if kms_connect_callback is None: try: @@ -168,6 +168,11 @@ async def _connect_kms( 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." + ) sock.settimeout(opts.socket_timeout) try: return await _async_wrap_socket_tls(sock, address, opts) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 4eac109c47..7e66213780 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -76,7 +76,7 @@ class KMSConnectContext: :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, else the default KMS - connect timeout. Never ``None``. + connect timeout. .. note:: ``timeoutMS`` does not constrain KMS requests for explicit encryption, so ``timeout`` is always the default there. Automatic @@ -88,7 +88,7 @@ class KMSConnectContext: host: str port: int - timeout: Optional[float] + timeout: float # A callback that opens a connection to a KMS host. @@ -100,10 +100,8 @@ class KMSConnectContext: _MAX_CONNECT_HEADER = 8192 -def _remaining(deadline: Optional[float]) -> Optional[float]: - """Seconds left before ``deadline``, or None when there is no deadline.""" - if deadline is None: - return None +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") @@ -194,14 +192,27 @@ def relay(src: socket.socket, dst: socket.socket) -> None: pass src.close() - for pair in ((relay_side, proxy), (proxy, relay_side)): - threading.Thread(target=relay, args=pair, daemon=True).start() + started = [] + try: + for pair in ((relay_side, proxy), (proxy, relay_side)): + thread = threading.Thread(target=relay, args=pair, daemon=True) + thread.start() + started.append(thread) + 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) -> socket.socket: # One deadline for all three phases; a timeout per phase would let the # total run to several times the caller's budget. - deadline = None if context.timeout is None else time.monotonic() + context.timeout + deadline = time.monotonic() + context.timeout sock = socket.create_connection((self.host, self.port), timeout=_remaining(deadline)) try: if self.ssl_context is not None: diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 6b7c89bb46..723c3dfc22 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -132,8 +132,8 @@ def _close_rejected_kms_socket(obj: Any) -> None: def _connect_kms( address: _Address, opts: PoolOptions, - kms_connect_callback: Optional[KMSConnectCallback] = None, - timeout: Optional[float] = None, + kms_connect_callback: Optional[KMSConnectCallback], + timeout: float, ) -> Union[socket.socket, _sslConn]: if kms_connect_callback is None: try: @@ -167,6 +167,11 @@ def _connect_kms( 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." + ) sock.settimeout(opts.socket_timeout) try: return _wrap_socket_tls(sock, address, opts) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index d56f88fafc..c74195c1cc 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -39,6 +39,7 @@ from collections.abc import Mapping from threading import Thread from typing import Any, Optional +from unittest import mock import pytest @@ -69,6 +70,7 @@ AsyncClientEncryption, QueryType, _connect_kms, + _EncryptionIO, ) from pymongo.asynchronous.helpers import anext from pymongo.asynchronous.mongo_client import AsyncMongoClient @@ -262,7 +264,7 @@ class TestKmsConnectCallbackUnit(AsyncPyMongoTestCase): def _pool_options(): return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=None) - async def test_non_socket_return_is_not_retried(self): + async def test_non_socket_return_raises_configuration_error(self): async def callback(context): return "not-a-socket" @@ -327,11 +329,16 @@ def serve(): 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) - async def callback(context): + 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()) @@ -555,13 +562,55 @@ def stub_proxy(): async def test_remaining_raises_once_the_deadline_passes(self): from pymongo.encryption_options import _remaining - self.assertIsNone(_remaining(None)) - left = _remaining(time.monotonic() + 5) - assert left is not None - self.assertGreater(left, 0) + 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): + # A connected UDP socket passes isinstance and getpeername, but TLS + # then 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_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") @@ -2343,6 +2392,13 @@ async def tls_callback(self, 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 diff --git a/test/test_encryption.py b/test/test_encryption.py index 96f7555b3e..d256ecd11a 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -39,6 +39,7 @@ from collections.abc import Mapping from threading import Thread from typing import Any, Optional +from unittest import mock import pytest @@ -94,6 +95,7 @@ ClientEncryption, QueryType, _connect_kms, + _EncryptionIO, ) from pymongo.synchronous.helpers import next from pymongo.synchronous.mongo_client import MongoClient @@ -262,7 +264,7 @@ class TestKmsConnectCallbackUnit(PyMongoTestCase): def _pool_options(): return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=None) - def test_non_socket_return_is_not_retried(self): + def test_non_socket_return_raises_configuration_error(self): def callback(context): return "not-a-socket" @@ -327,11 +329,16 @@ def serve(): 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 callback(context): + 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()) @@ -555,13 +562,55 @@ def stub_proxy(): def test_remaining_raises_once_the_deadline_passes(self): from pymongo.encryption_options import _remaining - self.assertIsNone(_remaining(None)) - left = _remaining(time.monotonic() + 5) - assert left is not None - self.assertGreater(left, 0) + 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): + # A connected UDP socket passes isinstance and getpeername, but TLS + # then 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_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") @@ -2335,6 +2384,13 @@ def tls_callback(self, 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 From 26e3cb42bfa34232f1b0fa010772b92f5dd68a29 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 21 Aug 2026 13:29:53 -0500 Subject: [PATCH 33/50] PYTHON-5805 Close leaks found in the third review pass Close the proxy socket when _bridge fails at socketpair, not only at thread start. Shield the executor future so a cancelled await still closes the socket its thread goes on to open. Also test the public error type: _wrap_encryption_errors turns the ConfigurationError into an EncryptionError, so callers see the latter. --- pymongo/encryption_options.py | 20 +++++++++-- test/asynchronous/test_encryption.py | 54 ++++++++++++++++++++++++++++ test/test_encryption.py | 53 +++++++++++++++++++++++++++ 3 files changed, 125 insertions(+), 2 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 7e66213780..5e2a575f29 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -100,6 +100,12 @@ class KMSConnectContext: _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() @@ -225,7 +231,11 @@ def __call__(self, context: KMSConnectContext) -> socket.socket: raise if self.ssl_context is None: return sock - return self._bridge(sock) + try: + return self._bridge(sock) + except BaseException: + sock.close() + raise class AsyncHTTPProxyKMSConnect(HTTPProxyKMSConnect): @@ -241,7 +251,13 @@ class AsyncHTTPProxyKMSConnect(HTTPProxyKMSConnect): async def __call__(self, context: KMSConnectContext) -> socket.socket: # type: ignore[override] # run_in_executor, as auth_oidc.py does for user callbacks. connect = functools.partial(super().__call__, context) - return await asyncio.get_running_loop().run_in_executor(None, connect) + 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: diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index c74195c1cc..6860749b64 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -71,6 +71,7 @@ QueryType, _connect_kms, _EncryptionIO, + _wrap_encryption_errors, ) from pymongo.asynchronous.helpers import anext from pymongo.asynchronous.mongo_client import AsyncMongoClient @@ -79,6 +80,7 @@ _HAVE_PYMONGOCRYPT, AsyncHTTPProxyKMSConnect, AutoEncryptionOpts, + HTTPProxyKMSConnect, KMSConnectContext, RangeOpts, TextOpts, @@ -611,6 +613,58 @@ def fail(self): await io.kms_request(StubKmsContext()) self.assertEqual(len(calls), 1) + async def test_contract_violation_surfaces_as_encryption_error(self): + # Public operations run under _wrap_encryption_errors, so 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") diff --git a/test/test_encryption.py b/test/test_encryption.py index d256ecd11a..2325f1e0f7 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -96,6 +96,7 @@ QueryType, _connect_kms, _EncryptionIO, + _wrap_encryption_errors, ) from pymongo.synchronous.helpers import next from pymongo.synchronous.mongo_client import MongoClient @@ -611,6 +612,58 @@ def fail(self): io.kms_request(StubKmsContext()) self.assertEqual(len(calls), 1) + def test_contract_violation_surfaces_as_encryption_error(self): + # Public operations run under _wrap_encryption_errors, so 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") From 10d360049d18a2ef95cc431f5b5ce7cb1d97c252 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 8 Sep 2026 09:28:05 -0500 Subject: [PATCH 34/50] fix handling of csot --- pymongo/encryption_options.py | 9 +++--- test/asynchronous/test_encryption.py | 46 ++++++++++++++++++++++++++++ test/test_encryption.py | 46 ++++++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 4 deletions(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 5e2a575f29..d4b288f2a3 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -179,6 +179,10 @@ def _bridge(self, proxy: socket.socket) -> socket.socket: :class:`AsyncHTTPProxyKMSConnect`, because the event loop cannot read an :class:`ssl.SSLSocket`. """ + # The caller's timeout covered the CONNECT handshake, not the KMS + # request that follows through the tunnel; let the driver's own + # timeout on driver_side govern the relay instead. + proxy.settimeout(None) driver_side, relay_side = socket.socketpair() def relay(src: socket.socket, dst: socket.socket) -> None: @@ -198,12 +202,9 @@ def relay(src: socket.socket, dst: socket.socket) -> None: pass src.close() - started = [] try: for pair in ((relay_side, proxy), (proxy, relay_side)): - thread = threading.Thread(target=relay, args=pair, daemon=True) - thread.start() - started.append(thread) + 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): diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 6860749b64..68ecf525e4 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -435,6 +435,52 @@ def stub_proxy(): sock.sendall(b"ping") self.assertEqual(sock.recv(64), b"echo:ping") + async def test_bridge_does_not_inherit_the_connect_deadline(self): + # The CONNECT deadline can be much shorter than the KMS request that + # follows through the tunnel; the relay must outlast it. + 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(0.4) + 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=0.2) + + sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + async def test_non_coroutine_callback_is_rejected(self): # The async API needs a coroutine function; a plain def must not be # awaited and retried. diff --git a/test/test_encryption.py b/test/test_encryption.py index 2325f1e0f7..07fdd6881d 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -434,6 +434,52 @@ def stub_proxy(): sock.sendall(b"ping") self.assertEqual(sock.recv(64), b"echo:ping") + def test_bridge_does_not_inherit_the_connect_deadline(self): + # The CONNECT deadline can be much shorter than the KMS request that + # follows through the tunnel; the relay must outlast it. + 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(0.4) + 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=0.2) + + sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + def test_non_coroutine_callback_is_rejected(self): # The async API needs a coroutine function; a plain def must not be # awaited and retried. From 6df62d35749b524e6edb89e0b29746522c4695b5 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 8 Sep 2026 12:43:31 -0500 Subject: [PATCH 35/50] PYTHON-5805 Move KMS connect callback tests to test_kms_connect.py --- test/asynchronous/test_encryption.py | 669 +----------------------- test/asynchronous/test_kms_connect.py | 714 ++++++++++++++++++++++++++ test/test_encryption.py | 669 +----------------------- test/test_kms_connect.py | 713 +++++++++++++++++++++++++ 4 files changed, 1431 insertions(+), 1334 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 cbf01de934..bc300e4a05 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -262,493 +262,7 @@ async def callback(context): context.host = "evil.example.com" # type: ignore[misc] -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_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) - - 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 tunnelled 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) - sock.sendall(b"ping") - self.assertEqual(sock.recv(64), b"echo:ping") - - async def test_bridge_does_not_inherit_the_connect_deadline(self): - # The CONNECT deadline can be much shorter than the KMS request that - # follows through the tunnel; the relay must outlast it. - 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(0.4) - 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=0.2) - - sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) - self.addCleanup(sock.close) - sock.settimeout(10) - sock.sendall(b"ping") - self.assertEqual(sock.recv(64), b"echo:ping") - - async def test_non_coroutine_callback_is_rejected(self): - # The async API needs a coroutine function; a plain def must not be - # awaited and retried. - if _IS_SYNC: - raise unittest.SkipTest("a regular function is correct for the sync API") - - left, right = socket.socketpair() - self.addCleanup(right.close) - - def callback(context): - return left - - with self.assertRaisesRegex(ConfigurationError, "coroutine function"): - await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) - - 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 tunnelled bytes; reading past the - # header would silently 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) - self.assertEqual(sock.recv(64), b"early-bytes") - - async def test_unconnected_socket_from_callback_is_rejected(self): - # The contract says connected; an unconnected socket would otherwise - # 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): - # A connected UDP socket passes isinstance and getpeername, but TLS - # then 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): - # Public operations run under _wrap_encryption_errors, so 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) - - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") - 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) - - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") - 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] - ) +# KMS connect callback unit and prose tests moved to test_kms_connect.py. class TestClientOptions(AsyncPyMongoTestCase): @@ -2521,186 +2035,7 @@ async def test_invalid_hostname_in_kms_certificate(self): await self.client_encrypted.create_data_key("aws", master_key=key) -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", -} - - -# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-connect-callback -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 - 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 - 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) - - self.assertGreaterEqual(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) +# KMS connect callback unit and prose tests moved to test_kms_connect.py. # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-tls-options-tests diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py new file mode 100644 index 0000000000..992a7d34f9 --- /dev/null +++ b/test/asynchronous/test_kms_connect.py @@ -0,0 +1,714 @@ +"""Tests for the KMS connect callback and HTTP proxy support. + +.. note:: This module is compatible with both the synchronous and + asynchronous PyMongo test suites. +""" + +from __future__ import annotations + +import asyncio +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 + +from bson.binary import Binary +from pymongo.asynchronous.encryption import ( + AsyncClientEncryption, + _connect_kms, + _EncryptionIO, + _wrap_encryption_errors, +) +from pymongo.encryption_options import ( + _HAVE_PYMONGOCRYPT, + AsyncHTTPProxyKMSConnect, + AutoEncryptionOpts, + HTTPProxyKMSConnect, + KMSConnectContext, +) +from pymongo.errors import ConfigurationError, EncryptionError +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_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) + + 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 tunnelled 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) + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + + async def test_bridge_does_not_inherit_the_connect_deadline(self): + # The CONNECT deadline can be much shorter than the KMS request that + # follows through the tunnel; the relay must outlast it. + 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(0.4) + 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=0.2) + + sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + + async def test_non_coroutine_callback_is_rejected(self): + # The async API needs a coroutine function; a plain def must not be + # awaited and retried. + if _IS_SYNC: + raise unittest.SkipTest("a regular function is correct for the sync API") + + left, right = socket.socketpair() + self.addCleanup(right.close) + + def callback(context): + return left + + with self.assertRaisesRegex(ConfigurationError, "coroutine function"): + await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + 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 tunnelled bytes; reading past the + # header would silently 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) + self.assertEqual(sock.recv(64), b"early-bytes") + + async def test_unconnected_socket_from_callback_is_rejected(self): + # The contract says connected; an unconnected socket would otherwise + # 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): + # A connected UDP socket passes isinstance and getpeername, but TLS + # then 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): + # Public operations run under _wrap_encryption_errors, so 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) + + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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) + + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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 + 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 + 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) + + self.assertGreaterEqual(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 9fc6a7c2ce..d2ec35e8a0 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -261,493 +261,7 @@ def callback(context): context.host = "evil.example.com" # type: ignore[misc] -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_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) - - 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 tunnelled 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) - sock.sendall(b"ping") - self.assertEqual(sock.recv(64), b"echo:ping") - - def test_bridge_does_not_inherit_the_connect_deadline(self): - # The CONNECT deadline can be much shorter than the KMS request that - # follows through the tunnel; the relay must outlast it. - 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(0.4) - 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=0.2) - - sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) - self.addCleanup(sock.close) - sock.settimeout(10) - sock.sendall(b"ping") - self.assertEqual(sock.recv(64), b"echo:ping") - - def test_non_coroutine_callback_is_rejected(self): - # The async API needs a coroutine function; a plain def must not be - # awaited and retried. - if _IS_SYNC: - raise unittest.SkipTest("a regular function is correct for the sync API") - - left, right = socket.socketpair() - self.addCleanup(right.close) - - def callback(context): - return left - - with self.assertRaisesRegex(ConfigurationError, "coroutine function"): - _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) - - 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 tunnelled bytes; reading past the - # header would silently 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) - self.assertEqual(sock.recv(64), b"early-bytes") - - def test_unconnected_socket_from_callback_is_rejected(self): - # The contract says connected; an unconnected socket would otherwise - # 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): - # A connected UDP socket passes isinstance and getpeername, but TLS - # then 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): - # Public operations run under _wrap_encryption_errors, so 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) - - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") - 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) - - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") - 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] - ) +# KMS connect callback unit and prose tests moved to test_kms_connect.py. class TestClientOptions(PyMongoTestCase): @@ -2512,186 +2026,7 @@ def test_invalid_hostname_in_kms_certificate(self): self.client_encrypted.create_data_key("aws", master_key=key) -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", -} - - -# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-connect-callback -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 - 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 - 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) - - self.assertGreaterEqual(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) +# KMS connect callback unit and prose tests moved to test_kms_connect.py. # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-tls-options-tests diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py new file mode 100644 index 0000000000..fbbacf5fc6 --- /dev/null +++ b/test/test_kms_connect.py @@ -0,0 +1,713 @@ +"""Tests for the KMS connect callback and HTTP proxy support. + +.. note:: This module is compatible with both the synchronous and + synchronous PyMongo test suites. +""" + +from __future__ import annotations + +import asyncio +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 + +from bson.binary import Binary +from pymongo.encryption_options import ( + _HAVE_PYMONGOCRYPT, + AutoEncryptionOpts, + HTTPProxyKMSConnect, + KMSConnectContext, +) +from pymongo.errors import ConfigurationError, EncryptionError +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_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) + + 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 tunnelled 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) + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + + def test_bridge_does_not_inherit_the_connect_deadline(self): + # The CONNECT deadline can be much shorter than the KMS request that + # follows through the tunnel; the relay must outlast it. + 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(0.4) + 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=0.2) + + sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + + def test_non_coroutine_callback_is_rejected(self): + # The async API needs a coroutine function; a plain def must not be + # awaited and retried. + if _IS_SYNC: + raise unittest.SkipTest("a regular function is correct for the sync API") + + left, right = socket.socketpair() + self.addCleanup(right.close) + + def callback(context): + return left + + with self.assertRaisesRegex(ConfigurationError, "coroutine function"): + _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + + 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 tunnelled bytes; reading past the + # header would silently 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) + self.assertEqual(sock.recv(64), b"early-bytes") + + def test_unconnected_socket_from_callback_is_rejected(self): + # The contract says connected; an unconnected socket would otherwise + # 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): + # A connected UDP socket passes isinstance and getpeername, but TLS + # then 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): + # Public operations run under _wrap_encryption_errors, so 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) + + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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) + + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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 + 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 + 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) + + self.assertGreaterEqual(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) From 8d4e86423664a5b549898895fbd29b4a3f635940 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 8 Sep 2026 12:55:49 -0500 Subject: [PATCH 36/50] PYTHON-5805 Address review notes on KMS proxy docs and changelog --- doc/changelog.rst | 26 +++++++++++++++----------- pymongo/asynchronous/encryption.py | 9 +++++++-- pymongo/encryption_options.py | 19 +++++++++---------- pymongo/synchronous/encryption.py | 9 +++++++-- test/asynchronous/test_encryption.py | 4 ++-- test/test_encryption.py | 4 ++-- 6 files changed, 42 insertions(+), 29 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 0a07096d4d..fab57910b3 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -1,6 +1,21 @@ Changelog ========= +Changes in Version 4.19.0 (unreleased) +-------------------------------------- + +- 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. + Changes in Version 4.18.0 (2026/09/03) -------------------------------------- @@ -96,17 +111,6 @@ PyMongo 4.18 brings a number of changes including: :meth:`~pymongo.synchronous.database.Database.aggregate`, and :meth:`~pymongo.asynchronous.collection.AsyncCollection.list_search_indexes` and :meth:`~pymongo.synchronous.collection.Collection.list_search_indexes`. -- 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. Changes in Version 4.17.0 (2026/04/20) -------------------------------------- diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index b6a5aed4a6..5ae620b515 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -145,6 +145,11 @@ async def _connect_kms( 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) @@ -667,14 +672,14 @@ def __init__( :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`; the + and returns a connected, unwrapped :class:`socket.socket`. The driver then performs the KMS TLS handshake over it. For an ordinary HTTP proxy, pass :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect`. Defaults to ``None``, meaning the driver connects to KMS hosts directly. - .. versionchanged:: 4.18 + .. versionchanged:: 4.19 Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index a25615edde..85a0854f9e 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -84,7 +84,7 @@ class KMSConnectContext: encryption passes the remaining budget. This deviates from the Client Side Operations Timeout specification; see PYTHON-6037. - .. versionadded:: 4.18 + .. versionadded:: 4.19 """ host: str @@ -144,7 +144,7 @@ class HTTPProxyKMSConnect: :param ssl_context: Optional :class:`ssl.SSLContext` for connecting to the proxy over TLS. Defaults to ``None``, meaning a plain connection. - .. versionadded:: 4.18 + .. versionadded:: 4.19 """ def __init__(self, host: str, port: int, ssl_context: Optional[ssl.SSLContext] = None): @@ -176,13 +176,13 @@ 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. Threads rather than tasks, even in + plain end of a pair. Uses threads rather than tasks, even in :class:`AsyncHTTPProxyKMSConnect`, because the event loop cannot read an :class:`ssl.SSLSocket`. """ - # The caller's timeout covered the CONNECT handshake, not the KMS - # request that follows through the tunnel; let the driver's own - # timeout on driver_side govern the relay instead. + # Clear the CONNECT-phase timeout before relaying: the KMS request that + # follows through the tunnel is governed by the driver's own timeout on + # driver_side, not the already-elapsed connect budget. proxy.settimeout(None) driver_side, relay_side = socket.socketpair() @@ -196,7 +196,7 @@ def relay(src: socket.socket, dst: socket.socket) -> None: except OSError: pass finally: - # EOF the peer instead of closing a socket it may be reading. + # Send EOF to the peer instead of closing a socket it may be reading. try: dst.shutdown(socket.SHUT_RDWR) except OSError: @@ -247,11 +247,10 @@ class AsyncHTTPProxyKMSConnect(HTTPProxyKMSConnect): callable and runs the blocking connect in a thread so the event loop stays free. - .. versionadded:: 4.18 + .. versionadded:: 4.19 """ async def __call__(self, context: KMSConnectContext) -> socket.socket: # type: ignore[override] - # run_in_executor, as auth_oidc.py does for user callbacks. connect = functools.partial(super().__call__, context) future = asyncio.get_running_loop().run_in_executor(None, connect) try: @@ -432,7 +431,7 @@ def __init__( :class:`AsyncHTTPProxyKMSConnect`. Defaults to ``None``, meaning the driver connects to KMS hosts directly. - .. versionchanged:: 4.18 + .. versionchanged:: 4.19 Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 2417599ba5..e87ab29aba 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -144,6 +144,11 @@ def _connect_kms( 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) @@ -664,14 +669,14 @@ def __init__( :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`; the + and returns a connected, unwrapped :class:`socket.socket`. The driver then performs the KMS TLS handshake over it. For an ordinary HTTP proxy, pass :class:`~pymongo.encryption_options.HTTPProxyKMSConnect`. Defaults to ``None``, meaning the driver connects to KMS hosts directly. - .. versionchanged:: 4.18 + .. versionchanged:: 4.19 Added the `kms_connect_callback` parameter. .. versionchanged:: 4.12 Added the `key_expiration_ms` parameter. diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index bc300e4a05..eeedf92647 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -262,7 +262,7 @@ async def callback(context): context.host = "evil.example.com" # type: ignore[misc] -# KMS connect callback unit and prose tests moved to test_kms_connect.py. +# KMS connect callback unit and prose tests (new in this PR) live in test_kms_connect.py. class TestClientOptions(AsyncPyMongoTestCase): @@ -2035,7 +2035,7 @@ 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 moved to test_kms_connect.py. +# KMS connect callback unit and prose tests (new in this PR) live in test_kms_connect.py. # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-tls-options-tests diff --git a/test/test_encryption.py b/test/test_encryption.py index d2ec35e8a0..18ef994950 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -261,7 +261,7 @@ def callback(context): context.host = "evil.example.com" # type: ignore[misc] -# KMS connect callback unit and prose tests moved to test_kms_connect.py. +# KMS connect callback unit and prose tests (new in this PR) live in test_kms_connect.py. class TestClientOptions(PyMongoTestCase): @@ -2026,7 +2026,7 @@ 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 moved to test_kms_connect.py. +# KMS connect callback unit and prose tests (new in this PR) live in test_kms_connect.py. # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-tls-options-tests From 91ed31d4fd3c56952cc5fddb7659cac9959bb943 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 8 Sep 2026 12:59:04 -0500 Subject: [PATCH 37/50] PYTHON-5805 Shorten KMS callback test pointer comment --- test/asynchronous/test_encryption.py | 4 ++-- test/test_encryption.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index eeedf92647..f5f3d206a6 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -262,7 +262,7 @@ async def callback(context): context.host = "evil.example.com" # type: ignore[misc] -# KMS connect callback unit and prose tests (new in this PR) live in test_kms_connect.py. +# KMS connect callback unit and prose tests live in test_kms_connect.py. class TestClientOptions(AsyncPyMongoTestCase): @@ -2035,7 +2035,7 @@ 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 (new in this PR) live in test_kms_connect.py. +# 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 diff --git a/test/test_encryption.py b/test/test_encryption.py index 18ef994950..d9fec9be62 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -261,7 +261,7 @@ def callback(context): context.host = "evil.example.com" # type: ignore[misc] -# KMS connect callback unit and prose tests (new in this PR) live in test_kms_connect.py. +# KMS connect callback unit and prose tests live in test_kms_connect.py. class TestClientOptions(PyMongoTestCase): @@ -2026,7 +2026,7 @@ 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 (new in this PR) live in test_kms_connect.py. +# 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 From c6190ceb257c004b4d5238470788c50e3a296871 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 8 Sep 2026 13:57:40 -0500 Subject: [PATCH 38/50] PYTHON-5805 Move test_init_kms_connect_callback to test_kms_connect.py --- test/asynchronous/test_encryption.py | 24 ------------------------ test/asynchronous/test_kms_connect.py | 23 +++++++++++++++++++++++ test/test_encryption.py | 24 ------------------------ test/test_kms_connect.py | 23 +++++++++++++++++++++++ 4 files changed, 46 insertions(+), 48 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index f5f3d206a6..f4fea08251 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -19,7 +19,6 @@ import asyncio import base64 import copy -import dataclasses import http.client import inspect import json @@ -82,7 +81,6 @@ AsyncHTTPProxyKMSConnect, AutoEncryptionOpts, HTTPProxyKMSConnect, - KMSConnectContext, RangeOpts, StringOpts, TextOpts, @@ -239,28 +237,6 @@ async def test_init_kms_tls_options(self): self.assertEqual(ctx.check_hostname, True) self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED) - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") - 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] - # KMS connect callback unit and prose tests live in test_kms_connect.py. diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index 992a7d34f9..3f11456741 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio +import dataclasses import http.client import socket import ssl @@ -61,6 +62,28 @@ class TestKmsConnectCallbackUnit(AsyncPyMongoTestCase): def _pool_options(): return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=None) + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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" diff --git a/test/test_encryption.py b/test/test_encryption.py index d9fec9be62..6da4d74a6e 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -19,7 +19,6 @@ import asyncio import base64 import copy -import dataclasses import http.client import inspect import json @@ -70,7 +69,6 @@ _HAVE_PYMONGOCRYPT, AutoEncryptionOpts, HTTPProxyKMSConnect, - KMSConnectContext, RangeOpts, StringOpts, TextOpts, @@ -238,28 +236,6 @@ def test_init_kms_tls_options(self): self.assertEqual(ctx.check_hostname, True) self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED) - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") - 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] - # KMS connect callback unit and prose tests live in test_kms_connect.py. diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index fbbacf5fc6..8861042c52 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio +import dataclasses import http.client import socket import ssl @@ -60,6 +61,28 @@ class TestKmsConnectCallbackUnit(PyMongoTestCase): def _pool_options(): return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=None) + @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") + 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" From f9153327eb2cc3b8cd8bd6e8484487482e31f011 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Tue, 8 Sep 2026 19:44:08 -0500 Subject: [PATCH 39/50] PYTHON-5805 Do not verify the local KMS proxy cert in prose tests --- test/asynchronous/test_kms_connect.py | 2 ++ test/test_kms_connect.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index 3f11456741..62c16a3063 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -580,6 +580,7 @@ async def tls_callback(self, context): self.callback_calls.append(context) ctx = ssl.create_default_context(cafile=CA_PEM) ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE callback = AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) return await callback(context) @@ -595,6 +596,7 @@ def _proxy_request(self, method, path, tls=False): if tls: ctx = ssl.create_default_context(cafile=CA_PEM) ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE conn = http.client.HTTPSConnection( f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=ctx ) diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index 8861042c52..8e57f585ec 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -579,6 +579,7 @@ def tls_callback(self, context): self.callback_calls.append(context) ctx = ssl.create_default_context(cafile=CA_PEM) ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE callback = HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) return callback(context) @@ -594,6 +595,7 @@ def _proxy_request(self, method, path, tls=False): if tls: ctx = ssl.create_default_context(cafile=CA_PEM) ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE conn = http.client.HTTPSConnection( f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=ctx ) From 2a0b1562ddd2c28a2759a499717ca92623d9551c Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 05:11:30 -0500 Subject: [PATCH 40/50] PYTHON-5805 Use standard changelog version placeholder --- doc/changelog.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index b4265f8c22..135ddae3b9 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -1,7 +1,7 @@ Changelog ========= -Changes in Version 4.19.0 (unreleased) +Changes in Version 4.19.0 (2026/XX/XX) -------------------------------------- - Added support for routing Key Management Service (KMS) requests for From a60ed6752c82bfbfd39076ac52b6c4f111bf16a6 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 05:17:49 -0500 Subject: [PATCH 41/50] PYTHON-5805 Remove unused imports from encryption tests --- test/asynchronous/test_encryption.py | 18 +----------------- test/test_encryption.py | 17 +---------------- 2 files changed, 2 insertions(+), 33 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index f4fea08251..65ee55c5b0 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -16,30 +16,24 @@ from __future__ import annotations -import asyncio import base64 import copy import http.client import inspect import json import os -import pathlib import re import socket import socketserver 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 @@ -57,7 +51,7 @@ sys.path[0:0] = [""] -from bson import BSON, DatetimeMS, Decimal128, encode, json_util +from bson import DatetimeMS, Decimal128, encode, json_util from bson.binary import UUID_SUBTYPE, Binary, UuidRepresentation from bson.codec_options import CodecOptions from bson.errors import BSONError @@ -69,18 +63,13 @@ 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, @@ -99,18 +88,14 @@ 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, ) -from test.asynchronous.test_bulk import AsyncBulkTestBase from test.asynchronous.unified_format import generate_test_classes, get_test_path from test.helpers_shared import ( ALL_KMS_PROVIDERS, AWS_CREDS, - AWS_TEMP_CREDS, AZURE_CREDS, CA_PEM, CLIENT_PEM, @@ -124,7 +109,6 @@ OvertCommandListener, TopologyEventListener, async_wait_until, - camel_to_snake_args, is_greenthread_patched, ) from test.version import Version diff --git a/test/test_encryption.py b/test/test_encryption.py index 6da4d74a6e..0f341ded7a 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -16,30 +16,24 @@ from __future__ import annotations -import asyncio import base64 import copy import http.client import inspect import json import os -import pathlib import re import socket import socketserver 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 @@ -57,7 +51,7 @@ sys.path[0:0] = [""] -from bson import BSON, DatetimeMS, Decimal128, encode, json_util +from bson import DatetimeMS, Decimal128, encode, json_util from bson.binary import UUID_SUBTYPE, Binary, UuidRepresentation from bson.codec_options import CodecOptions from bson.errors import BSONError @@ -68,7 +62,6 @@ from pymongo.encryption_options import ( _HAVE_PYMONGOCRYPT, AutoEncryptionOpts, - HTTPProxyKMSConnect, RangeOpts, StringOpts, TextOpts, @@ -87,16 +80,11 @@ 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, - _connect_kms, - _EncryptionIO, - _wrap_encryption_errors, ) from pymongo.synchronous.helpers import next from pymongo.synchronous.mongo_client import MongoClient @@ -107,7 +95,6 @@ from test.helpers_shared import ( ALL_KMS_PROVIDERS, AWS_CREDS, - AWS_TEMP_CREDS, AZURE_CREDS, CA_PEM, CLIENT_PEM, @@ -116,13 +103,11 @@ KMIP_CREDS, LOCAL_MASTER_KEY, ) -from test.test_bulk import BulkTestBase from test.unified_format import generate_test_classes, get_test_path from test.utils_shared import ( AllowListEventListener, OvertCommandListener, TopologyEventListener, - camel_to_snake_args, is_greenthread_patched, wait_until, ) From 6f7061bc87151619b42a8925a7cfa12c55121da8 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 05:22:19 -0500 Subject: [PATCH 42/50] Revert "PYTHON-5805 Remove unused imports from encryption tests" This reverts commit a60ed6752c82bfbfd39076ac52b6c4f111bf16a6. --- test/asynchronous/test_encryption.py | 18 +++++++++++++++++- test/test_encryption.py | 17 ++++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 65ee55c5b0..f4fea08251 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -16,24 +16,30 @@ from __future__ import annotations +import asyncio import base64 import copy import http.client import inspect import json import os +import pathlib import re import socket import socketserver 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 @@ -51,7 +57,7 @@ sys.path[0:0] = [""] -from bson import DatetimeMS, Decimal128, encode, json_util +from bson import BSON, DatetimeMS, Decimal128, encode, json_util from bson.binary import UUID_SUBTYPE, Binary, UuidRepresentation from bson.codec_options import CodecOptions from bson.errors import BSONError @@ -63,13 +69,18 @@ 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, @@ -88,14 +99,18 @@ 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, ) +from test.asynchronous.test_bulk import AsyncBulkTestBase from test.asynchronous.unified_format import generate_test_classes, get_test_path from test.helpers_shared import ( ALL_KMS_PROVIDERS, AWS_CREDS, + AWS_TEMP_CREDS, AZURE_CREDS, CA_PEM, CLIENT_PEM, @@ -109,6 +124,7 @@ OvertCommandListener, TopologyEventListener, async_wait_until, + camel_to_snake_args, is_greenthread_patched, ) from test.version import Version diff --git a/test/test_encryption.py b/test/test_encryption.py index 0f341ded7a..6da4d74a6e 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -16,24 +16,30 @@ from __future__ import annotations +import asyncio import base64 import copy import http.client import inspect import json import os +import pathlib import re import socket import socketserver 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 @@ -51,7 +57,7 @@ sys.path[0:0] = [""] -from bson import DatetimeMS, Decimal128, encode, json_util +from bson import BSON, DatetimeMS, Decimal128, encode, json_util from bson.binary import UUID_SUBTYPE, Binary, UuidRepresentation from bson.codec_options import CodecOptions from bson.errors import BSONError @@ -62,6 +68,7 @@ from pymongo.encryption_options import ( _HAVE_PYMONGOCRYPT, AutoEncryptionOpts, + HTTPProxyKMSConnect, RangeOpts, StringOpts, TextOpts, @@ -80,11 +87,16 @@ 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, + _connect_kms, + _EncryptionIO, + _wrap_encryption_errors, ) from pymongo.synchronous.helpers import next from pymongo.synchronous.mongo_client import MongoClient @@ -95,6 +107,7 @@ from test.helpers_shared import ( ALL_KMS_PROVIDERS, AWS_CREDS, + AWS_TEMP_CREDS, AZURE_CREDS, CA_PEM, CLIENT_PEM, @@ -103,11 +116,13 @@ KMIP_CREDS, LOCAL_MASTER_KEY, ) +from test.test_bulk import BulkTestBase from test.unified_format import generate_test_classes, get_test_path from test.utils_shared import ( AllowListEventListener, OvertCommandListener, TopologyEventListener, + camel_to_snake_args, is_greenthread_patched, wait_until, ) From 569a7751b63a77487b0e1eb189be874ca8695c66 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 05:27:22 -0500 Subject: [PATCH 43/50] PYTHON-5805 Drop redundant module compatibility note --- test/asynchronous/test_kms_connect.py | 6 +----- test/test_kms_connect.py | 6 +----- 2 files changed, 2 insertions(+), 10 deletions(-) diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index 62c16a3063..b24b907872 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -1,8 +1,4 @@ -"""Tests for the KMS connect callback and HTTP proxy support. - -.. note:: This module is compatible with both the synchronous and - asynchronous PyMongo test suites. -""" +"""Tests for the KMS connect callback and HTTP proxy support.""" from __future__ import annotations diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index 8e57f585ec..9f1dc64afd 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -1,8 +1,4 @@ -"""Tests for the KMS connect callback and HTTP proxy support. - -.. note:: This module is compatible with both the synchronous and - synchronous PyMongo test suites. -""" +"""Tests for the KMS connect callback and HTTP proxy support.""" from __future__ import annotations From b8b67e541b58415f5007ffc40857a0c61b7de7b3 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 05:35:44 -0500 Subject: [PATCH 44/50] PYTHON-5805 Rely on encryption marker instead of pymongocrypt skip --- test/asynchronous/test_kms_connect.py | 4 ---- test/test_kms_connect.py | 4 ---- 2 files changed, 8 deletions(-) diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index b24b907872..5aa1a5db7c 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -24,7 +24,6 @@ _wrap_encryption_errors, ) from pymongo.encryption_options import ( - _HAVE_PYMONGOCRYPT, AsyncHTTPProxyKMSConnect, AutoEncryptionOpts, HTTPProxyKMSConnect, @@ -58,7 +57,6 @@ class TestKmsConnectCallbackUnit(AsyncPyMongoTestCase): def _pool_options(): return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=None) - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") async def test_init_kms_connect_callback(self): opts = AutoEncryptionOpts({}, "k.d") self.assertIsNone(opts._kms_connect_callback) @@ -533,7 +531,6 @@ async def callback(context): with self.assertRaises(OSError): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") async def test_client_encryption_accepts_callback(self): async def callback(context): raise AssertionError("not called") @@ -549,7 +546,6 @@ async def callback(context): self.addAsyncCleanup(encryption.close) self.assertIs(encryption._io_callbacks.opts._kms_connect_callback, callback) - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") async def test_client_encryption_rejects_non_callable(self): client = self.simple_client() with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index 9f1dc64afd..b0489964b1 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -18,7 +18,6 @@ from bson.binary import Binary from pymongo.encryption_options import ( - _HAVE_PYMONGOCRYPT, AutoEncryptionOpts, HTTPProxyKMSConnect, KMSConnectContext, @@ -57,7 +56,6 @@ class TestKmsConnectCallbackUnit(PyMongoTestCase): def _pool_options(): return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=None) - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") def test_init_kms_connect_callback(self): opts = AutoEncryptionOpts({}, "k.d") self.assertIsNone(opts._kms_connect_callback) @@ -532,7 +530,6 @@ def callback(context): with self.assertRaises(OSError): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") def test_client_encryption_accepts_callback(self): def callback(context): raise AssertionError("not called") @@ -548,7 +545,6 @@ def callback(context): self.addCleanup(encryption.close) self.assertIs(encryption._io_callbacks.opts._kms_connect_callback, callback) - @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed") def test_client_encryption_rejects_non_callable(self): client = self.simple_client() with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): From 969b885250ec12aa302e97a53f5a8e9ced7c2bae Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 06:15:47 -0500 Subject: [PATCH 45/50] PYTHON-5805 Parse CONNECT status line and accept any 2xx --- pymongo/encryption_options.py | 7 ++++++- test/asynchronous/test_kms_connect.py | 11 +++++++++++ test/test_kms_connect.py | 11 +++++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 85a0854f9e..1da8fc5a02 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -169,7 +169,12 @@ def _tunnel(self, sock: socket.socket, context: KMSConnectContext) -> None: 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] - if not status.startswith(b"HTTP/1.1 200"): + # A CONNECT is successful for any 2xx status, e.g. "HTTP/1.0 200" or + # "HTTP/1.1 201", so reject malformed lines and non-2xx statuses rather + # than matching a single prefix. + parts = status.split(b" ", 2) + valid = len(parts) >= 2 and parts[0].startswith(b"HTTP/") 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: diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index 5aa1a5db7c..71619cb3a5 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -202,6 +202,17 @@ def run_stub(reply): 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 malformed status line must be rejected, not accepted as a 2xx. + host, port = run_stub(b"HTTP/1.1 2000 Evil\r\n\r\n") + with self.assertRaisesRegex(OSError, "refused CONNECT"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + 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) diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index b0489964b1..ddf20cfb7a 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -201,6 +201,17 @@ def run_stub(reply): 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 malformed status line must be rejected, not accepted as a 2xx. + host, port = run_stub(b"HTTP/1.1 2000 Evil\r\n\r\n") + with self.assertRaisesRegex(OSError, "refused CONNECT"): + HTTPProxyKMSConnect(host, port)(context) + 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) From b6f5406b898a85e4390c79bb7ec137fbe591ac15 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 06:36:10 -0500 Subject: [PATCH 46/50] PYTHON-5805 Reject bad callbacks and guard CONNECT status and injection --- pymongo/asynchronous/encryption.py | 24 +++++++++++++++--------- pymongo/encryption_options.py | 6 ++++++ pymongo/synchronous/encryption.py | 22 ++++++++++++++-------- test/asynchronous/test_kms_connect.py | 19 ++++++++++++++----- test/test_kms_connect.py | 19 ++++++++++++++----- 5 files changed, 63 insertions(+), 27 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 5ae620b515..2be9c4e5d0 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -157,22 +157,26 @@ async def _connect_kms( _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) # TLS targets address, not the peer, so verification follows the KMS host. + # For the async API, a plain callable would run its blocking connect on 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." + ) result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) - # The synchronous module takes a regular function and awaits nothing. - if not _IS_SYNC and not inspect.isawaitable(result): - _close_rejected_kms_socket(result) - raise ConfigurationError( - "kms_connect_callback must be a coroutine function for the async " - f"API, but returned {type(result)}." - ) sock = await result 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." + f"socket.socket, not {type(sock)}; consider AsyncHTTPProxyKMSConnect." ) # wrap_socket refuses a non-blocking socket, so normalize the mode here. try: @@ -187,7 +191,9 @@ async def _connect_kms( raise ConfigurationError( "kms_connect_callback must return a stream socket, not a datagram one." ) - sock.settimeout(opts.socket_timeout) + # The callback established the tunnel and may have consumed much of the + # CSOT budget, so recompute the remaining time for the TLS handshake. + sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) try: return await _async_wrap_socket_tls(sock, address, opts) except Exception as exc: diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 1da8fc5a02..7c387b5929 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -223,6 +223,12 @@ def relay(src: socket.socket, dst: socket.socket) -> None: return driver_side def __call__(self, context: KMSConnectContext) -> socket.socket: + # A configurable KMS host could inject CR/LF into the CONNECT request + # or Host header, so reject it before opening the proxy connection. + 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; a timeout per phase would let the # total run to several times the caller's budget. deadline = time.monotonic() + context.timeout diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index e87ab29aba..7a10525c70 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -156,16 +156,20 @@ def _connect_kms( _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) # TLS targets address, not the peer, so verification follows the KMS host. + # For the async API, a plain callable would run its blocking connect on 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." + ) result = kms_connect_callback( KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) ) - # The synchronous module takes a regular function and awaits nothing. - if not _IS_SYNC and not inspect.isawaitable(result): - _close_rejected_kms_socket(result) - raise ConfigurationError( - "kms_connect_callback must be a coroutine function for the async " - f"API, but returned {type(result)}." - ) sock = result if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket): _close_rejected_kms_socket(sock) @@ -186,7 +190,9 @@ def _connect_kms( raise ConfigurationError( "kms_connect_callback must return a stream socket, not a datagram one." ) - sock.settimeout(opts.socket_timeout) + # The callback established the tunnel and may have consumed much of the + # CSOT budget, so recompute the remaining time for the TLS handshake. + sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) try: return _wrap_socket_tls(sock, address, opts) except Exception as exc: diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index 71619cb3a5..e3145338af 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -213,6 +213,14 @@ def 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): + # The host is configurable, so reject CR/LF before it reaches the + # CONNECT request or Host header. + 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_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) @@ -305,19 +313,20 @@ def stub_proxy(): self.assertEqual(sock.recv(64), b"echo:ping") async def test_non_coroutine_callback_is_rejected(self): - # The async API needs a coroutine function; a plain def must not be - # awaited and retried. + # The async API needs a coroutine function; a plain def must be + # rejected before it runs its blocking connect on the event loop. if _IS_SYNC: raise unittest.SkipTest("a regular function is correct for the sync API") - left, right = socket.socketpair() - self.addCleanup(right.close) + entered = [] def callback(context): - return left + 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() diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index ddf20cfb7a..59a08db062 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -212,6 +212,14 @@ def run_stub(reply): with self.assertRaisesRegex(OSError, "refused CONNECT"): HTTPProxyKMSConnect(host, port)(context) + def test_control_characters_in_kms_host_are_rejected(self): + # The host is configurable, so reject CR/LF before it reaches the + # CONNECT request or Host header. + 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_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) @@ -304,19 +312,20 @@ def stub_proxy(): self.assertEqual(sock.recv(64), b"echo:ping") def test_non_coroutine_callback_is_rejected(self): - # The async API needs a coroutine function; a plain def must not be - # awaited and retried. + # The async API needs a coroutine function; a plain def must be + # rejected before it runs its blocking connect on the event loop. if _IS_SYNC: raise unittest.SkipTest("a regular function is correct for the sync API") - left, right = socket.socketpair() - self.addCleanup(right.close) + entered = [] def callback(context): - return left + 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() From 66e943ce80d0fed972c0a33ba75c57b3ee180614 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 07:54:21 -0500 Subject: [PATCH 47/50] PYTHON-5805 Enforce KMS budget across proxy tunnel and TLS --- pymongo/asynchronous/encryption.py | 8 +++- pymongo/encryption_options.py | 56 +++++++++++++++++++++------ pymongo/synchronous/encryption.py | 8 +++- test/asynchronous/test_kms_connect.py | 10 +++-- test/test_kms_connect.py | 10 +++-- 5 files changed, 69 insertions(+), 23 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 2be9c4e5d0..73e91fe66d 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -192,12 +192,16 @@ async def _connect_kms( "kms_connect_callback must return a stream socket, not a datagram one." ) # The callback established the tunnel and may have consumed much of the - # CSOT budget, so recompute the remaining time for the TLS handshake. + # CSOT budget, so recompute the remaining time for the TLS handshake and + # again for the KMS request that follows: wrapping resets the socket + # timeout to the pre-callback capture. sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) try: - return await _async_wrap_socket_tls(sock, address, opts) + conn = await _async_wrap_socket_tls(sock, address, opts) 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] diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 7c387b5929..1a30b56a56 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -152,16 +152,18 @@ def __init__(self, host: str, port: int, ssl_context: Optional[ssl.SSLContext] = self.port = port self.ssl_context = ssl_context - def _tunnel(self, sock: socket.socket, context: KMSConnectContext) -> None: + 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()) # A byte at a time: a bulk read could consume tunnelled bytes sent in # the same segment as the response, and the driver reads those from - # this same socket. + # this same socket. Reapply the remaining budget before each read so a + # proxy that trickles bytes cannot keep this loop alive past 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}") @@ -170,10 +172,14 @@ def _tunnel(self, sock: socket.socket, context: KMSConnectContext) -> None: 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", so reject malformed lines and non-2xx statuses rather - # than matching a single prefix. + # "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 parts[1].isdigit() + 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}") @@ -222,7 +228,9 @@ def relay(src: socket.socket, dst: socket.socket) -> None: raise return driver_side - def __call__(self, context: KMSConnectContext) -> socket.socket: + def __call__( + self, context: KMSConnectContext, deadline: Optional[float] = None + ) -> socket.socket: # A configurable KMS host could inject CR/LF into the CONNECT request # or Host header, so reject it before opening the proxy connection. if "\r" in context.host or "\n" in context.host: @@ -230,15 +238,18 @@ def __call__(self, context: KMSConnectContext) -> socket.socket: f"KMS host must not contain control characters: {context.host!r}" ) # One deadline for all three phases; a timeout per phase would let the - # total run to several times the caller's budget. - deadline = time.monotonic() + context.timeout - sock = socket.create_connection((self.host, self.port), timeout=_remaining(deadline)) + # total run to several times the caller's budget. The async wrapper + # computes the deadline before scheduling so queueing behind other + # executor work is charged against it. + 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) + self._tunnel(sock, context, deadline) except BaseException: sock.close() raise @@ -250,6 +261,26 @@ def __call__(self, context: KMSConnectContext) -> socket.socket: sock.close() raise + def _connect_proxy(self, deadline: float) -> socket.socket: + # Resolve and try each address with the budget recomputed per attempt, + # instead of create_connection applying 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: + sock.settimeout(_remaining(deadline)) + 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. @@ -262,7 +293,10 @@ class AsyncHTTPProxyKMSConnect(HTTPProxyKMSConnect): """ async def __call__(self, context: KMSConnectContext) -> socket.socket: # type: ignore[override] - connect = functools.partial(super().__call__, context) + # 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) diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 7a10525c70..b84acad81e 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -191,12 +191,16 @@ def _connect_kms( "kms_connect_callback must return a stream socket, not a datagram one." ) # The callback established the tunnel and may have consumed much of the - # CSOT budget, so recompute the remaining time for the TLS handshake. + # CSOT budget, so recompute the remaining time for the TLS handshake and + # again for the KMS request that follows: wrapping resets the socket + # timeout to the pre-callback capture. sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) try: - return _wrap_socket_tls(sock, address, opts) + conn = _wrap_socket_tls(sock, address, opts) 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] diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index e3145338af..2dd7a52aa2 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -208,10 +208,12 @@ def run_stub(reply): self.addCleanup(sock.close) self.assertIsInstance(sock, socket.socket) - # A malformed status line must be rejected, not accepted as a 2xx. - host, port = run_stub(b"HTTP/1.1 2000 Evil\r\n\r\n") - with self.assertRaisesRegex(OSError, "refused CONNECT"): - await AsyncHTTPProxyKMSConnect(host, port)(context) + # A status code must be exactly three digits; both a four-digit code + # and a zero-padded code are malformed. + 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): # The host is configurable, so reject CR/LF before it reaches the diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index 59a08db062..33cb68db47 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -207,10 +207,12 @@ def run_stub(reply): self.addCleanup(sock.close) self.assertIsInstance(sock, socket.socket) - # A malformed status line must be rejected, not accepted as a 2xx. - host, port = run_stub(b"HTTP/1.1 2000 Evil\r\n\r\n") - with self.assertRaisesRegex(OSError, "refused CONNECT"): - HTTPProxyKMSConnect(host, port)(context) + # A status code must be exactly three digits; both a four-digit code + # and a zero-padded code are malformed. + 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): # The host is configurable, so reject CR/LF before it reaches the From 5fbb9c10154914b140811be2efc58424b7a36185 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 08:17:41 -0500 Subject: [PATCH 48/50] PYTHON-5805 Harden proxy relay teardown and cancellation cleanup --- pymongo/asynchronous/encryption.py | 6 ++++++ pymongo/encryption_options.py | 13 +++++++++++-- pymongo/synchronous/encryption.py | 7 +++++++ 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 73e91fe66d..5248171ebd 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -198,6 +198,12 @@ async def _connect_kms( sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) try: conn = await _async_wrap_socket_tls(sock, address, opts) + except asyncio.CancelledError: + # The executor thread may still be wrapping the socket, so close the + # callback's socket here; otherwise it would leak and, for a TLS proxy, + # leave the bridge relay threads running. + 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)) diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index 1a30b56a56..f774cde181 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -198,19 +198,22 @@ def _bridge(self, proxy: socket.socket) -> socket.socket: driver_side, relay_side = socket.socketpair() def relay(src: socket.socket, dst: socket.socket) -> None: + # These run in daemon threads: any error, including the ValueError + # an SSLSocket.shutdown can raise during the teardown race, ends the + # relay quietly. try: while True: buf = src.recv(16384) if not buf: break dst.sendall(buf) - except OSError: + 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: + except (OSError, ValueError): pass src.close() @@ -270,7 +273,13 @@ def _connect_proxy(self, deadline: float) -> socket.socket: ): sock = socket.socket(family, socktype, proto) try: + # _remaining raises socket.timeout once the deadline passes; + # propagate it rather than disguise it as a connect error below. sock.settimeout(_remaining(deadline)) + except socket.timeout: + sock.close() + raise + try: sock.connect(sockaddr) except OSError as exc: last_error = exc diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index b84acad81e..539a9d377a 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -16,6 +16,7 @@ from __future__ import annotations +import asyncio import contextlib import functools import inspect @@ -197,6 +198,12 @@ def _connect_kms( sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) try: conn = _wrap_socket_tls(sock, address, opts) + except asyncio.CancelledError: + # The executor thread may still be wrapping the socket, so close the + # callback's socket here; otherwise it would leak and, for a TLS proxy, + # leave the bridge relay threads running. + 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)) From b143c99876ee4049863e4f951e08380b1877c079 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 08:44:01 -0500 Subject: [PATCH 49/50] PYTHON-5805 Make KMS TLS handshake cancellation-safe --- pymongo/pool_shared.py | 22 +++++++++++---- test/asynchronous/test_kms_connect.py | 39 +++++++++++++++++++++++---- test/test_kms_connect.py | 39 +++++++++++++++++++++++---- 3 files changed, 85 insertions(+), 15 deletions(-) diff --git a/pymongo/pool_shared.py b/pymongo/pool_shared.py index 5e286fc957..09404bb228 100644 --- a/pymongo/pool_shared.py +++ b/pymongo/pool_shared.py @@ -304,6 +304,12 @@ async def _async_create_connection(address: _Address, options: PoolOptions) -> s raise OSError("getaddrinfo failed") +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]: @@ -329,13 +335,19 @@ async def _async_wrap_socket_tls( # 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 diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index 2dd7a52aa2..493c893f28 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -223,6 +223,21 @@ async def test_control_characters_in_kms_host_are_rejected(self): with self.assertRaisesRegex(ConfigurationError, "control characters"): await callback(context) + async def test_cancelled_tls_wrap_closes_late_socket(self): + # Cancelling a TLS wrap can leave the executor producing an SSLSocket + # after the task is gone; the done callback must close that result. + 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) @@ -265,8 +280,13 @@ def stub_proxy(): sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) self.addCleanup(sock.close) sock.settimeout(10) - sock.sendall(b"ping") - self.assertEqual(sock.recv(64), b"echo:ping") + 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 CONNECT deadline can be much shorter than the KMS request that @@ -311,8 +331,13 @@ def stub_proxy(): sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) self.addCleanup(sock.close) sock.settimeout(10) - sock.sendall(b"ping") - self.assertEqual(sock.recv(64), b"echo:ping") + 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): # The async API needs a coroutine function; a plain def must be @@ -376,7 +401,11 @@ def stub_proxy(): sock = await AsyncHTTPProxyKMSConnect(host, port)(context) self.addCleanup(sock.close) sock.settimeout(10) - self.assertEqual(sock.recv(64), b"early-bytes") + 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): # The contract says connected; an unconnected socket would otherwise diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index 33cb68db47..f1536ee13b 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -222,6 +222,21 @@ def test_control_characters_in_kms_host_are_rejected(self): with self.assertRaisesRegex(ConfigurationError, "control characters"): callback(context) + def test_cancelled_tls_wrap_closes_late_socket(self): + # Cancelling a TLS wrap can leave the executor producing an SSLSocket + # after the task is gone; the done callback must close that result. + 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) @@ -264,8 +279,13 @@ def stub_proxy(): sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) self.addCleanup(sock.close) sock.settimeout(10) - sock.sendall(b"ping") - self.assertEqual(sock.recv(64), b"echo:ping") + 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 CONNECT deadline can be much shorter than the KMS request that @@ -310,8 +330,13 @@ def stub_proxy(): sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) self.addCleanup(sock.close) sock.settimeout(10) - sock.sendall(b"ping") - self.assertEqual(sock.recv(64), b"echo:ping") + 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): # The async API needs a coroutine function; a plain def must be @@ -375,7 +400,11 @@ def stub_proxy(): sock = HTTPProxyKMSConnect(host, port)(context) self.addCleanup(sock.close) sock.settimeout(10) - self.assertEqual(sock.recv(64), b"early-bytes") + 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): # The contract says connected; an unconnected socket would otherwise From 97d1ba2d0f3ff6e7975199389e3c8fee183166c4 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Wed, 9 Sep 2026 11:31:15 -0500 Subject: [PATCH 50/50] PYTHON-5805 Address review notes on prose verification and docs --- pymongo/asynchronous/encryption.py | 6 ++++-- pymongo/encryption_options.py | 2 +- pymongo/synchronous/encryption.py | 6 ++++-- test/asynchronous/test_kms_connect.py | 10 +++++++++- test/test_kms_connect.py | 10 +++++++++- 5 files changed, 27 insertions(+), 7 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 5248171ebd..242440f64f 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -689,8 +689,10 @@ def __init__( 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`. The - driver then performs the KMS TLS handshake over it. For an ordinary - HTTP proxy, pass + driver then performs the KMS TLS handshake over it. The callback + must be a coroutine function for the asynchronous API; a plain + callable is rejected before it can block the event loop. For an + ordinary HTTP proxy, pass :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect`. Defaults to ``None``, meaning the driver connects to KMS hosts directly. diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index f774cde181..e10de60313 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -121,7 +121,7 @@ class HTTPProxyKMSConnect: Pass an instance as ``kms_connect_callback`` to reach KMS hosts through a forward proxy that speaks HTTP ``CONNECT``:: - from pymongo.encryption_options import HTTPProxyKMSConnect + from pymongo.encryption_options import AutoEncryptionOpts, HTTPProxyKMSConnect opts = AutoEncryptionOpts( kms_providers={"aws": aws_creds}, diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 539a9d377a..2f61b51764 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -687,8 +687,10 @@ def __init__( 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`. The - driver then performs the KMS TLS handshake over it. For an ordinary - HTTP proxy, pass + driver then performs the KMS TLS handshake over it. The callback + must be a coroutine function for the synchronous API; a plain + callable is rejected before it can block the event loop. For an + ordinary HTTP proxy, pass :class:`~pymongo.encryption_options.HTTPProxyKMSConnect`. Defaults to ``None``, meaning the driver connects to KMS hosts directly. diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index 493c893f28..176e9911cd 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -623,6 +623,9 @@ 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 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 callback = AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) return await callback(context) @@ -639,6 +642,9 @@ 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 @@ -725,7 +731,9 @@ async def test_03_auto_encryption_through_proxy(self): raw = await self.client.db.coll.find_one({"_id": 1}) self.assertIsInstance(raw["encrypted_string"], Binary) - self.assertGreaterEqual(await self.connect_count(), 1) + # 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): diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index f1536ee13b..f60d195ce6 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -622,6 +622,9 @@ 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 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 callback = HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) return callback(context) @@ -638,6 +641,9 @@ 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 @@ -724,7 +730,9 @@ def test_03_auto_encryption_through_proxy(self): raw = self.client.db.coll.find_one({"_id": 1}) self.assertIsInstance(raw["encrypted_string"], Binary) - self.assertGreaterEqual(self.connect_count(), 1) + # 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):