PYTHON-5805 CSFLE/QE Support for HTTP Proxies

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 and the driver performs the KMS TLS handshake over it, so
+  verification still targets the KMS host rather than the proxy. For an ordinary
+  HTTP proxy, pass :class:`~pymongo.encryption_options.HTTPProxyKMSConnect` or
+  :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect` instead of
+  writing a callback.
 
 Changes in Version 4.17.0 (2026/04/20)
 --------------------------------------

pymongo/asynchronous/encryption.py

@@ -19,7 +19,9 @@
 import asyncio
 import contextlib
 import enum
+import inspect
 import socket
+import ssl
 import time as time  # noqa: PLC0414 # needed in sync version
 import uuid
 import weakref
@@ -63,7 +65,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 +86,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 +117,53 @@
 _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument)
 
 
-async def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]:
+def _close_rejected_kms_socket(obj: Any) -> None:
+    """Close a rejected kms_connect_callback return value, 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,
+    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))
+
+    # 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)
+    )
+    # The synchronous module takes a regular function and awaits nothing.
+    if not _IS_SYNC and not inspect.isawaitable(result):
+        _close_rejected_kms_socket(result)
+        raise ConfigurationError(
+            "kms_connect_callback must be a coroutine function for the async "
+            f"API, but returned {type(result)}."
+        )
+    sock = await result
+    if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket):
+        _close_rejected_kms_socket(sock)
+        raise ConfigurationError(
+            "kms_connect_callback must return a connected, unwrapped "
+            f"socket.socket, not {type(sock)}; consider HTTPProxyKMSConnect."
+        )
+    # wrap_socket refuses a non-blocking socket, so normalize the mode here.
+    sock.settimeout(opts.socket_timeout)
     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 +246,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 +287,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):
@@ -596,6 +652,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.
 
@@ -665,7 +722,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. 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.
         .. versionchanged:: 4.12
            Added the `key_expiration_ms` parameter.
         .. versionchanged:: 4.0
@@ -709,6 +777,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,8 +19,14 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
-from typing import TYPE_CHECKING, Any, Optional, TypedDict
+import asyncio
+import functools
+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
 
 from pymongo.uri_parser_shared import _parse_kms_tls_options
 
@@ -54,6 +60,152 @@ 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
+    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 left in the timeout budget, else the default KMS
+        connect timeout. Never ``None``.
+
+    .. 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
+    """
+
+    host: str
+    port: int
+    timeout: Optional[float]
+
+
+# A callback that opens a connection to a KMS host.
+AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]]
+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 over TLS, pass an :class:`ssl.SSLContext`. It applies
+    only to the proxy connection; KMS TLS is still negotiated end to end::
+
+      import ssl
+
+      proxy_tls = ssl.create_default_context(cafile="proxy-ca.pem")
+      callback = HTTPProxyKMSConnect("proxy.example.com", 8443, proxy_tls)
+
+    Use :class:`AsyncHTTPProxyKMSConnect` with the asynchronous API.
+
+    :param host: Hostname of the proxy.
+    :param port: Port of the proxy.
+    :param ssl_context: Optional :class:`ssl.SSLContext` for connecting to the
+        proxy over TLS. Defaults to ``None``, meaning a plain connection.
+
+    .. versionadded:: 4.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.
+
+        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()
+
+        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]
+        # 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)
+
+
 class AutoEncryptionOpts:
     """Options to configure automatic client-side field level encryption."""
 
@@ -74,6 +226,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 +364,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`. 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.
         .. versionchanged:: 4.12
            Added the `key_expiration_ms` parameter.
         .. versionchanged:: 4.2
@@ -258,6 +424,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: