diff --git a/doc/changelog.rst b/doc/changelog.rst index 9ea2f05c2e..5f76f3f49e 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -16,14 +16,18 @@ PyMongo 4.19 brings a number of changes including: interpreter remain daemon threads and shutdown behavior is unchanged. Note that because these threads are non-daemon, a subinterpreter may block on teardown until any in-flight monitor work completes. -- Added the ``kms_connect_callback`` option on +- Added support for routing Key Management Service (KMS) requests for + Client-Side Field Level Encryption and Queryable Encryption through an HTTP + proxy, using the new ``kms_connect_callback`` option on :class:`~pymongo.encryption_options.AutoEncryptionOpts`, :class:`~pymongo.encryption.ClientEncryption`, and - :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`: a callable - that opens the connection to a Key Management Service (KMS) host for - Client-Side Field Level Encryption and Queryable Encryption, e.g. to route - KMS requests through a proxy. The driver performs the KMS TLS handshake over - the callback's connection, so verification still targets the KMS host. + :class:`~pymongo.asynchronous.encryption.AsyncClientEncryption`. The callback + opens the connection and the driver performs the KMS TLS handshake over it, so + verification still targets the KMS host rather than the proxy. For an ordinary + HTTP proxy, pass :class:`~pymongo.encryption_options.HTTPProxyKMSConnect` or + :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect` instead of + writing a callback. Both helpers accept custom ``CONNECT`` headers, e.g. + ``Proxy-Authorization`` for proxies that require authentication. - Added the ``srv_host_validator`` keyword argument to :class:`~pymongo.synchronous.mongo_client.MongoClient` and :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient`, an alternative to diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 0eea43d898..bb0488d359 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -199,7 +199,7 @@ async def _connect_kms( _close_rejected_kms_socket(sock) raise ConfigurationError( "kms_connect_callback must return a connected, unwrapped " - f"socket.socket, not {type(sock)}." + f"socket.socket, not {type(sock)}; consider AsyncHTTPProxyKMSConnect." ) # wrap_socket refuses a non-blocking socket, so normalize the mode here. try: @@ -715,7 +715,9 @@ def __init__( which the driver performs the KMS TLS handshake. The callback must be a coroutine function for the asynchronous API; a plain callable is rejected before it can block the event loop. When a CSOT timeout is active, the driver stops waiting at the - deadline and closes any socket the callback yields later. + deadline and closes any socket the callback yields later. For an + ordinary HTTP proxy, pass + :class:`~pymongo.encryption_options.AsyncHTTPProxyKMSConnect`. Defaults to ``None``, meaning the driver connects to KMS hosts directly. diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py index f77e461704..91f5e584eb 100644 --- a/pymongo/encryption_options.py +++ b/pymongo/encryption_options.py @@ -19,7 +19,13 @@ from __future__ import annotations +import asyncio +import contextlib +import functools import socket +import ssl +import threading +import time import warnings from collections.abc import Awaitable, Mapping from dataclasses import dataclass @@ -65,6 +71,9 @@ class KMSConnectContext: :class:`socket.socket`. The driver performs the KMS TLS handshake over it, verifying against ``host`` rather than the peer actually reached. + Prefer :class:`HTTPProxyKMSConnect` or :class:`AsyncHTTPProxyKMSConnect` + over writing a callback. + :param host: Hostname of the KMS server, and the TLS verification target. :param port: Port of the KMS server. :param timeout: Seconds left in the timeout budget, or the default KMS @@ -87,6 +96,247 @@ class KMSConnectContext: AsyncKMSConnectCallback = Callable[[KMSConnectContext], Awaitable[socket.socket]] KMSConnectCallback = Callable[[KMSConnectContext], socket.socket] +# Cap the CONNECT response header so a silent proxy cannot grow the buffer without bound. +_MAX_CONNECT_HEADER = 8192 + + +def _close_completed_socket(future: asyncio.Future[socket.socket]) -> None: + """Close a socket produced after its awaiting task was cancelled.""" + if not future.cancelled() and future.exception() is None: + # The callback may have returned a non-socket; close best effort. + close = getattr(future.result(), "close", None) + if callable(close): + with contextlib.suppress(Exception): + close() + + +def _remaining(deadline: float) -> float: + """Seconds left before ``deadline``.""" + left = deadline - time.monotonic() + if left <= 0: + raise socket.timeout("timed out connecting through the proxy") + return left + + +class HTTPProxyKMSConnect: + """Route KMS connections through an HTTP proxy, for the synchronous API. + + Pass an instance as ``kms_connect_callback`` to reach KMS hosts through a + forward proxy that speaks HTTP ``CONNECT``:: + + from pymongo.encryption_options import AutoEncryptionOpts, HTTPProxyKMSConnect + + opts = AutoEncryptionOpts( + kms_providers={"aws": aws_creds}, + key_vault_namespace="keyvault.datakeys", + kms_connect_callback=HTTPProxyKMSConnect("proxy.example.com", 8080), + ) + + To reach the proxy over TLS, pass an :class:`ssl.SSLContext`. It applies + only to the proxy connection; KMS TLS is still negotiated end to end:: + + import ssl + + proxy_tls = ssl.create_default_context(cafile="proxy-ca.pem") + callback = HTTPProxyKMSConnect("proxy.example.com", 8443, proxy_tls) + + Use :class:`AsyncHTTPProxyKMSConnect` with the asynchronous API. + + :param host: Hostname of the proxy. + :param port: Port of the proxy. + :param ssl_context: Optional :class:`ssl.SSLContext` for connecting to the + proxy over TLS. Defaults to ``None``, meaning a plain connection. + :param headers: Optional mapping of extra ``CONNECT`` request headers, + e.g. ``{"Proxy-Authorization": "Basic ..."}`` for a proxy that + requires authentication. ``Host`` is always set from the KMS address. + When sending credentials, also pass ``ssl_context`` so they are not + sent to the proxy in cleartext. + + .. versionadded:: 4.19 + """ + + def __init__( + self, + host: str, + port: int, + ssl_context: Optional[ssl.SSLContext] = None, + headers: Optional[Mapping[str, str]] = None, + ): + self.host = host + self.port = port + self.ssl_context = ssl_context + self.headers = dict(headers) if headers else {} + for name, value in self.headers.items(): + if not isinstance(name, str) or not isinstance(value, str): + raise TypeError("proxy header names and values must be strings") + # Header fields become CONNECT request lines; a colon or control + # character in a name, or CR/LF in a value, would corrupt or + # inject request lines. + if not name or ":" in name or "\r" in name or "\n" in name: + raise ConfigurationError(f"invalid proxy header name: {name!r}") + if name.lower() == "host": + raise ConfigurationError("the Host CONNECT header is set from the KMS address") + if "\r" in value or "\n" in value: + raise ConfigurationError(f"invalid proxy header value for {name!r}") + + def _tunnel(self, sock: socket.socket, context: KMSConnectContext, deadline: float) -> None: + # An IPv6 literal needs brackets to be a valid HTTP authority. + host = f"[{context.host}]" if ":" in context.host else context.host + target = f"{host}:{context.port}" + lines = [f"CONNECT {target} HTTP/1.1", f"Host: {target}"] + lines.extend(f"{name}: {value}" for name, value in self.headers.items()) + sock.sendall(("\r\n".join(lines) + "\r\n\r\n").encode()) + # Read a byte at a time: a bulk read could consume tunneled bytes from + # this same socket. Reapply the budget before each read so a trickling + # proxy cannot outlive the deadline. + response = bytearray() + while not response.endswith(b"\r\n\r\n"): + sock.settimeout(_remaining(deadline)) + chunk = sock.recv(1) + if not chunk: + raise OSError(f"proxy closed the connection while tunneling to {target}") + response += chunk + if len(response) > _MAX_CONNECT_HEADER: + raise OSError(f"proxy sent an oversized CONNECT response for {target}") + status = bytes(response).split(b"\r\n", 1)[0] + # A CONNECT is successful for any 2xx status, e.g. "HTTP/1.0 200" or + # "HTTP/1.1 201"; require a three-digit code and reject malformed lines. + parts = status.split(b" ", 2) + valid = ( + len(parts) >= 2 + and parts[0].startswith(b"HTTP/") + and len(parts[1]) == 3 + and parts[1].isdigit() + ) + if not valid or not 200 <= int(parts[1]) < 300: + raise OSError(f"proxy refused CONNECT to {target}: {status!r}") + + def _bridge(self, proxy: socket.socket) -> socket.socket: + """Relay a TLS proxy connection through a socketpair. + + Python cannot layer TLS over an :class:`ssl.SSLSocket`, so return the + plain end of a pair, using threads rather than tasks even in + :class:`AsyncHTTPProxyKMSConnect`: the event loop cannot read an + :class:`ssl.SSLSocket`. + """ + # Clear the CONNECT-phase timeout; the tunneled KMS request is governed + # by the driver's own timeout, not the elapsed connect budget. + proxy.settimeout(None) + driver_side, relay_side = socket.socketpair() + + def relay(src: socket.socket, dst: socket.socket) -> None: + # Daemon threads: any error, including the ValueError an + # SSLSocket.shutdown can raise in the teardown race, ends the relay. + try: + while True: + buf = src.recv(16384) + if not buf: + break + dst.sendall(buf) + except (OSError, ValueError): + pass + finally: + # Send EOF to the peer instead of closing a socket it may be reading. + try: + dst.shutdown(socket.SHUT_RDWR) + except (OSError, ValueError): + pass + src.close() + + try: + for pair in ((relay_side, proxy), (proxy, relay_side)): + threading.Thread(target=relay, args=pair, daemon=True).start() + except BaseException: + # Unblock any thread that did start, then drop every socket. + for sock in (proxy, relay_side, driver_side): + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + sock.close() + raise + return driver_side + + def __call__( + self, context: KMSConnectContext, deadline: Optional[float] = None + ) -> socket.socket: + # A configurable KMS host could inject CR/LF into the CONNECT request. + if "\r" in context.host or "\n" in context.host: + raise ConfigurationError( + f"KMS host must not contain control characters: {context.host!r}" + ) + # One deadline for all three phases; per-phase timeouts would multiply + # the caller's budget. + if deadline is None: + deadline = time.monotonic() + context.timeout + sock = self._connect_proxy(deadline) + try: + if self.ssl_context is not None: + sock.settimeout(_remaining(deadline)) + sock = self.ssl_context.wrap_socket(sock, server_hostname=self.host) + sock.settimeout(_remaining(deadline)) + self._tunnel(sock, context, deadline) + except BaseException: + sock.close() + raise + if self.ssl_context is None: + return sock + try: + return self._bridge(sock) + except BaseException: + sock.close() + raise + + def _connect_proxy(self, deadline: float) -> socket.socket: + # Recompute the budget per address, rather than socket.create_connection, + # which applies the timeout to every address. + last_error: Optional[OSError] = None + for family, socktype, proto, _, sockaddr in socket.getaddrinfo( + self.host, self.port, type=socket.SOCK_STREAM + ): + sock = socket.socket(family, socktype, proto) + try: + # Propagate the timeout from _remaining rather than report a + # connect error. + sock.settimeout(_remaining(deadline)) + except socket.timeout: + sock.close() + raise + try: + sock.connect(sockaddr) + except OSError as exc: + last_error = exc + sock.close() + continue + return sock + raise OSError( + f"could not connect to proxy {self.host}:{self.port}: {last_error}" + ) from last_error + + +class AsyncHTTPProxyKMSConnect(HTTPProxyKMSConnect): + """Route KMS connections through an HTTP proxy, for the asynchronous API. + + Behaves exactly like :class:`HTTPProxyKMSConnect`, but is a coroutine + callable and runs the blocking connect in a thread so the event loop stays + free. + + .. versionadded:: 4.19 + """ + + async def __call__(self, context: KMSConnectContext) -> socket.socket: # type: ignore[override] + # Capture the deadline before scheduling so time spent queued behind + # other executor work counts against the KMS budget. + deadline = time.monotonic() + context.timeout + connect = functools.partial(super().__call__, context, deadline) + future = asyncio.get_running_loop().run_in_executor(None, connect) + try: + return await asyncio.shield(future) + except asyncio.CancelledError: + # The thread runs on regardless, so close the socket it returns. + future.add_done_callback(_close_completed_socket) + raise + class AutoEncryptionOpts: """Options to configure automatic client-side field level encryption.""" @@ -253,8 +503,10 @@ def __init__( the KMS TLS handshake. Must be a coroutine function for :class:`~pymongo.asynchronous.mongo_client.AsyncMongoClient` and a regular function for - :class:`~pymongo.synchronous.mongo_client.MongoClient`. Defaults - to ``None``, meaning the driver connects to KMS hosts directly. + :class:`~pymongo.synchronous.mongo_client.MongoClient`. For an + ordinary proxy, pass :class:`HTTPProxyKMSConnect` or + :class:`AsyncHTTPProxyKMSConnect`. Defaults to ``None``, meaning + the driver connects to KMS hosts directly. .. versionchanged:: 4.19 Added the `kms_connect_callback` parameter. diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index b89ad933f1..79a8566177 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -199,7 +199,7 @@ def _connect_kms( _close_rejected_kms_socket(sock) raise ConfigurationError( "kms_connect_callback must return a connected, unwrapped " - f"socket.socket, not {type(sock)}." + f"socket.socket, not {type(sock)}; consider HTTPProxyKMSConnect." ) # wrap_socket refuses a non-blocking socket, so normalize the mode here. try: @@ -713,7 +713,9 @@ def __init__( which the driver performs the KMS TLS handshake. The callback must be a regular function; the async API requires a coroutine function and rejects plain callables before they can block the event loop. When a CSOT timeout is active, the driver stops waiting at the - deadline and closes any socket the callback yields later. + deadline and closes any socket the callback yields later. For an + ordinary HTTP proxy, pass + :class:`~pymongo.encryption_options.HTTPProxyKMSConnect`. Defaults to ``None``, meaning the driver connects to KMS hosts directly. diff --git a/test/asynchronous/test_kms_connect.py b/test/asynchronous/test_kms_connect.py index 994ec5ca21..c3dcba8a52 100644 --- a/test/asynchronous/test_kms_connect.py +++ b/test/asynchronous/test_kms_connect.py @@ -1,20 +1,23 @@ -"""Tests for the KMS connect callback.""" +"""Tests for the KMS connect callback and HTTP proxy support.""" from __future__ import annotations import asyncio import dataclasses +import http.client import socket import ssl import threading import time import unittest from asyncio.trsock import TransportSocket +from typing import Any from unittest import mock import pytest import pymongo +from bson.binary import Binary from pymongo.asynchronous.encryption import ( AsyncClientEncryption, _connect_kms, @@ -22,20 +25,31 @@ _wrap_encryption_errors, ) from pymongo.encryption_options import ( + AsyncHTTPProxyKMSConnect, AutoEncryptionOpts, + HTTPProxyKMSConnect, KMSConnectContext, ) from pymongo.errors import ConfigurationError, 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 CLIENT_PEM +from test.asynchronous.test_encryption import OPTS, AsyncEncryptionIntegrationTest +from test.helpers_shared import AWS_CREDS, CA_PEM, CLIENT_PEM _IS_SYNC = False pytestmark = pytest.mark.encryption +KMS_PROXY_HOST = "127.0.0.1" +KMS_PROXY_PORT = 9004 +KMS_TLS_PROXY_PORT = 9005 + +AWS_MASTER_KEY = { + "region": "us-east-1", + "key": "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0", +} + class TestKmsConnectCallbackUnit(AsyncPyMongoTestCase): """Contract checks for kms_connect_callback that need no KMS server.""" @@ -156,6 +170,158 @@ async def callback(context): with self.assertRaisesRegex(ConfigurationError, "TransportSocket"): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + async def test_http_proxy_helper_tunnels_and_reports_refusal(self): + # Covers the CONNECT handshake without KMS credentials. + accepted = [] + + def stub(listener, reply): + try: + conn, _ = listener.accept() + except OSError: + return + accepted.append(conn.recv(4096)) + conn.sendall(reply) + conn.close() + + def run_stub(reply): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + threading.Thread(target=stub, args=(listener, reply), daemon=True).start() + return listener.getsockname() + + host, port = run_stub(b"HTTP/1.1 200 Connection Established\r\n\r\n") + callback = AsyncHTTPProxyKMSConnect(host, port) + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = await callback(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + self.assertEqual(accepted[0].split(b"\r\n")[0], b"CONNECT kms.example.com:443 HTTP/1.1") + + host, port = run_stub(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n") + with self.assertRaisesRegex(OSError, "refused CONNECT"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + + # Any 2xx status is a successful tunnel, not just HTTP/1.1 200. + host, port = run_stub(b"HTTP/1.0 200 Connection Established\r\n\r\n") + sock = await AsyncHTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + + # A status code must be exactly three digits, with no zero padding. + for reply in (b"HTTP/1.1 2000 Evil\r\n\r\n", b"HTTP/1.1 00200 Evil\r\n\r\n"): + host, port = run_stub(reply) + with self.assertRaisesRegex(OSError, "refused CONNECT"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + + async def test_control_characters_in_kms_host_are_rejected(self): + # Reject CR/LF in the configurable host before it reaches CONNECT. + callback = AsyncHTTPProxyKMSConnect("proxy.example.com", 8080) + context = KMSConnectContext(host="kms.example.com\r\nX-Injected: 1", port=443, timeout=10) + with self.assertRaisesRegex(ConfigurationError, "control characters"): + await callback(context) + + async def test_http_proxy_helper_sends_custom_headers(self): + # Extra CONNECT headers reach the proxy verbatim. + accepted = [] + + def stub(listener): + try: + conn, _ = listener.accept() + except OSError: + return + request = b"" + while b"\r\n\r\n" not in request: + chunk = conn.recv(4096) + if not chunk: + return + request += chunk + accepted.append(request) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + conn.close() + + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + threading.Thread(target=stub, args=(listener,), daemon=True).start() + + host, port = listener.getsockname() + headers = {"Proxy-Authorization": "Basic dXNlcjpwYXNz", "X-Trace-Id": "abc123"} + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = await AsyncHTTPProxyKMSConnect(host, port, headers=headers)(context) + self.addCleanup(sock.close) + request = accepted[0] + self.assertEqual(request.split(b"\r\n")[0], b"CONNECT kms.example.com:443 HTTP/1.1") + self.assertIn(b"\r\nProxy-Authorization: Basic dXNlcjpwYXNz\r\n", request) + self.assertIn(b"\r\nX-Trace-Id: abc123\r\n", request) + self.assertEqual(request.count(b"\r\nHost: "), 1) + + async def test_http_proxy_helper_authenticates_to_the_proxy(self): + # The motivating case: 407 without credentials, 200 with them. + def stub(listener): + for _ in range(2): + try: + conn, _ = listener.accept() + except OSError: + return + request = b"" + while b"\r\n\r\n" not in request: + chunk = conn.recv(4096) + if not chunk: + return + request += chunk + if b"\r\nProxy-Authorization: Basic dXNlcjpwYXNz\r\n" in request: + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + else: + conn.sendall(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n") + conn.close() + + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(2) + self.addCleanup(listener.close) + threading.Thread(target=stub, args=(listener,), daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "refused CONNECT"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + headers = {"Proxy-Authorization": "Basic dXNlcjpwYXNz"} + sock = await AsyncHTTPProxyKMSConnect(host, port, headers=headers)(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + + async def test_http_proxy_helper_rejects_bad_headers(self): + for headers in [ + {"Bad\r\nName": "x"}, + {"X-Ok": "ok\r\nInjected: 1"}, + {"Host": "evil.example.com"}, + {"host": "evil.example.com"}, + {"": "x"}, + {"Bad:Name": "x"}, + ]: + with self.assertRaisesRegex(ConfigurationError, "proxy header|Host CONNECT header"): + AsyncHTTPProxyKMSConnect("proxy.example.com", 8080, headers=headers) + + for headers in [{1: "x"}, {"X-Ok": 1}, {None: "x"}, {"X-Ok": None}]: + with self.assertRaisesRegex(TypeError, "must be strings"): + AsyncHTTPProxyKMSConnect("proxy.example.com", 8080, headers=headers) + + async def test_http_proxy_helper_accepts_legal_header_values(self): + # Colons and spaces are legal in values (e.g. auth schemes); only + # CR/LF would let a value inject a request line. + callback = AsyncHTTPProxyKMSConnect( + "proxy.example.com", + 8080, + headers={"Proxy-Authorization": "Basic dXNlcjpwYXNz", "X-Token": "a: b"}, + ) + self.assertEqual( + callback.headers, + {"Proxy-Authorization": "Basic dXNlcjpwYXNz", "X-Token": "a: b"}, + ) + 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. @@ -171,6 +337,105 @@ async def test_cancelled_tls_wrap_closes_late_socket(self): self.assertEqual(left.fileno(), -1) self.addCleanup(right.close) + async def test_tls_proxy_helper_bridges_the_tunnel(self): + # Covers the TLS-proxy path and the socketpair relay without KMS creds. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + request = b"" + while b"\r\n\r\n" not in request: + chunk = tls.recv(4096) + if not chunk: + return + request += chunk + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + # The tunneled peer speaks only after the client does, as a TLS server would. + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + client_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + + sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + else: + await asyncio.get_running_loop().run_in_executor(None, sock.sendall, b"ping") + data = await asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"echo:ping") + + async def test_bridge_does_not_inherit_the_connect_deadline(self): + # The relay must outlast the much shorter CONNECT deadline. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + request = b"" + while b"\r\n\r\n" not in request: + chunk = tls.recv(4096) + if not chunk: + return + request += chunk + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + # Outlast the connect-phase deadline before the KMS side speaks. + time.sleep(2.0) + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + client_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=1.0) + + sock = await AsyncHTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + else: + await asyncio.get_running_loop().run_in_executor(None, sock.sendall, b"ping") + data = await asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"echo:ping") + async def test_non_coroutine_callback_is_rejected(self): # A plain def must be rejected before it blocks the event loop. if _IS_SYNC: @@ -186,6 +451,57 @@ def callback(context): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) self.assertEqual(entered, [], "invalid callback must not be entered") + async def test_proxy_closing_before_connect_reply_raises(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + # Read the CONNECT request, then hang up without replying. + conn.recv(4096) + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "proxy closed the connection"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + + async def test_tunnel_keeps_bytes_sent_with_the_connect_reply(self): + # A proxy may coalesce its 200 with tunneled bytes; reading past the header would drop them. + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + conn.recv(4096) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\nearly-bytes") + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = await AsyncHTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + data = sock.recv(64) + else: + data = await asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"early-bytes") + 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() @@ -197,6 +513,64 @@ async def callback(context): with self.assertRaisesRegex(ConfigurationError, "already connected"): await _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + async def test_ipv6_host_is_bracketed_in_connect(self): + accepted = [] + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + accepted.append(conn.recv(4096)) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="::1", port=443, timeout=10) + sock = await AsyncHTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + self.assertEqual(accepted[0].split(b"\r\n")[0], b"CONNECT [::1]:443 HTTP/1.1") + + async def test_oversized_connect_response_is_rejected(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + conn.recv(4096) + # Never sends the terminator. + while True: + conn.sendall(b"x" * 1024) + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "oversized CONNECT response"): + await AsyncHTTPProxyKMSConnect(host, port)(context) + + async def test_remaining_raises_once_the_deadline_passes(self): + from pymongo.encryption_options import _remaining + + self.assertGreater(_remaining(time.monotonic() + 5), 0) + with self.assertRaises(socket.timeout): + _remaining(time.monotonic() - 1) + 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) @@ -248,6 +622,50 @@ async def test_contract_violation_surfaces_as_encryption_error(self): raise ConfigurationError("kms_connect_callback must return ...") self.assertIsInstance(caught.exception.__cause__, ConfigurationError) + async def test_bridge_failure_closes_the_proxy_socket(self): + # A failure inside _bridge must not strand the connected proxy socket. + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + tls.recv(4096) + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + captured = [] + + def failing_bridge(self, proxy): + captured.append(proxy) + raise OSError("no file descriptors") + + host, port = listener.getsockname() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + + with mock.patch.object(HTTPProxyKMSConnect, "_bridge", failing_bridge): + with self.assertRaisesRegex(OSError, "no file descriptors"): + await AsyncHTTPProxyKMSConnect(host, port, ctx)(context) + + self.assertEqual(captured[0].fileno(), -1, "proxy socket was left open") + async def test_network_error_from_callback_propagates(self): async def callback(context): raise OSError("proxy unreachable") @@ -391,3 +809,183 @@ async def test_client_encryption_rejects_non_callable(self): OPTS, kms_connect_callback="not-callable", # type: ignore[arg-type] ) + + +class TestKmsConnectCallbackProse(AsyncEncryptionIntegrationTest): + @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") + async def asyncSetUp(self): + await super().asyncSetUp() + self.callback_calls: list[Any] = [] + + async def plain_callback(self, context): + self.callback_calls.append(context) + return await AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) + + async def tls_callback(self, context): + self.callback_calls.append(context) + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + # PYTHON-5040 tracks re-enabling verification: the evergreen-tools CA + # lacks an Authority Key Identifier newer OpenSSL requires. + ctx.verify_mode = ssl.CERT_NONE + callback = AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) + return await callback(context) + + async def proxy_request(self, method, path, tls=False): + """Call the proxy's control endpoints and return the body.""" + if _IS_SYNC: + return self._proxy_request(method, path, tls) + return await asyncio.get_running_loop().run_in_executor( + None, self._proxy_request, method, path, tls + ) + + def _proxy_request(self, method, path, tls=False): + if tls: + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + # PYTHON-5040 tracks re-enabling verification once the test CA cert + # is fixed; the evergreen-tools CA lacks an Authority Key Identifier + # that newer OpenSSL requires, so verification fails on Windows 3.14. + ctx.verify_mode = ssl.CERT_NONE + conn = http.client.HTTPSConnection( + f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=ctx + ) + else: + conn = http.client.HTTPConnection(f"{KMS_PROXY_HOST}:{KMS_PROXY_PORT}") + try: + conn.request(method, path) + return conn.getresponse().read().decode() + finally: + conn.close() + + async def connect_count(self, tls=False): + body = await self.proxy_request("GET", "/metrics", tls=tls) + # One "key value" per line; the server also emits connect_target. + for line in body.splitlines(): + key, _, value = line.partition(" ") + if key == "connect_count": + return int(value) + raise AssertionError(f"no connect_count in metrics body: {body!r}") + + async def test_01_plain_http_proxy(self): + await self.proxy_request("POST", "/reset") + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(await self.connect_count(), 1) + + async def test_02_https_proxy(self): + await self.proxy_request("POST", "/reset", tls=True) + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.tls_callback, + ) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(await self.connect_count(tls=True), 1) + + async def test_03_auto_encryption_through_proxy(self): + await self.client.keyvault.datakeys.drop() + await self.client.db.coll.drop() + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + data_key_id = await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + schema = { + "bsonType": "object", + "properties": { + "encrypted_string": { + "encrypt": { + "keyId": [data_key_id], + "bsonType": "string", + "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic", + } + } + }, + } + + await self.proxy_request("POST", "/reset") + opts = AutoEncryptionOpts( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + schema_map={"db.coll": schema}, + kms_connect_callback=self.plain_callback, + ) + client_encrypted = await self.async_rs_or_single_client(auto_encryption_opts=opts) + + await client_encrypted.db.coll.insert_one({"_id": 1, "encrypted_string": "hello"}) + decrypted = await client_encrypted.db.coll.find_one({"_id": 1}) + self.assertEqual(decrypted["encrypted_string"], "hello") + + raw = await self.client.db.coll.find_one({"_id": 1}) + self.assertIsInstance(raw["encrypted_string"], Binary) + + # The decrypt reuses the cached key, so exactly one KMS request follows + # the reset. + self.assertEqual(await self.connect_count(), 1) + + async def test_04_callback_error(self): + async def failing_callback(context): + raise OSError("proxy is on fire") + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=failing_callback, + ) + with self.assertRaisesRegex(EncryptionError, "proxy is on fire"): + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + + @unittest.skip( + "PYTHON-6037 ClientEncryption does not support timeoutMS, so the " + "callback always receives the default KMS connect timeout" + ) + async def test_05_callback_receives_timeout(self): + key_vault_client = await self.async_rs_or_single_client(timeoutMS=1000) + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + key_vault_client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + + self.assertTrue(self.callback_calls, "callback was never invoked") + for context in self.callback_calls: + # Checks only the spec's non-zero requirement, which cannot fail. + self.assertIsNotNone(context.timeout) + self.assertGreater(context.timeout, 0) + + async def test_06_retry_after_network_error(self): + state = {"calls": 0} + + async def flaky_callback(context): + state["calls"] += 1 + if state["calls"] == 1: + raise OSError("first attempt fails") + return await AsyncHTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=flaky_callback, + ) + await encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(state["calls"], 2) diff --git a/test/test_kms_connect.py b/test/test_kms_connect.py index 25c96191ab..0fc6a6f447 100644 --- a/test/test_kms_connect.py +++ b/test/test_kms_connect.py @@ -1,22 +1,26 @@ -"""Tests for the KMS connect callback.""" +"""Tests for the KMS connect callback and HTTP proxy support.""" from __future__ import annotations import asyncio import dataclasses +import http.client import socket import ssl import threading import time import unittest from asyncio.trsock import TransportSocket +from typing import Any from unittest import mock import pytest import pymongo +from bson.binary import Binary from pymongo.encryption_options import ( AutoEncryptionOpts, + HTTPProxyKMSConnect, KMSConnectContext, ) from pymongo.errors import ConfigurationError, EncryptionError, NetworkTimeout @@ -29,13 +33,22 @@ _wrap_encryption_errors, ) from test import PyMongoTestCase -from test.helpers_shared import CLIENT_PEM -from test.test_encryption import OPTS +from test.helpers_shared import AWS_CREDS, CA_PEM, CLIENT_PEM +from test.test_encryption import OPTS, EncryptionIntegrationTest _IS_SYNC = True pytestmark = pytest.mark.encryption +KMS_PROXY_HOST = "127.0.0.1" +KMS_PROXY_PORT = 9004 +KMS_TLS_PROXY_PORT = 9005 + +AWS_MASTER_KEY = { + "region": "us-east-1", + "key": "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0", +} + class TestKmsConnectCallbackUnit(PyMongoTestCase): """Contract checks for kms_connect_callback that need no KMS server.""" @@ -156,6 +169,158 @@ def callback(context): with self.assertRaisesRegex(ConfigurationError, "TransportSocket"): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + def test_http_proxy_helper_tunnels_and_reports_refusal(self): + # Covers the CONNECT handshake without KMS credentials. + accepted = [] + + def stub(listener, reply): + try: + conn, _ = listener.accept() + except OSError: + return + accepted.append(conn.recv(4096)) + conn.sendall(reply) + conn.close() + + def run_stub(reply): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + threading.Thread(target=stub, args=(listener, reply), daemon=True).start() + return listener.getsockname() + + host, port = run_stub(b"HTTP/1.1 200 Connection Established\r\n\r\n") + callback = HTTPProxyKMSConnect(host, port) + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = callback(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + self.assertEqual(accepted[0].split(b"\r\n")[0], b"CONNECT kms.example.com:443 HTTP/1.1") + + host, port = run_stub(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n") + with self.assertRaisesRegex(OSError, "refused CONNECT"): + HTTPProxyKMSConnect(host, port)(context) + + # Any 2xx status is a successful tunnel, not just HTTP/1.1 200. + host, port = run_stub(b"HTTP/1.0 200 Connection Established\r\n\r\n") + sock = HTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + + # A status code must be exactly three digits, with no zero padding. + for reply in (b"HTTP/1.1 2000 Evil\r\n\r\n", b"HTTP/1.1 00200 Evil\r\n\r\n"): + host, port = run_stub(reply) + with self.assertRaisesRegex(OSError, "refused CONNECT"): + HTTPProxyKMSConnect(host, port)(context) + + def test_control_characters_in_kms_host_are_rejected(self): + # Reject CR/LF in the configurable host before it reaches CONNECT. + callback = HTTPProxyKMSConnect("proxy.example.com", 8080) + context = KMSConnectContext(host="kms.example.com\r\nX-Injected: 1", port=443, timeout=10) + with self.assertRaisesRegex(ConfigurationError, "control characters"): + callback(context) + + def test_http_proxy_helper_sends_custom_headers(self): + # Extra CONNECT headers reach the proxy verbatim. + accepted = [] + + def stub(listener): + try: + conn, _ = listener.accept() + except OSError: + return + request = b"" + while b"\r\n\r\n" not in request: + chunk = conn.recv(4096) + if not chunk: + return + request += chunk + accepted.append(request) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + conn.close() + + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + threading.Thread(target=stub, args=(listener,), daemon=True).start() + + host, port = listener.getsockname() + headers = {"Proxy-Authorization": "Basic dXNlcjpwYXNz", "X-Trace-Id": "abc123"} + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = HTTPProxyKMSConnect(host, port, headers=headers)(context) + self.addCleanup(sock.close) + request = accepted[0] + self.assertEqual(request.split(b"\r\n")[0], b"CONNECT kms.example.com:443 HTTP/1.1") + self.assertIn(b"\r\nProxy-Authorization: Basic dXNlcjpwYXNz\r\n", request) + self.assertIn(b"\r\nX-Trace-Id: abc123\r\n", request) + self.assertEqual(request.count(b"\r\nHost: "), 1) + + def test_http_proxy_helper_authenticates_to_the_proxy(self): + # The motivating case: 407 without credentials, 200 with them. + def stub(listener): + for _ in range(2): + try: + conn, _ = listener.accept() + except OSError: + return + request = b"" + while b"\r\n\r\n" not in request: + chunk = conn.recv(4096) + if not chunk: + return + request += chunk + if b"\r\nProxy-Authorization: Basic dXNlcjpwYXNz\r\n" in request: + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + else: + conn.sendall(b"HTTP/1.1 407 Proxy Authentication Required\r\n\r\n") + conn.close() + + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(2) + self.addCleanup(listener.close) + threading.Thread(target=stub, args=(listener,), daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "refused CONNECT"): + HTTPProxyKMSConnect(host, port)(context) + headers = {"Proxy-Authorization": "Basic dXNlcjpwYXNz"} + sock = HTTPProxyKMSConnect(host, port, headers=headers)(context) + self.addCleanup(sock.close) + self.assertIsInstance(sock, socket.socket) + + def test_http_proxy_helper_rejects_bad_headers(self): + for headers in [ + {"Bad\r\nName": "x"}, + {"X-Ok": "ok\r\nInjected: 1"}, + {"Host": "evil.example.com"}, + {"host": "evil.example.com"}, + {"": "x"}, + {"Bad:Name": "x"}, + ]: + with self.assertRaisesRegex(ConfigurationError, "proxy header|Host CONNECT header"): + HTTPProxyKMSConnect("proxy.example.com", 8080, headers=headers) + + for headers in [{1: "x"}, {"X-Ok": 1}, {None: "x"}, {"X-Ok": None}]: + with self.assertRaisesRegex(TypeError, "must be strings"): + HTTPProxyKMSConnect("proxy.example.com", 8080, headers=headers) + + def test_http_proxy_helper_accepts_legal_header_values(self): + # Colons and spaces are legal in values (e.g. auth schemes); only + # CR/LF would let a value inject a request line. + callback = HTTPProxyKMSConnect( + "proxy.example.com", + 8080, + headers={"Proxy-Authorization": "Basic dXNlcjpwYXNz", "X-Token": "a: b"}, + ) + self.assertEqual( + callback.headers, + {"Proxy-Authorization": "Basic dXNlcjpwYXNz", "X-Token": "a: b"}, + ) + 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. @@ -171,6 +336,105 @@ def test_cancelled_tls_wrap_closes_late_socket(self): self.assertEqual(left.fileno(), -1) self.addCleanup(right.close) + def test_tls_proxy_helper_bridges_the_tunnel(self): + # Covers the TLS-proxy path and the socketpair relay without KMS creds. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + request = b"" + while b"\r\n\r\n" not in request: + chunk = tls.recv(4096) + if not chunk: + return + request += chunk + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + # The tunneled peer speaks only after the client does, as a TLS server would. + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + client_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + + sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + else: + asyncio.get_running_loop().run_in_executor(None, sock.sendall, b"ping") + data = asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"echo:ping") + + def test_bridge_does_not_inherit_the_connect_deadline(self): + # The relay must outlast the much shorter CONNECT deadline. + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + request = b"" + while b"\r\n\r\n" not in request: + chunk = tls.recv(4096) + if not chunk: + return + request += chunk + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + # Outlast the connect-phase deadline before the KMS side speaks. + time.sleep(2.0) + tls.sendall(b"echo:" + tls.recv(64)) + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + client_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + client_ctx.check_hostname = False + client_ctx.verify_mode = ssl.CERT_NONE + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=1.0) + + sock = HTTPProxyKMSConnect(host, port, client_ctx)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + sock.sendall(b"ping") + self.assertEqual(sock.recv(64), b"echo:ping") + else: + asyncio.get_running_loop().run_in_executor(None, sock.sendall, b"ping") + data = asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"echo:ping") + def test_non_coroutine_callback_is_rejected(self): # A plain def must be rejected before it blocks the event loop. if _IS_SYNC: @@ -186,6 +450,57 @@ def callback(context): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) self.assertEqual(entered, [], "invalid callback must not be entered") + def test_proxy_closing_before_connect_reply_raises(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + # Read the CONNECT request, then hang up without replying. + conn.recv(4096) + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "proxy closed the connection"): + HTTPProxyKMSConnect(host, port)(context) + + def test_tunnel_keeps_bytes_sent_with_the_connect_reply(self): + # A proxy may coalesce its 200 with tunneled bytes; reading past the header would drop them. + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + conn.recv(4096) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\nearly-bytes") + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + sock = HTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + sock.settimeout(10) + if _IS_SYNC: + data = sock.recv(64) + else: + data = asyncio.get_running_loop().run_in_executor(None, sock.recv, 64) + self.assertEqual(data, b"early-bytes") + 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() @@ -197,6 +512,64 @@ def callback(context): with self.assertRaisesRegex(ConfigurationError, "already connected"): _connect_kms(("kms.example.com", 443), self._pool_options(), callback, 10.0) + def test_ipv6_host_is_bracketed_in_connect(self): + accepted = [] + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + try: + conn, _ = listener.accept() + accepted.append(conn.recv(4096)) + conn.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + conn.close() + except OSError: + pass + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="::1", port=443, timeout=10) + sock = HTTPProxyKMSConnect(host, port)(context) + self.addCleanup(sock.close) + self.assertEqual(accepted[0].split(b"\r\n")[0], b"CONNECT [::1]:443 HTTP/1.1") + + def test_oversized_connect_response_is_rejected(self): + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + conn.recv(4096) + # Never sends the terminator. + while True: + conn.sendall(b"x" * 1024) + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + host, port = listener.getsockname() + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + with self.assertRaisesRegex(OSError, "oversized CONNECT response"): + HTTPProxyKMSConnect(host, port)(context) + + def test_remaining_raises_once_the_deadline_passes(self): + from pymongo.encryption_options import _remaining + + self.assertGreater(_remaining(time.monotonic() + 5), 0) + with self.assertRaises(socket.timeout): + _remaining(time.monotonic() - 1) + 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) @@ -248,6 +621,50 @@ def test_contract_violation_surfaces_as_encryption_error(self): raise ConfigurationError("kms_connect_callback must return ...") self.assertIsInstance(caught.exception.__cause__, ConfigurationError) + def test_bridge_failure_closes_the_proxy_socket(self): + # A failure inside _bridge must not strand the connected proxy socket. + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + self.addCleanup(listener.close) + + server_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_ctx.load_cert_chain(CLIENT_PEM) + + def stub_proxy(): + conn = None + try: + conn, _ = listener.accept() + tls = server_ctx.wrap_socket(conn, server_side=True) + tls.recv(4096) + tls.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n") + tls.close() + except OSError: + pass + finally: + if conn is not None: + conn.close() + + threading.Thread(target=stub_proxy, daemon=True).start() + + captured = [] + + def failing_bridge(self, proxy): + captured.append(proxy) + raise OSError("no file descriptors") + + host, port = listener.getsockname() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + context = KMSConnectContext(host="kms.example.com", port=443, timeout=10) + + with mock.patch.object(HTTPProxyKMSConnect, "_bridge", failing_bridge): + with self.assertRaisesRegex(OSError, "no file descriptors"): + HTTPProxyKMSConnect(host, port, ctx)(context) + + self.assertEqual(captured[0].fileno(), -1, "proxy socket was left open") + def test_network_error_from_callback_propagates(self): def callback(context): raise OSError("proxy unreachable") @@ -389,3 +806,183 @@ def test_client_encryption_rejects_non_callable(self): OPTS, kms_connect_callback="not-callable", # type: ignore[arg-type] ) + + +class TestKmsConnectCallbackProse(EncryptionIntegrationTest): + @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") + def setUp(self): + super().setUp() + self.callback_calls: list[Any] = [] + + def plain_callback(self, context): + self.callback_calls.append(context) + return HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) + + def tls_callback(self, context): + self.callback_calls.append(context) + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + # PYTHON-5040 tracks re-enabling verification: the evergreen-tools CA + # lacks an Authority Key Identifier newer OpenSSL requires. + ctx.verify_mode = ssl.CERT_NONE + callback = HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_TLS_PROXY_PORT, ctx) + return callback(context) + + def proxy_request(self, method, path, tls=False): + """Call the proxy's control endpoints and return the body.""" + if _IS_SYNC: + return self._proxy_request(method, path, tls) + return asyncio.get_running_loop().run_in_executor( + None, self._proxy_request, method, path, tls + ) + + def _proxy_request(self, method, path, tls=False): + if tls: + ctx = ssl.create_default_context(cafile=CA_PEM) + ctx.check_hostname = False + # PYTHON-5040 tracks re-enabling verification once the test CA cert + # is fixed; the evergreen-tools CA lacks an Authority Key Identifier + # that newer OpenSSL requires, so verification fails on Windows 3.14. + ctx.verify_mode = ssl.CERT_NONE + conn = http.client.HTTPSConnection( + f"{KMS_PROXY_HOST}:{KMS_TLS_PROXY_PORT}", context=ctx + ) + else: + conn = http.client.HTTPConnection(f"{KMS_PROXY_HOST}:{KMS_PROXY_PORT}") + try: + conn.request(method, path) + return conn.getresponse().read().decode() + finally: + conn.close() + + def connect_count(self, tls=False): + body = self.proxy_request("GET", "/metrics", tls=tls) + # One "key value" per line; the server also emits connect_target. + for line in body.splitlines(): + key, _, value = line.partition(" ") + if key == "connect_count": + return int(value) + raise AssertionError(f"no connect_count in metrics body: {body!r}") + + def test_01_plain_http_proxy(self): + self.proxy_request("POST", "/reset") + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(self.connect_count(), 1) + + def test_02_https_proxy(self): + self.proxy_request("POST", "/reset", tls=True) + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.tls_callback, + ) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(self.connect_count(tls=True), 1) + + def test_03_auto_encryption_through_proxy(self): + self.client.keyvault.datakeys.drop() + self.client.db.coll.drop() + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + data_key_id = encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + schema = { + "bsonType": "object", + "properties": { + "encrypted_string": { + "encrypt": { + "keyId": [data_key_id], + "bsonType": "string", + "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic", + } + } + }, + } + + self.proxy_request("POST", "/reset") + opts = AutoEncryptionOpts( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + schema_map={"db.coll": schema}, + kms_connect_callback=self.plain_callback, + ) + client_encrypted = self.rs_or_single_client(auto_encryption_opts=opts) + + client_encrypted.db.coll.insert_one({"_id": 1, "encrypted_string": "hello"}) + decrypted = client_encrypted.db.coll.find_one({"_id": 1}) + self.assertEqual(decrypted["encrypted_string"], "hello") + + raw = self.client.db.coll.find_one({"_id": 1}) + self.assertIsInstance(raw["encrypted_string"], Binary) + + # The decrypt reuses the cached key, so exactly one KMS request follows + # the reset. + self.assertEqual(self.connect_count(), 1) + + def test_04_callback_error(self): + def failing_callback(context): + raise OSError("proxy is on fire") + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=failing_callback, + ) + with self.assertRaisesRegex(EncryptionError, "proxy is on fire"): + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + + @unittest.skip( + "PYTHON-6037 ClientEncryption does not support timeoutMS, so the " + "callback always receives the default KMS connect timeout" + ) + def test_05_callback_receives_timeout(self): + key_vault_client = self.rs_or_single_client(timeoutMS=1000) + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + key_vault_client, + OPTS, + kms_connect_callback=self.plain_callback, + ) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + + self.assertTrue(self.callback_calls, "callback was never invoked") + for context in self.callback_calls: + # Checks only the spec's non-zero requirement, which cannot fail. + self.assertIsNotNone(context.timeout) + self.assertGreater(context.timeout, 0) + + def test_06_retry_after_network_error(self): + state = {"calls": 0} + + def flaky_callback(context): + state["calls"] += 1 + if state["calls"] == 1: + raise OSError("first attempt fails") + return HTTPProxyKMSConnect(KMS_PROXY_HOST, KMS_PROXY_PORT)(context) + + encryption = self.create_client_encryption( + {"aws": AWS_CREDS}, + "keyvault.datakeys", + self.client, + OPTS, + kms_connect_callback=flaky_callback, + ) + encryption.create_data_key("aws", master_key=AWS_MASTER_KEY) + self.assertGreaterEqual(state["calls"], 2) diff --git a/tools/synchro.py b/tools/synchro.py index bd5edebda3..fc99bb2747 100644 --- a/tools/synchro.py +++ b/tools/synchro.py @@ -73,6 +73,7 @@ "AsyncClientEncryption": "ClientEncryption", "AsyncMongoCryptCallback": "MongoCryptCallback", "AsyncKMSConnectCallback": "KMSConnectCallback", + "AsyncHTTPProxyKMSConnect": "HTTPProxyKMSConnect", "AsyncExplicitEncrypter": "ExplicitEncrypter", "AsyncAutoEncrypter": "AutoEncrypter", "AsyncContextManager": "ContextManager",