PYTHON-5805 CSFLE/QE Support for HTTP Proxies

doc/changelog.rst

@@ -16,6 +16,14 @@ PyMongo 4.19 brings a number of changes including:
   interpreter remain daemon threads and shutdown behavior is unchanged.
   Note that because these threads are non-daemon, a subinterpreter may block
   on teardown until any in-flight monitor work completes.
+- Added the ``kms_connect_callback`` option on
+  :class:`~pymongo.encryption_options.AutoEncryptionOpts`,
+  :class:`~pymongo.encryption.ClientEncryption`, and
+  :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`: a callable
+  that opens the connection to a Key Management Service (KMS) host for
+  Client-Side Field Level Encryption and Queryable Encryption, e.g. to route
+  KMS requests through a proxy. The driver performs the KMS TLS handshake over
+  the callback's connection, so verification still targets the KMS host.
 - Added the ``srv_host_validator`` keyword argument to
   :class:`~pymongo.synchronous.mongo_client.MongoClient` and
   :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient`, an alternative to

pymongo/asynchronous/encryption.py

@@ -17,8 +17,11 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import functools
+import inspect
 import socket
+import ssl
 import time as time  # noqa: PLC0414 # needed in sync version
 import uuid
 import weakref
@@ -60,7 +63,9 @@
 from pymongo.common import CONNECT_TIMEOUT
 from pymongo.daemon import _spawn_daemon
 from pymongo.encryption_options import (
+    AsyncKMSConnectCallback,
     AutoEncryptionOpts,
+    KMSConnectContext,
     RangeOpts,
     StringOpts,
     # Re-exported for backwards compatibility: TextOpts is deprecated but must
@@ -90,6 +95,8 @@
 from pymongo.pool_options import PoolOptions
 from pymongo.pool_shared import (
     _async_configured_socket,
+    _async_wrap_socket_tls,
+    _close_late_socket,
     _raise_connection_failure,
 )
 from pymongo.read_concern import ReadConcern
@@ -120,11 +127,108 @@
 _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument)
 
 
-async def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]:
+def _close_rejected_kms_socket(obj: Any) -> None:
+    """Close a callback return value that failed validation, best effort.
+
+    ``_connect_kms`` raises on a contract violation instead of returning the
+    value, so no caller ever takes ownership of it. Close it here, tolerating
+    non-socket values and ``close()`` failures.
+    """
+    close = getattr(obj, "close", None)
+    if callable(close):
+        with contextlib.suppress(Exception):
+            close()
+
+
+async def _connect_kms(
+    address: _Address,
+    opts: PoolOptions,
+    kms_connect_callback: Optional[AsyncKMSConnectCallback],
+    timeout: float,
+) -> Union[socket.socket, _sslConn]:
+    """Connect to a KMS host and perform the TLS handshake over the socket.
+
+    Uses ``kms_connect_callback`` when one is provided, otherwise connects
+    directly, and always verifies against ``address`` (the KMS host).
+    """
+    if kms_connect_callback is None:
+        try:
+            return await _async_configured_socket(address, opts)
+        except Exception as exc:
+            _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
+
+    # TLS targets address, not the peer, so verification follows the KMS host.
+    # A plain callable would block the event loop before we could reject it,
+    # so check the callback first.
+    if not _IS_SYNC:
+        callback_any: Any = kms_connect_callback
+        is_coro = inspect.iscoroutinefunction(callback_any)
+        if not is_coro and callable(callback_any):
+            is_coro = inspect.iscoroutinefunction(callback_any.__call__)
+        if not is_coro:
+            raise ConfigurationError(
+                "kms_connect_callback must be a coroutine function for the async API."
+            )
+    # Any: the sync version's callback returns a plain socket.
+    result: Any = kms_connect_callback(
+        KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout)
+    )
+    remaining = _csot.remaining()
+    if remaining is None or _IS_SYNC:
+        # The synchronous API cannot interrupt a callback that has started
+        # running; honoring the deadline is the callback's contract there.
+        sock = await result
+    else:
+        # CSOT is cooperative: a callback that ignores the timeout could block
+        # past the deadline. Shield the task so stopping the wait does not
+        # cancel it mid-flight, and close any socket it yields later.
+        task = asyncio.ensure_future(result)
+        try:
+            sock = await asyncio.wait_for(asyncio.shield(task), remaining)
+        except asyncio.CancelledError:
+            task.add_done_callback(_close_late_socket)
+            raise
+        except asyncio.TimeoutError:
+            task.add_done_callback(_close_late_socket)
+            _raise_connection_failure(
+                address,
+                socket.timeout("timed out"),
+                timeout_details=_get_timeout_details(opts),
+            )
+    if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket):
+        _close_rejected_kms_socket(sock)
+        raise ConfigurationError(
+            "kms_connect_callback must return a connected, unwrapped "
+            f"socket.socket, not {type(sock)}."
+        )
+    # wrap_socket refuses a non-blocking socket, so normalize the mode here.
+    try:
+        sock.getpeername()
+    except OSError:
+        _close_rejected_kms_socket(sock)
+        raise ConfigurationError(
+            "kms_connect_callback must return an already connected socket."
+        ) from None
+    if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM:
+        _close_rejected_kms_socket(sock)
+        raise ConfigurationError(
+            "kms_connect_callback must return a stream socket, not a datagram one."
+        )
+    # The callback may have consumed much of the CSOT budget, and wrapping
+    # resets the socket timeout, so recompute the remaining time here and for
+    # the KMS request that follows.
+    sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001))
     try:
-        return await _async_configured_socket(address, opts)
+        conn = await _async_wrap_socket_tls(sock, address, opts)
+    except asyncio.CancelledError:
+        # The executor may still be wrapping the socket; close it so a TLS
+        # proxy's relay threads wind down instead of leaking.
+        sock.close()
+        raise
     except Exception as exc:
         _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
+    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001))
+    return conn
 
 
 class _EncryptionIO(AsyncMongoCryptCallback):  # type: ignore[misc]
@@ -179,22 +283,28 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 False,  # disable_ocsp_endpoint_check
                 _IS_SYNC,
             )
-        # CSOT: set timeout for socket creation.
-        connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)
-        opts = PoolOptions(
-            connect_timeout=connect_timeout,
-            socket_timeout=connect_timeout,
-            ssl_context=ctx,
-        )
         address = parse_host(endpoint, _HTTPS_PORT)
         if address[0].endswith(".sock"):
             raise ConfigurationError(f"Invalid KMS endpoint {endpoint!r}")
         sleep_u = kms_context.usleep
         if sleep_u:
             sleep_sec = float(sleep_u) / 1e6
             await asyncio.sleep(sleep_sec)
+        # Set the connect timeout after the retry backoff so the budget
+        # reflects the sleep.
+        connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)
+        opts = PoolOptions(
+            connect_timeout=connect_timeout,
+            socket_timeout=connect_timeout,
+            ssl_context=ctx,
+        )
         try:
-            conn = await _connect_kms(address, opts)
+            conn = await _connect_kms(
+                address,
+                opts,
+                self.opts._kms_connect_callback,
+                connect_timeout,
+            )
             try:
                 await async_socket_sendall(conn, message)
                 while kms_context.bytes_needed > 0:
@@ -230,6 +340,8 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 conn.close()
         except MongoCryptError:
             raise  # Propagate MongoCryptError errors directly.
+        except ConfigurationError:
+            raise  # A callback contract violation is not transient.
         except Exception as exc:
             remaining = _csot.remaining()
             if isinstance(exc, NetworkTimeout) or (remaining is not None and remaining <= 0):
@@ -526,6 +638,7 @@ def __init__(
         codec_options: CodecOptions[_DocumentTypeArg],
         kms_tls_options: Optional[Mapping[str, Any]] = None,
         key_expiration_ms: Optional[int] = None,
+        kms_connect_callback: Optional[AsyncKMSConnectCallback] = None,
     ) -> None:
         """Explicit client-side field level encryption.
 
@@ -595,7 +708,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, e.g. to route KMS requests through a proxy. It receives
+            a :class:`~pymongo.encryption_options.KMSConnectContext` and
+            returns a connected, unwrapped :class:`socket.socket`, over which
+            the driver performs the KMS TLS handshake.
+            Must be a coroutine function for the asynchronous API.
+            The callback is responsible for honoring ``context.timeout``.
+            Defaults to ``None``, meaning the driver connects to KMS hosts
+            directly.
+
+        .. versionchanged:: 4.19
+           Added the `kms_connect_callback` parameter.
         .. versionchanged:: 4.12
            Added the `key_expiration_ms` parameter.
         .. versionchanged:: 4.0
@@ -639,6 +763,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(

pymongo/encryption_options.py

@@ -19,9 +19,11 @@
 
 from __future__ import annotations
 
+import socket
 import warnings
-from collections.abc import Mapping
-from typing import TYPE_CHECKING, Any, Optional, TypedDict
+from collections.abc import Awaitable, Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict
 
 from pymongo.uri_parser_shared import _parse_kms_tls_options
 
@@ -55,6 +57,38 @@ def check_min_pymongocrypt() -> None:
         )
 
 
+@dataclass(frozen=True)
+class KMSConnectContext:
+    """Information about a pending KMS connection.
+
+    Passed to ``kms_connect_callback``, which must return a plain, unwrapped
+    :class:`socket.socket`. The driver performs the KMS TLS handshake over it,
+    verifying against ``host`` rather than the peer actually reached.
+
+    :param host: Hostname of the KMS server, and the TLS verification target.
+    :param port: Port of the KMS server.
+    :param timeout: Seconds allowed for the connection: the default KMS
+        connect timeout, capped by the remaining time of an active operation
+        timeout (``timeoutMS``).
+
+    .. note:: ``timeoutMS`` does not constrain KMS requests for explicit
+       encryption, so ``timeout`` is always the default there. Automatic
+       encryption passes the remaining budget. This deviates from the Client
+       Side Operations Timeout specification; see PYTHON-6037.
+
+    .. versionadded:: 4.19
+    """
+
+    host: str
+    port: int
+    timeout: float
+
+
+# A callback that opens a connection to a KMS host.
+AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]]
+KMSConnectCallback = Callable[[KMSConnectContext], socket.socket]
+
+
 class AutoEncryptionOpts:
     """Options to configure automatic client-side field level encryption."""
 
@@ -75,6 +109,7 @@ def __init__(
         bypass_query_analysis: bool = False,
         encrypted_fields_map: Optional[Mapping[str, Any]] = None,
         key_expiration_ms: Optional[int] = None,
+        kms_connect_callback: Optional[Callable[[KMSConnectContext], Any]] = None,
     ) -> None:
         """Options to configure automatic client-side field level encryption.
 
@@ -212,7 +247,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:`KMSConnectContext` and returns a connected,
+            unwrapped :class:`socket.socket`, over which the driver performs
+            the KMS TLS handshake. Must be a coroutine function for
+            :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient` and a
+            regular function for
+            :class:`~pymongo.synchronous.mongo_client.MongoClient`. Defaults
+            to ``None``, meaning the driver connects to KMS hosts directly.
+
+        .. versionchanged:: 4.19
+           Added the `kms_connect_callback` parameter.
         .. versionchanged:: 4.12
            Added the `key_expiration_ms` parameter.
         .. versionchanged:: 4.2
@@ -259,6 +305,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: