packages/mosaic/framework/tools/lease-broker/daemon.py

@@ -358,7 +358,16 @@ def read_frame(connection: socket.socket, deadline: float) -> dict[str, object]:
         if not chunk:
             break
         data.extend(chunk)
-    if len(data) > MAX_FRAME or not data.endswith(b"\n") or data.count(b"\n") != 1:
+    if len(data) > MAX_FRAME:
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise BrokerFailure("MALFORMED_REQUEST")
+            connection.settimeout(remaining)
+            if not connection.recv(4096):
+                break
+        raise BrokerFailure("MALFORMED_REQUEST")
+    if not data.endswith(b"\n") or data.count(b"\n") != 1:
         raise BrokerFailure("MALFORMED_REQUEST")
     try:
         value = json.loads(data)