diff --git a/bridge/cutter/rea_bridge.py b/bridge/cutter/rea_bridge.py index 9dfb4d875..df359ecae 100644 --- a/bridge/cutter/rea_bridge.py +++ b/bridge/cutter/rea_bridge.py @@ -783,14 +783,18 @@ def _within_response_limit(value: Any) -> bool: current = pending.pop() if isinstance(current, str): size += 2 + if len(current) > MAX_FRAME_BYTES - size: + return False if current.isascii(): size += len(current) - size += current.count('"') + current.count("\\") - size += sum(5 for character in current if ord(character) < 0x20) else: - size += 4 * len(current) - size += current.count('"') + current.count("\\") - size += sum(5 for character in current if ord(character) < 0x20) + # Count UTF-8 bytes without allocating a copy of the whole output. + size += sum( + len(current[start:start + 65536].encode("utf-8")) + for start in range(0, len(current), 65536) + ) + size += current.count('"') + current.count("\\") + size += sum(5 for character in current if ord(character) < 0x20) elif isinstance(current, dict): identity = id(current) if identity in seen: diff --git a/tests/fixtures/cutter_bridge_protocol_test.py b/tests/fixtures/cutter_bridge_protocol_test.py index b1abdc275..164a43509 100644 --- a/tests/fixtures/cutter_bridge_protocol_test.py +++ b/tests/fixtures/cutter_bridge_protocol_test.py @@ -320,6 +320,34 @@ class HostVersionTest(unittest.TestCase): self.assertIsNone(bridge._RequestDispatcher().cutter_version) +class ResponseEncodingTest(unittest.TestCase): + def setUp(self): + self.server = bridge._BridgeServer.__new__(bridge._BridgeServer) + self.server.token = "test-token" + self.server.session_id = str(uuid.uuid4()) + + def test_large_mixed_unicode_output_fits_the_authenticated_frame(self): + for suffix in ("a", "é", "雪", "😀"): + with self.subTest(suffix=suffix): + output = "a" * (5 * 1024 * 1024 - 1) + suffix + value = {"ok": True, "output": output, "execution_state": "complete"} + encoded = self.server._encode_authenticated_response(value, "a" * 64, "b" * 64) + self.assertIsNotNone(encoded) + self.assertLessEqual(len(encoded), bridge.MAX_FRAME_BYTES) + frame = json.loads(encoded) + self.assertEqual(json.loads(frame["response"]), value) + expected = hmac.new( + self.server.token.encode(), + f"rea-cutter-response-v2:{self.server.session_id}:{'a' * 64}:{'b' * 64}:{frame['response']}".encode(), + hashlib.sha256, + ).hexdigest() + self.assertTrue(hmac.compare_digest(frame["proof"], expected)) + self.assertIsNone(self.server._encode_authenticated_response({ + "ok": True, "output": "😀" * (5 * 1024 * 1024), + "execution_state": "complete", + }, "a" * 64, "b" * 64)) + + class ProtocolProducerTest(unittest.TestCase): def test_rejects_legacy_peers_before_ui_dispatch(self): with socket.create_connection(("127.0.0.1", self.server.port), timeout=2) as connection: