fix: harden DGR-002 protocol bounds

tests/test_native_shard_protocol.py

@@ -26,7 +26,9 @@ from google.protobuf import descriptor_pb2
 from meshnet_node.activation_compression import CompressionPolicy
 from meshnet_node.native_protocol import (
     DEFAULT_MAX_CHUNK_BYTES,
+    DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
     DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
+    DEFAULT_MAX_TENSORS_PER_BUNDLE,
     HIDDEN_STATES,
     PayloadCorrupt,
     ProtocolError,
@@ -39,6 +41,7 @@ from meshnet_node.native_protocol import (
     negotiate_flow_control,
     pb,
     plan_prefill_chunks,
+    validate_session_message_size,
 )
 from meshnet_node.native_protocol import conformance
 
@@ -233,10 +236,161 @@ def test_bundle_from_a_newer_layout_is_refused():
     bundle = encode_bundle([])
     bundle.bundle_version = 99
 
-    with pytest.raises(ProtocolError, match="newer"):
+    with pytest.raises(ProtocolError, match="not supported"):
         decode_bundle(bundle)
 
 
+def test_bundle_with_unspecified_version_is_refused():
+    with pytest.raises(ProtocolError, match="version 0"):
+        decode_bundle(pb.TensorBundle())
+
+
+def test_tensor_with_unspecified_compression_is_refused():
+    tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
+    tensor.compression = pb.COMPRESSION_UNSPECIFIED
+
+    with pytest.raises(ProtocolError, match="unspecified"):
+        decode_tensor(tensor)
+
+
+def test_tensor_with_unspecified_checksum_is_refused():
+    tensor = encode_tensor(HIDDEN_STATES, b"\xaa" * 32, [1, 16, 1], pb.DTYPE_BFLOAT16)
+    tensor.checksum.algorithm = pb.CHECKSUM_ALGORITHM_UNSPECIFIED
+
+    with pytest.raises(ProtocolError, match="unspecified"):
+        decode_tensor(tensor)
+
+
+def test_declared_tensor_size_cannot_exceed_negotiated_chunk_bound():
+    tensor = pb.NamedTensor(
+        name=HIDDEN_STATES,
+        shape=[1],
+        dtype=pb.DTYPE_UINT8,
+        byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN,
+        total_bytes=DEFAULT_MAX_CHUNK_BYTES + 1,
+        compression=pb.COMPRESSION_NONE,
+    )
+
+    with pytest.raises(ProtocolError, match="negotiated chunk bound"):
+        decode_tensor(tensor)
+
+
+def test_stricter_session_negotiation_overrides_global_defaults():
+    tensor = encode_tensor(
+        HIDDEN_STATES,
+        b"\xaa" * 32,
+        [32],
+        pb.DTYPE_UINT8,
+        max_fragment_bytes=32,
+    )
+
+    with pytest.raises(ProtocolError, match="16-byte negotiated"):
+        decode_tensor(tensor, max_chunk_bytes=16, max_fragment_bytes=16)
+    with pytest.raises(ProtocolError, match="16-byte"):
+        decode_bundle(
+            encode_bundle([tensor]),
+            max_chunk_bytes=16,
+            max_fragment_bytes=16,
+        )
+
+
+def test_fragment_cannot_exceed_protocol_fragment_bound():
+    tensor = encode_tensor(
+        HIDDEN_STATES,
+        b"\xaa" * (1024 * 1024 + 1),
+        [1024 * 1024 + 1],
+        pb.DTYPE_UINT8,
+        max_fragment_bytes=1024 * 1024 + 1,
+    )
+
+    with pytest.raises(PayloadCorrupt, match="fragment larger"):
+        decode_tensor(tensor)
+
+
+def test_fragment_count_is_bounded_before_sorting_or_allocation():
+    tensor = encode_tensor(
+        HIDDEN_STATES,
+        b"\xaa" * (DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1),
+        [DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1],
+        pb.DTYPE_UINT8,
+        max_fragment_bytes=1,
+        max_fragments=DEFAULT_MAX_FRAGMENTS_PER_TENSOR + 1,
+    )
+
+    with pytest.raises(PayloadCorrupt, match="fragments.*exceeding"):
+        decode_tensor(tensor)
+
+
+def test_tensor_count_is_bounded_before_bundle_decoding():
+    tensors = [
+        pb.NamedTensor(name=f"t-{index}")
+        for index in range(DEFAULT_MAX_TENSORS_PER_BUNDLE + 1)
+    ]
+
+    with pytest.raises(ProtocolError, match="tensors.*exceeding"):
+        encode_bundle(tensors)
+
+
+@pytest.mark.parametrize(
+    ("shape", "match"),
+    [
+        ([1] * 9, "rank"),
+        ([-1], "dimension outside"),
+        ([1 << 31], "dimension outside"),
+    ],
+)
+def test_shape_rank_and_dimensions_are_bounded_before_payload_work(shape, match):
+    with pytest.raises(ProtocolError, match=match):
+        encode_tensor(HIDDEN_STATES, b"\x00", shape, pb.DTYPE_UINT8)
+
+
+def test_complete_session_message_includes_envelope_in_byte_bound():
+    small = pb.SessionRequest(chunk=pb.ActivationChunk(bundle=encode_bundle([])))
+    assert validate_session_message_size(small, max_chunk_bytes=64) == small.ByteSize()
+
+    oversized = pb.SessionRequest(
+        chunk=pb.ActivationChunk(
+            envelope=pb.Envelope(work_id="w" * 128),
+            bundle=encode_bundle([]),
+        )
+    )
+    with pytest.raises(ProtocolError, match="serialized session message"):
+        validate_session_message_size(oversized, max_chunk_bytes=64)
+
+
+def test_bundle_serialized_size_is_bounded():
+    payload = b"\xaa" * (DEFAULT_MAX_CHUNK_BYTES // 2)
+    first = encode_tensor("first", payload, [len(payload)], pb.DTYPE_UINT8)
+    second = encode_tensor("second", payload, [len(payload)], pb.DTYPE_UINT8)
+
+    with pytest.raises(ProtocolError, match="serialized tensor bundle"):
+        encode_bundle([first, second])
+
+
+def test_compressed_tensor_cannot_expand_past_declared_size():
+    pytest.importorskip("zstandard")
+    payload = b"\x00" * 65536
+    always = CompressionPolicy(min_input_bytes=0, min_savings_bytes=0, min_savings_ratio=0.0)
+    tensor = encode_tensor(
+        HIDDEN_STATES, payload, [len(payload)], pb.DTYPE_UINT8, policy=always
+    )
+    assert tensor.compression == pb.COMPRESSION_ZSTD
+    # Simulate a hostile frame that advertises a small output while expanding to
+    # the original 64 KiB. Shape and total_bytes agree, so only the bounded
+    # decompressor protects the receiver.
+    tensor.shape[:] = [1024]
+    tensor.total_bytes = 1024
+
+    with pytest.raises(PayloadCorrupt, match="bounded zstd"):
+        decode_tensor(tensor)
+
+
+def test_generated_runtime_floors_match_committed_stubs():
+    metadata = (REPO_ROOT / "packages/node/pyproject.toml").read_text()
+    assert '"grpcio>=1.82.1"' in metadata
+    assert '"protobuf>=7.35.0"' in metadata
+
+
 # --- Bounded prefill chunking and the decode fast path ----------------------
 
 
@@ -310,10 +464,21 @@ def test_flow_control_negotiation_takes_the_strictest_bound():
     # A sender cannot talk a worker into unbounded queues by proposing a large
     # window: neither peer can raise the other's ceiling.
     assert settled.max_inflight_chunks == default_flow_control().max_inflight_chunks
+    assert settled.credits_granted <= settled.max_inflight_chunks
     assert settled.max_chunk_bytes == DEFAULT_MAX_CHUNK_BYTES
     assert settled.max_prefill_chunk_tokens == DEFAULT_MAX_PREFILL_CHUNK_TOKENS
 
 
+def test_initial_credits_never_exceed_negotiated_inflight_limit():
+    proposed = pb.FlowControl(credits_granted=64, max_inflight_chunks=64)
+    limits = pb.FlowControl(credits_granted=64, max_inflight_chunks=8)
+
+    settled = negotiate_flow_control(proposed, limits)
+
+    assert settled.credits_granted == 8
+    assert settled.max_inflight_chunks == 8
+
+
 # --- Compatibility ----------------------------------------------------------