diff --git a/pymongo/_kms_connect_shared.py b/pymongo/_kms_connect_shared.py new file mode 100644 index 0000000000..c160657da9 --- /dev/null +++ b/pymongo/_kms_connect_shared.py @@ -0,0 +1,75 @@ +# Copyright 2026-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""KMS connection support shared by the synchronous and asynchronous APIs. + +The per-API connection logic lives in ``pymongo.asynchronous._kms_connect`` +and its generated synchronous mirror. +""" + +from __future__ import annotations + +import contextlib +import socket +from collections.abc import Awaitable +from dataclasses import dataclass +from typing import Any, Callable + + +@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, + verifying against ``host`` rather than the peer actually reached. + + :param host: Hostname of the KMS server, and the TLS verification target. + :param port: Port of the KMS server. + :param timeout: Seconds allowed for the connection: the default KMS + connect timeout, capped by the remaining ``timeoutMS`` budget. + + .. 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.19 + """ + + host: str + port: int + timeout: float + + +# A callback that opens a connection to a KMS host. +AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]] +KMSConnectCallback = Callable[[KMSConnectContext], socket.socket] + + +def _close_rejected_kms_socket(obj: Any) -> None: + """Close a callback return value that failed validation, best effort. + + ``_connect_kms`` raises on a contract violation, so no caller takes + ownership. Tolerates non-socket values and ``close()`` failures. + """ + close = getattr(obj, "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + +# Sphinx documents this class under pymongo.encryption_options, so the +# definition must claim that module name. +KMSConnectContext.__module__ = "pymongo.encryption_options" diff --git a/pymongo/asynchronous/_kms_connect.py b/pymongo/asynchronous/_kms_connect.py new file mode 100644 index 0000000000..168d63425d --- /dev/null +++ b/pymongo/asynchronous/_kms_connect.py @@ -0,0 +1,145 @@ +# Copyright 2026-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""KMS connection for the asynchronous API. + +Connects to a KMS host, directly or through a ``kms_connect_callback``, and +performs the KMS TLS handshake over the connection. The helpers shared by +both APIs live in ``pymongo._kms_connect_shared``. +""" + +from __future__ import annotations + +import asyncio +import inspect +import socket +import ssl +from typing import TYPE_CHECKING, Any, Optional, Union, cast + +from pymongo import _csot +from pymongo._kms_connect_shared import ( + AsyncKMSConnectCallback, + KMSConnectContext, + _close_rejected_kms_socket, +) +from pymongo.common import CONNECT_TIMEOUT +from pymongo.errors import ConfigurationError +from pymongo.helpers_shared import _get_timeout_details +from pymongo.pool_options import PoolOptions +from pymongo.pool_shared import ( + _async_configured_socket, + _async_wrap_socket_tls, + _close_late_socket, + _raise_connection_failure, +) + +if TYPE_CHECKING: + from pymongo.pyopenssl_context import _sslConn + from pymongo.typings import _Address + + +_IS_SYNC = False + +_KMS_CONNECT_TIMEOUT = CONNECT_TIMEOUT # CDRIVER-3262 redefined this value to CONNECT_TIMEOUT + + +async def _connect_kms( + address: _Address, + opts: PoolOptions, + kms_connect_callback: Optional[AsyncKMSConnectCallback], + timeout: float, +) -> Union[socket.socket, _sslConn]: + """Connect to a KMS host and perform the TLS handshake over the socket. + + Uses ``kms_connect_callback`` when one is provided, otherwise connects + directly, and always verifies against ``address`` (the KMS host). + """ + 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. Reject plain callables up front: invoking one would block the + # event loop. + if not _IS_SYNC: + callback_any: Any = kms_connect_callback + is_coro = inspect.iscoroutinefunction(callback_any) + if not is_coro and callable(callback_any): + is_coro = inspect.iscoroutinefunction(callback_any.__call__) + if not is_coro: + raise ConfigurationError( + "kms_connect_callback must be a coroutine function for the async API." + ) + # Any: the sync version's callback returns a plain socket. + result: Any = kms_connect_callback( + KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) + ) + remaining = _csot.remaining() + if remaining is None or _IS_SYNC: + # The synchronous API cannot interrupt a running callback; honoring + # the deadline is the callback's contract there. + sock = await result + else: + # CSOT is cooperative: a callback that ignores the timeout can + # outlive the deadline. Shield the task so cancelling the wait does + # not cancel it mid-flight, and close the socket it yields later. + task = asyncio.ensure_future(result) + try: + sock = await asyncio.wait_for(asyncio.shield(task), remaining) + except asyncio.CancelledError: + task.add_done_callback(_close_late_socket) + raise + except asyncio.TimeoutError: + task.add_done_callback(_close_late_socket) + _raise_connection_failure( + address, + socket.timeout("timed out"), + timeout_details=_get_timeout_details(opts), + ) + 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)}." + ) + # wrap_socket refuses a non-blocking socket, so normalize the mode here. + try: + sock.getpeername() + except OSError: + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return an already connected socket." + ) from None + if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM: + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return a stream socket, not a datagram one." + ) + # The callback may have consumed much of the CSOT budget, and wrapping + # resets the socket timeout, so recompute the remaining time for the KMS + # request that follows. + sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) + try: + conn = await _async_wrap_socket_tls(sock, address, opts) + except asyncio.CancelledError: + # The executor may still be wrapping the socket; close it so a TLS + # proxy's relay threads wind down instead of leaking. + sock.close() + raise + except Exception as exc: + _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) + conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) + return conn diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 4ae63716e3..9c2d72604c 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -17,11 +17,8 @@ from __future__ import annotations import asyncio -import contextlib import functools -import inspect import socket -import ssl import time as time # noqa: PLC0414 # needed in sync version import uuid import weakref @@ -56,16 +53,15 @@ from bson.codec_options import CodecOptions from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument, _inflate_bson from pymongo import _csot, _op_id +from pymongo._kms_connect_shared import AsyncKMSConnectCallback +from pymongo.asynchronous._kms_connect import _KMS_CONNECT_TIMEOUT, _connect_kms from pymongo.asynchronous.collection import AsyncCollection from pymongo.asynchronous.cursor import AsyncCursor from pymongo.asynchronous.database import AsyncDatabase from pymongo.asynchronous.mongo_client import AsyncMongoClient -from pymongo.common import CONNECT_TIMEOUT from pymongo.daemon import _spawn_daemon from pymongo.encryption_options import ( - AsyncKMSConnectCallback, AutoEncryptionOpts, - KMSConnectContext, RangeOpts, StringOpts, # Re-exported for backwards compatibility: TextOpts is deprecated but must @@ -94,9 +90,6 @@ from pymongo.operations import UpdateOne from pymongo.pool_options import PoolOptions from pymongo.pool_shared import ( - _async_configured_socket, - _async_wrap_socket_tls, - _close_late_socket, _raise_connection_failure, ) from pymongo.read_concern import ReadConcern @@ -109,14 +102,10 @@ if TYPE_CHECKING: from pymongocrypt.mongocrypt import MongoCryptKmsContext - from pymongo.pyopenssl_context import _sslConn - from pymongo.typings import _Address - _IS_SYNC = False _HTTPS_PORT = 443 -_KMS_CONNECT_TIMEOUT = CONNECT_TIMEOUT # CDRIVER-3262 redefined this value to CONNECT_TIMEOUT _MONGOCRYPTD_TIMEOUT_MS = 10000 _DATA_KEY_OPTS: CodecOptions[dict[str, Any]] = CodecOptions( @@ -127,110 +116,6 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -def _close_rejected_kms_socket(obj: Any) -> None: - """Close a callback return value that failed validation, best effort. - - ``_connect_kms`` raises on a contract violation instead of returning the - value, so no caller ever takes ownership of it. Close it here, tolerating - non-socket values and ``close()`` failures. - """ - 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], - timeout: float, -) -> Union[socket.socket, _sslConn]: - """Connect to a KMS host and perform the TLS handshake over the socket. - - Uses ``kms_connect_callback`` when one is provided, otherwise connects - directly, and always verifies against ``address`` (the KMS host). - """ - 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. - # A plain callable would block the event loop before we could reject it, - # so check the callback first. - if not _IS_SYNC: - callback_any: Any = kms_connect_callback - is_coro = inspect.iscoroutinefunction(callback_any) - if not is_coro and callable(callback_any): - is_coro = inspect.iscoroutinefunction(callback_any.__call__) - if not is_coro: - raise ConfigurationError( - "kms_connect_callback must be a coroutine function for the async API." - ) - # Any: the sync version's callback returns a plain socket. - result: Any = kms_connect_callback( - KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) - ) - remaining = _csot.remaining() - if remaining is None or _IS_SYNC: - # The synchronous API cannot interrupt a callback that has started - # running; honoring the deadline is the callback's contract there. - sock = await result - else: - # CSOT is cooperative: a callback that ignores the timeout could block - # past the deadline. Shield the task so stopping the wait does not - # cancel it mid-flight, and close any socket it yields later. - task = asyncio.ensure_future(result) - try: - sock = await asyncio.wait_for(asyncio.shield(task), remaining) - except asyncio.CancelledError: - task.add_done_callback(_close_late_socket) - raise - except asyncio.TimeoutError: - task.add_done_callback(_close_late_socket) - _raise_connection_failure( - address, - socket.timeout("timed out"), - timeout_details=_get_timeout_details(opts), - ) - 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)}." - ) - # wrap_socket refuses a non-blocking socket, so normalize the mode here. - try: - sock.getpeername() - except OSError: - _close_rejected_kms_socket(sock) - raise ConfigurationError( - "kms_connect_callback must return an already connected socket." - ) from None - if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM: - _close_rejected_kms_socket(sock) - raise ConfigurationError( - "kms_connect_callback must return a stream socket, not a datagram one." - ) - # The callback may have consumed much of the CSOT budget, and wrapping - # resets the socket timeout, so recompute the remaining time here and for - # the KMS request that follows. - sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) - try: - conn = await _async_wrap_socket_tls(sock, address, opts) - except asyncio.CancelledError: - # The executor may still be wrapping the socket; close it so a TLS - # proxy's relay threads wind down instead of leaking. - sock.close() - raise - except Exception as exc: - _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) - conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) - return conn - - class _EncryptionIO(AsyncMongoCryptCallback): # type: ignore[misc] def __init__( self, diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index b1d13700da..852ce37210 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -19,10 +19,8 @@ from __future__ import annotations -import socket import warnings -from collections.abc import Awaitable, Mapping -from dataclasses import dataclass +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict from pymongo.uri_parser_shared import _parse_kms_tls_options @@ -37,6 +35,11 @@ except ImportError: _HAVE_PYMONGOCRYPT = False from bson import int64 +from pymongo._kms_connect_shared import ( # noqa: F401 + AsyncKMSConnectCallback, + KMSConnectCallback, + KMSConnectContext, +) from pymongo.common import check_for_min_version, validate_is_mapping from pymongo.errors import ConfigurationError @@ -57,38 +60,6 @@ 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, - verifying against ``host`` rather than the peer actually reached. - - :param host: Hostname of the KMS server, and the TLS verification target. - :param port: Port of the KMS server. - :param timeout: Seconds allowed for the connection: the default KMS - connect timeout, capped by the remaining time of an active operation - timeout (``timeoutMS``). - - .. 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.19 - """ - - host: str - port: int - timeout: float - - -# A callback that opens a connection to a KMS host. -AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]] -KMSConnectCallback = Callable[[KMSConnectContext], socket.socket] - - class AutoEncryptionOpts: """Options to configure automatic client-side field level encryption.""" diff --git a/pymongo/synchronous/_kms_connect.py b/pymongo/synchronous/_kms_connect.py new file mode 100644 index 0000000000..12a8da93dc --- /dev/null +++ b/pymongo/synchronous/_kms_connect.py @@ -0,0 +1,145 @@ +# Copyright 2026-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""KMS connection for the synchronous API. + +Connects to a KMS host, directly or through a ``kms_connect_callback``, and +performs the KMS TLS handshake over the connection. The helpers shared by +both APIs live in ``pymongo._kms_connect_shared``. +""" + +from __future__ import annotations + +import asyncio +import inspect +import socket +import ssl +from typing import TYPE_CHECKING, Any, Optional, Union, cast + +from pymongo import _csot +from pymongo._kms_connect_shared import ( + KMSConnectCallback, + KMSConnectContext, + _close_rejected_kms_socket, +) +from pymongo.common import CONNECT_TIMEOUT +from pymongo.errors import ConfigurationError +from pymongo.helpers_shared import _get_timeout_details +from pymongo.pool_options import PoolOptions +from pymongo.pool_shared import ( + _close_late_socket, + _configured_socket, + _raise_connection_failure, + _wrap_socket_tls, +) + +if TYPE_CHECKING: + from pymongo.pyopenssl_context import _sslConn + from pymongo.typings import _Address + + +_IS_SYNC = True + +_KMS_CONNECT_TIMEOUT = CONNECT_TIMEOUT # CDRIVER-3262 redefined this value to CONNECT_TIMEOUT + + +def _connect_kms( + address: _Address, + opts: PoolOptions, + kms_connect_callback: Optional[KMSConnectCallback], + timeout: float, +) -> Union[socket.socket, _sslConn]: + """Connect to a KMS host and perform the TLS handshake over the socket. + + Uses ``kms_connect_callback`` when one is provided, otherwise connects + directly, and always verifies against ``address`` (the KMS host). + """ + if kms_connect_callback is None: + try: + return _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. Reject plain callables up front: invoking one would block the + # event loop. + if not _IS_SYNC: + callback_any: Any = kms_connect_callback + is_coro = inspect.iscoroutinefunction(callback_any) + if not is_coro and callable(callback_any): + is_coro = inspect.iscoroutinefunction(callback_any.__call__) + if not is_coro: + raise ConfigurationError( + "kms_connect_callback must be a coroutine function for the async API." + ) + # Any: the sync version's callback returns a plain socket. + result: Any = kms_connect_callback( + KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) + ) + remaining = _csot.remaining() + if remaining is None or _IS_SYNC: + # The synchronous API cannot interrupt a running callback; honoring + # the deadline is the callback's contract there. + sock = result + else: + # CSOT is cooperative: a callback that ignores the timeout can + # outlive the deadline. Shield the task so cancelling the wait does + # not cancel it mid-flight, and close the socket it yields later. + task = asyncio.ensure_future(result) + try: + sock = asyncio.wait_for(asyncio.shield(task), remaining) + except asyncio.CancelledError: + task.add_done_callback(_close_late_socket) + raise + except asyncio.TimeoutError: + task.add_done_callback(_close_late_socket) + _raise_connection_failure( + address, + socket.timeout("timed out"), + timeout_details=_get_timeout_details(opts), + ) + 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)}." + ) + # wrap_socket refuses a non-blocking socket, so normalize the mode here. + try: + sock.getpeername() + except OSError: + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return an already connected socket." + ) from None + if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM: + _close_rejected_kms_socket(sock) + raise ConfigurationError( + "kms_connect_callback must return a stream socket, not a datagram one." + ) + # The callback may have consumed much of the CSOT budget, and wrapping + # resets the socket timeout, so recompute the remaining time for the KMS + # request that follows. + sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) + try: + conn = _wrap_socket_tls(sock, address, opts) + except asyncio.CancelledError: + # The executor may still be wrapping the socket; close it so a TLS + # proxy's relay threads wind down instead of leaking. + sock.close() + raise + except Exception as exc: + _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) + conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) + return conn diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index ac2f710382..d73aec4b68 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -16,12 +16,8 @@ from __future__ import annotations -import asyncio -import contextlib import functools -import inspect import socket -import ssl import time as time # noqa: PLC0414 # needed in sync version import uuid import weakref @@ -56,12 +52,10 @@ from bson.codec_options import CodecOptions from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument, _inflate_bson from pymongo import _csot, _op_id -from pymongo.common import CONNECT_TIMEOUT +from pymongo._kms_connect_shared import KMSConnectCallback from pymongo.daemon import _spawn_daemon from pymongo.encryption_options import ( AutoEncryptionOpts, - KMSConnectCallback, - KMSConnectContext, RangeOpts, StringOpts, # Re-exported for backwards compatibility: TextOpts is deprecated but must @@ -90,14 +84,12 @@ from pymongo.operations import UpdateOne from pymongo.pool_options import PoolOptions from pymongo.pool_shared import ( - _close_late_socket, - _configured_socket, _raise_connection_failure, - _wrap_socket_tls, ) from pymongo.read_concern import ReadConcern from pymongo.results import DeleteResult from pymongo.ssl_support import BLOCKING_IO_ERRORS, get_ssl_context +from pymongo.synchronous._kms_connect import _KMS_CONNECT_TIMEOUT, _connect_kms from pymongo.synchronous.collection import Collection from pymongo.synchronous.cursor import Cursor from pymongo.synchronous.database import Database @@ -109,14 +101,10 @@ if TYPE_CHECKING: from pymongocrypt.mongocrypt import MongoCryptKmsContext - from pymongo.pyopenssl_context import _sslConn - from pymongo.typings import _Address - _IS_SYNC = True _HTTPS_PORT = 443 -_KMS_CONNECT_TIMEOUT = CONNECT_TIMEOUT # CDRIVER-3262 redefined this value to CONNECT_TIMEOUT _MONGOCRYPTD_TIMEOUT_MS = 10000 _DATA_KEY_OPTS: CodecOptions[dict[str, Any]] = CodecOptions( @@ -127,110 +115,6 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -def _close_rejected_kms_socket(obj: Any) -> None: - """Close a callback return value that failed validation, best effort. - - ``_connect_kms`` raises on a contract violation instead of returning the - value, so no caller ever takes ownership of it. Close it here, tolerating - non-socket values and ``close()`` failures. - """ - close = getattr(obj, "close", None) - if callable(close): - with contextlib.suppress(Exception): - close() - - -def _connect_kms( - address: _Address, - opts: PoolOptions, - kms_connect_callback: Optional[KMSConnectCallback], - timeout: float, -) -> Union[socket.socket, _sslConn]: - """Connect to a KMS host and perform the TLS handshake over the socket. - - Uses ``kms_connect_callback`` when one is provided, otherwise connects - directly, and always verifies against ``address`` (the KMS host). - """ - if kms_connect_callback is None: - try: - return _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. - # A plain callable would block the event loop before we could reject it, - # so check the callback first. - if not _IS_SYNC: - callback_any: Any = kms_connect_callback - is_coro = inspect.iscoroutinefunction(callback_any) - if not is_coro and callable(callback_any): - is_coro = inspect.iscoroutinefunction(callback_any.__call__) - if not is_coro: - raise ConfigurationError( - "kms_connect_callback must be a coroutine function for the async API." - ) - # Any: the sync version's callback returns a plain socket. - result: Any = kms_connect_callback( - KMSConnectContext(host=address[0], port=cast(int, address[1]), timeout=timeout) - ) - remaining = _csot.remaining() - if remaining is None or _IS_SYNC: - # The synchronous API cannot interrupt a callback that has started - # running; honoring the deadline is the callback's contract there. - sock = result - else: - # CSOT is cooperative: a callback that ignores the timeout could block - # past the deadline. Shield the task so stopping the wait does not - # cancel it mid-flight, and close any socket it yields later. - task = asyncio.ensure_future(result) - try: - sock = asyncio.wait_for(asyncio.shield(task), remaining) - except asyncio.CancelledError: - task.add_done_callback(_close_late_socket) - raise - except asyncio.TimeoutError: - task.add_done_callback(_close_late_socket) - _raise_connection_failure( - address, - socket.timeout("timed out"), - timeout_details=_get_timeout_details(opts), - ) - 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)}." - ) - # wrap_socket refuses a non-blocking socket, so normalize the mode here. - try: - sock.getpeername() - except OSError: - _close_rejected_kms_socket(sock) - raise ConfigurationError( - "kms_connect_callback must return an already connected socket." - ) from None - if sock.getsockopt(socket.SOL_SOCKET, socket.SO_TYPE) != socket.SOCK_STREAM: - _close_rejected_kms_socket(sock) - raise ConfigurationError( - "kms_connect_callback must return a stream socket, not a datagram one." - ) - # The callback may have consumed much of the CSOT budget, and wrapping - # resets the socket timeout, so recompute the remaining time here and for - # the KMS request that follows. - sock.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) - try: - conn = _wrap_socket_tls(sock, address, opts) - except asyncio.CancelledError: - # The executor may still be wrapping the socket; close it so a TLS - # proxy's relay threads wind down instead of leaking. - sock.close() - raise - except Exception as exc: - _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) - conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)) - return conn - - class _EncryptionIO(MongoCryptCallback): # type: ignore[misc] def __init__( self, diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py deleted file mode 100644 index 6a577b90ee..0000000000 --- a/test/asynchronous/test_kms_connect.py +++ /dev/null @@ -1,449 +0,0 @@ -"""Tests for the KMS connect callback.""" - -from __future__ import annotations - -import asyncio -import dataclasses -import os -import socket -import ssl -import threading -import time -import unittest -from asyncio.trsock import TransportSocket -from unittest import mock - -import pytest - -import pymongo -from pymongo.asynchronous.encryption import ( - AsyncClientEncryption, - _connect_kms, - _EncryptionIO, - _wrap_encryption_errors, -) -from pymongo.encryption_options import ( - AutoEncryptionOpts, - KMSConnectContext, -) -from pymongo.errors import ConfigurationError, ConnectionFailure, EncryptionError, NetworkTimeout -from pymongo.pool_options import PoolOptions -from pymongo.ssl_support import get_ssl_context -from test.asynchronous import AsyncPyMongoTestCase -from test.asynchronous.test_encryption import OPTS -from test.helpers_shared import CA_PEM, CERT_PATH, CLIENT_PEM - -_IS_SYNC = False - -pytestmark = pytest.mark.encryption - -_KMS_ADDRESS = ("kms.example.com", 443) - - -def _tls_server_context(cert=CLIENT_PEM): - ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) - ctx.load_cert_chain(cert) - return ctx - - -def _client_tls_context(verify=False): - # verify=False matches the driver's test mode: the local certs don't verify. - if verify: - return get_ssl_context(None, None, CA_PEM, None, False, False, False, _IS_SYNC) - return get_ssl_context(None, None, None, None, True, True, False, _IS_SYNC) - - -async def _run_blocking(func, *args): - """Run a blocking callable off the event loop (inline in the sync version).""" - if _IS_SYNC: - return func(*args) - return await asyncio.get_running_loop().run_in_executor(None, func, *args) - - -def _callback_returning(value): - """A kms_connect_callback that always produces ``value``.""" - - async def callback(context): - return value - - return callback - - -class TestKmsConnectCallbackUnit(AsyncPyMongoTestCase): - """Contract checks for kms_connect_callback that need no KMS server.""" - - @staticmethod - def _pool_options(ssl_context=None): - return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=ssl_context) - - def _listen(self, backlog=1): - listener = socket.socket() - listener.bind(("127.0.0.1", 0)) - listener.listen(backlog) - self.addCleanup(listener.close) - return listener - - def _socketpair(self): - left, right = socket.socketpair() - self.addCleanup(left.close) - self.addCleanup(right.close) - return left, right - - async def test_init_kms_connect_callback(self): - opts = AutoEncryptionOpts({}, "k.d") - self.assertIsNone(opts._kms_connect_callback) - - async def callback(context): - raise AssertionError("not called") - - opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) - self.assertIs(opts._kms_connect_callback, callback) - - for bad in [1, "not-callable", object()]: - with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): - AutoEncryptionOpts({}, "k.d", kms_connect_callback=bad) # type: ignore[arg-type] - - context = KMSConnectContext(host="kms.example.com", port=443, timeout=9.5) - self.assertEqual(context.host, "kms.example.com") - self.assertEqual(context.port, 443) - self.assertEqual(context.timeout, 9.5) - with self.assertRaises(dataclasses.FrozenInstanceError): - context.host = "evil.example.com" # type: ignore[misc] - - async def test_non_socket_return_raises_configuration_error(self): - with self.assertRaisesRegex(ConfigurationError, "must return a connected"): - await _connect_kms( - _KMS_ADDRESS, self._pool_options(), _callback_returning("not-a-socket"), 10.0 - ) - - async def test_already_wrapped_socket_is_rejected(self): - # ssl.SSLSocket passes isinstance but cannot be TLS-wrapped again. - ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) - ctx.check_hostname = False - ctx.verify_mode = ssl.CERT_NONE - left, _right = self._socketpair() - # No peer needed to produce a genuine ssl.SSLSocket. - wrapped = ctx.wrap_socket(left, do_handshake_on_connect=False, server_hostname="x") - self.addCleanup(wrapped.close) - - with self.assertRaisesRegex(ConfigurationError, "unwrapped"): - await _connect_kms( - _KMS_ADDRESS, self._pool_options(), _callback_returning(wrapped), 10.0 - ) - - async def test_context_receives_host_port_and_timeout(self): - received = [] - left, _right = self._socketpair() - - async def callback(context): - received.append(context) - return left - - # ssl_context=None returns the socket unchanged, so a plain socket is accepted. - conn = await _connect_kms(_KMS_ADDRESS, self._pool_options(), callback, 12.5) - self.assertIs(conn, left) - - self.assertEqual(len(received), 1) - self.assertEqual(received[0].host, "kms.example.com") - self.assertEqual(received[0].port, 443) - self.assertEqual(received[0].timeout, 12.5) - - async def test_non_blocking_socket_from_callback_is_accepted(self): - # Without the driver normalizing the mode, this raises ValueError. - server_ctx = _tls_server_context() - listener = self._listen() - - def serve(): - try: - conn, _ = listener.accept() - server_ctx.wrap_socket(conn, server_side=True).close() - except OSError: - pass - - threading.Thread(target=serve, daemon=True).start() - - options = self._pool_options(_client_tls_context()) - - def connect(): - sock = socket.create_connection(listener.getsockname(), timeout=10) - sock.setblocking(False) - return sock - - async def callback(context): - return await _run_blocking(connect) - - conn = await _connect_kms(listener.getsockname(), options, callback, 10.0) - self.addCleanup(conn.close) - self.assertIsNotNone(conn.gettimeout()) - - async def test_tls_verification_targets_the_kms_host(self): - # The handshake must verify against the KMS address, not the peer the - # callback connected to. The server cert covers 127.0.0.1 (the peer) - # and localhost, but not the KMS hostname used below, so only - # address-based verification produces this outcome. - server_ctx = _tls_server_context(os.path.join(CERT_PATH, "server.pem")) - listener = self._listen(2) - - def serve(): - for _ in range(2): - try: - conn, _ = listener.accept() - server_ctx.wrap_socket(conn, server_side=True).close() - except (OSError, ssl.SSLError): - # The mismatched-name attempt fails mid-handshake. - pass - - threading.Thread(target=serve, daemon=True).start() - - # Full verification: trusted CA, invalid certs and hostnames rejected. - options = self._pool_options(_client_tls_context(verify=True)) - - created = [] - - def connect(): - sock = socket.create_connection(listener.getsockname(), timeout=10) - created.append(sock) - return sock - - async def callback(context): - return await _run_blocking(connect) - - port = listener.getsockname()[1] - # The cert covers localhost: verifying against the KMS address succeeds. - conn = await _connect_kms(("localhost", port), options, callback, 10.0) - self.addCleanup(conn.close) - # TLS-wrapped in either SSL flavor: a new object, not the plain socket. - self.assertIsNot(conn, created[0]) - # The cert does not cover this name: verification must fail even - # though the peer (127.0.0.1) presents a cert valid for itself. - with self.assertRaises(ConnectionFailure): - await _connect_kms(("kms.example.com", port), options, callback, 10.0) - - async def test_asyncio_transport_socket_is_rejected(self): - # get_extra_info("socket") is a TransportSocket, not a socket.socket. - left, _right = self._socketpair() - - with self.assertRaisesRegex(ConfigurationError, "TransportSocket"): - await _connect_kms( - _KMS_ADDRESS, self._pool_options(), _callback_returning(TransportSocket(left)), 10.0 - ) - - async def test_cancelled_tls_wrap_closes_late_socket(self): - # A cancelled wrap can leave the executor producing an SSLSocket; the - # done callback must close it. - if _IS_SYNC: - raise unittest.SkipTest("the cancel-safe wrap is an async path") - from pymongo.pool_shared import _close_late_socket - - left, _right = self._socketpair() - future = asyncio.get_running_loop().create_future() - future.set_result(left) - self.assertNotEqual(left.fileno(), -1) - _close_late_socket(future) - self.assertEqual(left.fileno(), -1) - - async def test_non_coroutine_callback_is_rejected(self): - # A plain def must be rejected before it blocks the event loop. - if _IS_SYNC: - raise unittest.SkipTest("a regular function is correct for the sync API") - - entered = [] - - def callback(context): - entered.append(context) - return None - - with self.assertRaisesRegex(ConfigurationError, "coroutine function"): - await _connect_kms(_KMS_ADDRESS, self._pool_options(), callback, 10.0) - self.assertEqual(entered, [], "invalid callback must not be entered") - - async def test_unconnected_socket_from_callback_is_rejected(self): - # An unconnected socket would fail later as a transient error and be retried. - bare = socket.socket() - self.addCleanup(bare.close) - - with self.assertRaisesRegex(ConfigurationError, "already connected"): - await _connect_kms(_KMS_ADDRESS, self._pool_options(), _callback_returning(bare), 10.0) - - async def test_datagram_socket_from_callback_is_rejected(self): - # TLS on a connected UDP socket raises NotImplementedError, which would be retried. - left = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - right = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - self.addCleanup(left.close) - self.addCleanup(right.close) - right.bind(("127.0.0.1", 0)) - left.connect(right.getsockname()) - - with self.assertRaisesRegex(ConfigurationError, "stream socket"): - await _connect_kms(_KMS_ADDRESS, self._pool_options(), _callback_returning(left), 10.0) - - async def test_kms_request_does_not_retry_a_contract_violation(self): - # _connect_kms has no retry loop; the no-retry guarantee is in - # kms_request, so exercise that instead. - calls = [] - - async def callback(context): - calls.append(context) - return "not-a-socket" - - opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) - io = _EncryptionIO(None, mock.MagicMock(), None, opts) - - class StubKmsContext: - endpoint = "kms.example.com:443" - message = b"request" - kms_provider = "aws" - usleep = 0 - bytes_needed = 1 - - def feed(self, data): - raise AssertionError("should not reach the socket") - - def fail(self): - raise AssertionError("a contract violation must not be retried") - - with self.assertRaises(ConfigurationError): - await io.kms_request(StubKmsContext()) - self.assertEqual(len(calls), 1) - - async def test_contract_violation_surfaces_as_encryption_error(self): - # Callers see EncryptionError with ConfigurationError as its cause. - with self.assertRaises(EncryptionError) as caught: - with _wrap_encryption_errors(): - raise ConfigurationError("kms_connect_callback must return ...") - self.assertIsInstance(caught.exception.__cause__, ConfigurationError) - - async def test_network_error_from_callback_propagates(self): - async def callback(context): - raise OSError("proxy unreachable") - - # Not a ConfigurationError, so kms_request retries it. - with self.assertRaises(OSError): - await _connect_kms(_KMS_ADDRESS, self._pool_options(), callback, 10.0) - - async def test_csot_deadline_stops_a_hung_callback(self): - # A callback that ignores the timeout cannot block past the CSOT - # deadline, and a socket it yields later must be closed. - if _IS_SYNC: - raise unittest.SkipTest("the sync API cannot interrupt a callback") - - left, _right = self._socketpair() - - async def hung_callback(context): - await asyncio.sleep(0.5) - return left - - with self.assertRaises(NetworkTimeout): - with pymongo.timeout(0.1): - await _connect_kms(_KMS_ADDRESS, self._pool_options(), hung_callback, 10.0) - self.assertNotEqual(left.fileno(), -1) - # Let the shielded callback finish; the driver closes the late result. - await asyncio.sleep(0.75) - self.assertEqual(left.fileno(), -1) - - async def test_cancelling_kms_connect_closes_the_callback_socket(self): - # Cancelling during the TLS handshake must close the callback's socket, - # so a TLS proxy's relay threads wind down. - if _IS_SYNC: - raise unittest.SkipTest("cancellation is an async-only behavior") - - server_ctx = _tls_server_context() - listener = self._listen() - gate = threading.Event() - eof = threading.Event() - - def stub_server(): - conn = None - try: - conn, _ = listener.accept() - # The cancel may land before or after the executor starts the - # handshake. Peek for the ClientHello without consuming it, or - # for EOF if the driver closed it, before wrap_socket detaches conn. - while True: - data = conn.recv(4096, socket.MSG_PEEK) - if not data: - eof.set() - return - if data[:1] == b"\x16": # TLS handshake record - break - # Hold the handshake open until the test has cancelled. - if not gate.wait(5): - return - tls = server_ctx.wrap_socket(conn, server_side=True, do_handshake_on_connect=False) - try: - tls.do_handshake() - # A discarded connection may end in a reset rather than a - # clean EOF; either proves the driver closed it. - while tls.recv(4096): - pass - except OSError: - pass - eof.set() - tls.close() - except OSError: - pass - finally: - if conn is not None: - conn.close() - - threading.Thread(target=stub_server, daemon=True).start() - - options = self._pool_options(_client_tls_context()) - socks = [] - - def connect(): - sock = socket.create_connection(listener.getsockname(), timeout=10) - socks.append(sock) - return sock - - async def callback(context): - return await _run_blocking(connect) - - # The sync version returns a socket instead of a coroutine, so both - # error codes are needed depending on the version being checked. - pending = _connect_kms(listener.getsockname(), options, callback, 10.0) - task = asyncio.ensure_future(pending) # type: ignore[type-var,arg-type] - for _ in range(100): - if socks: - break - await asyncio.sleep(0.01) - self.assertTrue(socks, "callback was never invoked") - # Bias the cancel to land mid-handshake; the stub handles the earlier - # window too. - await asyncio.sleep(0.1) - task.cancel() - with self.assertRaises(asyncio.CancelledError): - await task - # The late SSLSocket (or raw socket) must be closed; the stub sees EOF. - gate.set() - for _ in range(50): - if eof.is_set(): - break - await asyncio.sleep(0.1) - self.assertTrue(eof.is_set(), "driver never closed the callback socket") - - async def test_client_encryption_accepts_callback(self): - async def callback(context): - raise AssertionError("not called") - - client = self.simple_client() - encryption = AsyncClientEncryption( - {"local": {"key": b"\x00" * 96}}, - "keyvault.datakeys", - client, - OPTS, - kms_connect_callback=callback, - ) - self.addAsyncCleanup(encryption.close) - self.assertIs(encryption._io_callbacks.opts._kms_connect_callback, callback) - - async def test_client_encryption_rejects_non_callable(self): - client = self.simple_client() - with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): - AsyncClientEncryption( - {"local": {"key": b"\x00" * 96}}, - "keyvault.datakeys", - client, - OPTS, - kms_connect_callback="not-callable", # type: ignore[arg-type] - ) diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index 8bd59ca689..8209651776 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -1,4 +1,30 @@ -"""Tests for the KMS connect callback.""" +# Copyright 2026-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the KMS connect callback. + +Not processed by synchro: the tests are written once and parameterized over +the asynchronous and synchronous APIs with the ``[async]``/``[sync]`` ids. +Every test runs as a coroutine on a pytest-asyncio loop; the ``[sync]`` +variants call the blocking synchronous APIs inline, which is harmless for +these self-contained tests. The ``api`` parameter provides the per-API +accessors. + +The tests must run single threaded under thread based parallelization such +as pytest-run-parallel: pytest-asyncio does not support concurrent replicas, +and the tests spin up real sockets and threads that replicas would race on. +""" from __future__ import annotations @@ -8,148 +34,242 @@ import socket import ssl import threading -import time -import unittest from asyncio.trsock import TransportSocket +from collections.abc import Callable +from contextlib import contextmanager +from typing import Any from unittest import mock import pytest import pymongo -from pymongo.encryption_options import ( - AutoEncryptionOpts, - KMSConnectContext, -) +from bson.codec_options import CodecOptions +from pymongo.encryption_options import _HAVE_PYMONGOCRYPT, AutoEncryptionOpts, KMSConnectContext from pymongo.errors import ConfigurationError, ConnectionFailure, EncryptionError, NetworkTimeout from pymongo.pool_options import PoolOptions from pymongo.ssl_support import get_ssl_context -from pymongo.synchronous.encryption import ( - ClientEncryption, - _connect_kms, - _EncryptionIO, - _wrap_encryption_errors, -) -from test import PyMongoTestCase from test.helpers_shared import CA_PEM, CERT_PATH, CLIENT_PEM -from test.test_encryption import OPTS -_IS_SYNC = True - -pytestmark = pytest.mark.encryption +pytestmark = [pytest.mark.encryption, pytest.mark.asyncio] _KMS_ADDRESS = ("kms.example.com", 443) +OPTS = CodecOptions() -def _tls_server_context(cert=CLIENT_PEM): - ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) - ctx.load_cert_chain(cert) - return ctx +class Facade: + """The per-API accessors, shared by the tests through the ``api`` parameter.""" -def _client_tls_context(verify=False): - # verify=False matches the driver's test mode: the local certs don't verify. - if verify: - return get_ssl_context(None, None, CA_PEM, None, False, False, False, _IS_SYNC) - return get_ssl_context(None, None, None, None, True, True, False, _IS_SYNC) + def __init__(self, is_async: bool) -> None: + self.is_async = is_async + async def maybe_await(self, result: Any) -> Any: + """Await ``result`` in the asynchronous API (a no-op otherwise).""" + if self.is_async: + return await result + return result -def _run_blocking(func, *args): - """Run a blocking callable off the event loop (inline in the sync version).""" - if _IS_SYNC: + async def offload(self, func: Callable[..., Any], *args: Any) -> Any: + """Run a blocking callable off the event loop (inline when synchronous).""" + if self.is_async: + return await asyncio.get_running_loop().run_in_executor(None, func, *args) return func(*args) - return asyncio.get_running_loop().run_in_executor(None, func, *args) + def encryption(self): + """The API's ``encryption`` module.""" + if self.is_async: + from pymongo.asynchronous import encryption + else: + from pymongo.synchronous import encryption + return encryption + + def kms_connect(self): + """The API's ``_kms_connect`` module.""" + if self.is_async: + from pymongo.asynchronous import _kms_connect + else: + from pymongo.synchronous import _kms_connect + return _kms_connect + + async def connect(self, address, pool_options, callback, timeout): + """``_connect_kms`` for this API.""" + module = self.kms_connect() + if self.is_async: + return await module._connect_kms(address, pool_options, callback, timeout) + return module._connect_kms(address, pool_options, callback, timeout) + + def callback(self, func): + """Adapt a non-blocking ``func(context)`` to the API's callback form.""" + if self.is_async: + + async def callback(context): + return func(context) + + return callback + + return func + + def blocking_callback(self, func): + """Adapt a blocking ``func(context)``, offloaded in the async API.""" + if self.is_async: + + async def callback(context): + return await asyncio.get_running_loop().run_in_executor(None, func, context) + + return callback + + return func + + def callback_returning(self, value): + """A kms_connect_callback that always produces ``value``.""" + return self.callback(lambda context: value) + + def client_tls_context(self, verify=False): + # verify=False matches the driver's test mode: the local certs don't verify. + if verify: + return get_ssl_context(None, None, CA_PEM, None, False, False, False, not self.is_async) + return get_ssl_context(None, None, None, None, True, True, False, not self.is_async) + + def client_encryption(self, kms_providers, key_vault_namespace, client, kms_connect_callback): + """A ClientEncryption for this API using a local key provider.""" + if self.is_async: + from pymongo.asynchronous.encryption import AsyncClientEncryption + + return AsyncClientEncryption( + kms_providers, + key_vault_namespace, + client, + OPTS, + kms_connect_callback=kms_connect_callback, + ) + from pymongo.synchronous.encryption import ClientEncryption -def _callback_returning(value): - """A kms_connect_callback that always produces ``value``.""" + return ClientEncryption( + kms_providers, + key_vault_namespace, + client, + OPTS, + kms_connect_callback=kms_connect_callback, + ) - def callback(context): - return value + def simple_client(self): + """A lazily-connecting client for this API.""" + if self.is_async: + from pymongo.asynchronous.mongo_client import AsyncMongoClient - return callback + return AsyncMongoClient() + from pymongo import MongoClient + return MongoClient() -class TestKmsConnectCallbackUnit(PyMongoTestCase): - """Contract checks for kms_connect_callback that need no KMS server.""" - @staticmethod - def _pool_options(ssl_context=None): - return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=ssl_context) +ASYNC = Facade(is_async=True) +SYNC = Facade(is_async=False) - def _listen(self, backlog=1): - listener = socket.socket() - listener.bind(("127.0.0.1", 0)) - listener.listen(backlog) - self.addCleanup(listener.close) - return listener +both_apis = pytest.mark.parametrize("api", [ASYNC, SYNC], ids=["async", "sync"]) +async_only = pytest.mark.parametrize("api", [ASYNC], ids=["async"]) - def _socketpair(self): - left, right = socket.socketpair() - self.addCleanup(left.close) - self.addCleanup(right.close) - return left, right - def test_init_kms_connect_callback(self): - opts = AutoEncryptionOpts({}, "k.d") - self.assertIsNone(opts._kms_connect_callback) +def _pool_options(ssl_context=None): + return PoolOptions(connect_timeout=10, socket_timeout=10, ssl_context=ssl_context) - def callback(context): - raise AssertionError("not called") - opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) - self.assertIs(opts._kms_connect_callback, callback) +def _tls_server_context(cert=CLIENT_PEM): + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ctx.load_cert_chain(cert) + return ctx - for bad in [1, "not-callable", object()]: - with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): - AutoEncryptionOpts({}, "k.d", kms_connect_callback=bad) # type: ignore[arg-type] - context = KMSConnectContext(host="kms.example.com", port=443, timeout=9.5) - self.assertEqual(context.host, "kms.example.com") - self.assertEqual(context.port, 443) - self.assertEqual(context.timeout, 9.5) - with self.assertRaises(dataclasses.FrozenInstanceError): - context.host = "evil.example.com" # type: ignore[misc] +@contextmanager +def _listen(backlog=1): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(backlog) + try: + yield listener + finally: + listener.close() + + +@contextmanager +def _socketpair(): + left, right = socket.socketpair() + try: + yield left, right + finally: + left.close() + right.close() + + +@both_apis +async def test_init_kms_connect_callback(api): + opts = AutoEncryptionOpts({}, "k.d") + assert opts._kms_connect_callback is None + + def action(context): + raise AssertionError("not called") + + callback = api.callback(action) + opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) + assert opts._kms_connect_callback is callback + + for bad in [1, "not-callable", object()]: + with pytest.raises(TypeError, match="kms_connect_callback must be callable"): + AutoEncryptionOpts({}, "k.d", kms_connect_callback=bad) # type: ignore[arg-type] + + context = KMSConnectContext(host="kms.example.com", port=443, timeout=9.5) + assert context.host == "kms.example.com" + assert context.port == 443 + assert context.timeout == 9.5 + with pytest.raises(dataclasses.FrozenInstanceError): + context.host = "evil.example.com" # type: ignore[misc] + + +@both_apis +async def test_non_socket_return_raises_configuration_error(api): + with pytest.raises(ConfigurationError, match="must return a connected"): + await api.connect( + _KMS_ADDRESS, _pool_options(), api.callback_returning("not-a-socket"), 10.0 + ) - def test_non_socket_return_raises_configuration_error(self): - with self.assertRaisesRegex(ConfigurationError, "must return a connected"): - _connect_kms( - _KMS_ADDRESS, self._pool_options(), _callback_returning("not-a-socket"), 10.0 - ) - def test_already_wrapped_socket_is_rejected(self): - # ssl.SSLSocket passes isinstance but cannot be TLS-wrapped again. - ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) - ctx.check_hostname = False - ctx.verify_mode = ssl.CERT_NONE - left, _right = self._socketpair() - # No peer needed to produce a genuine ssl.SSLSocket. +@both_apis +async def test_already_wrapped_socket_is_rejected(api): + # ssl.SSLSocket passes isinstance but cannot be TLS-wrapped again. + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + with _socketpair() as (left, _right): + # No peer is needed to produce a genuine ssl.SSLSocket. wrapped = ctx.wrap_socket(left, do_handshake_on_connect=False, server_hostname="x") - self.addCleanup(wrapped.close) + with pytest.raises(ConfigurationError, match="unwrapped"): + await api.connect(_KMS_ADDRESS, _pool_options(), api.callback_returning(wrapped), 10.0) - with self.assertRaisesRegex(ConfigurationError, "unwrapped"): - _connect_kms(_KMS_ADDRESS, self._pool_options(), _callback_returning(wrapped), 10.0) - def test_context_receives_host_port_and_timeout(self): - received = [] - left, _right = self._socketpair() +@both_apis +async def test_context_receives_host_port_and_timeout(api): + received = [] + with _socketpair() as (left, _right): - def callback(context): + def action(context): received.append(context) return left # ssl_context=None returns the socket unchanged, so a plain socket is accepted. - conn = _connect_kms(_KMS_ADDRESS, self._pool_options(), callback, 12.5) - self.assertIs(conn, left) + conn = await api.connect(_KMS_ADDRESS, _pool_options(), api.callback(action), 12.5) + assert conn is left + + assert len(received) == 1 + assert received[0].host == "kms.example.com" + assert received[0].port == 443 + assert received[0].timeout == 12.5 - self.assertEqual(len(received), 1) - self.assertEqual(received[0].host, "kms.example.com") - self.assertEqual(received[0].port, 443) - self.assertEqual(received[0].timeout, 12.5) - def test_non_blocking_socket_from_callback_is_accepted(self): - # Without the driver normalizing the mode, this raises ValueError. - server_ctx = _tls_server_context() - listener = self._listen() +@both_apis +async def test_non_blocking_socket_from_callback_is_accepted(api): + # Without the driver normalizing the mode, this raises ValueError. + server_ctx = _tls_server_context() + with _listen() as listener: def serve(): try: @@ -160,27 +280,32 @@ def serve(): threading.Thread(target=serve, daemon=True).start() - options = self._pool_options(_client_tls_context()) + options = _pool_options(api.client_tls_context()) def connect(): sock = socket.create_connection(listener.getsockname(), timeout=10) sock.setblocking(False) return sock - def callback(context): - return _run_blocking(connect) + conn = await api.connect( + listener.getsockname(), + options, + api.blocking_callback(lambda context: connect()), + 10.0, + ) + try: + assert conn.gettimeout() is not None + finally: + conn.close() - conn = _connect_kms(listener.getsockname(), options, callback, 10.0) - self.addCleanup(conn.close) - self.assertIsNotNone(conn.gettimeout()) - def test_tls_verification_targets_the_kms_host(self): - # The handshake must verify against the KMS address, not the peer the - # callback connected to. The server cert covers 127.0.0.1 (the peer) - # and localhost, but not the KMS hostname used below, so only - # address-based verification produces this outcome. - server_ctx = _tls_server_context(os.path.join(CERT_PATH, "server.pem")) - listener = self._listen(2) +@both_apis +async def test_tls_verification_targets_the_kms_host(api): + # The handshake must verify against the KMS address, not the peer the + # callback connected to: the cert covers 127.0.0.1 and localhost, but + # not the KMS hostname used below. + server_ctx = _tls_server_context(os.path.join(CERT_PATH, "server.pem")) + with _listen(2) as listener: def serve(): for _ in range(2): @@ -194,7 +319,7 @@ def serve(): threading.Thread(target=serve, daemon=True).start() # Full verification: trusted CA, invalid certs and hostnames rejected. - options = self._pool_options(_client_tls_context(verify=True)) + options = _pool_options(api.client_tls_context(verify=True)) created = [] @@ -203,190 +328,207 @@ def connect(): created.append(sock) return sock - def callback(context): - return _run_blocking(connect) - port = listener.getsockname()[1] # The cert covers localhost: verifying against the KMS address succeeds. - conn = _connect_kms(("localhost", port), options, callback, 10.0) - self.addCleanup(conn.close) - # TLS-wrapped in either SSL flavor: a new object, not the plain socket. - self.assertIsNot(conn, created[0]) - # The cert does not cover this name: verification must fail even - # though the peer (127.0.0.1) presents a cert valid for itself. - with self.assertRaises(ConnectionFailure): - _connect_kms(("kms.example.com", port), options, callback, 10.0) - - def test_asyncio_transport_socket_is_rejected(self): - # get_extra_info("socket") is a TransportSocket, not a socket.socket. - left, _right = self._socketpair() - - with self.assertRaisesRegex(ConfigurationError, "TransportSocket"): - _connect_kms( - _KMS_ADDRESS, self._pool_options(), _callback_returning(TransportSocket(left)), 10.0 + conn = await api.connect( + ("localhost", port), + options, + api.blocking_callback(lambda context: connect()), + 10.0, + ) + try: + # TLS-wrapped in either API: a new object, not the plain socket. + assert conn is not created[0] + finally: + conn.close() + # The cert does not cover this name, so verification fails even + # though the peer's cert is valid for itself. + with pytest.raises(ConnectionFailure): + await api.connect( + ("kms.example.com", port), + options, + api.blocking_callback(lambda context: connect()), + 10.0, ) - def test_cancelled_tls_wrap_closes_late_socket(self): - # A cancelled wrap can leave the executor producing an SSLSocket; the - # done callback must close it. - if _IS_SYNC: - raise unittest.SkipTest("the cancel-safe wrap is an async path") - from pymongo.pool_shared import _close_late_socket - left, _right = self._socketpair() +@both_apis +async def test_asyncio_transport_socket_is_rejected(api): + # get_extra_info("socket") is a TransportSocket, not a socket.socket. + with _socketpair() as (left, _right): + with pytest.raises(ConfigurationError, match="TransportSocket"): + await api.connect( + _KMS_ADDRESS, + _pool_options(), + api.callback_returning(TransportSocket(left)), + 10.0, + ) + + +@async_only +async def test_cancelled_tls_wrap_closes_late_socket(api): + # A cancelled wrap can leave the executor producing an SSLSocket. The + # done callback must close it. + from pymongo.pool_shared import _close_late_socket + + with _socketpair() as (left, _right): future = asyncio.get_running_loop().create_future() future.set_result(left) - self.assertNotEqual(left.fileno(), -1) + assert left.fileno() != -1 _close_late_socket(future) - self.assertEqual(left.fileno(), -1) - - def test_non_coroutine_callback_is_rejected(self): - # A plain def must be rejected before it blocks the event loop. - if _IS_SYNC: - raise unittest.SkipTest("a regular function is correct for the sync API") - - entered = [] - - def callback(context): - entered.append(context) - return None - - with self.assertRaisesRegex(ConfigurationError, "coroutine function"): - _connect_kms(_KMS_ADDRESS, self._pool_options(), callback, 10.0) - self.assertEqual(entered, [], "invalid callback must not be entered") - - def test_unconnected_socket_from_callback_is_rejected(self): - # An unconnected socket would fail later as a transient error and be retried. - bare = socket.socket() - self.addCleanup(bare.close) - - with self.assertRaisesRegex(ConfigurationError, "already connected"): - _connect_kms(_KMS_ADDRESS, self._pool_options(), _callback_returning(bare), 10.0) - - def test_datagram_socket_from_callback_is_rejected(self): - # TLS on a connected UDP socket raises NotImplementedError, which would be retried. - left = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - right = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - self.addCleanup(left.close) - self.addCleanup(right.close) - right.bind(("127.0.0.1", 0)) - left.connect(right.getsockname()) - - with self.assertRaisesRegex(ConfigurationError, "stream socket"): - _connect_kms(_KMS_ADDRESS, self._pool_options(), _callback_returning(left), 10.0) - - def test_kms_request_does_not_retry_a_contract_violation(self): - # _connect_kms has no retry loop; the no-retry guarantee is in - # kms_request, so exercise that instead. - calls = [] - - def callback(context): - calls.append(context) - return "not-a-socket" - - opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=callback) - io = _EncryptionIO(None, mock.MagicMock(), None, opts) - - class StubKmsContext: - endpoint = "kms.example.com:443" - message = b"request" - kms_provider = "aws" - usleep = 0 - bytes_needed = 1 - - def feed(self, data): - raise AssertionError("should not reach the socket") - - def fail(self): - raise AssertionError("a contract violation must not be retried") - - with self.assertRaises(ConfigurationError): - io.kms_request(StubKmsContext()) - self.assertEqual(len(calls), 1) - - def test_contract_violation_surfaces_as_encryption_error(self): - # Callers see EncryptionError with ConfigurationError as its cause. - with self.assertRaises(EncryptionError) as caught: - with _wrap_encryption_errors(): - raise ConfigurationError("kms_connect_callback must return ...") - self.assertIsInstance(caught.exception.__cause__, ConfigurationError) - - def test_network_error_from_callback_propagates(self): - def callback(context): - raise OSError("proxy unreachable") - - # Not a ConfigurationError, so kms_request retries it. - with self.assertRaises(OSError): - _connect_kms(_KMS_ADDRESS, self._pool_options(), callback, 10.0) - - def test_csot_deadline_stops_a_hung_callback(self): - # A callback that ignores the timeout cannot block past the CSOT - # deadline, and a socket it yields later must be closed. - if _IS_SYNC: - raise unittest.SkipTest("the sync API cannot interrupt a callback") - - left, _right = self._socketpair() - - def hung_callback(context): - time.sleep(0.5) + assert left.fileno() == -1 + + +@async_only +async def test_non_coroutine_callback_is_rejected(api): + # A plain def must be rejected before it blocks the event loop. + entered = [] + + def callback(context): + entered.append(context) + return None + + with pytest.raises(ConfigurationError, match="coroutine function"): + await api.connect(_KMS_ADDRESS, _pool_options(), callback, 10.0) + assert entered == [], "invalid callback must not be entered" + + +@both_apis +async def test_unconnected_socket_from_callback_is_rejected(api): + # An unconnected socket would fail later as a transient error and be retried. + with socket.socket() as bare: + with pytest.raises(ConfigurationError, match="already connected"): + await api.connect(_KMS_ADDRESS, _pool_options(), api.callback_returning(bare), 10.0) + + +@both_apis +async def test_datagram_socket_from_callback_is_rejected(api): + # TLS on a connected UDP socket raises NotImplementedError, which would be retried. + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as left: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as right: + right.bind(("127.0.0.1", 0)) + left.connect(right.getsockname()) + + with pytest.raises(ConfigurationError, match="stream socket"): + await api.connect(_KMS_ADDRESS, _pool_options(), api.callback_returning(left), 10.0) + + +@both_apis +async def test_kms_request_does_not_retry_a_contract_violation(api): + # _connect_kms has no retry loop; the no-retry guarantee is in + # kms_request, so exercise that instead. + calls = [] + + def action(context): + calls.append(context) + return "not-a-socket" + + opts = AutoEncryptionOpts({}, "k.d", kms_connect_callback=api.callback(action)) + io = api.encryption()._EncryptionIO(None, mock.MagicMock(), None, opts) + + class StubKmsContext: + endpoint = "kms.example.com:443" + message = b"request" + kms_provider = "aws" + usleep = 0 + bytes_needed = 1 + + def feed(self, data): + raise AssertionError("should not reach the socket") + + def fail(self): + raise AssertionError("a contract violation must not be retried") + + with pytest.raises(ConfigurationError): + await api.maybe_await(io.kms_request(StubKmsContext())) + assert len(calls) == 1 + + +@both_apis +async def test_contract_violation_surfaces_as_encryption_error(api): + # Callers see EncryptionError with ConfigurationError as its cause. + with pytest.raises(EncryptionError) as exc_info: + with api.encryption()._wrap_encryption_errors(): + raise ConfigurationError("kms_connect_callback must return ...") + assert isinstance(exc_info.value.__cause__, ConfigurationError) + + +@both_apis +async def test_network_error_from_callback_propagates(api): + def action(context): + raise OSError("proxy unreachable") + + # Not a ConfigurationError, so kms_request retries it. + with pytest.raises(OSError): + await api.connect(_KMS_ADDRESS, _pool_options(), api.callback(action), 10.0) + + +@async_only +async def test_csot_deadline_stops_a_hung_callback(api): + # A callback that ignores the timeout cannot block past the CSOT + # deadline, and a socket it yields later must be closed. + with _socketpair() as (left, _right): + + async def hung_callback(context): + await asyncio.sleep(0.5) return left - with self.assertRaises(NetworkTimeout): + with pytest.raises(NetworkTimeout): with pymongo.timeout(0.1): - _connect_kms(_KMS_ADDRESS, self._pool_options(), hung_callback, 10.0) - self.assertNotEqual(left.fileno(), -1) - # Let the shielded callback finish; the driver closes the late result. - time.sleep(0.75) - self.assertEqual(left.fileno(), -1) - - def test_cancelling_kms_connect_closes_the_callback_socket(self): - # Cancelling during the TLS handshake must close the callback's socket, - # so a TLS proxy's relay threads wind down. - if _IS_SYNC: - raise unittest.SkipTest("cancellation is an async-only behavior") - - server_ctx = _tls_server_context() - listener = self._listen() - gate = threading.Event() - eof = threading.Event() - - def stub_server(): - conn = None - try: - conn, _ = listener.accept() - # The cancel may land before or after the executor starts the - # handshake. Peek for the ClientHello without consuming it, or - # for EOF if the driver closed it, before wrap_socket detaches conn. - while True: - data = conn.recv(4096, socket.MSG_PEEK) - if not data: - eof.set() - return - if data[:1] == b"\x16": # TLS handshake record - break - # Hold the handshake open until the test has cancelled. - if not gate.wait(5): + await api.connect(_KMS_ADDRESS, _pool_options(), hung_callback, 10.0) + assert left.fileno() != -1 + # Let the shielded callback finish. The driver closes the late result. + await asyncio.sleep(0.75) + assert left.fileno() == -1 + + +@async_only +async def test_cancelling_kms_connect_closes_the_callback_socket(api): + # Cancelling during the TLS handshake must close the callback's socket, + # so a TLS proxy's relay threads wind down. + server_ctx = _tls_server_context() + gate = threading.Event() + eof = threading.Event() + + def stub_server(listener): + conn = None + try: + conn, _ = listener.accept() + # The cancel may land before or after the executor starts the + # handshake: peek for a ClientHello without consuming it, or for + # EOF if the driver closed the socket, before wrap_socket + # detaches conn. + while True: + data = conn.recv(4096, socket.MSG_PEEK) + if not data: + eof.set() return - tls = server_ctx.wrap_socket(conn, server_side=True, do_handshake_on_connect=False) - try: - tls.do_handshake() - # A discarded connection may end in a reset rather than a - # clean EOF; either proves the driver closed it. - while tls.recv(4096): - pass - except OSError: + if data[:1] == b"\x16": # TLS handshake record + break + # Hold the handshake open until the test has cancelled. + if not gate.wait(5): + return + tls = server_ctx.wrap_socket(conn, server_side=True, do_handshake_on_connect=False) + try: + tls.do_handshake() + # A discarded connection may reset instead of reaching clean + # EOF; either proves the driver closed it. + while tls.recv(4096): pass - eof.set() - tls.close() except OSError: pass - finally: - if conn is not None: - conn.close() - - threading.Thread(target=stub_server, daemon=True).start() - - options = self._pool_options(_client_tls_context()) + eof.set() + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + with _listen() as listener: + threading.Thread(target=stub_server, args=(listener,), daemon=True).start() + + options = _pool_options(api.client_tls_context()) socks = [] def connect(): @@ -394,54 +536,63 @@ def connect(): socks.append(sock) return sock - def callback(context): - return _run_blocking(connect) - - # The sync version returns a socket instead of a coroutine, so both - # error codes are needed depending on the version being checked. - pending = _connect_kms(listener.getsockname(), options, callback, 10.0) - task = asyncio.ensure_future(pending) # type: ignore[type-var,arg-type] + # Schedule the connect now and cancel it once the callback socket + # exists, so the cancellation lands mid-handshake. + task = asyncio.ensure_future( + api.connect( + listener.getsockname(), + options, + api.blocking_callback(lambda context: connect()), + 10.0, + ) + ) for _ in range(100): if socks: break - time.sleep(0.01) - self.assertTrue(socks, "callback was never invoked") - # Bias the cancel to land mid-handshake; the stub handles the earlier + await asyncio.sleep(0.01) + assert socks, "callback was never invoked" + # Bias the cancel to land mid-handshake. The stub handles the earlier # window too. - time.sleep(0.1) + await asyncio.sleep(0.1) task.cancel() - with self.assertRaises(asyncio.CancelledError): - task - # The late SSLSocket (or raw socket) must be closed; the stub sees EOF. + with pytest.raises(asyncio.CancelledError): + await task + # The late SSLSocket (or raw socket) must be closed. The stub sees EOF. gate.set() for _ in range(50): if eof.is_set(): break - time.sleep(0.1) - self.assertTrue(eof.is_set(), "driver never closed the callback socket") - - def test_client_encryption_accepts_callback(self): - def callback(context): - raise AssertionError("not called") - - client = self.simple_client() - encryption = ClientEncryption( + await asyncio.sleep(0.1) + assert eof.is_set(), "driver never closed the callback socket" + + +@pytest.mark.skipif(not _HAVE_PYMONGOCRYPT, reason="pymongocrypt is not installed") +@both_apis +async def test_client_encryption_accepts_callback(api): + def action(context): + raise AssertionError("not called") + + callback = api.callback(action) + client = api.simple_client() + encryption = api.client_encryption( + {"local": {"key": b"\x00" * 96}}, "keyvault.datakeys", client, callback + ) + try: + assert encryption._io_callbacks.opts._kms_connect_callback is callback + finally: + await api.maybe_await(encryption.close()) + await api.maybe_await(client.close()) + + +@pytest.mark.skipif(not _HAVE_PYMONGOCRYPT, reason="pymongocrypt is not installed") +@both_apis +async def test_client_encryption_rejects_non_callable(api): + client = api.simple_client() + with pytest.raises(TypeError, match="kms_connect_callback must be callable"): + api.client_encryption( {"local": {"key": b"\x00" * 96}}, "keyvault.datakeys", client, - OPTS, - kms_connect_callback=callback, + "not-callable", ) - self.addCleanup(encryption.close) - self.assertIs(encryption._io_callbacks.opts._kms_connect_callback, callback) - - def test_client_encryption_rejects_non_callable(self): - client = self.simple_client() - with self.assertRaisesRegex(TypeError, "kms_connect_callback must be callable"): - ClientEncryption( - {"local": {"key": b"\x00" * 96}}, - "keyvault.datakeys", - client, - OPTS, - kms_connect_callback="not-callable", # type: ignore[arg-type] - ) + await api.maybe_await(client.close())