PYTHON-6147 Add CSFLE HTTP proxy KMS connect helpers

doc/changelog.rst

@@ -20,14 +20,18 @@ 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
+- 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`: 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.
+  :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. Both helpers accept custom ``CONNECT`` headers, e.g.
+  ``Proxy-Authorization`` for proxies that require authentication.
 - 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/_kms_connect.py

@@ -113,7 +113,7 @@ 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)}."
+            f"socket.socket, not {type(sock)}; consider AsyncHTTPProxyKMSConnect."
         )
     # wrap_socket refuses a non-blocking socket, so normalize the mode here.
     try:

pymongo/asynchronous/encryption.py

@@ -600,6 +600,10 @@ def __init__(
             the driver performs the KMS TLS handshake.
             Must be a coroutine function for the asynchronous API.
             The callback is responsible for honoring ``context.timeout``.
+            When a CSOT timeout is active, the driver stops waiting at the
+            deadline and closes any socket the callback yields later. For an
+            ordinary HTTP proxy, pass
+            :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect`.
             Defaults to ``None``, meaning the driver connects to KMS hosts
             directly.
 

pymongo/encryption_options.py

@@ -36,7 +36,9 @@
     _HAVE_PYMONGOCRYPT = False
 from bson import int64
 from pymongo._kms_connect_shared import (  # noqa: F401
+    AsyncHTTPProxyKMSConnect,
     AsyncKMSConnectCallback,
+    HTTPProxyKMSConnect,
     KMSConnectCallback,
     KMSConnectContext,
 )

pymongo/synchronous/_kms_connect.py

@@ -113,7 +113,7 @@ def _connect_kms(
         _close_rejected_kms_socket(sock)
         raise ConfigurationError(
             "kms_connect_callback must return a connected, unwrapped "
-            f"socket.socket, not {type(sock)}."
+            f"socket.socket, not {type(sock)}; consider HTTPProxyKMSConnect."
         )
     # wrap_socket refuses a non-blocking socket, so normalize the mode here.
     try:

pymongo/synchronous/encryption.py

@@ -597,6 +597,10 @@ def __init__(
             the driver performs the KMS TLS handshake.
             Must be a regular function.
             The callback is responsible for honoring ``context.timeout``.
+            When a CSOT timeout is active, the driver stops waiting at the
+            deadline and closes any socket the callback yields later. For an
+            ordinary HTTP proxy, pass
+            :class:`~pymongo.encryption_options.HTTPProxyKMSConnect`.
             Defaults to ``None``, meaning the driver connects to KMS hosts
             directly.
 

test/asynchronous/test_encryption.py

@@ -226,7 +226,8 @@ async def test_init_kms_tls_options(self):
         self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
 
 
-# KMS connect callback tests live in test_kms_connect.py.
+# KMS connect callback unit tests live in test_kms_connect.py. The prose
+# tests are in test_kms_connect_prose.py.
 
 
 class TestClientOptions(AsyncPyMongoTestCase):
@@ -2007,7 +2008,8 @@ async def test_invalid_hostname_in_kms_certificate(self):
             await self.client_encrypted.create_data_key("aws", master_key=key)
 
 
-# KMS connect callback tests live in test_kms_connect.py.
+# KMS connect callback unit tests live in test_kms_connect.py. The prose
+# tests are in test_kms_connect_prose.py.
 
 
 # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-tls-options-tests

test/asynchronous/test_kms_connect_prose.py

@@ -0,0 +1,218 @@
+"""Integration tests for the KMS connect callback and HTTP proxy support.
+
+The unit tests live in ``test/test_kms_connect.py`` (a file that synchro
+does not process). This module adds the integration tests, which run real
+KMS traffic through a local proxy and are executed against both APIs via
+the generated synchronous mirror.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import http.client
+import ssl
+import unittest
+from typing import Any
+
+import pytest
+
+from bson.binary import Binary
+from pymongo.encryption_options import AsyncHTTPProxyKMSConnect, AutoEncryptionOpts
+from pymongo.errors import EncryptionError
+from test.asynchronous.test_encryption import OPTS, AsyncEncryptionIntegrationTest
+from test.helpers_shared import AWS_CREDS, CA_PEM
+
+_IS_SYNC = False
+
+pytestmark = pytest.mark.encryption
+
+KMS_PROXY_HOST = "127.0.0.1"
+KMS_PROXY_PORT = 9004
+KMS_TLS_PROXY_PORT = 9005
+
+AWS_MASTER_KEY = {
+    "region": "us-east-1",
+    "key": "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0",
+}
+
+
+class TestKmsConnectCallbackProse(AsyncEncryptionIntegrationTest):
+    @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set")
+    async def asyncSetUp(self):
+        await super().asyncSetUp()
+        self.callback_calls: list[Any] = []
+
+    async def plain_callback(self, context):
+        self.callback_calls.append(context)
+        return await AsyncHTTPProxyKMSConnect(f"http://{KMS_PROXY_HOST}:{KMS_PROXY_PORT}")(context)
+
+    def _proxy_tls_context(self):
+        ctx = ssl.create_default_context(cafile=CA_PEM)
+        ctx.check_hostname = False
+        # PYTHON-5040 tracks re-enabling verification once the test CA cert
+        # is fixed. The evergreen-tools CA lacks an Authority Key Identifier
+        # that newer OpenSSL requires, so verification fails on Windows 3.14.
+        ctx.verify_mode = ssl.CERT_NONE
+        return ctx
+
+    async def tls_callback(self, context):
+        self.callback_calls.append(context)
+        callback = AsyncHTTPProxyKMSConnect(
+            f"https://{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", self._proxy_tls_context()
+        )
+        return await callback(context)
+
+    async def proxy_request(self, method, path, tls=False):
+        """Call the proxy's control endpoints and return the body."""
+        if _IS_SYNC:
+            return self._proxy_request(method, path, tls)
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self._proxy_request, method, path, tls
+        )
+
+    def _proxy_request(self, method, path, tls=False):
+        if tls:
+            conn = http.client.HTTPSConnection(
+                f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=self._proxy_tls_context()
+            )
+        else:
+            conn = http.client.HTTPConnection(f"{KMS_PROXY_HOST}:{KMS_PROXY_PORT}")
+        try:
+            conn.request(method, path)
+            return conn.getresponse().read().decode()
+        finally:
+            conn.close()
+
+    async def connect_count(self, tls=False):
+        body = await self.proxy_request("GET", "/metrics", tls=tls)
+        # One "key value" per line. The server also emits connect_target.
+        for line in body.splitlines():
+            key, _, value = line.partition(" ")
+            if key == "connect_count":
+                return int(value)
+        raise AssertionError(f"no connect_count in metrics body: {body!r}")
+
+    async def test_01_plain_http_proxy(self):
+        await self.proxy_request("POST", "/reset")
+        encryption = self.create_client_encryption(
+            {"aws": AWS_CREDS},
+            "keyvault.datakeys",
+            self.client,
+            OPTS,
+            kms_connect_callback=self.plain_callback,
+        )
+        await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY)
+        self.assertGreaterEqual(await self.connect_count(), 1)
+
+    async def test_02_https_proxy(self):
+        await self.proxy_request("POST", "/reset", tls=True)
+        encryption = self.create_client_encryption(
+            {"aws": AWS_CREDS},
+            "keyvault.datakeys",
+            self.client,
+            OPTS,
+            kms_connect_callback=self.tls_callback,
+        )
+        await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY)
+        self.assertGreaterEqual(await self.connect_count(tls=True), 1)
+
+    async def test_03_auto_encryption_through_proxy(self):
+        await self.client.keyvault.datakeys.drop()
+        await self.client.db.coll.drop()
+
+        encryption = self.create_client_encryption(
+            {"aws": AWS_CREDS},
+            "keyvault.datakeys",
+            self.client,
+            OPTS,
+            kms_connect_callback=self.plain_callback,
+        )
+        data_key_id = await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY)
+        schema = {
+            "bsonType": "object",
+            "properties": {
+                "encrypted_string": {
+                    "encrypt": {
+                        "keyId": [data_key_id],
+                        "bsonType": "string",
+                        "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic",
+                    }
+                }
+            },
+        }
+
+        await self.proxy_request("POST", "/reset")
+        opts = AutoEncryptionOpts(
+            {"aws": AWS_CREDS},
+            "keyvault.datakeys",
+            schema_map={"db.coll": schema},
+            kms_connect_callback=self.plain_callback,
+        )
+        client_encrypted = await self.async_rs_or_single_client(auto_encryption_opts=opts)
+
+        await client_encrypted.db.coll.insert_one({"_id": 1, "encrypted_string": "hello"})
+        decrypted = await client_encrypted.db.coll.find_one({"_id": 1})
+        self.assertEqual(decrypted["encrypted_string"], "hello")
+
+        raw = await self.client.db.coll.find_one({"_id": 1})
+        self.assertIsInstance(raw["encrypted_string"], Binary)
+
+        # The decrypt reuses the cached key, so exactly one KMS request follows
+        # the reset.
+        self.assertEqual(await self.connect_count(), 1)
+
+    async def test_04_callback_error(self):
+        async def failing_callback(context):
+            raise OSError("proxy is on fire")
+
+        encryption = self.create_client_encryption(
+            {"aws": AWS_CREDS},
+            "keyvault.datakeys",
+            self.client,
+            OPTS,
+            kms_connect_callback=failing_callback,
+        )
+        with self.assertRaisesRegex(EncryptionError, "proxy is on fire"):
+            await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY)
+
+    @unittest.skip(
+        "PYTHON-6037 ClientEncryption does not support timeoutMS, so the "
+        "callback always receives the default KMS connect timeout"
+    )
+    async def test_05_callback_receives_timeout(self):
+        key_vault_client = await self.async_rs_or_single_client(timeoutMS=1000)
+        encryption = self.create_client_encryption(
+            {"aws": AWS_CREDS},
+            "keyvault.datakeys",
+            key_vault_client,
+            OPTS,
+            kms_connect_callback=self.plain_callback,
+        )
+        await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY)
+
+        self.assertTrue(self.callback_calls, "callback was never invoked")
+        for context in self.callback_calls:
+            # Checks only the spec's non-zero requirement, which cannot fail.
+            self.assertIsNotNone(context.timeout)
+            self.assertGreater(context.timeout, 0)
+
+    async def test_06_retry_after_network_error(self):
+        state = {"calls": 0}
+
+        async def flaky_callback(context):
+            state["calls"] += 1
+            if state["calls"] == 1:
+                raise OSError("first attempt fails")
+            return await AsyncHTTPProxyKMSConnect(f"http://{KMS_PROXY_HOST}:{KMS_PROXY_PORT}")(
+                context
+            )
+
+        encryption = self.create_client_encryption(
+            {"aws": AWS_CREDS},
+            "keyvault.datakeys",
+            self.client,
+            OPTS,
+            kms_connect_callback=flaky_callback,
+        )
+        await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY)
+        self.assertGreaterEqual(state["calls"], 2)

test/test_encryption.py

@@ -226,7 +226,8 @@ def test_init_kms_tls_options(self):
         self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
 
 
-# KMS connect callback tests live in test_kms_connect.py.
+# KMS connect callback unit tests live in test_kms_connect.py. The prose
+# tests are in test_kms_connect_prose.py.
 
 
 class TestClientOptions(PyMongoTestCase):
@@ -1999,7 +2000,8 @@ def test_invalid_hostname_in_kms_certificate(self):
             self.client_encrypted.create_data_key("aws", master_key=key)
 
 
-# KMS connect callback tests live in test_kms_connect.py.
+# KMS connect callback unit tests live in test_kms_connect.py. The prose
+# tests are in test_kms_connect_prose.py.
 
 
 # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#kms-tls-options-tests