Fix byte-based JSON-RPC framing and request-id race (Python client)

keyext.client.python/keyapi/rpc.py

@@ -6,7 +6,6 @@
 
 from keyapi import KEY_DATA_CLASSES, KEY_DATA_CLASSES_REV
 
-JSON_RPC_REQ_FORMAT = "Content-Length: {json_string_len}\r\n\r\n{json_string}"
 LEN_HEADER = "Content-Length: "
 TYPE_HEADER = "Content-Type: "
 
@@ -47,46 +46,72 @@ def __init__(self, stdin, stdout):
         self.write_lock = threading.Lock()
 
     @staticmethod
-    def __add_header(json_string):
+    def __add_header(content_bytes):
         '''
-        Adds a header for the given json string
+        Prepends the JSON-RPC framing header to an already UTF-8 encoded body.
 
-        :param str json_string: The string
-        :return: the string with the header
+        The ``Content-Length`` value is the number of *bytes* of the body (per
+        the LSP base protocol), not the number of characters. Counting
+        characters truncates every message that contains a non-ASCII glyph
+        (common in KeY terms), so the header must be computed from the encoded
+        bytes.
+
+        :param bytes content_bytes: the UTF-8 encoded JSON body
+        :return: the framed message as bytes
         '''
-        return JSON_RPC_REQ_FORMAT.format(json_string_len=len(json_string), json_string=json_string)
+        header = "%s%d\r\n\r\n" % (LEN_HEADER, len(content_bytes))
+        return header.encode("ascii") + content_bytes
 
     def send_request(self, message):
         '''
         Sends the given message.
 
         :param dict message: The message to send.
         '''
-        json_string = json.dumps(message , cls=MyEncoder)
-        jsonrpc_req = self.__add_header(json_string)
+        json_string = json.dumps(message, cls=MyEncoder)
+        content_bytes = json_string.encode("utf-8")
+        jsonrpc_req = self.__add_header(content_bytes)
         with self.write_lock:
             self.stdout.write(jsonrpc_req)
             self.stdout.flush()
 
+    def _read_exactly(self, count):
+        '''
+        Reads exactly ``count`` bytes, looping until they have all arrived.
+
+        ``read(n)`` on a socket stream may legally return fewer than ``n`` bytes,
+        so a single ``read`` is not enough to consume a framed message.
+
+        :return: the bytes, or ``None`` if EOF is reached first
+        '''
+        chunks = []
+        remaining = count
+        while remaining > 0:
+            chunk = self.stdin.read(remaining)
+            if not chunk:
+                return None
+            chunks.append(chunk)
+            remaining -= len(chunk)
+        return b"".join(chunks)
+
     def recv_response(self) -> object:
         '''
-        Recives a message.
+        Receives a message. Expects the input stream to be binary.
 
-        :return: a message
+        :return: a message, or ``None`` when the stream has reached EOF
         '''
-        with (self.read_lock):
+        with self.read_lock:
             message_size = None
             while True:
                 # read header
                 line = self.stdin.readline()
                 if not line:
                     # server quit
                     return None
-                # line = line.decode("utf-8")
-                if not line.endswith("\r\n"):
+                if not line.endswith(b"\r\n"):
                     raise ResponseError(ErrorCodes.ParseError, "Bad header: missing newline")
-                # remove the "\r\n"
-                line = line[:-2]
+                # remove the "\r\n" and decode the (ASCII) header line
+                line = line[:-2].decode("ascii")
                 if line == "":
                     # done with the headers
                     break
@@ -104,8 +129,13 @@ def recv_response(self) -> object:
             if not message_size:
                 raise ResponseError(ErrorCodes.ParseError, "Bad header: missing size")
 
-            jsonrpc_res = self.stdin.read(message_size)  # .decode("utf-8")
-            return json.loads(jsonrpc_res, object_hook=object_decoder)
+            # Content-Length counts bytes, so read bytes (not characters) and
+            # only then decode as UTF-8.
+            body = self._read_exactly(message_size)
+            if body is None:
+                # EOF in the middle of a message
+                return None
+            return json.loads(body.decode("utf-8"), object_hook=object_decoder)
 
 
 def object_decoder(obj):
@@ -130,6 +160,7 @@ def __init__(self, json_rpc_endpoint: JsonRpcEndpoint, method_callbacks=None, no
         self.event_dict = {}
         self.response_dict = {}
         self.next_id = 0
+        self._id_lock = threading.Lock()
         self._timeout = timeout
         self.shutdown_flag = False
 
@@ -196,8 +227,12 @@ def send_message(self, method_name, params, id=None):
         self.json_rpc_endpoint.send_request(message_dict)
 
     def call_method(self, method_name, args):
-        current_id = self.next_id
-        self.next_id += 1
+        # Allocate the request id atomically: concurrent callers must never
+        # share an id, otherwise their entries in event_dict/response_dict
+        # collide and a response gets delivered to the wrong caller.
+        with self._id_lock:
+            current_id = self.next_id
+            self.next_id += 1
         cond = threading.Condition()
         self.event_dict[current_id] = cond
 

keyext.client.python/keyapi/server.py

@@ -7,8 +7,10 @@ class NetKeY(object):
     def __init__(self, target):
         self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.socket.connect(target)
-        self.inStream = self.socket.makefile("r", newline="\r\n")
-        self.outStream = self.socket.makefile("w", newline="\r\n")
+        # Binary streams: the JSON-RPC framing counts bytes, not characters, so
+        # the transport must not do any text decoding or newline translation.
+        self.inStream = self.socket.makefile("rb")
+        self.outStream = self.socket.makefile("wb")
 
         self.rpc_endpoint = JsonRpcEndpoint(self.inStream, self.outStream)
         self.endpoint = LspEndpoint(self.rpc_endpoint)

keyext.client.python/tests/test_rpc.py

@@ -0,0 +1,156 @@
+# This file is part of KeY - https://key-project.org
+# KeY is licensed under the GNU General Public License Version 2
+# SPDX-License-Identifier: GPL-2.0-only
+"""Regression tests for the JSON-RPC transport (keyapi.rpc).
+
+Run with::
+
+    python3 -m unittest discover -s tests -t .
+
+from the ``keyext.client.python`` directory.
+"""
+import io
+import queue
+import threading
+import unittest
+
+from keyapi.rpc import JsonRpcEndpoint, LspEndpoint
+
+# Logic symbols / umlauts as they appear in KeY terms. Each of these is more
+# than one byte in UTF-8, so byte length and character length differ.
+UNICODE = "∀x; (x = x) — Grüße ∃∈≤"
+
+# Direct handle to the name-mangled static framing helper.
+_add_header = JsonRpcEndpoint._JsonRpcEndpoint__add_header
+
+
+def _frame_server_side(json_text):
+    """Frame a JSON body the way the Gson-based server does: raw UTF-8 on the
+    wire with a byte-counted Content-Length. This is the input the client must
+    be able to read back."""
+    body = json_text.encode("utf-8")
+    return ("Content-Length: %d\r\n\r\n" % len(body)).encode("ascii") + body
+
+
+class _DripBytesIO(io.BytesIO):
+    """A BytesIO that returns at most ``chunk`` bytes per ``read`` call, to
+    mimic the short reads a real socket may hand back."""
+
+    def __init__(self, data, chunk=1):
+        super().__init__(data)
+        self._chunk = chunk
+
+    def read(self, size=-1):
+        if size is None or size < 0:
+            return super().read(size)
+        return super().read(min(size, self._chunk))
+
+
+class _EchoEndpoint:
+    """In-memory transport for LspEndpoint: every request is echoed back as a
+    response whose ``result`` is the request id."""
+
+    def __init__(self):
+        self.sent = []
+        self._lock = threading.Lock()
+        self._responses = queue.Queue()
+
+    def send_request(self, message):
+        with self._lock:
+            self.sent.append(message)
+        if message.get("id") is not None and "method" in message:
+            self._responses.put(
+                {"jsonrpc": "2.0", "id": message["id"], "result": message["id"]})
+
+    def recv_response(self):
+        return self._responses.get()
+
+
+class WriteFramingTest(unittest.TestCase):
+    def test_add_header_uses_byte_length(self):
+        # #6: Content-Length must count UTF-8 bytes, not characters. "∀é" is
+        # 2 characters but 5 bytes (3 + 2); a char count would advertise 2.
+        body = "∀é".encode("utf-8")
+        self.assertEqual(len(body), 5)
+        framed = _add_header(body)
+        header, sep, rest = framed.partition(b"\r\n\r\n")
+        self.assertEqual(header, b"Content-Length: 5")
+        self.assertEqual(rest, body)
+
+    def test_send_request_advertises_body_byte_length(self):
+        out = io.BytesIO()
+        JsonRpcEndpoint(io.BytesIO(), out).send_request(
+            {"jsonrpc": "2.0", "id": 1, "params": [UNICODE]})
+        header, body = out.getvalue().split(b"\r\n\r\n", 1)
+        advertised = int(header[len(b"Content-Length: "):])
+        self.assertEqual(advertised, len(body))
+
+
+class ReadFramingTest(unittest.TestCase):
+    def test_read_raw_utf8_body(self):
+        # #6: a server message carrying raw multi-byte UTF-8 must decode
+        # correctly. The old code read message_size *characters* from a text
+        # stream and corrupted anything past the first multi-byte glyph.
+        text = '{"jsonrpc": "2.0", "id": 2, "result": "' + UNICODE + '"}'
+        ep = JsonRpcEndpoint(io.BytesIO(_frame_server_side(text)), io.BytesIO())
+        msg = ep.recv_response()
+        self.assertEqual(msg["result"], UNICODE)
+        self.assertEqual(msg["id"], 2)
+
+    def test_read_back_to_back_raw_utf8(self):
+        # #6: framing two multi-byte messages must keep them aligned; reading
+        # the wrong number of bytes for the first bleeds into the second.
+        t1 = '{"jsonrpc": "2.0", "id": 1, "result": "first ∀∃"}'
+        t2 = '{"jsonrpc": "2.0", "id": 2, "result": "second é≤"}'
+        buf = _frame_server_side(t1) + _frame_server_side(t2)
+        ep = JsonRpcEndpoint(io.BytesIO(buf), io.BytesIO())
+        self.assertEqual(ep.recv_response()["result"], "first ∀∃")
+        self.assertEqual(ep.recv_response()["result"], "second é≤")
+        self.assertIsNone(ep.recv_response())  # clean EOF
+
+    def test_partial_reads_are_assembled(self):
+        # #6: read() may return fewer bytes than requested; the body must still
+        # be reassembled in full before being decoded.
+        text = '{"jsonrpc": "2.0", "id": 7, "result": "' + UNICODE + '"}'
+        ep = JsonRpcEndpoint(_DripBytesIO(_frame_server_side(text), chunk=1),
+                             io.BytesIO())
+        self.assertEqual(ep.recv_response()["result"], UNICODE)
+
+
+class ConcurrencyTest(unittest.TestCase):
+    def test_unique_ids_under_concurrency(self):
+        # #7: concurrent call_method invocations must each get a distinct id and
+        # all complete without a TimeoutError caused by a collided id.
+        n = 32
+        endpoint = LspEndpoint(_EchoEndpoint(), timeout=10)
+        endpoint.daemon = True
+        endpoint.start()
+
+        results = [None] * n
+        errors = []
+        barrier = threading.Barrier(n)
+
+        def worker(i):
+            try:
+                barrier.wait()
+                results[i] = endpoint.call_method("ping", [i])
+            except Exception as exc:  # noqa: BLE001 - recorded for assertion
+                errors.append(exc)
+
+        threads = [threading.Thread(target=worker, args=(i,)) for i in range(n)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join(timeout=15)
+
+        endpoint.stop()
+        self.assertEqual(errors, [])
+        # Every worker got a response back (the echoed id), none timed out.
+        self.assertTrue(all(r is not None for r in results))
+        sent_ids = [m["id"] for m in endpoint.json_rpc_endpoint.sent]
+        self.assertEqual(len(sent_ids), n)
+        self.assertEqual(len(set(sent_ids)), n, "request ids must be unique")
+
+
+if __name__ == "__main__":
+    unittest.main()