From 493fc2ab3e237c2155fde4400002ed2aafe9b2be Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Fri, 10 Jan 2025 13:05:57 -0600 Subject: [PATCH] PYTHON-5014 Fix handling of async socket errors in kms request (#2054) --- pymongo/asynchronous/encryption.py | 9 ++++++++- pymongo/synchronous/encryption.py | 9 ++++++++- test/asynchronous/test_encryption.py | 12 ++++++++---- test/test_encryption.py | 12 ++++++++---- 4 files changed, 32 insertions(+), 10 deletions(-) diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py index 1cf165e6a..98ab68527 100644 --- a/pymongo/asynchronous/encryption.py +++ b/pymongo/asynchronous/encryption.py @@ -219,7 +219,14 @@ class _EncryptionIO(AsyncMongoCryptCallback): # type: ignore[misc] # 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)) + # Async raises an OSError instead of returning empty bytes. + if isinstance(exc, OSError): + msg_prefix = "KMS connection closed" + else: + msg_prefix = None + _raise_connection_failure( + address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts) + ) finally: conn.close() except MongoCryptError: diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py index ef4985505..d41169861 100644 --- a/pymongo/synchronous/encryption.py +++ b/pymongo/synchronous/encryption.py @@ -219,7 +219,14 @@ class _EncryptionIO(MongoCryptCallback): # type: ignore[misc] # 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)) + # Async raises an OSError instead of returning empty bytes. + if isinstance(exc, OSError): + msg_prefix = "KMS connection closed" + else: + msg_prefix = None + _raise_connection_failure( + address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts) + ) finally: conn.close() except MongoCryptError: diff --git a/test/asynchronous/test_encryption.py b/test/asynchronous/test_encryption.py index 10c4c8a56..ef53d8ccd 100644 --- a/test/asynchronous/test_encryption.py +++ b/test/asynchronous/test_encryption.py @@ -2162,7 +2162,8 @@ class TestKmsTLSOptions(AsyncEncryptionIntegrationTest): # 127.0.0.1:9001: ('Certificate does not contain any `subjectAltName`s.',) key["endpoint"] = "127.0.0.1:9001" with self.assertRaisesRegex( - EncryptionError, "IP address mismatch|wronghost|IPAddressMismatch|Certificate" + EncryptionError, + "IP address mismatch|wronghost|IPAddressMismatch|Certificate|SSL handshake failed", ): await self.client_encryption_invalid_hostname.create_data_key("aws", key) @@ -2179,7 +2180,8 @@ class TestKmsTLSOptions(AsyncEncryptionIntegrationTest): await self.client_encryption_expired.create_data_key("azure", key) # Invalid cert hostname error. with self.assertRaisesRegex( - EncryptionError, "IP address mismatch|wronghost|IPAddressMismatch|Certificate" + EncryptionError, + "IP address mismatch|wronghost|IPAddressMismatch|Certificate|SSL handshake failed", ): await self.client_encryption_invalid_hostname.create_data_key("azure", key) @@ -2196,7 +2198,8 @@ class TestKmsTLSOptions(AsyncEncryptionIntegrationTest): await self.client_encryption_expired.create_data_key("gcp", key) # Invalid cert hostname error. with self.assertRaisesRegex( - EncryptionError, "IP address mismatch|wronghost|IPAddressMismatch|Certificate" + EncryptionError, + "IP address mismatch|wronghost|IPAddressMismatch|Certificate|SSL handshake failed", ): await self.client_encryption_invalid_hostname.create_data_key("gcp", key) @@ -2210,7 +2213,8 @@ class TestKmsTLSOptions(AsyncEncryptionIntegrationTest): await self.client_encryption_expired.create_data_key("kmip") # Invalid cert hostname error. with self.assertRaisesRegex( - EncryptionError, "IP address mismatch|wronghost|IPAddressMismatch|Certificate" + EncryptionError, + "IP address mismatch|wronghost|IPAddressMismatch|Certificate|SSL handshake failed", ): await self.client_encryption_invalid_hostname.create_data_key("kmip") diff --git a/test/test_encryption.py b/test/test_encryption.py index 7b5aa776e..726463c41 100644 --- a/test/test_encryption.py +++ b/test/test_encryption.py @@ -2154,7 +2154,8 @@ class TestKmsTLSOptions(EncryptionIntegrationTest): # 127.0.0.1:9001: ('Certificate does not contain any `subjectAltName`s.',) key["endpoint"] = "127.0.0.1:9001" with self.assertRaisesRegex( - EncryptionError, "IP address mismatch|wronghost|IPAddressMismatch|Certificate" + EncryptionError, + "IP address mismatch|wronghost|IPAddressMismatch|Certificate|SSL handshake failed", ): self.client_encryption_invalid_hostname.create_data_key("aws", key) @@ -2171,7 +2172,8 @@ class TestKmsTLSOptions(EncryptionIntegrationTest): self.client_encryption_expired.create_data_key("azure", key) # Invalid cert hostname error. with self.assertRaisesRegex( - EncryptionError, "IP address mismatch|wronghost|IPAddressMismatch|Certificate" + EncryptionError, + "IP address mismatch|wronghost|IPAddressMismatch|Certificate|SSL handshake failed", ): self.client_encryption_invalid_hostname.create_data_key("azure", key) @@ -2188,7 +2190,8 @@ class TestKmsTLSOptions(EncryptionIntegrationTest): self.client_encryption_expired.create_data_key("gcp", key) # Invalid cert hostname error. with self.assertRaisesRegex( - EncryptionError, "IP address mismatch|wronghost|IPAddressMismatch|Certificate" + EncryptionError, + "IP address mismatch|wronghost|IPAddressMismatch|Certificate|SSL handshake failed", ): self.client_encryption_invalid_hostname.create_data_key("gcp", key) @@ -2202,7 +2205,8 @@ class TestKmsTLSOptions(EncryptionIntegrationTest): self.client_encryption_expired.create_data_key("kmip") # Invalid cert hostname error. with self.assertRaisesRegex( - EncryptionError, "IP address mismatch|wronghost|IPAddressMismatch|Certificate" + EncryptionError, + "IP address mismatch|wronghost|IPAddressMismatch|Certificate|SSL handshake failed", ): self.client_encryption_invalid_hostname.create_data_key("kmip")