From c37efd53eb7e89aa36d67431598f3fd1be1fd7fd Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Mon, 17 Aug 2026 12:42:43 -0500 Subject: [PATCH 01/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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")