fix: reject malformed binary frames and decode empty metadata as {}

hivemind/binary.py

@@ -6,9 +6,28 @@
 
 from __future__ import annotations
 
+try:
+    import ujson as json  # type: ignore[import]
+except ImportError:
+    import json  # type: ignore[no-redef]
+
 # Message type constant (must match client.py)
 MSG_BINARY = 12
 
+# HIVEMIND-WIRE-1 §4.1: the frame-format version this client implements.
+FRAME_FORMAT_VERSION = 1
+
+# HIVEMIND-WIRE-1 §4.2: assigned message-type codes (0-10 and 12); 11 is
+# reserved (removed THIRDPRTY), 13-31 are unassigned. A receiver MUST
+# reject any other code as malformed rather than mapping it to a type.
+ASSIGNED_MSG_TYPES = frozenset(range(11)) | {MSG_BINARY}
+
+
+class MalformedBinaryFrame(ValueError):
+    """A WIRE-1 §4 binary frame that cannot be decoded: truncated, a
+    metadata length past the end of the frame, an unassigned message-type
+    code, or a metadata block that is not valid JSON."""
+
 # Binary payload sub-types
 BIN_UNDEFINED = 0
 BIN_RAW_AUDIO = 1
@@ -101,7 +120,7 @@ def encode(msg_type: int, bin_type: int, metadata: bytes,
     Returns:
         Encoded binary frame.
     """
-    version = 1
+    version = FRAME_FORMAT_VERSION
     compressed = False
 
     # Calculate total content bits
@@ -136,28 +155,58 @@ def encode(msg_type: int, bin_type: int, metadata: bytes,
 
 
 def decode(data: bytes) -> dict:
-    """Decode a HiveMind binary frame.
+    """Decode a HiveMind binary frame (HIVEMIND-WIRE-1 §4).
 
     Returns:
         Dict with keys: ``msg_type``, ``bin_type``, ``versioned``,
         ``protocol_version``, ``compressed``, ``metadata``, ``payload``.
+        ``metadata`` is the decoded JSON object; a zero-length metadata
+        block decodes as ``{}`` (§4.1).
+
+    Raises:
+        MalformedBinaryFrame: the frame-format version is not the one
+            this client implements (§4.1), the message-type code is not
+            one of the codes §4.2 assigns, the metadata length claims
+            more bytes than remain in the frame, the metadata block is
+            not valid UTF-8 JSON, or the frame is truncated.
     """
-    r = BitReader(data)
-
-    # Skip leading zeros until pad marker (1 bit)
-    while r.read_bits(1) == 0:
-        pass
-
-    versioned = bool(r.read_bits(1))
-    protocol_version = r.read_bits(8) if versioned else 0
-    msg_type = r.read_bits(5)
-    compressed = bool(r.read_bits(1))
-    meta_length = r.read_bits(8)
-    metadata = r.read_bytes(meta_length)
-
-    bin_type = r.read_bits(4) if msg_type == MSG_BINARY else 0
-
-    payload = r.read_bytes(r.remaining_bytes())
+    try:
+        r = BitReader(data)
+
+        # Skip leading zeros until pad marker (1 bit)
+        while r.read_bits(1) == 0:
+            pass
+
+        versioned = bool(r.read_bits(1))
+        protocol_version = r.read_bits(8) if versioned else 0
+        if versioned and protocol_version != FRAME_FORMAT_VERSION:
+            raise MalformedBinaryFrame(
+                "malformed binary frame: unimplemented frame-format "
+                "version %d" % protocol_version)
+        msg_type = r.read_bits(5)
+        if msg_type not in ASSIGNED_MSG_TYPES:
+            raise MalformedBinaryFrame(
+                "malformed binary frame: unassigned WIRE-1 message-type "
+                "code %d" % msg_type)
+        compressed = bool(r.read_bits(1))
+        meta_length = r.read_bits(8)
+        if meta_length > r.remaining_bytes():
+            raise MalformedBinaryFrame(
+                "malformed binary frame: metadata_len claims %d bytes "
+                "but only %d remain" % (meta_length, r.remaining_bytes()))
+        raw_metadata = r.read_bytes(meta_length)
+        if meta_length:
+            metadata = json.loads(raw_metadata.decode("utf-8"))
+        else:
+            metadata = {}
+
+        bin_type = r.read_bits(4) if msg_type == MSG_BINARY else 0
+
+        payload = r.read_bytes(r.remaining_bytes())
+    except MalformedBinaryFrame:
+        raise
+    except (IndexError, ValueError, UnicodeDecodeError) as e:
+        raise MalformedBinaryFrame("malformed binary frame: %s" % e) from e
 
     return {
         "msg_type": msg_type,

hivemind/client.py

@@ -43,7 +43,11 @@ def b2a_base64(data: bytes) -> bytes:
     CIPHER_CHACHA20,
     _norm_cipher,
 )
-from hivemind.binary import encode as binary_encode, decode as binary_decode
+from hivemind.binary import (
+    encode as binary_encode,
+    decode as binary_decode,
+    MalformedBinaryFrame,
+)
 from hivemind.noise import (
     PROTOCOL_V3,
     NoiseHandshake,
@@ -592,7 +596,11 @@ async def _dispatch_message(self, raw) -> None:
             if isinstance(frame, bytes):
                 # HIVEMIND-WIRE-1 binary frame
                 if self.on_binary is not None:
-                    decoded = binary_decode(frame)
+                    try:
+                        decoded = binary_decode(frame)
+                    except MalformedBinaryFrame as e:
+                        print("dropping malformed binary frame:", e)
+                        return
                     self.on_binary(decoded["bin_type"], decoded["payload"])
                 return
             try:

test/test_binary.py

@@ -1,6 +1,43 @@
 """Binary protocol unit tests."""
+import os
+import random
 import unittest
-from hivemind.binary import encode, decode, MSG_BINARY, BIN_RAW_AUDIO
+
+from hivemind.binary import (
+    encode,
+    decode,
+    MalformedBinaryFrame,
+    MSG_BINARY,
+    BIN_RAW_AUDIO,
+    ASSIGNED_MSG_TYPES,
+    FRAME_FORMAT_VERSION,
+    BitWriter,
+)
+
+
+def _encode_with_version(msg_type, bin_type, metadata, payload, version):
+    """Build a versioned frame like ``encode()``, but with an arbitrary
+    frame-format version instead of FRAME_FORMAT_VERSION, so tests can
+    exercise HIVEMIND-WIRE-1 §4.1 version rejection."""
+    content_bits = 1 + 1 + 8 + 5 + 1 + 8 + len(metadata) * 8
+    if msg_type == MSG_BINARY:
+        content_bits += 4
+    content_bits += len(payload) * 8
+    leading_zeros = (8 - (content_bits % 8)) % 8
+
+    w = BitWriter()
+    w.write_bits(0, leading_zeros)
+    w.write_bits(1, 1)  # pad marker
+    w.write_bits(1, 1)  # versioned flag
+    w.write_bits(version, 8)
+    w.write_bits(msg_type, 5)
+    w.write_bits(0, 1)  # compressed flag
+    w.write_bits(len(metadata), 8)
+    w.write_bytes(metadata)
+    if msg_type == MSG_BINARY:
+        w.write_bits(bin_type, 4)
+    w.write_bytes(payload)
+    return w.finish()
 
 
 class TestBinaryCodec(unittest.TestCase):
@@ -13,7 +50,7 @@ def test_bus_message_roundtrip(self) -> None:
         frame = encode(1, 0, meta, payload)
         result = decode(frame)
         self.assertEqual(result["msg_type"], 1)
-        self.assertEqual(result["metadata"], meta)
+        self.assertEqual(result["metadata"], {"source": "test"})
         self.assertEqual(result["payload"], payload)
         self.assertTrue(result["versioned"])
         self.assertEqual(result["protocol_version"], 1)
@@ -29,15 +66,17 @@ def test_binary_with_bin_type(self) -> None:
         self.assertEqual(result["payload"], payload)
 
     def test_empty_metadata(self) -> None:
-        """Empty metadata must roundtrip correctly."""
+        """HIVEMIND-WIRE-1 §4.1: a zero-length metadata block decodes as
+        the empty object."""
         frame = encode(1, 0, b'', b'test payload')
         result = decode(frame)
-        self.assertEqual(result["metadata"], b'')
+        self.assertEqual(result["metadata"], {})
         self.assertEqual(result["payload"], b'test payload')
 
-    def test_all_message_types(self) -> None:
-        """All 14 message types (0-13) must roundtrip."""
-        for mt in range(14):
+    def test_all_assigned_message_types_roundtrip(self) -> None:
+        """Every HIVEMIND-WIRE-1 §4.2 assigned message type must roundtrip."""
+        self.assertEqual(ASSIGNED_MSG_TYPES, frozenset(range(11)) | {12})
+        for mt in sorted(ASSIGNED_MSG_TYPES):
             bt = BIN_RAW_AUDIO if mt == MSG_BINARY else 0
             frame = encode(mt, bt, b'{}', b'x')
             result = decode(frame)
@@ -51,5 +90,79 @@ def test_large_payload(self) -> None:
         self.assertEqual(result["payload"], payload)
 
 
+class TestMalformedFrames(unittest.TestCase):
+    """HIVEMIND-WIRE-1 §4.2: a receiver MUST reject a frame carrying an
+    unassigned message-type code as malformed."""
+
+    def test_unimplemented_frame_format_version_2_rejected(self) -> None:
+        frame = _encode_with_version(1, 0, b'{}', b'x',
+                                      FRAME_FORMAT_VERSION + 1)
+        with self.assertRaisesRegex(MalformedBinaryFrame, "2"):
+            decode(frame)
+
+    def test_unimplemented_frame_format_version_0_rejected(self) -> None:
+        frame = _encode_with_version(1, 0, b'{}', b'x', 0)
+        with self.assertRaisesRegex(MalformedBinaryFrame, "0"):
+            decode(frame)
+
+    def test_unimplemented_frame_format_version_255_rejected(self) -> None:
+        frame = _encode_with_version(1, 0, b'{}', b'x', 255)
+        with self.assertRaisesRegex(MalformedBinaryFrame, "255"):
+            decode(frame)
+
+    def test_reserved_code_11_rejected(self) -> None:
+        frame = encode(11, 0, b'{}', b'x')
+        with self.assertRaises(MalformedBinaryFrame):
+            decode(frame)
+
+    def test_unassigned_code_13_rejected(self) -> None:
+        frame = encode(13, 0, b'{}', b'x')
+        with self.assertRaises(MalformedBinaryFrame):
+            decode(frame)
+
+    def test_unassigned_code_31_rejected(self) -> None:
+        frame = encode(31, 0, b'{}', b'x')
+        with self.assertRaises(MalformedBinaryFrame):
+            decode(frame)
+
+    def test_metadata_length_exceeds_remaining_bytes(self) -> None:
+        # A well-formed frame with metadata_len=1 but zero metadata bytes
+        # and no payload: claims one byte that does not exist.
+        frame = encode(1, 0, b'', b'')
+        # Corrupt the metadata-length field (last full byte before the
+        # zero-length metadata/payload) to claim 255 bytes.
+        corrupted = bytearray(frame)
+        # metadata_len is the 8 bits right after msg_type(5)+compressed(1)
+        # bits following the header; simplest reliable corruption is to
+        # encode a frame with 1 byte of real metadata, then truncate the
+        # trailing metadata byte off the wire representation.
+        good = bytearray(encode(1, 0, b'x', b''))
+        truncated = bytes(good[:-1])
+        with self.assertRaises(MalformedBinaryFrame):
+            decode(truncated)
+
+    def test_truncated_frame(self) -> None:
+        frame = encode(1, 0, b'{}', b'payload bytes here')
+        with self.assertRaises(MalformedBinaryFrame):
+            decode(frame[:1])
+
+    def test_metadata_not_valid_json(self) -> None:
+        frame = encode(1, 0, b'not json', b'x')
+        with self.assertRaises(MalformedBinaryFrame):
+            decode(frame)
+
+    def test_fuzz_random_bytes_only_raise_malformed_or_decode_clean(self) -> None:
+        """Feeding random bytes must never escape with anything other than
+        MalformedBinaryFrame, or a clean decode."""
+        rng = random.Random(1234)
+        for _ in range(200):
+            length = rng.randint(0, 64)
+            data = bytes(rng.randrange(256) for _ in range(length))
+            try:
+                decode(data)
+            except MalformedBinaryFrame:
+                pass
+
+
 if __name__ == "__main__":
     unittest.main()