Skip to content

Commit e413334

Browse files
ShaneHarveyShaneHarvey
committed
PYTHON-2560 Add prose test
1 parent 5c79502 commit e413334

File tree

6 files changed

+180
-8
lines changed

6 files changed

+180
-8
lines changed

‎.evergreen/run-tests.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ if [ -n "$TEST_ENCRYPTION" ] || [ -n "$TEST_FLE_AZURE_AUTO" ] || [ -n "$TEST_FLE
158158

159159
# TODO: Test with 'pip install pymongocrypt'
160160
if [ ! -d "libmongocrypt_git" ]; then
161-
git clone https://github.com/mongodb/libmongocrypt.git libmongocrypt_git
161+
git clone https://github.com/ShaneHarvey/libmongocrypt.git --branch PYTHON-4992 libmongocrypt_git
162162
fi
163163
python -m pip install -U setuptools
164164
python -m pip install ./libmongocrypt_git/bindings/python

‎.evergreen/scripts/prepare-resources.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,6 @@ if [ "$PROJECT" = "drivers-tools" ]; then
77
# If this was a patch build, doing a fresh clone would not actually test the patch
88
cp -R $PROJECT_DIRECTORY/ $DRIVERS_TOOLS
99
else
10-
git clone https://github.com/mongodb-labs/drivers-evergreen-tools.git $DRIVERS_TOOLS
10+
git clone https://github.com/ShaneHarvey/drivers-evergreen-tools.git --branch DRIVERS-1541 $DRIVERS_TOOLS
1111
fi
1212
echo "{ \"releases\": { \"default\": \"$MONGODB_BINARIES\" }}" >$MONGO_ORCHESTRATION_HOME/orchestration.config

‎pymongo/asynchronous/encryption.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@
7272
EncryptedCollectionError,
7373
EncryptionError,
7474
InvalidOperation,
75+
NetworkTimeout,
7576
ServerSelectionTimeoutError,
7677
)
7778
from pymongo.network_layer import BLOCKING_IO_ERRORS, async_sendall
@@ -165,8 +166,8 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
165166
None, # crlfile
166167
False, # allow_invalid_certificates
167168
False, # allow_invalid_hostnames
168-
False,
169-
) # disable_ocsp_endpoint_check
169+
False, # disable_ocsp_endpoint_check
170+
)
170171
# CSOT: set timeout for socket creation.
171172
connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)
172173
opts = PoolOptions(
@@ -207,7 +208,9 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
207208
raise # Propagate MongoCryptError errors directly.
208209
except Exception as exc:
209210
remaining = _csot.remaining()
210-
if remaining is not None and remaining <= 0:
211+
if isinstance(exc, (socket.timeout, NetworkTimeout)) or (
212+
remaining is not None and remaining <= 0
213+
):
211214
# Wrap I/O errors in PyMongo exceptions.
212215
_raise_connection_failure((host, port), exc)
213216
# Mark this attempt as failed and defer to libmongocrypt to retry.

‎pymongo/synchronous/encryption.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@
6767
EncryptedCollectionError,
6868
EncryptionError,
6969
InvalidOperation,
70+
NetworkTimeout,
7071
ServerSelectionTimeoutError,
7172
)
7273
from pymongo.network_layer import BLOCKING_IO_ERRORS, sendall
@@ -165,8 +166,8 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
165166
None, # crlfile
166167
False, # allow_invalid_certificates
167168
False, # allow_invalid_hostnames
168-
False,
169-
) # disable_ocsp_endpoint_check
169+
False, # disable_ocsp_endpoint_check
170+
)
170171
# CSOT: set timeout for socket creation.
171172
connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)
172173
opts = PoolOptions(
@@ -207,7 +208,9 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
207208
raise # Propagate MongoCryptError errors directly.
208209
except Exception as exc:
209210
remaining = _csot.remaining()
210-
if remaining is not None and remaining <= 0:
211+
if isinstance(exc, (socket.timeout, NetworkTimeout)) or (
212+
remaining is not None and remaining <= 0
213+
):
211214
# Wrap I/O errors in PyMongo exceptions.
212215
_raise_connection_failure((host, port), exc)
213216
# Mark this attempt as failed and defer to libmongocrypt to retry.

‎test/asynchronous/test_encryption.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
import base64
1919
import copy
20+
import http.client
21+
import json
2022
import os
2123
import pathlib
2224
import re
@@ -91,6 +93,7 @@
9193
WriteError,
9294
)
9395
from pymongo.operations import InsertOne, ReplaceOne, UpdateOne
96+
from pymongo.ssl_support import get_ssl_context
9497
from pymongo.write_concern import WriteConcern
9598

9699
_IS_SYNC = False
@@ -2853,6 +2856,86 @@ async def test_accepts_trim_factor_0(self):
28532856
assert len(payload) > len(self.payload_defaults)
28542857

28552858

2859+
# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#24-kms-retry-tests
2860+
class TestKmsRetryProse(AsyncEncryptionIntegrationTest):
2861+
@unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set")
2862+
async def asyncSetUp(self):
2863+
await super().asyncSetUp()
2864+
# 1, create client with only tlsCAFile.
2865+
providers: dict = copy.deepcopy(ALL_KMS_PROVIDERS)
2866+
providers["azure"]["identityPlatformEndpoint"] = "127.0.0.1:9003"
2867+
providers["gcp"]["endpoint"] = "127.0.0.1:9003"
2868+
kms_tls_opts = {
2869+
p: {"tlsCAFile": CA_PEM, "tlsCertificateKeyFile": CLIENT_PEM} for p in providers
2870+
}
2871+
self.client_encryption = self.create_client_encryption(
2872+
providers, "keyvault.datakeys", self.client, OPTS, kms_tls_options=kms_tls_opts
2873+
)
2874+
2875+
async def http_post(self, path, data=None):
2876+
# Note, the connection to the mock server needs to be closed after
2877+
# each request because the server is single threaded.
2878+
ctx = get_ssl_context(
2879+
CLIENT_PEM, # certfile
2880+
None, # passphrase
2881+
CA_PEM, # ca_certs
2882+
None, # crlfile
2883+
False, # allow_invalid_certificates
2884+
False, # allow_invalid_hostnames
2885+
False, # disable_ocsp_endpoint_check
2886+
)
2887+
conn = http.client.HTTPSConnection("127.0.0.1:9003", context=ctx)
2888+
try:
2889+
if data is not None:
2890+
headers = {"Content-type": "application/json"}
2891+
body = json.dumps(data)
2892+
else:
2893+
headers = {}
2894+
body = None
2895+
conn.request("POST", path, body, headers)
2896+
res = conn.getresponse()
2897+
res.read()
2898+
finally:
2899+
conn.close()
2900+
2901+
async def _test(self, provider, master_key):
2902+
await self.http_post("/reset")
2903+
# Case 1: createDataKey and encrypt with TCP retry
2904+
await self.http_post("/set_failpoint/network", {"count": 1})
2905+
key_id = await self.client_encryption.create_data_key(provider, master_key=master_key)
2906+
await self.http_post("/set_failpoint/network", {"count": 1})
2907+
await self.client_encryption.encrypt(
2908+
123, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id
2909+
)
2910+
2911+
# Case 2: createDataKey and encrypt with HTTP retry
2912+
await self.http_post("/set_failpoint/http", {"count": 1})
2913+
key_id = await self.client_encryption.create_data_key(provider, master_key=master_key)
2914+
await self.http_post("/set_failpoint/http", {"count": 1})
2915+
await self.client_encryption.encrypt(
2916+
123, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id
2917+
)
2918+
2919+
# Case 3: createDataKey fails after too many retries
2920+
await self.http_post("/set_failpoint/network", {"count": 4})
2921+
with self.assertRaisesRegex(EncryptionError, "KMS request failed after"):
2922+
await self.client_encryption.create_data_key(provider, master_key=master_key)
2923+
2924+
async def test_kms_retry(self):
2925+
await self._test("aws", {"region": "foo", "key": "bar", "endpoint": "127.0.0.1:9003"})
2926+
await self._test("azure", {"keyVaultEndpoint": "127.0.0.1:9003", "keyName": "foo"})
2927+
await self._test(
2928+
"gcp",
2929+
{
2930+
"projectId": "foo",
2931+
"location": "bar",
2932+
"keyRing": "baz",
2933+
"keyName": "qux",
2934+
"endpoint": "127.0.0.1:9003",
2935+
},
2936+
)
2937+
2938+
28562939
# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#automatic-data-encryption-keys
28572940
class TestAutomaticDecryptionKeys(AsyncEncryptionIntegrationTest):
28582941
@async_client_context.require_no_standalone

‎test/test_encryption.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
import base64
1919
import copy
20+
import http.client
21+
import json
2022
import os
2123
import pathlib
2224
import re
@@ -88,6 +90,7 @@
8890
WriteError,
8991
)
9092
from pymongo.operations import InsertOne, ReplaceOne, UpdateOne
93+
from pymongo.ssl_support import get_ssl_context
9194
from pymongo.synchronous import encryption
9295
from pymongo.synchronous.encryption import Algorithm, ClientEncryption, QueryType
9396
from pymongo.synchronous.mongo_client import MongoClient
@@ -2835,6 +2838,86 @@ def test_accepts_trim_factor_0(self):
28352838
assert len(payload) > len(self.payload_defaults)
28362839

28372840

2841+
# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#24-kms-retry-tests
2842+
class TestKmsRetryProse(EncryptionIntegrationTest):
2843+
@unittest.skipUnless(any(AWS_CREDS.values()), "AWS environment credentials are not set")
2844+
def setUp(self):
2845+
super().setUp()
2846+
# 1, create client with only tlsCAFile.
2847+
providers: dict = copy.deepcopy(ALL_KMS_PROVIDERS)
2848+
providers["azure"]["identityPlatformEndpoint"] = "127.0.0.1:9003"
2849+
providers["gcp"]["endpoint"] = "127.0.0.1:9003"
2850+
kms_tls_opts = {
2851+
p: {"tlsCAFile": CA_PEM, "tlsCertificateKeyFile": CLIENT_PEM} for p in providers
2852+
}
2853+
self.client_encryption = self.create_client_encryption(
2854+
providers, "keyvault.datakeys", self.client, OPTS, kms_tls_options=kms_tls_opts
2855+
)
2856+
2857+
def http_post(self, path, data=None):
2858+
# Note, the connection to the mock server needs to be closed after
2859+
# each request because the server is single threaded.
2860+
ctx = get_ssl_context(
2861+
CLIENT_PEM, # certfile
2862+
None, # passphrase
2863+
CA_PEM, # ca_certs
2864+
None, # crlfile
2865+
False, # allow_invalid_certificates
2866+
False, # allow_invalid_hostnames
2867+
False, # disable_ocsp_endpoint_check
2868+
)
2869+
conn = http.client.HTTPSConnection("127.0.0.1:9003", context=ctx)
2870+
try:
2871+
if data is not None:
2872+
headers = {"Content-type": "application/json"}
2873+
body = json.dumps(data)
2874+
else:
2875+
headers = {}
2876+
body = None
2877+
conn.request("POST", path, body, headers)
2878+
res = conn.getresponse()
2879+
res.read()
2880+
finally:
2881+
conn.close()
2882+
2883+
def _test(self, provider, master_key):
2884+
self.http_post("/reset")
2885+
# Case 1: createDataKey and encrypt with TCP retry
2886+
self.http_post("/set_failpoint/network", {"count": 1})
2887+
key_id = self.client_encryption.create_data_key(provider, master_key=master_key)
2888+
self.http_post("/set_failpoint/network", {"count": 1})
2889+
self.client_encryption.encrypt(
2890+
123, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id
2891+
)
2892+
2893+
# Case 2: createDataKey and encrypt with HTTP retry
2894+
self.http_post("/set_failpoint/http", {"count": 1})
2895+
key_id = self.client_encryption.create_data_key(provider, master_key=master_key)
2896+
self.http_post("/set_failpoint/http", {"count": 1})
2897+
self.client_encryption.encrypt(
2898+
123, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, key_id
2899+
)
2900+
2901+
# Case 3: createDataKey fails after too many retries
2902+
self.http_post("/set_failpoint/network", {"count": 4})
2903+
with self.assertRaisesRegex(EncryptionError, "KMS request failed after"):
2904+
self.client_encryption.create_data_key(provider, master_key=master_key)
2905+
2906+
def test_kms_retry(self):
2907+
self._test("aws", {"region": "foo", "key": "bar", "endpoint": "127.0.0.1:9003"})
2908+
self._test("azure", {"keyVaultEndpoint": "127.0.0.1:9003", "keyName": "foo"})
2909+
self._test(
2910+
"gcp",
2911+
{
2912+
"projectId": "foo",
2913+
"location": "bar",
2914+
"keyRing": "baz",
2915+
"keyName": "qux",
2916+
"endpoint": "127.0.0.1:9003",
2917+
},
2918+
)
2919+
2920+
28382921
# https://github.com/mongodb/specifications/blob/master/source/client-side-encryption/tests/README.md#automatic-data-encryption-keys
28392922
class TestAutomaticDecryptionKeys(EncryptionIntegrationTest):
28402923
@client_context.require_no_standalone

0 commit comments

Comments
 (0)