wip

Author: blink1073

pymongo/asynchronous/encryption.py

@@ -76,7 +76,7 @@
     NetworkTimeout,
     ServerSelectionTimeoutError,
 )
-from pymongo.network_layer import async_receive_kms, async_sendall
+from pymongo.network_layer import PyMongoKMSProtocol, async_receive_kms, async_sendall
 from pymongo.operations import UpdateOne
 from pymongo.pool_options import PoolOptions
 from pymongo.pool_shared import (
@@ -187,16 +187,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)
+            interface = await _configured_protocol_interface(address, opts, PyMongoKMSProtocol)
             conn = AsyncBaseConnection(interface, opts)
             try:
                 await async_sendall(interface.get_conn, message)
-                # CSOT: update timeout.
-                interface.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                data = await async_receive_kms(conn, kms_context.bytes_needed)
-                if not data:
-                    raise OSError("KMS connection closed")
-                kms_context.feed(bytes)
+                while kms_context.bytes_needed > 0:
+                    # CSOT: update timeout.
+                    interface.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    data = await async_receive_kms(conn, kms_context.bytes_needed)
+                    if not data:
+                        raise OSError("KMS connection closed")
+                    kms_context.feed(data)
             except MongoCryptError:
                 raise  # Propagate MongoCryptError errors directly.
             except Exception as exc:
@@ -212,7 +213,7 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts)
                 )
             finally:
-                conn.close_conn()
+                conn.close_conn(None)
         except MongoCryptError:
             raise  # Propagate MongoCryptError errors directly.
         except Exception as exc:

pymongo/network_layer.py

@@ -257,6 +257,7 @@ def __init__(self, timeout: Optional[float] = None):
         self.transport: Transport = None  # type: ignore[assignment]
         self._timeout = timeout
         self._closed = asyncio.get_running_loop().create_future()
+        self._connection_lost = False
 
     def settimeout(self, timeout: float | None) -> None:
         self._timeout = timeout
@@ -309,7 +310,6 @@ def __init__(self, timeout: Optional[float] = None):
         self._expecting_compression = False
         self._message_size = 0
         self._op_code = 0
-        self._connection_lost = False
         self._read_waiter: Optional[Future[Any]] = None
         self._is_compressed = False
         self._compressor_id: Optional[int] = None
@@ -551,7 +551,8 @@ def buffer_updated(self, nbytes: int) -> None:
             buffer = self._buffers.popleft()
             buffer_remaining = len(buffer.buffer) - buffer.start_index
             # if we didn't exhaust the buffer, read the partial data and put it back.
-            if buffer_remaining <= n_remaining:
+            if buffer_remaining > n_remaining:
+                n_remaining = 0
                 output_buf[out_index : n_remaining + out_index] = buffer.buffer[
                     buffer.start_index : buffer.start_index + n_remaining
                 ]
@@ -562,6 +563,7 @@ def buffer_updated(self, nbytes: int) -> None:
                 output_buf[out_index : out_index + buffer_remaining] = buffer.buffer[
                     buffer.start_index :
                 ]
+                n_remaining -= buffer_remaining
                 buffer.start_index = 0
                 self._buffer_pool.append(buffer)
         waiter.set_result(output_buf)

pymongo/pool_shared.py

@@ -36,7 +36,12 @@
     NetworkTimeout,
     _CertificateError,
 )
-from pymongo.network_layer import AsyncNetworkingInterface, NetworkingInterface, PyMongoProtocol
+from pymongo.network_layer import (
+    AsyncNetworkingInterface,
+    NetworkingInterface,
+    PyMongoBaseProtocol,
+    PyMongoProtocol,
+)
 from pymongo.pool_options import PoolOptions
 from pymongo.ssl_support import PYSSLError, SSLError, _has_sni
 
@@ -326,7 +331,7 @@ async def _async_configured_socket(
 
 
 async def _configured_protocol_interface(
-    address: _Address, options: PoolOptions
+    address: _Address, options: PoolOptions, protocol_kls: PyMongoBaseProtocol = PyMongoProtocol
 ) -> AsyncNetworkingInterface:
     """Given (host, port) and PoolOptions, return a configured AsyncNetworkingInterface.
 
@@ -341,7 +346,7 @@ async def _configured_protocol_interface(
     if ssl_context is None:
         return AsyncNetworkingInterface(
             await asyncio.get_running_loop().create_connection(
-                lambda: PyMongoProtocol(timeout=timeout), sock=sock
+                lambda: protocol_kls(timeout=timeout), sock=sock
             )
         )
 
@@ -350,7 +355,7 @@ async def _configured_protocol_interface(
         # We have to pass hostname / ip address to wrap_socket
         # to use SSLContext.check_hostname.
         transport, protocol = await asyncio.get_running_loop().create_connection(  # type: ignore[call-overload]
-            lambda: PyMongoProtocol(timeout=timeout),
+            lambda: protocol_kls(timeout=timeout),
             sock=sock,
             server_hostname=host,
             ssl=ssl_context,

pymongo/synchronous/encryption.py

@@ -70,7 +70,7 @@
     NetworkTimeout,
     ServerSelectionTimeoutError,
 )
-from pymongo.network_layer import receive_kms, sendall
+from pymongo.network_layer import PyMongoKMSProtocol, receive_kms, sendall
 from pymongo.operations import UpdateOne
 from pymongo.pool_options import PoolOptions
 from pymongo.pool_shared import (
@@ -186,16 +186,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)
+            interface = _configured_socket_interface(address, opts, PyMongoKMSProtocol)
             conn = BaseConnection(interface, opts)
             try:
                 sendall(interface.get_conn, message)
-                # CSOT: update timeout.
-                interface.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                data = receive_kms(conn, kms_context.bytes_needed)
-                if not data:
-                    raise OSError("KMS connection closed")
-                kms_context.feed(bytes)
+                while kms_context.bytes_needed > 0:
+                    # CSOT: update timeout.
+                    interface.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    data = receive_kms(conn, kms_context.bytes_needed)
+                    if not data:
+                        raise OSError("KMS connection closed")
+                    kms_context.feed(data)
             except MongoCryptError:
                 raise  # Propagate MongoCryptError errors directly.
             except Exception as exc:
@@ -211,7 +212,7 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts)
                 )
             finally:
-                conn.close_conn()
+                conn.close_conn(None)
         except MongoCryptError:
             raise  # Propagate MongoCryptError errors directly.
         except Exception as exc: