PYTHON-5805 CSFLE/QE Support for HTTP Proxies - #24
Draft
blink1073 wants to merge 44 commits into
Draft
Conversation
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.
… 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.
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.
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.
…pwire 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.
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.
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.
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.
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.
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.
The worked CONNECT example moved to HTTPProxyKMSConnect when the helper landed, so the cross-reference to KMSConnectContext was stale.
# Conflicts: # doc/changelog.rst
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.
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.
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.
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.
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.
# Conflicts: # pymongo/asynchronous/encryption.py # pymongo/encryption_options.py # pymongo/synchronous/encryption.py # test/asynchronous/test_encryption.py # test/test_encryption.py
blink1073
commented
Sep 8, 2026
# Conflicts: # doc/changelog.rst
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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, plus aKMSConnectContext.HTTPProxyKMSConnect/AsyncHTTPProxyKMSConnecthelpers.pymongo/pool_shared.py, leaving existing connection paths unchanged.Test Plan
Ran the encryption prose suite against the drivers-evergreen-tools proxies: 10 passed, 2 skipped.
just lint,just typing, andjust docsclean.Checklist
Checklist for Author
Checklist for Reviewer