always allow partial reads

Author: blink1073

pymongo/asynchronous/encryption.py

@@ -191,15 +191,13 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
             conn = AsyncBaseConnection(interface, opts)
             try:
                 await async_sendall(interface.get_conn, message)
-                first = True
                 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, first)
+                    data = await async_receive_kms(conn, kms_context.bytes_needed)
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
-                    first = False
             except MongoCryptError:
                 raise  # Propagate MongoCryptError errors directly.
             except Exception as exc:

pymongo/network_layer.py

@@ -499,13 +499,13 @@ def connection_made(self, transport: BaseTransport) -> None:
         """
         self.transport = transport  # type: ignore[assignment]
 
-    async def read(self, bytes_needed: int, first=False) -> bytes:
-        """Read the requested bytes from this connection."""
-        if self._bytes_ready >= bytes_needed or (self._bytes_ready > 0 and first):
+    async def read(self, bytes_needed: int) -> bytes:
+        """Read up to the requested bytes from this connection."""
+        if self._bytes_ready > 0:
             # Wait for other listeners first.
             if len(self._pending_listeners):
                 await asyncio.gather(*self._pending_listeners)
-            return self._read(bytes_needed, first)
+            return self._read(bytes_needed)
         if self.transport:
             try:
                 self.transport.resume_reading()
@@ -514,7 +514,7 @@ async def read(self, bytes_needed: int, first=False) -> bytes:
                 raise OSError("connection is already closed") from None
         if self.transport and self.transport.is_closing():
             raise OSError("connection is already closed")
-        self._pending_reads.append((bytes_needed, first))
+        self._pending_reads.append(bytes_needed)
         read_waiter = asyncio.get_running_loop().create_future()
         self._pending_listeners.append(read_waiter)
         return await read_waiter
@@ -548,19 +548,15 @@ def buffer_updated(self, nbytes: int) -> None:
         if not len(self._pending_reads):
             return
 
-        bytes_needed, first = self._pending_reads.popleft()
-        if not first and (bytes_needed == 0 or self._bytes_ready < bytes_needed):
-            self._pending_reads.appendleft((bytes_needed, first))
-            return
-
-        data = self._read(bytes_needed, first)
+        bytes_needed = self._pending_reads.popleft()
+        data = self._read(bytes_needed)
         waiter = self._pending_listeners.popleft()
         waiter.set_result(data)
 
-    def _read(self, bytes_needed, first=False):
+    def _read(self, bytes_needed):
         """Read bytes from the buffer."""
         # Send the bytes to the listener.
-        if first and self._bytes_ready < bytes_needed:
+        if self._bytes_ready < bytes_needed:
             bytes_needed = self._bytes_ready
         self._bytes_ready -= bytes_needed
 
@@ -596,13 +592,13 @@ async def async_sendall(conn: PyMongoBaseProtocol, buf: bytes) -> None:
         raise socket.timeout("timed out") from exc
 
 
-async def async_receive_kms(conn: AsyncBaseConnection, bytes_needed: int, first=False) -> bytes:
+async def async_receive_kms(conn: AsyncBaseConnection, bytes_needed: int) -> bytes:
     """Receive raw bytes from the kms connection."""
 
     def callback(result: Any) -> bytes:
         return result
 
-    return await _async_receive_data(conn, callback, bytes_needed, first)
+    return await _async_receive_data(conn, callback, bytes_needed)
 
 
 async def _async_receive_data(

pymongo/synchronous/encryption.py

@@ -190,15 +190,13 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
             conn = BaseConnection(interface, opts)
             try:
                 sendall(interface.get_conn, message)
-                first = True
                 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, first)
+                    data = receive_kms(conn, kms_context.bytes_needed)
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
-                    first = False
             except MongoCryptError:
                 raise  # Propagate MongoCryptError errors directly.
             except Exception as exc: