Skip to content

Commit 56d8e6a

Browse files
blink1073blink1073
authored
PYTHON-6154 Extract KMS connect helpers into dedicated modules (#3110)
1 parent 99ec73e commit 56d8e6a

8 files changed

Lines changed: 844 additions & 1037 deletions

File tree

‎pymongo/_kms_connect_shared.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
# Copyright 2026-present MongoDB, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""KMS connection support shared by the synchronous and asynchronous APIs.
16+
17+
The per-API connection logic lives in ``pymongo.asynchronous._kms_connect``
18+
and its generated synchronous mirror.
19+
"""
20+
21+
from __future__ import annotations
22+
23+
import contextlib
24+
import socket
25+
from collections.abc import Awaitable
26+
from dataclasses import dataclass
27+
from typing import Any, Callable
28+
29+
30+
@dataclass(frozen=True)
31+
class KMSConnectContext:
32+
"""Information about a pending KMS connection.
33+
34+
Passed to ``kms_connect_callback``, which must return a plain, unwrapped
35+
:class:`socket.socket`. The driver performs the KMS TLS handshake over it,
36+
verifying against ``host`` rather than the peer actually reached.
37+
38+
:param host: Hostname of the KMS server, and the TLS verification target.
39+
:param port: Port of the KMS server.
40+
:param timeout: Seconds allowed for the connection: the default KMS
41+
connect timeout, capped by the remaining ``timeoutMS`` budget.
42+
43+
.. note:: ``timeoutMS`` does not constrain KMS requests for explicit
44+
encryption, so ``timeout`` is always the default there. Automatic
45+
encryption passes the remaining budget. This deviates from the Client
46+
Side Operations Timeout specification; see PYTHON-6037.
47+
48+
.. versionadded:: 4.19
49+
"""
50+
51+
host: str
52+
port: int
53+
timeout: float
54+
55+
56+
# A callback that opens a connection to a KMS host.
57+
AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]]
58+
KMSConnectCallback = Callable[[KMSConnectContext], socket.socket]
59+
60+
61+
def _close_rejected_kms_socket(obj: Any) -> None:
62+
"""Close a callback return value that failed validation, best effort.
63+
64+
``_connect_kms`` raises on a contract violation, so no caller takes
65+
ownership. Tolerates non-socket values and ``close()`` failures.
66+
"""
67+
close = getattr(obj, "close", None)
68+
if callable(close):
69+
with contextlib.suppress(Exception):
70+
close()
71+
72+
73+
# Sphinx documents this class under pymongo.encryption_options, so the
74+
# definition must claim that module name.
75+
KMSConnectContext.__module__ = "pymongo.encryption_options"

‎pymongo/asynchronous/_kms_connect.py‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
# Copyright 2026-present MongoDB, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""KMS connection for the asynchronous API.
16+
17+
Connects to a KMS host, directly or through a ``kms_connect_callback``, and
18+
performs the KMS TLS handshake over the connection. The helpers shared by
19+
both APIs live in ``pymongo._kms_connect_shared``.
20+
"""
21+
22+
from __future__ import annotations
23+
24+
import asyncio
25+
import inspect
26+
import socket
27+
import ssl
28+
from typing import TYPE_CHECKING, Any, Optional, Union, cast
29+
30+
from pymongo import _csot
31+
from pymongo._kms_connect_shared import (
32+
AsyncKMSConnectCallback,
33+
KMSConnectContext,
34+
_close_rejected_kms_socket,
35+
)
36+
from pymongo.common import CONNECT_TIMEOUT
37+
from pymongo.errors import ConfigurationError
38+
from pymongo.helpers_shared import _get_timeout_details
39+
from pymongo.pool_options import PoolOptions
40+
from pymongo.pool_shared import (
41+
_async_configured_socket,
42+
_async_wrap_socket_tls,
43+
_close_late_socket,
44+
_raise_connection_failure,
45+
)
46+
47+
if TYPE_CHECKING:
48+
from pymongo.pyopenssl_context import _sslConn
49+
from pymongo.typings import _Address
50+
51+
52+
_IS_SYNC = False
53+
54+
_KMS_CONNECT_TIMEOUT = CONNECT_TIMEOUT # CDRIVER-3262 redefined this value to CONNECT_TIMEOUT
55+
56+
57+
async def _connect_kms(
58+
address: _Address,
59+
opts: PoolOptions,
60+
kms_connect_callback: Optional[AsyncKMSConnectCallback],
61+
timeout: float,
62+
) -> Union[socket.socket, _sslConn]:
63+
"""Connect to a KMS host and perform the TLS handshake over the socket.
64+
65+
Uses ``kms_connect_callback`` when one is provided, otherwise connects
66+
directly, and always verifies against ``address`` (the KMS host).
67+
"""
68+
if kms_connect_callback is None:
69+
try:
70+
return await _async_configured_socket(address, opts)
71+
except Exception as exc:
72+
_raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
73+
74+
# TLS targets ``address``, not the peer, so verification follows the KMS
75+
# host. Reject plain callables up front: invoking one would block the
76+
# event loop.
77+
if not _IS_SYNC:
78+
callback_any: Any = kms_connect_callback
79+
is_coro = inspect.iscoroutinefunction(callback_any)
80+
if not is_coro and callable(callback_any):
81+
is_coro = inspect.iscoroutinefunction(callback_any.__call__)
82+
if not is_coro:
83+
raise ConfigurationError(
84+
"kms_connect_callback must be a coroutine function for the async API."
85+
)
86+
# Any: the sync version's callback returns a plain socket.
87+
result: Any = kms_connect_callback(
88+
KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout)
89+
)
90+
remaining = _csot.remaining()
91+
if remaining is None or _IS_SYNC:
92+
# The synchronous API cannot interrupt a running callback; honoring
93+
# the deadline is the callback's contract there.
94+
sock = await result
95+
else:
96+
# CSOT is cooperative: a callback that ignores the timeout can
97+
# outlive the deadline. Shield the task so cancelling the wait does
98+
# not cancel it mid-flight, and close the socket it yields later.
99+
task = asyncio.ensure_future(result)
100+
try:
101+
sock = await asyncio.wait_for(asyncio.shield(task), remaining)
102+
except asyncio.CancelledError:
103+
task.add_done_callback(_close_late_socket)
104+
raise
105+
except asyncio.TimeoutError:
106+
task.add_done_callback(_close_late_socket)
107+
_raise_connection_failure(
108+
address,
109+
socket.timeout("timed out"),
110+
timeout_details=_get_timeout_details(opts),
111+
)
112+
if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket):
113+
_close_rejected_kms_socket(sock)
114+
raise ConfigurationError(
115+
"kms_connect_callback must return a connected, unwrapped "
116+
f"socket.socket, not {type(sock)}."
117+
)
118+
# wrap_socket refuses a non-blocking socket, so normalize the mode here.
119+
try:
120+
sock.getpeername()
121+
except OSError:
122+
_close_rejected_kms_socket(sock)
123+
raise ConfigurationError(
124+
"kms_connect_callback must return an already connected socket."
125+
) from None
126+
if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM:
127+
_close_rejected_kms_socket(sock)
128+
raise ConfigurationError(
129+
"kms_connect_callback must return a stream socket, not a datagram one."
130+
)
131+
# The callback may have consumed much of the CSOT budget, and wrapping
132+
# resets the socket timeout, so recompute the remaining time for the KMS
133+
# request that follows.
134+
sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001))
135+
try:
136+
conn = await _async_wrap_socket_tls(sock, address, opts)
137+
except asyncio.CancelledError:
138+
# The executor may still be wrapping the socket; close it so a TLS
139+
# proxy's relay threads wind down instead of leaking.
140+
sock.close()
141+
raise
142+
except Exception as exc:
143+
_raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
144+
conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001))
145+
return conn

‎pymongo/asynchronous/encryption.py‎

Lines changed: 2 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,8 @@
1717
from __future__ import annotations
1818

1919
import asyncio
20-
import contextlib
2120
import functools
22-
import inspect
2321
import socket
24-
import ssl
2522
import time as time # noqa: PLC0414 # needed in sync version
2623
import uuid
2724
import weakref
@@ -56,16 +53,15 @@
5653
from bson.codec_options import CodecOptions
5754
from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument, _inflate_bson
5855
from pymongo import _csot, _op_id
56+
from pymongo._kms_connect_shared import AsyncKMSConnectCallback
57+
from pymongo.asynchronous._kms_connect import _KMS_CONNECT_TIMEOUT, _connect_kms
5958
from pymongo.asynchronous.collection import AsyncCollection
6059
from pymongo.asynchronous.cursor import AsyncCursor
6160
from pymongo.asynchronous.database import AsyncDatabase
6261
from pymongo.asynchronous.mongo_client import AsyncMongoClient
63-
from pymongo.common import CONNECT_TIMEOUT
6462
from pymongo.daemon import _spawn_daemon
6563
from pymongo.encryption_options import (
66-
AsyncKMSConnectCallback,
6764
AutoEncryptionOpts,
68-
KMSConnectContext,
6965
RangeOpts,
7066
StringOpts,
7167
# Re-exported for backwards compatibility: TextOpts is deprecated but must
@@ -94,9 +90,6 @@
9490
from pymongo.operations import UpdateOne
9591
from pymongo.pool_options import PoolOptions
9692
from pymongo.pool_shared import (
97-
_async_configured_socket,
98-
_async_wrap_socket_tls,
99-
_close_late_socket,
10093
_raise_connection_failure,
10194
)
10295
from pymongo.read_concern import ReadConcern
@@ -109,14 +102,10 @@
109102
if TYPE_CHECKING:
110103
from pymongocrypt.mongocrypt import MongoCryptKmsContext
111104

112-
from pymongo.pyopenssl_context import _sslConn
113-
from pymongo.typings import _Address
114-
115105

116106
_IS_SYNC = False
117107

118108
_HTTPS_PORT = 443
119-
_KMS_CONNECT_TIMEOUT = CONNECT_TIMEOUT # CDRIVER-3262 redefined this value to CONNECT_TIMEOUT
120109
_MONGOCRYPTD_TIMEOUT_MS = 10000
121110

122111
_DATA_KEY_OPTS: CodecOptions[dict[str, Any]] = CodecOptions(
@@ -127,110 +116,6 @@
127116
_KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument)
128117

129118

130-
def _close_rejected_kms_socket(obj: Any) -> None:
131-
"""Close a callback return value that failed validation, best effort.
132-
133-
``_connect_kms`` raises on a contract violation instead of returning the
134-
value, so no caller ever takes ownership of it. Close it here, tolerating
135-
non-socket values and ``close()`` failures.
136-
"""
137-
close = getattr(obj, "close", None)
138-
if callable(close):
139-
with contextlib.suppress(Exception):
140-
close()
141-
142-
143-
async def _connect_kms(
144-
address: _Address,
145-
opts: PoolOptions,
146-
kms_connect_callback: Optional[AsyncKMSConnectCallback],
147-
timeout: float,
148-
) -> Union[socket.socket, _sslConn]:
149-
"""Connect to a KMS host and perform the TLS handshake over the socket.
150-
151-
Uses ``kms_connect_callback`` when one is provided, otherwise connects
152-
directly, and always verifies against ``address`` (the KMS host).
153-
"""
154-
if kms_connect_callback is None:
155-
try:
156-
return await _async_configured_socket(address, opts)
157-
except Exception as exc:
158-
_raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
159-
160-
# TLS targets address, not the peer, so verification follows the KMS host.
161-
# A plain callable would block the event loop before we could reject it,
162-
# so check the callback first.
163-
if not _IS_SYNC:
164-
callback_any: Any = kms_connect_callback
165-
is_coro = inspect.iscoroutinefunction(callback_any)
166-
if not is_coro and callable(callback_any):
167-
is_coro = inspect.iscoroutinefunction(callback_any.__call__)
168-
if not is_coro:
169-
raise ConfigurationError(
170-
"kms_connect_callback must be a coroutine function for the async API."
171-
)
172-
# Any: the sync version's callback returns a plain socket.
173-
result: Any = kms_connect_callback(
174-
KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout)
175-
)
176-
remaining = _csot.remaining()
177-
if remaining is None or _IS_SYNC:
178-
# The synchronous API cannot interrupt a callback that has started
179-
# running; honoring the deadline is the callback's contract there.
180-
sock = await result
181-
else:
182-
# CSOT is cooperative: a callback that ignores the timeout could block
183-
# past the deadline. Shield the task so stopping the wait does not
184-
# cancel it mid-flight, and close any socket it yields later.
185-
task = asyncio.ensure_future(result)
186-
try:
187-
sock = await asyncio.wait_for(asyncio.shield(task), remaining)
188-
except asyncio.CancelledError:
189-
task.add_done_callback(_close_late_socket)
190-
raise
191-
except asyncio.TimeoutError:
192-
task.add_done_callback(_close_late_socket)
193-
_raise_connection_failure(
194-
address,
195-
socket.timeout("timed out"),
196-
timeout_details=_get_timeout_details(opts),
197-
)
198-
if not isinstance(sock, socket.socket) or isinstance(sock, ssl.SSLSocket):
199-
_close_rejected_kms_socket(sock)
200-
raise ConfigurationError(
201-
"kms_connect_callback must return a connected, unwrapped "
202-
f"socket.socket, not {type(sock)}."
203-
)
204-
# wrap_socket refuses a non-blocking socket, so normalize the mode here.
205-
try:
206-
sock.getpeername()
207-
except OSError:
208-
_close_rejected_kms_socket(sock)
209-
raise ConfigurationError(
210-
"kms_connect_callback must return an already connected socket."
211-
) from None
212-
if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM:
213-
_close_rejected_kms_socket(sock)
214-
raise ConfigurationError(
215-
"kms_connect_callback must return a stream socket, not a datagram one."
216-
)
217-
# The callback may have consumed much of the CSOT budget, and wrapping
218-
# resets the socket timeout, so recompute the remaining time here and for
219-
# the KMS request that follows.
220-
sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001))
221-
try:
222-
conn = await _async_wrap_socket_tls(sock, address, opts)
223-
except asyncio.CancelledError:
224-
# The executor may still be wrapping the socket; close it so a TLS
225-
# proxy's relay threads wind down instead of leaking.
226-
sock.close()
227-
raise
228-
except Exception as exc:
229-
_raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
230-
conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001))
231-
return conn
232-
233-
234119
class _EncryptionIO(AsyncMongoCryptCallback): # type: ignore[misc]
235120
def __init__(
236121
self,

0 commit comments

Comments
 (0)