Conversation
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Pull request overview
Adds HTTP proxy tunneling for CSFLE and Queryable Encryption KMS traffic while preserving end-to-end KMS TLS verification.
Changes:
- Adds synchronous/asynchronous KMS connection callbacks and HTTP proxy helpers.
- Separates socket connection from TLS wrapping.
- Adds unit, integration, generated mirror tests, and documentation.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
pymongo/encryption_options.py |
Defines callback APIs, context, and proxy helpers. |
pymongo/asynchronous/encryption.py |
Integrates callbacks into asynchronous KMS requests. |
pymongo/synchronous/encryption.py |
Provides the generated synchronous integration. |
pymongo/pool_shared.py |
Extracts TLS wrapping from connection creation. |
test/asynchronous/test_encryption.py |
Tests asynchronous callbacks and proxy tunneling. |
test/test_encryption.py |
Provides synchronous encryption tests. |
test/asynchronous/test_pooling.py |
Tests asynchronous TLS wrapping. |
test/test_pooling.py |
Provides synchronous TLS-wrapping tests. |
tools/synchro.py |
Adds mappings for new asynchronous symbols. |
doc/changelog.rst |
Documents HTTP proxy support. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.
Suppressed comments (3)
Previously missed (3) — in code that hasn't changed since the last review.
pymongo/encryption_options.py:198
- Thread startup is not failure-atomic. If
socket.socketpair()or eitherThread.start()fails (for example under descriptor or thread exhaustion),_bridgeraises after the caller's close guard and can leave the TLS proxy socket, socketpair, and possibly the first relay thread alive. Add cleanup for every partially created resource before propagating startup failures.
for pair in ((relay_side, proxy), (proxy, relay_side)):
threading.Thread(target=relay, args=pair, daemon=True).start()
test/asynchronous/test_encryption.py:2356
- This async helper performs synchronous
http.clientconnect, request, and response reads on the event-loop thread, with no connection timeout. A slow or unavailable proxy can block the entire async test loop indefinitely. Run the complete blocking transaction in a worker thread (or use async networking) while preserving the generated synchronous variant.
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()
test/asynchronous/test_encryption.py:333
- This coroutine callback calls blocking
socket.create_connectiondirectly, which can stall the event loop during DNS resolution or connection setup. Use the event loop's nonblocking socket connection API or offload the connect before returning the deliberately nonblocking socket.
async def callback(context):
sock = socket.create_connection(listener.getsockname(), timeout=10)
sock.setblocking(False)
return sock
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
pymongo/encryption_options.py:244
- A task cancellation here cancels the asyncio Future but cannot stop an already-running executor thread. If
super().__call__subsequently succeeds, its connected tunnel socket is discarded without being closed; for HTTPS proxies this can also leave both relay threads alive. Retain/shield the worker future and arrange to close its eventual socket result when the awaiting task is cancelled.
return await asyncio.get_running_loop().run_in_executor(None, connect)
pymongo/encryption_options.py:228
- If
socket.socketpair()fails before_bridgereaches its cleanup block (for example under file-descriptor exhaustion), this already-connected TLS proxy socket is leaked. Ensure every_bridgefailure closessock.
return self._bridge(sock)
There was a problem hiding this comment.
🟡 Changes recommended
HTTP CONNECT status parsing must be corrected before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
test/asynchronous/test_encryption.py:19
- The KMS callback tests were moved to
test_kms_connect.py, but their newly added imports remain here unused (asyncio,threading,time,TransportSocket,mock, the KMS helpers/private encryption symbols,PoolOptions, andget_ssl_context). Remove them from this asynchronous source test and regenerate the synchronous mirror to avoid stale dependencies and import-time overhead.
import asyncio
- Files reviewed: 12/12 changed files
- Comments generated: 1
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Critical callback-validation and request-injection issues must be fixed before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
pymongo/asynchronous/encryption.py:192
- This restores the original
opts.socket_timeoutafter the callback, so time spent establishing the proxy tunnel is not deducted before the KMS TLS handshake. Under automatic encryption withtimeoutMS, a callback can consume nearly the entire remaining budget and the handshake then receives that budget again, allowing the operation to overrun its deadline. Track a deadline across callback execution and TLS wrapping, or recompute the remaining CSOT budget before the handshake.
pymongo/asynchronous/encryption.py:175
- The async error directs users to
HTTPProxyKMSConnect, but that synchronous helper is incompatible with this API and can block the event loop before being rejected. NameAsyncHTTPProxyKMSConnectin the async source; the synchro mapping will produce the synchronous helper name in the generated module.
raise ConfigurationError(
"kms_connect_callback must return a connected, unwrapped "
f"socket.socket, not {type(sock)}; consider HTTPProxyKMSConnect."
- Files reviewed: 12/12 changed files
- Comments generated: 2
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Five moderate timeout/deadline and HTTP status-validation issues must be addressed before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
pymongo/encryption_options.py:176
- HTTP status codes are exactly three digits, but this accepts malformed codes such as
00200as success and can callint()on thousands of digits, raisingValueErrorinstead of the intended proxyOSError. Require a three-digit status code before converting it.
pymongo/encryption_options.py:235
socket.create_connectionapplies this timeout separately to every address returned by DNS. For a dual-stack or multi-address proxy, each unreachable address can consume the entire remaining budget, so this call can exceed the KMS/CSOT deadline by a multiple. Resolve and try addresses with_remaining(deadline)recomputed for each attempt instead of delegating the whole loop tocreate_connection.
sock = socket.create_connection((self.host, self.port), timeout=_remaining(deadline))
- Files reviewed: 12/12 changed files
- Comments generated: 3
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Async TLS-handshake cancellation can leak sockets and proxy threads; async tests also contain blocking socket operations.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (3)
test/asynchronous/test_kms_connect.py:315
sendallandrecvare blocking socket calls executed directly by this coroutine; if the delayed proxy response never arrives, the event loop is blocked for up to 10 seconds. Offload them withrun_in_executor(while retaining the sync branch for synchro generation).
sock.settimeout(10)
sock.sendall(b"ping")
self.assertEqual(sock.recv(64), b"echo:ping")
test/asynchronous/test_kms_connect.py:269
- These synchronous socket operations run inside an async test and can block the event-loop thread for up to the configured 10-second timeout when the relay stalls. Run both operations in an executor (as
proxy_requestdoes below) so failures do not freeze the async test loop.
sock.settimeout(10)
sock.sendall(b"ping")
self.assertEqual(sock.recv(64), b"echo:ping")
test/asynchronous/test_kms_connect.py:379
- This blocking
recvis invoked on the event-loop thread and can stall the entire async test for up to 10 seconds if the expected coalesced bytes are delayed. Await an executor-backed receive instead, following the pattern used byproxy_request.
sock.settimeout(10)
self.assertEqual(sock.recv(64), b"early-bytes")
- Files reviewed: 12/12 changed files
- Comments generated: 1
- Review effort level: Balanced
There was a problem hiding this comment.
🔵 Needs a closer look
The HTTPS prose tests disable required certificate verification, and some new API documentation is incomplete.
Review details
Suppressed comments (5)
Previously missed (5) — in code that hasn't changed since the last review.
test/asynchronous/test_kms_connect.py:626
- This disables both certificate and hostname verification for the HTTPS proxy, so prose Case 2 can pass with an untrusted or wrong-host proxy certificate. The specification requires the callback's proxy TLS connection to be verified with
ca.pem; keep the verification defaults established bycreate_default_context.
test/asynchronous/test_kms_connect.py:642 - The HTTPS control requests also disable the
ca.pemcertificate and hostname checks, contrary to the prose setup for Case 2. Leaving the default context verification enabled ensures reset/metrics calls fail when the proxy presents an invalid identity.
test/asynchronous/test_kms_connect.py:728 - This only proves that some KMS connection occurred. Prose Case 3 also expects exactly one request after the reset, verifying that the subsequent decrypt reuses the cached key; asserting equality would cover that required behavior.
pymongo/asynchronous/encryption.py:696 - This async API rejects ordinary callables before invoking them, but the new parameter documentation says any callable is accepted. Document the coroutine requirement so users do not pass the synchronous helper or a blocking function; phrasing it with
async defalso lets synchro generate the corresponding sync wording.
pymongo/encryption_options.py:124 - The example imports only
HTTPProxyKMSConnectbut immediately constructsAutoEncryptionOpts, so it cannot run as shown. Import both public classes in the snippet.
- Files reviewed: 12/12 changed files
- Comments generated: 0 new
- Review effort level: Balanced
c40ae84 to
d6ea9c2
Compare
_extract _async_wrap_socket_tls / _wrap_socket_tls so TLS can be applied to sockets obtained from a KMS connect callback, and shield the executor handshake so cancellation cannot orphan the wrapped socket.
Add KMSConnectContext and the kms_connect_callback option to AutoEncryptionOpts and ClientEncryption so callers can route KMS connections through an HTTP proxy. The driver performs the KMS TLS handshake over the returned socket, so verification still targets the KMS host, and CSOT deadlines cover the callback. For ordinary proxies, callers can pass the new HTTPProxyKMSConnect or AsyncHTTPProxyKMSConnect helper instead of writing a callback. Enforce the CSOT deadline across the proxy tunnel, relay, and TLS handshake, and make the async callback contract strict: coroutine functions for the async API, plain callables rejected.
Cover the callback contract, proxy tunnel and relay, CONNECT status handling, CSOT deadline enforcement, and cancellation safety.
d6ea9c2 to
47e250c
Compare
PYTHON-5805
Changes in this PR
Lets CSFLE and Queryable Encryption route KMS traffic through an HTTP proxy via a user-supplied connect callback, while still performing the KMS TLS handshake end to end against the KMS host.
kms_connect_callbacktoAutoEncryptionOpts,ClientEncryption, andAsyncClientEncryption, which accept aKMSConnectContextto with the metadata about the connection.HTTPProxyKMSConnect/AsyncHTTPProxyKMSConnecthelpers for the simple case. They also server as an example for more complicated cases.pymongo/pool_shared.py, leaving existing connection paths unchanged.Test Plan
Encryption tests: all of the new tests passed and the expected skips were applied.
Checklist
Checklist for Author
Checklist for Reviewer