From 5c795028c9f98258ef48ef11f9dcd6e4216456ef Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Mon, 25 Nov 2024 17:29:57 -0800 Subject: [PATCH 1/9] PYTHON-2560 Retry KMS requests on transient errors --- pymongo/asynchronous/encryption.py | 20 ++++++++++++-------- pymongo/synchronous/encryption.py | 20 ++++++++++++-------- 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 4802c3f54e..a4258159cc 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -72,7 +72,6 @@ EncryptedCollectionError, EncryptionError, InvalidOperation, - PyMongoError, ServerSelectionTimeoutError, ) from pymongo.network_layer import BLOCKING_IO_ERRORS, async_sendall @@ -176,6 +175,9 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: ssl_context=ctx, ) host, port = parse_host(endpoint, _HTTPS_PORT) + sleep_usec = kms_context.usleep + if sleep_usec: + await asyncio.sleep(float(sleep_usec) / 1e6) try: conn = await _configured_socket((host, port), opts) try: @@ -201,13 +203,15 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: raise socket.timeout("timed out") from None finally: conn.close() - except (PyMongoError, MongoCryptError): - raise # Propagate pymongo errors directly. - except asyncio.CancelledError: - raise - except Exception as error: - # Wrap I/O errors in PyMongo exceptions. - _raise_connection_failure((host, port), error) + except MongoCryptError: + raise # Propagate MongoCryptError errors directly. + except Exception as exc: + remaining = _csot.remaining() + if remaining is not None and remaining <= 0: + # Wrap I/O errors in PyMongo exceptions. + _raise_connection_failure((host, port), exc) + # Mark this attempt as failed and defer to libmongocrypt to retry. + kms_context.fail() async def collection_info(self, database: str, filter: bytes) -> Optional[bytes]: """Get the collection info for a namespace. diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 09d0c0f2fd..f6bdbb913d 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -67,7 +67,6 @@ EncryptedCollectionError, EncryptionError, InvalidOperation, - PyMongoError, ServerSelectionTimeoutError, ) from pymongo.network_layer import BLOCKING_IO_ERRORS, sendall @@ -176,6 +175,9 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: ssl_context=ctx, ) host, port = parse_host(endpoint, _HTTPS_PORT) + sleep_usec = kms_context.usleep + if sleep_usec: + asyncio.sleep(float(sleep_usec) / 1e6) try: conn = _configured_socket((host, port), opts) try: @@ -201,13 +203,15 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: raise socket.timeout("timed out") from None finally: conn.close() - except (PyMongoError, MongoCryptError): - raise # Propagate pymongo errors directly. - except asyncio.CancelledError: - raise - except Exception as error: - # Wrap I/O errors in PyMongo exceptions. - _raise_connection_failure((host, port), error) + except MongoCryptError: + raise # Propagate MongoCryptError errors directly. + except Exception as exc: + remaining = _csot.remaining() + if remaining is not None and remaining <= 0: + # Wrap I/O errors in PyMongo exceptions. + _raise_connection_failure((host, port), exc) + # Mark this attempt as failed and defer to libmongocrypt to retry. + kms_context.fail() def collection_info(self, database: str, filter: bytes) -> Optional[bytes]: """Get the collection info for a namespace. From e413334ac6692d03fdf3e4d06742ccee193bb832 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Mon, 2 Dec 2024 16:45:04 -0800 Subject: [PATCH 2/9] PYTHON-2560 Add prose test --- .evergreen/run-tests.sh | 2 +- .evergreen/scripts/prepare-resources.sh | 2 +- pymongo/asynchronous/encryption.py | 9 ++- pymongo/synchronous/encryption.py | 9 ++- test/asynchronous/test_encryption.py | 83 +++++++++++++++++++++++++ test/test_encryption.py | 83 +++++++++++++++++++++++++ 6 files changed, 180 insertions(+), 8 deletions(-) diff --git a/.evergreen/run-tests.sh b/.evergreen/run-tests.sh index 95fe10a6c3..307abe741f 100755 --- a/.evergreen/run-tests.sh +++ b/.evergreen/run-tests.sh @@ -158,7 +158,7 @@ if [ -n "$TEST_ENCRYPTION" ] || [ -n "$TEST_FLE_AZURE_AUTO" ] || [ -n "$TEST_FLE # TODO: Test with 'pip install pymongocrypt' if [ ! -d "libmongocrypt_git" ]; then - git clone https://github.com/mongodb/libmongocrypt.git libmongocrypt_git + git clone https://github.com/ShaneHarvey/libmongocrypt.git --branch PYTHON-4992 libmongocrypt_git fi python -m pip install -U setuptools python -m pip install ./libmongocrypt_git/bindings/python diff --git a/.evergreen/scripts/prepare-resources.sh b/.evergreen/scripts/prepare-resources.sh index 33394b55ff..e6d211bac2 100755 --- a/.evergreen/scripts/prepare-resources.sh +++ b/.evergreen/scripts/prepare-resources.sh @@ -7,6 +7,6 @@ if [ "$PROJECT" = "drivers-tools" ]; then # If this was a patch build, doing a fresh clone would not actually test the patch cp -R $PROJECT_DIRECTORY/ $DRIVERS_TOOLS else - git clone https://github.com/mongodb-labs/drivers-evergreen-tools.git $DRIVERS_TOOLS + git clone https://github.com/ShaneHarvey/drivers-evergreen-tools.git --branch DRIVERS-1541 $DRIVERS_TOOLS fi echo "{ \"releases\": { \"default\": \"$MONGODB_BINARIES\" }}" >$MONGO_ORCHESTRATION_HOME/orchestration.config diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index a4258159cc..eff6194698 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -72,6 +72,7 @@ EncryptedCollectionError, EncryptionError, InvalidOperation, + NetworkTimeout, ServerSelectionTimeoutError, ) from pymongo.network_layer import BLOCKING_IO_ERRORS, async_sendall @@ -165,8 +166,8 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: None, # crlfile False, # allow_invalid_certificates False, # allow_invalid_hostnames - False, - ) # disable_ocsp_endpoint_check + False, # disable_ocsp_endpoint_check + ) # CSOT: set timeout for socket creation. connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001) opts = PoolOptions( @@ -207,7 +208,9 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: raise # Propagate MongoCryptError errors directly. except Exception as exc: remaining = _csot.remaining() - if remaining is not None and remaining <= 0: + if isinstance(exc, (socket.timeout, NetworkTimeout)) or ( + remaining is not None and remaining <= 0 + ): # Wrap I/O errors in PyMongo exceptions. _raise_connection_failure((host, port), exc) # Mark this attempt as failed and defer to libmongocrypt to retry. diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index f6bdbb913d..627b9deb29 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -67,6 +67,7 @@ EncryptedCollectionError, EncryptionError, InvalidOperation, + NetworkTimeout, ServerSelectionTimeoutError, ) from pymongo.network_layer import BLOCKING_IO_ERRORS, sendall @@ -165,8 +166,8 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: None, # crlfile False, # allow_invalid_certificates False, # allow_invalid_hostnames - False, - ) # disable_ocsp_endpoint_check + False, # disable_ocsp_endpoint_check + ) # CSOT: set timeout for socket creation. connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001) opts = PoolOptions( @@ -207,7 +208,9 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: raise # Propagate MongoCryptError errors directly. except Exception as exc: remaining = _csot.remaining() - if remaining is not None and remaining <= 0: + if isinstance(exc, (socket.timeout, NetworkTimeout)) or ( + remaining is not None and remaining <= 0 + ): # Wrap I/O errors in PyMongo exceptions. _raise_connection_failure((host, port), exc) # Mark this attempt as failed and defer to libmongocrypt to retry. diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 21cd5e2666..1039b13268 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -17,6 +17,8 @@ import base64 import copy +import http.client +import json import os import pathlib import re @@ -91,6 +93,7 @@ WriteError, ) from pymongo.operations import InsertOne, ReplaceOne, UpdateOne +from pymongo.ssl_support import get_ssl_context from pymongo.write_concern import WriteConcern _IS_SYNC = False @@ -2853,6 +2856,86 @@ async def test_accepts_trim_factor_0(self): assert len(payload) > len(self.payload_defaults) +# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#24-kms-retry-tests +class TestKmsRetryProse(AsyncEncryptionIntegrationTest): + @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") + async def asyncSetUp(self): + await super().asyncSetUp() + # 1, create client with only tlsCAFile. + providers: dict = copy.deepcopy(ALL_KMS_PROVIDERS) + providers["azure"]["identityPlatformEndpoint"] = "127.0.0.1:9003" + providers["gcp"]["endpoint"] = "127.0.0.1:9003" + kms_tls_opts = { + p: {"tlsCAFile": CA_PEM, "tlsCertificateKeyFile": CLIENT_PEM} for p in providers + } + self.client_encryption = self.create_client_encryption( + providers, "keyvault.datakeys", self.client, OPTS, kms_tls_options=kms_tls_opts + ) + + async def http_post(self, path, data=None): + # Note, the connection to the mock server needs to be closed after + # each request because the server is single threaded. + ctx = get_ssl_context( + CLIENT_PEM, # certfile + None, # passphrase + CA_PEM, # ca_certs + None, # crlfile + False, # allow_invalid_certificates + False, # allow_invalid_hostnames + False, # disable_ocsp_endpoint_check + ) + conn = http.client.HTTPSConnection("127.0.0.1:9003", context=ctx) + try: + if data is not None: + headers = {"Content-type": "application/json"} + body = json.dumps(data) + else: + headers = {} + body = None + conn.request("POST", path, body, headers) + res = conn.getresponse() + res.read() + finally: + conn.close() + + async def _test(self, provider, master_key): + await self.http_post("/reset") + # Case 1: createDataKey and encrypt with TCP retry + await self.http_post("/set_failpoint/network", {"count": 1}) + key_id = await self.client_encryption.create_data_key(provider, master_key=master_key) + await self.http_post("/set_failpoint/network", {"count": 1}) + await self.client_encryption.encrypt( + 123, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id + ) + + # Case 2: createDataKey and encrypt with HTTP retry + await self.http_post("/set_failpoint/http", {"count": 1}) + key_id = await self.client_encryption.create_data_key(provider, master_key=master_key) + await self.http_post("/set_failpoint/http", {"count": 1}) + await self.client_encryption.encrypt( + 123, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id + ) + + # Case 3: createDataKey fails after too many retries + await self.http_post("/set_failpoint/network", {"count": 4}) + with self.assertRaisesRegex(EncryptionError, "KMS request failed after"): + await self.client_encryption.create_data_key(provider, master_key=master_key) + + async def test_kms_retry(self): + await self._test("aws", {"region": "foo", "key": "bar", "endpoint": "127.0.0.1:9003"}) + await self._test("azure", {"keyVaultEndpoint": "127.0.0.1:9003", "keyName": "foo"}) + await self._test( + "gcp", + { + "projectId": "foo", + "location": "bar", + "keyRing": "baz", + "keyName": "qux", + "endpoint": "127.0.0.1:9003", + }, + ) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#automatic-data-encryption-keys class TestAutomaticDecryptionKeys(AsyncEncryptionIntegrationTest): @async_client_context.require_no_standalone diff --git a/test/test_encryption.py b/test/test_encryption.py index 18e21fe6a7..3f7a8c8d53 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -17,6 +17,8 @@ import base64 import copy +import http.client +import json import os import pathlib import re @@ -88,6 +90,7 @@ WriteError, ) from pymongo.operations import InsertOne, ReplaceOne, UpdateOne +from pymongo.ssl_support import get_ssl_context from pymongo.synchronous import encryption from pymongo.synchronous.encryption import Algorithm, ClientEncryption, QueryType from pymongo.synchronous.mongo_client import MongoClient @@ -2835,6 +2838,86 @@ def test_accepts_trim_factor_0(self): assert len(payload) > len(self.payload_defaults) +# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#24-kms-retry-tests +class TestKmsRetryProse(EncryptionIntegrationTest): + @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") + def setUp(self): + super().setUp() + # 1, create client with only tlsCAFile. + providers: dict = copy.deepcopy(ALL_KMS_PROVIDERS) + providers["azure"]["identityPlatformEndpoint"] = "127.0.0.1:9003" + providers["gcp"]["endpoint"] = "127.0.0.1:9003" + kms_tls_opts = { + p: {"tlsCAFile": CA_PEM, "tlsCertificateKeyFile": CLIENT_PEM} for p in providers + } + self.client_encryption = self.create_client_encryption( + providers, "keyvault.datakeys", self.client, OPTS, kms_tls_options=kms_tls_opts + ) + + def http_post(self, path, data=None): + # Note, the connection to the mock server needs to be closed after + # each request because the server is single threaded. + ctx = get_ssl_context( + CLIENT_PEM, # certfile + None, # passphrase + CA_PEM, # ca_certs + None, # crlfile + False, # allow_invalid_certificates + False, # allow_invalid_hostnames + False, # disable_ocsp_endpoint_check + ) + conn = http.client.HTTPSConnection("127.0.0.1:9003", context=ctx) + try: + if data is not None: + headers = {"Content-type": "application/json"} + body = json.dumps(data) + else: + headers = {} + body = None + conn.request("POST", path, body, headers) + res = conn.getresponse() + res.read() + finally: + conn.close() + + def _test(self, provider, master_key): + self.http_post("/reset") + # Case 1: createDataKey and encrypt with TCP retry + self.http_post("/set_failpoint/network", {"count": 1}) + key_id = self.client_encryption.create_data_key(provider, master_key=master_key) + self.http_post("/set_failpoint/network", {"count": 1}) + self.client_encryption.encrypt( + 123, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id + ) + + # Case 2: createDataKey and encrypt with HTTP retry + self.http_post("/set_failpoint/http", {"count": 1}) + key_id = self.client_encryption.create_data_key(provider, master_key=master_key) + self.http_post("/set_failpoint/http", {"count": 1}) + self.client_encryption.encrypt( + 123, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id + ) + + # Case 3: createDataKey fails after too many retries + self.http_post("/set_failpoint/network", {"count": 4}) + with self.assertRaisesRegex(EncryptionError, "KMS request failed after"): + self.client_encryption.create_data_key(provider, master_key=master_key) + + def test_kms_retry(self): + self._test("aws", {"region": "foo", "key": "bar", "endpoint": "127.0.0.1:9003"}) + self._test("azure", {"keyVaultEndpoint": "127.0.0.1:9003", "keyName": "foo"}) + self._test( + "gcp", + { + "projectId": "foo", + "location": "bar", + "keyRing": "baz", + "keyName": "qux", + "endpoint": "127.0.0.1:9003", + }, + ) + + # https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#automatic-data-encryption-keys class TestAutomaticDecryptionKeys(EncryptionIntegrationTest): @client_context.require_no_standalone From 845568998c758524dbf7e4bf3617ab9ee8d0d2d0 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Tue, 3 Dec 2024 13:53:22 -0800 Subject: [PATCH 3/9] PYTHON-2560 Fix asyncio.sleep, synchro can't handle it --- pymongo/asynchronous/encryption.py | 8 +++++--- pymongo/synchronous/encryption.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index eff6194698..8d35e66189 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -19,6 +19,7 @@ import contextlib import enum import socket +import time as time # noqa: PLC0414 # needed in sync version import uuid import weakref from copy import deepcopy @@ -176,9 +177,10 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: ssl_context=ctx, ) host, port = parse_host(endpoint, _HTTPS_PORT) - sleep_usec = kms_context.usleep - if sleep_usec: - await asyncio.sleep(float(sleep_usec) / 1e6) + sleep_u = kms_context.usleep + if sleep_u: + sleep_sec = float(sleep_u) / 1e6 + await asyncio.sleep(sleep_sec) try: conn = await _configured_socket((host, port), opts) try: diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 627b9deb29..fe7f3f0465 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -19,6 +19,7 @@ import contextlib import enum import socket +import time as time # noqa: PLC0414 # needed in sync version import uuid import weakref from copy import deepcopy @@ -176,9 +177,10 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: ssl_context=ctx, ) host, port = parse_host(endpoint, _HTTPS_PORT) - sleep_usec = kms_context.usleep - if sleep_usec: - asyncio.sleep(float(sleep_usec) / 1e6) + sleep_u = kms_context.usleep + if sleep_u: + sleep_sec = float(sleep_u) / 1e6 + time.sleep(sleep_sec) try: conn = _configured_socket((host, port), opts) try: From 5f01703a049f0df316fcfc0c14de26b26c806678 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Tue, 3 Dec 2024 14:29:28 -0800 Subject: [PATCH 4/9] PYTHON-2560 Add more info to the KMS failed error --- pymongo/asynchronous/encryption.py | 11 +++++++---- pymongo/synchronous/encryption.py | 11 +++++++---- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 8d35e66189..878792132a 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -199,9 +199,6 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: if not data: raise OSError("KMS connection closed") kms_context.feed(data) - # Async raises an OSError instead of returning empty bytes - except OSError as err: - raise OSError("KMS connection closed") from err except BLOCKING_IO_ERRORS: raise socket.timeout("timed out") from None finally: @@ -216,7 +213,13 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: # Wrap I/O errors in PyMongo exceptions. _raise_connection_failure((host, port), exc) # Mark this attempt as failed and defer to libmongocrypt to retry. - kms_context.fail() + try: + kms_context.fail() + except MongoCryptError as final_err: + exc = MongoCryptError( + f"{final_err}, last attempt failed with: {exc}", final_err.code + ) + raise exc from final_err async def collection_info(self, database: str, filter: bytes) -> Optional[bytes]: """Get the collection info for a namespace. diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index fe7f3f0465..92baa7067f 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -199,9 +199,6 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: if not data: raise OSError("KMS connection closed") kms_context.feed(data) - # Async raises an OSError instead of returning empty bytes - except OSError as err: - raise OSError("KMS connection closed") from err except BLOCKING_IO_ERRORS: raise socket.timeout("timed out") from None finally: @@ -216,7 +213,13 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: # Wrap I/O errors in PyMongo exceptions. _raise_connection_failure((host, port), exc) # Mark this attempt as failed and defer to libmongocrypt to retry. - kms_context.fail() + try: + kms_context.fail() + except MongoCryptError as final_err: + exc = MongoCryptError( + f"{final_err}, last attempt failed with: {exc}", final_err.code + ) + raise exc from final_err def collection_info(self, database: str, filter: bytes) -> Optional[bytes]: """Get the collection info for a namespace. From a983726133da1b303b41ca797432ebb1da4d0515 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Tue, 3 Dec 2024 14:38:34 -0800 Subject: [PATCH 5/9] PYTHON-2560 Fix HTTPSConnection typechecking --- test/asynchronous/test_encryption.py | 2 +- test/test_encryption.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 1039b13268..0823158f10 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -2875,7 +2875,7 @@ async def asyncSetUp(self): async def http_post(self, path, data=None): # Note, the connection to the mock server needs to be closed after # each request because the server is single threaded. - ctx = get_ssl_context( + ctx: ssl.SSLContext = get_ssl_context( CLIENT_PEM, # certfile None, # passphrase CA_PEM, # ca_certs diff --git a/test/test_encryption.py b/test/test_encryption.py index 3f7a8c8d53..f0c6b74692 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -2857,7 +2857,7 @@ def setUp(self): def http_post(self, path, data=None): # Note, the connection to the mock server needs to be closed after # each request because the server is single threaded. - ctx = get_ssl_context( + ctx: ssl.SSLContext = get_ssl_context( CLIENT_PEM, # certfile None, # passphrase CA_PEM, # ca_certs From 6a85e84584a7b108f4102210d58788cf2ff2f942 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Tue, 3 Dec 2024 15:02:48 -0800 Subject: [PATCH 6/9] PYTHON-2560 Remove branch changes --- .evergreen/run-tests.sh | 2 +- .evergreen/scripts/prepare-resources.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.evergreen/run-tests.sh b/.evergreen/run-tests.sh index 307abe741f..95fe10a6c3 100755 --- a/.evergreen/run-tests.sh +++ b/.evergreen/run-tests.sh @@ -158,7 +158,7 @@ if [ -n "$TEST_ENCRYPTION" ] || [ -n "$TEST_FLE_AZURE_AUTO" ] || [ -n "$TEST_FLE # TODO: Test with 'pip install pymongocrypt' if [ ! -d "libmongocrypt_git" ]; then - git clone https://github.com/ShaneHarvey/libmongocrypt.git --branch PYTHON-4992 libmongocrypt_git + git clone https://github.com/mongodb/libmongocrypt.git libmongocrypt_git fi python -m pip install -U setuptools python -m pip install ./libmongocrypt_git/bindings/python diff --git a/.evergreen/scripts/prepare-resources.sh b/.evergreen/scripts/prepare-resources.sh index e6d211bac2..33394b55ff 100755 --- a/.evergreen/scripts/prepare-resources.sh +++ b/.evergreen/scripts/prepare-resources.sh @@ -7,6 +7,6 @@ if [ "$PROJECT" = "drivers-tools" ]; then # If this was a patch build, doing a fresh clone would not actually test the patch cp -R $PROJECT_DIRECTORY/ $DRIVERS_TOOLS else - git clone https://github.com/ShaneHarvey/drivers-evergreen-tools.git --branch DRIVERS-1541 $DRIVERS_TOOLS + git clone https://github.com/mongodb-labs/drivers-evergreen-tools.git $DRIVERS_TOOLS fi echo "{ \"releases\": { \"default\": \"$MONGODB_BINARIES\" }}" >$MONGO_ORCHESTRATION_HOME/orchestration.config From 7d539a24f0580b423029e5038f0f47956cb2c0a1 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Tue, 3 Dec 2024 15:09:42 -0800 Subject: [PATCH 7/9] PYTHON-2560 Fix TestCustomEndpoint, wrap error properly --- pymongo/asynchronous/encryption.py | 35 +++++++++++++++++++++--------- pymongo/synchronous/encryption.py | 35 +++++++++++++++++++++--------- 2 files changed, 50 insertions(+), 20 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 878792132a..48a2f191fd 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -64,7 +64,11 @@ from pymongo.asynchronous.cursor import AsyncCursor from pymongo.asynchronous.database import AsyncDatabase from pymongo.asynchronous.mongo_client import AsyncMongoClient -from pymongo.asynchronous.pool import _configured_socket, _raise_connection_failure +from pymongo.asynchronous.pool import ( + _configured_socket, + _get_timeout_details, + _raise_connection_failure, +) from pymongo.common import CONNECT_TIMEOUT from pymongo.daemon import _spawn_daemon from pymongo.encryption_options import AutoEncryptionOpts, RangeOpts @@ -89,6 +93,8 @@ if TYPE_CHECKING: from pymongocrypt.mongocrypt import MongoCryptKmsContext + from pymongo.typings import _Address + _IS_SYNC = False @@ -104,6 +110,13 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) +async def _connect_kms(address: _Address, opts: PoolOptions): + try: + return await _configured_socket(address, opts) + except Exception as exc: + _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) + + @contextlib.contextmanager def _wrap_encryption_errors() -> Iterator[None]: """Context manager to wrap encryption related errors.""" @@ -176,13 +189,13 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: socket_timeout=connect_timeout, ssl_context=ctx, ) - host, port = parse_host(endpoint, _HTTPS_PORT) + address = parse_host(endpoint, _HTTPS_PORT) sleep_u = kms_context.usleep if sleep_u: sleep_sec = float(sleep_u) / 1e6 await asyncio.sleep(sleep_sec) try: - conn = await _configured_socket((host, port), opts) + conn = await _connect_kms(address, opts) try: await async_sendall(conn, message) while kms_context.bytes_needed > 0: @@ -199,19 +212,21 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None: if not data: raise OSError("KMS connection closed") kms_context.feed(data) - except BLOCKING_IO_ERRORS: - raise socket.timeout("timed out") from None + except MongoCryptError: + raise # Propagate MongoCryptError errors directly. + except Exception as exc: + # Wrap I/O errors in PyMongo exceptions. + if isinstance(exc, BLOCKING_IO_ERRORS): + exc = socket.timeout("timed out") + _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) finally: conn.close() except MongoCryptError: raise # Propagate MongoCryptError errors directly. except Exception as exc: remaining = _csot.remaining() - if isinstance(exc, (socket.timeout, NetworkTimeout)) or ( - remaining is not None and remaining <= 0 - ): - # Wrap I/O errors in PyMongo exceptions. - _raise_connection_failure((host, port), exc) + if isinstance(exc, NetworkTimeout) or (remaining is not None and remaining <= 0): + raise # Mark this attempt as failed and defer to libmongocrypt to retry. try: kms_context.fail() diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index 92baa7067f..add5761337 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -81,7 +81,11 @@ from pymongo.synchronous.cursor import Cursor from pymongo.synchronous.database import Database from pymongo.synchronous.mongo_client import MongoClient -from pymongo.synchronous.pool import _configured_socket, _raise_connection_failure +from pymongo.synchronous.pool import ( + _configured_socket, + _get_timeout_details, + _raise_connection_failure, +) from pymongo.typings import _DocumentType, _DocumentTypeArg from pymongo.uri_parser import parse_host from pymongo.write_concern import WriteConcern @@ -89,6 +93,8 @@ if TYPE_CHECKING: from pymongocrypt.mongocrypt import MongoCryptKmsContext + from pymongo.typings import _Address + _IS_SYNC = True @@ -104,6 +110,13 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) +def _connect_kms(address: _Address, opts: PoolOptions): + try: + return _configured_socket(address, opts) + except Exception as exc: + _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) + + @contextlib.contextmanager def _wrap_encryption_errors() -> Iterator[None]: """Context manager to wrap encryption related errors.""" @@ -176,13 +189,13 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: socket_timeout=connect_timeout, ssl_context=ctx, ) - host, port = parse_host(endpoint, _HTTPS_PORT) + address = parse_host(endpoint, _HTTPS_PORT) sleep_u = kms_context.usleep if sleep_u: sleep_sec = float(sleep_u) / 1e6 time.sleep(sleep_sec) try: - conn = _configured_socket((host, port), opts) + conn = _connect_kms(address, opts) try: sendall(conn, message) while kms_context.bytes_needed > 0: @@ -199,19 +212,21 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None: if not data: raise OSError("KMS connection closed") kms_context.feed(data) - except BLOCKING_IO_ERRORS: - raise socket.timeout("timed out") from None + except MongoCryptError: + raise # Propagate MongoCryptError errors directly. + except Exception as exc: + # Wrap I/O errors in PyMongo exceptions. + if isinstance(exc, BLOCKING_IO_ERRORS): + exc = socket.timeout("timed out") + _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts)) finally: conn.close() except MongoCryptError: raise # Propagate MongoCryptError errors directly. except Exception as exc: remaining = _csot.remaining() - if isinstance(exc, (socket.timeout, NetworkTimeout)) or ( - remaining is not None and remaining <= 0 - ): - # Wrap I/O errors in PyMongo exceptions. - _raise_connection_failure((host, port), exc) + if isinstance(exc, NetworkTimeout) or (remaining is not None and remaining <= 0): + raise # Mark this attempt as failed and defer to libmongocrypt to retry. try: kms_context.fail() From 767904f90a5a74d196e8d7f10434e99ec4ec34b5 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Tue, 3 Dec 2024 15:44:25 -0800 Subject: [PATCH 8/9] PYTHON-2560 Fix test_04_aws_endpoint_invalid_port, cause is now MongoCryptError --- test/asynchronous/test_encryption.py | 3 +-- test/test_encryption.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 0823158f10..559b06ddf4 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -1369,9 +1369,8 @@ async def test_04_aws_endpoint_invalid_port(self): "key": ("arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0"), "endpoint": "kms.us-east-1.amazonaws.com:12345", } - with self.assertRaisesRegex(EncryptionError, "kms.us-east-1.amazonaws.com:12345") as ctx: + with self.assertRaisesRegex(EncryptionError, "kms.us-east-1.amazonaws.com:12345"): await self.client_encryption.create_data_key("aws", master_key=master_key) - self.assertIsInstance(ctx.exception.cause, AutoReconnect) @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") async def test_05_aws_endpoint_wrong_region(self): diff --git a/test/test_encryption.py b/test/test_encryption.py index f0c6b74692..7a9929b7fd 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -1363,9 +1363,8 @@ def test_04_aws_endpoint_invalid_port(self): "key": ("arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0"), "endpoint": "kms.us-east-1.amazonaws.com:12345", } - with self.assertRaisesRegex(EncryptionError, "kms.us-east-1.amazonaws.com:12345") as ctx: + with self.assertRaisesRegex(EncryptionError, "kms.us-east-1.amazonaws.com:12345"): self.client_encryption.create_data_key("aws", master_key=master_key) - self.assertIsInstance(ctx.exception.cause, AutoReconnect) @unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set") def test_05_aws_endpoint_wrong_region(self): From 20cd5d72936f19f271b24263e7867b3fd6758043 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Tue, 3 Dec 2024 15:46:30 -0800 Subject: [PATCH 9/9] PYTHON-2560 Fix typing --- pymongo/asynchronous/encryption.py | 3 ++- pymongo/synchronous/encryption.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 48a2f191fd..1cf165e6a2 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -93,6 +93,7 @@ if TYPE_CHECKING: from pymongocrypt.mongocrypt import MongoCryptKmsContext + from pymongo.pyopenssl_context import _sslConn from pymongo.typings import _Address @@ -110,7 +111,7 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -async def _connect_kms(address: _Address, opts: PoolOptions): +async def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]: try: return await _configured_socket(address, opts) except Exception as exc: diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index add5761337..ef49855059 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -93,6 +93,7 @@ if TYPE_CHECKING: from pymongocrypt.mongocrypt import MongoCryptKmsContext + from pymongo.pyopenssl_context import _sslConn from pymongo.typings import _Address @@ -110,7 +111,7 @@ _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument) -def _connect_kms(address: _Address, opts: PoolOptions): +def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]: try: return _configured_socket(address, opts) except Exception as exc: