fix: harden DGR-002 protocol bounds

packages/node/meshnet_node/native_protocol/codec.py

@@ -48,6 +48,10 @@ DEFAULT_MAX_CHUNK_BYTES = 4 * 1024 * 1024
 DEFAULT_MAX_FRAGMENT_BYTES = 1024 * 1024
 
 DEFAULT_MAX_INFLIGHT_CHUNKS = 8
+DEFAULT_MAX_FRAGMENTS_PER_TENSOR = 64
+DEFAULT_MAX_TENSORS_PER_BUNDLE = 64
+DEFAULT_MAX_TENSOR_RANK = 8
+DEFAULT_MAX_TENSOR_DIMENSION = (1 << 31) - 1
 
 # Canonical boundary tensor name for a dense transformer hidden state.
 HIDDEN_STATES = "hidden_states"
@@ -79,14 +83,28 @@ def itemsize(dtype: int) -> int:
         raise ProtocolError(f"unsupported dtype {dtype}") from None
 
 
-def expected_bytes(shape: Sequence[int], dtype: int) -> int:
+def expected_bytes(
+    shape: Sequence[int],
+    dtype: int,
+    *,
+    max_rank: int = DEFAULT_MAX_TENSOR_RANK,
+    max_dimension: int = DEFAULT_MAX_TENSOR_DIMENSION,
+    max_bytes: int | None = None,
+) -> int:
     """Byte count a tensor of `shape` and `dtype` must occupy."""
-    if any(dim < 0 for dim in shape):
-        raise ProtocolError(f"negative dimension in shape {list(shape)}")
+    if len(shape) > max_rank:
+        raise ProtocolError(f"tensor rank {len(shape)} exceeds limit {max_rank}")
+    if any(dim < 0 or dim > max_dimension for dim in shape):
+        raise ProtocolError(
+            f"dimension outside 0..{max_dimension} in shape {list(shape)}"
+        )
+    size = itemsize(dtype)
     count = 1
     for dim in shape:
         count *= dim
-    return count * itemsize(dtype)
+        if max_bytes is not None and count * size > max_bytes:
+            raise ProtocolError(f"tensor shape {list(shape)} exceeds byte limit {max_bytes}")
+    return count * size
 
 
 # --- CRC32C ----------------------------------------------------------------
@@ -136,16 +154,18 @@ def encode_tensor(
     dtype: int = pb.DTYPE_BFLOAT16,
     *,
     policy: CompressionPolicy | None = None,
+    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
     max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
+    max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
 ) -> pb.NamedTensor:
     """Build a NamedTensor, compressing and fragmenting as needed.
 
     `data` is the uncompressed little-endian payload. The checksum is taken over
     it *before* compression so it stays valid whichever framing a hop chooses.
     """
-    if max_fragment_bytes <= 0:
-        raise ProtocolError("max_fragment_bytes must be positive")
-    declared = expected_bytes(shape, dtype)
+    if max_chunk_bytes <= 0 or max_fragment_bytes <= 0 or max_fragments <= 0:
+        raise ProtocolError("tensor byte/count bounds must be positive")
+    declared = expected_bytes(shape, dtype, max_bytes=max_chunk_bytes)
     if len(data) != declared:
         raise ProtocolError(
             f"tensor {name!r} declares shape {list(shape)} ({declared} bytes) "
@@ -179,6 +199,10 @@ def encode_tensor(
         # A zero-element tensor is legal (e.g. an empty mask) and still needs a
         # fragment, so coverage checks have something to verify.
         slices = [b""]
+    if len(slices) > max_fragments:
+        raise ProtocolError(
+            f"tensor {name!r} needs {len(slices)} fragments, exceeding limit {max_fragments}"
+        )
 
     offset = 0
     for index, piece in enumerate(slices):
@@ -194,20 +218,57 @@ def encode_tensor(
     return tensor
 
 
-def decode_tensor(tensor: pb.NamedTensor) -> bytes:
+def decode_tensor(
+    tensor: pb.NamedTensor,
+    *,
+    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
+    max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
+    max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
+) -> bytes:
     """Reassemble, decompress and validate a NamedTensor's payload.
 
     Raises PayloadCorrupt rather than returning a payload it cannot fully
     account for: a hole in the fragments or a bad checksum means the activation
     is not what the sender computed, and continuing would corrupt the route.
     """
+    if max_chunk_bytes <= 0 or max_fragment_bytes <= 0 or max_fragments <= 0:
+        raise ProtocolError("negotiated byte/count bounds must be positive")
+    if tensor.total_bytes > max_chunk_bytes:
+        raise ProtocolError(
+            f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes, exceeding "
+            f"the {max_chunk_bytes}-byte negotiated chunk bound"
+        )
     if tensor.byte_order == pb.BYTE_ORDER_BIG_ENDIAN:
         raise ProtocolError(f"tensor {tensor.name!r} is big-endian; wire order is little-endian")
     if tensor.byte_order != pb.BYTE_ORDER_LITTLE_ENDIAN:
         raise ProtocolError(f"tensor {tensor.name!r} declares no byte order")
+    declared = expected_bytes(
+        tensor.shape, tensor.dtype, max_bytes=max_chunk_bytes
+    )
+    if declared != tensor.total_bytes:
+        raise PayloadCorrupt(
+            f"tensor {tensor.name!r} shape {list(tensor.shape)} implies {declared} bytes "
+            f"but declares {tensor.total_bytes}"
+        )
 
     if not tensor.fragments:
         raise PayloadCorrupt(f"tensor {tensor.name!r} carries no fragments")
+    if len(tensor.fragments) > max_fragments:
+        raise PayloadCorrupt(
+            f"tensor {tensor.name!r} carries {len(tensor.fragments)} fragments, "
+            f"exceeding limit {max_fragments}"
+        )
+    if any(len(fragment.payload) > max_fragment_bytes for fragment in tensor.fragments):
+        raise PayloadCorrupt(
+            f"tensor {tensor.name!r} carries a fragment larger than "
+            f"{max_fragment_bytes} bytes"
+        )
+    wire_bytes = sum(len(fragment.payload) for fragment in tensor.fragments)
+    if wire_bytes > max_chunk_bytes:
+        raise PayloadCorrupt(
+            f"tensor {tensor.name!r} wire body exceeds the "
+            f"{max_chunk_bytes}-byte negotiated chunk bound"
+        )
 
     fragments = sorted(tensor.fragments, key=lambda f: f.byte_offset)
     count = fragments[0].fragment_count
@@ -231,44 +292,90 @@ def decode_tensor(tensor: pb.NamedTensor) -> bytes:
         body.extend(fragment.payload)
 
     if tensor.compression == pb.COMPRESSION_ZSTD:
-        data = decompress_activation(bytes(body), "zstd").body
-    elif tensor.compression in (pb.COMPRESSION_NONE, pb.COMPRESSION_UNSPECIFIED):
+        try:
+            data = decompress_activation(
+                bytes(body), "zstd", max_output_bytes=tensor.total_bytes
+            ).body
+        except ValueError as exc:
+            raise PayloadCorrupt(
+                f"tensor {tensor.name!r} has invalid bounded zstd payload"
+            ) from exc
+    elif tensor.compression == pb.COMPRESSION_NONE:
         data = bytes(body)
     else:
-        raise ProtocolError(f"tensor {tensor.name!r} uses unsupported compression")
+        raise ProtocolError(
+            f"tensor {tensor.name!r} uses unspecified or unsupported compression"
+        )
 
     if len(data) != tensor.total_bytes:
         raise PayloadCorrupt(
             f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes but "
             f"reassembled {len(data)}"
         )
-    declared = expected_bytes(tensor.shape, tensor.dtype)
-    if declared != tensor.total_bytes:
-        raise PayloadCorrupt(
-            f"tensor {tensor.name!r} shape {list(tensor.shape)} implies {declared} bytes "
-            f"but declares {tensor.total_bytes}"
-        )
 
     algorithm = tensor.checksum.algorithm
     if algorithm == pb.CHECKSUM_ALGORITHM_CRC32C:
         if tensor.checksum.value != struct.pack(">I", crc32c(data)):
             raise PayloadCorrupt(f"tensor {tensor.name!r} failed its CRC32C check")
-    elif algorithm not in (pb.CHECKSUM_ALGORITHM_NONE, pb.CHECKSUM_ALGORITHM_UNSPECIFIED):
-        raise ProtocolError(f"tensor {tensor.name!r} uses unsupported checksum algorithm")
+    elif algorithm != pb.CHECKSUM_ALGORITHM_NONE:
+        raise ProtocolError(
+            f"tensor {tensor.name!r} uses unspecified or unsupported checksum algorithm"
+        )
 
     return data
 
 
-def encode_bundle(tensors: Iterable[pb.NamedTensor]) -> pb.TensorBundle:
-    return pb.TensorBundle(bundle_version=BUNDLE_VERSION, tensors=list(tensors))
-
-
-def decode_bundle(bundle: pb.TensorBundle) -> dict[str, bytes]:
-    """Validate every tensor in a bundle and return name -> payload."""
-    if bundle.bundle_version > BUNDLE_VERSION:
+def encode_bundle(
+    tensors: Iterable[pb.NamedTensor],
+    *,
+    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
+    max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
+) -> pb.TensorBundle:
+    if max_chunk_bytes <= 0 or max_tensors <= 0:
+        raise ProtocolError("bundle byte/count bounds must be positive")
+    tensor_list = list(tensors)
+    if len(tensor_list) > max_tensors:
         raise ProtocolError(
-            f"bundle version {bundle.bundle_version} is newer than this build "
-            f"understands ({BUNDLE_VERSION})"
+            f"bundle carries {len(tensor_list)} tensors, exceeding limit {max_tensors}"
+        )
+    bundle = pb.TensorBundle(bundle_version=BUNDLE_VERSION, tensors=tensor_list)
+    if bundle.ByteSize() > max_chunk_bytes:
+        raise ProtocolError(
+            f"serialized tensor bundle exceeds the {max_chunk_bytes}-byte "
+            "negotiated chunk bound"
+        )
+    return bundle
+
+
+def decode_bundle(
+    bundle: pb.TensorBundle,
+    *,
+    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
+    max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
+    max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
+    max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
+) -> dict[str, bytes]:
+    """Validate every tensor in a bundle and return name -> payload."""
+    if bundle.bundle_version != BUNDLE_VERSION:
+        raise ProtocolError(
+            f"bundle version {bundle.bundle_version} is not supported by this build "
+            f"({BUNDLE_VERSION})"
+        )
+    if (
+        max_chunk_bytes <= 0
+        or max_fragment_bytes <= 0
+        or max_fragments <= 0
+        or max_tensors <= 0
+    ):
+        raise ProtocolError("negotiated byte/count bounds must be positive")
+    if len(bundle.tensors) > max_tensors:
+        raise ProtocolError(
+            f"bundle carries {len(bundle.tensors)} tensors, exceeding limit {max_tensors}"
+        )
+    if bundle.ByteSize() > max_chunk_bytes:
+        raise ProtocolError(
+            f"serialized tensor bundle exceeds the {max_chunk_bytes}-byte "
+            "negotiated chunk bound"
         )
     payloads: dict[str, bytes] = {}
     for tensor in bundle.tensors:
@@ -276,10 +383,40 @@ def decode_bundle(bundle: pb.TensorBundle) -> dict[str, bytes]:
             raise ProtocolError("bundle carries an unnamed tensor")
         if tensor.name in payloads:
             raise ProtocolError(f"bundle carries duplicate tensor {tensor.name!r}")
-        payloads[tensor.name] = decode_tensor(tensor)
+        payloads[tensor.name] = decode_tensor(
+            tensor,
+            max_chunk_bytes=max_chunk_bytes,
+            max_fragment_bytes=max_fragment_bytes,
+            max_fragments=max_fragments,
+        )
     return payloads
 
 
+def validate_session_message_size(
+    message: pb.SessionRequest | pb.SessionResponse,
+    *,
+    max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
+) -> int:
+    """Reject an oversized complete stream frame, including protobuf overhead.
+
+    Bundle validation alone is insufficient because the envelope and oneof
+    framing are part of the same gRPC message. Senders call this immediately
+    before writing; receivers configure gRPC's receive limit to the same value
+    and call it again before semantic decoding.
+    """
+    if max_chunk_bytes <= 0:
+        raise ProtocolError("max_chunk_bytes must be positive")
+    if not isinstance(message, (pb.SessionRequest, pb.SessionResponse)):
+        raise ProtocolError("size validation requires a session request or response")
+    size = message.ByteSize()
+    if size > max_chunk_bytes:
+        raise ProtocolError(
+            f"serialized session message is {size} bytes, exceeding the "
+            f"{max_chunk_bytes}-byte negotiated chunk bound"
+        )
+    return size
+
+
 # --- Bounded prefill chunking ----------------------------------------------
 
 
@@ -363,15 +500,22 @@ def negotiate_flow_control(
         candidates = [v for v in (a, b) if v > 0]
         return min(candidates) if candidates else fallback
 
-    return pb.FlowControl(
-        credits_granted=_min(
-            proposed.credits_granted, limits.credits_granted, DEFAULT_MAX_INFLIGHT_CHUNKS
-        ),
-        max_inflight_chunks=_min(
-            proposed.max_inflight_chunks,
-            limits.max_inflight_chunks,
+    max_inflight_chunks = _min(
+        proposed.max_inflight_chunks,
+        limits.max_inflight_chunks,
+        DEFAULT_MAX_INFLIGHT_CHUNKS,
+    )
+    credits_granted = min(
+        _min(
+            proposed.credits_granted,
+            limits.credits_granted,
             DEFAULT_MAX_INFLIGHT_CHUNKS,
         ),
+        max_inflight_chunks,
+    )
+    return pb.FlowControl(
+        credits_granted=credits_granted,
+        max_inflight_chunks=max_inflight_chunks,
         max_chunk_bytes=_min(
             proposed.max_chunk_bytes, limits.max_chunk_bytes, DEFAULT_MAX_CHUNK_BYTES
         ),