Async KMS

Author: NoahStapp

pymongo/asynchronous/encryption.py

@@ -142,6 +142,32 @@ def __init__(
         self.opts = opts
         self._spawned = False
 
+    async def _async_kms_request(
+        self, kms_context: MongoCryptKmsContext, host, port, opts, message
+    ) -> None:
+        from pymongo.network_layer import _async_receive_data_socket
+
+        try:
+            conn = await _configured_socket((host, port), opts)
+            try:
+                await async_sendall(conn, message)
+                while kms_context.bytes_needed > 0:
+                    # CSOT: update timeout.
+                    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    data = await _async_receive_data_socket(conn, kms_context.bytes_needed)
+                    if not data:
+                        raise OSError("KMS connection closed")
+                    kms_context.feed(data)
+            except BLOCKING_IO_ERRORS:
+                raise socket.timeout("timed out") from None
+            finally:
+                conn.close()
+        except (PyMongoError, MongoCryptError):
+            raise  # Propagate pymongo errors directly.
+        except Exception as error:
+            # Wrap I/O errors in PyMongo exceptions.
+            _raise_connection_failure((host, port), error)
+
     async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         """Complete a KMS request.
 
@@ -174,20 +200,23 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         )
         host, port = parse_host(endpoint, _HTTPS_PORT)
         try:
-            conn = await _configured_socket((host, port), opts)
-            try:
-                await async_sendall(conn, message)
-                while kms_context.bytes_needed > 0:
-                    # CSOT: update timeout.
-                    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    data = conn.recv(kms_context.bytes_needed)
-                    if not data:
-                        raise OSError("KMS connection closed")
-                    kms_context.feed(data)
-            except BLOCKING_IO_ERRORS:
-                raise socket.timeout("timed out") from None
-            finally:
-                conn.close()
+            if _IS_SYNC:
+                conn = await _configured_socket((host, port), opts)
+                try:
+                    await async_sendall(conn, message)
+                    while kms_context.bytes_needed > 0:
+                        # CSOT: update timeout.
+                        conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                        data = conn.recv(kms_context.bytes_needed)
+                        if not data:
+                            raise OSError("KMS connection closed")
+                        kms_context.feed(data)
+                except BLOCKING_IO_ERRORS:
+                    raise socket.timeout("timed out") from None
+                finally:
+                    conn.close()
+            else:
+                return await self._async_kms_request(kms_context, host, port, opts, message)
         except (PyMongoError, MongoCryptError):
             raise  # Propagate pymongo errors directly.
         except Exception as error:

pymongo/network_layer.py

@@ -275,6 +275,40 @@ async def async_receive_data(
         sock.settimeout(sock_timeout)
 
 
+async def _async_receive_data_socket(
+    sock: socket.socket | _sslConn, length: int, deadline: Optional[float]
+) -> memoryview:
+    sock_timeout = sock.gettimeout()
+    timeout: Optional[Union[float, int]]
+    if deadline:
+        # When the timeout has expired perform one final check to
+        # see if the socket is readable. This helps avoid spurious
+        # timeouts on AWS Lambda and other FaaS environments.
+        timeout = max(deadline - time.monotonic(), 0)
+    else:
+        timeout = sock_timeout
+
+    sock.settimeout(0.0)
+    loop = asyncio.get_event_loop()
+    try:
+        if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
+            read_task = asyncio.create_task(_async_receive_ssl(sock, length, loop))  # type: ignore[arg-type]
+        else:
+            read_task = asyncio.create_task(_async_receive(sock, length, loop))  # type: ignore[arg-type]
+        tasks = [read_task]
+        done, pending = await asyncio.wait(
+            tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED
+        )
+        for task in pending:
+            task.cancel()
+        await asyncio.wait(pending)
+        if read_task in done:
+            return read_task.result()
+        raise socket.timeout("timed out")
+    finally:
+        sock.settimeout(sock_timeout)
+
+
 async def _async_receive(conn: socket.socket, length: int, loop: AbstractEventLoop) -> memoryview:
     mv = memoryview(bytearray(length))
     bytes_read = 0

pymongo/synchronous/encryption.py

@@ -142,6 +142,32 @@ def __init__(
         self.opts = opts
         self._spawned = False
 
+    def _async_kms_request(
+        self, kms_context: MongoCryptKmsContext, host, port, opts, message
+    ) -> None:
+        from pymongo.network_layer import _receive_data_socket
+
+        try:
+            conn = _configured_socket((host, port), opts)
+            try:
+                sendall(conn, message)
+                while kms_context.bytes_needed > 0:
+                    # CSOT: update timeout.
+                    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    data = _receive_data_socket(conn, kms_context.bytes_needed)
+                    if not data:
+                        raise OSError("KMS connection closed")
+                    kms_context.feed(data)
+            except BLOCKING_IO_ERRORS:
+                raise socket.timeout("timed out") from None
+            finally:
+                conn.close()
+        except (PyMongoError, MongoCryptError):
+            raise  # Propagate pymongo errors directly.
+        except Exception as error:
+            # Wrap I/O errors in PyMongo exceptions.
+            _raise_connection_failure((host, port), error)
+
     def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         """Complete a KMS request.
 
@@ -174,20 +200,23 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         )
         host, port = parse_host(endpoint, _HTTPS_PORT)
         try:
-            conn = _configured_socket((host, port), opts)
-            try:
-                sendall(conn, message)
-                while kms_context.bytes_needed > 0:
-                    # CSOT: update timeout.
-                    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    data = conn.recv(kms_context.bytes_needed)
-                    if not data:
-                        raise OSError("KMS connection closed")
-                    kms_context.feed(data)
-            except BLOCKING_IO_ERRORS:
-                raise socket.timeout("timed out") from None
-            finally:
-                conn.close()
+            if _IS_SYNC:
+                conn = _configured_socket((host, port), opts)
+                try:
+                    sendall(conn, message)
+                    while kms_context.bytes_needed > 0:
+                        # CSOT: update timeout.
+                        conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                        data = conn.recv(kms_context.bytes_needed)
+                        if not data:
+                            raise OSError("KMS connection closed")
+                        kms_context.feed(data)
+                except BLOCKING_IO_ERRORS:
+                    raise socket.timeout("timed out") from None
+                finally:
+                    conn.close()
+            else:
+                return self._async_kms_request(kms_context, host, port, opts, message)
         except (PyMongoError, MongoCryptError):
             raise  # Propagate pymongo errors directly.
         except Exception as error: