fix sync kms

Author: blink1073

pymongo/asynchronous/encryption.py

@@ -94,6 +94,7 @@
 if TYPE_CHECKING:
     from pymongocrypt.mongocrypt import MongoCryptKmsContext
 
+    from pymongo.typings import _Address
 
 _IS_SYNC = False
 
@@ -109,6 +110,14 @@
 _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument)
 
 
+async def _connect_kms(address: _Address, opts: PoolOptions) -> AsyncBaseConnection:
+    try:
+        interface = await _configured_protocol_interface(address, opts, PyMongoKMSProtocol)
+        return AsyncBaseConnection(interface, 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."""
@@ -187,13 +196,17 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
             sleep_sec = float(sleep_u) / 1e6
             await asyncio.sleep(sleep_sec)
         try:
-            interface = await _configured_protocol_interface(address, opts, PyMongoKMSProtocol)
-            conn = AsyncBaseConnection(interface, opts)
+            conn = await _connect_kms(address, opts)
             try:
-                await async_sendall(interface.get_conn, message)
+                await async_sendall(conn.conn.get_conn, message)
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
-                    interface.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    conn.set_conn_timeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    # if _IS_SYNC:
+                    #     # TODO: why can't we use receive_kms?
+                    #     data = conn.conn.sock.recv(kms_context.bytes_needed)
+                    # else:
+                    #     data = await async_receive_kms(conn, kms_context.bytes_needed)
                     data = await async_receive_kms(conn, kms_context.bytes_needed)
                     if not data:
                         raise OSError("KMS connection closed")

pymongo/network_layer.py

@@ -41,9 +41,9 @@
 from pymongo.socket_checker import _errno_from_exception
 
 if TYPE_CHECKING:
-    from pymongo.asynchronous.pool import AsyncBaseConnection
+    from pymongo.asynchronous.pool import AsyncBaseConnection, AsyncConnection
     from pymongo.pyopenssl_context import _sslConn
-    from pymongo.synchronous.pool import BaseConnection
+    from pymongo.synchronous.pool import BaseConnection, Connection
 
 _UNPACK_HEADER = struct.Struct("<iiii").unpack
 _UNPACK_COMPRESSION_HEADER = struct.Struct("<iiB").unpack
@@ -642,7 +642,7 @@ async def _async_receive_data(
 
 
 async def async_receive_message(
-    conn: AsyncBaseConnection,
+    conn: AsyncConnection,
     request_id: Optional[int],
     max_message_size: int = MAX_MESSAGE_SIZE,
 ) -> Union[_OpReply, _OpMsg]:
@@ -663,19 +663,11 @@ def callback(result: Any) -> _OpMsg | _OpReply:
 
 def receive_kms(conn: BaseConnection, bytes_needed: int) -> bytes:
     """Receive raw bytes from the kms connection."""
-    if _csot.get_timeout():
-        deadline = _csot.get_deadline()
-    else:
-        timeout = conn.conn.gettimeout()
-        if timeout:
-            deadline = time.monotonic() + timeout
-        else:
-            deadline = None
-    return receive_data(conn, bytes_needed, deadline)
+    return conn.conn.sock.recv(bytes_needed)
 
 
 def receive_message(
-    conn: BaseConnection, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
+    conn: Connection, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
 ) -> Union[_OpReply, _OpMsg]:
     """Receive a raw BSON message or raise socket.error."""
     if _csot.get_timeout():

pymongo/synchronous/encryption.py

@@ -93,6 +93,7 @@
 if TYPE_CHECKING:
     from pymongocrypt.mongocrypt import MongoCryptKmsContext
 
+    from pymongo.typings import _Address
 
 _IS_SYNC = True
 
@@ -108,6 +109,14 @@
 _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument)
 
 
+def _connect_kms(address: _Address, opts: PoolOptions) -> BaseConnection:
+    try:
+        interface = _configured_socket_interface(address, opts, PyMongoKMSProtocol)
+        return BaseConnection(interface, 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."""
@@ -186,13 +195,17 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
             sleep_sec = float(sleep_u) / 1e6
             time.sleep(sleep_sec)
         try:
-            interface = _configured_socket_interface(address, opts, PyMongoKMSProtocol)
-            conn = BaseConnection(interface, opts)
+            conn = _connect_kms(address, opts)
             try:
-                sendall(interface.get_conn, message)
+                sendall(conn.conn.get_conn, message)
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
-                    interface.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    conn.set_conn_timeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    # if _IS_SYNC:
+                    #     # TODO: why can't we use receive_kms?
+                    #     data = conn.conn.sock.recv(kms_context.bytes_needed)
+                    # else:
+                    #     data = receive_kms(conn, kms_context.bytes_needed)
                     data = receive_kms(conn, kms_context.bytes_needed)
                     if not data:
                         raise OSError("KMS connection closed")