fix: harden DGR-002 protocol bounds

This commit is contained in:
Dobromir Popov
2026-07-13 17:30:54 +03:00
parent 30dcf953fe
commit d904c40f66
10 changed files with 470 additions and 63 deletions

View File

@@ -38,7 +38,7 @@ New:
| Path | What | | Path | What |
|---|---| |---|---|
| `packages/node/native/proto/shard_runtime.proto` | The schema (sha256 `077ee349…`, see `protocol.json`) | | `packages/node/native/proto/shard_runtime.proto` | The schema (sha256 `9e211660…`, see `protocol.json`) |
| `packages/node/native/CMakeLists.txt` | C++ generation + build wiring + ctest | | `packages/node/native/CMakeLists.txt` | C++ generation + build wiring + ctest |
| `packages/node/native/tests/test_shard_protocol_conformance.cpp` | C++ conformance test | | `packages/node/native/tests/test_shard_protocol_conformance.cpp` | C++ conformance test |
| `packages/node/native/testdata/*.binpb` | Committed cross-language vectors | | `packages/node/native/testdata/*.binpb` | Committed cross-language vectors |
@@ -50,14 +50,19 @@ New:
| `scripts/generate_native_protocol.py` | Python generation, with `--check` | | `scripts/generate_native_protocol.py` | Python generation, with `--check` |
| `scripts/generate_protocol_goldens.py` | Vector generation, with `--check` | | `scripts/generate_protocol_goldens.py` | Vector generation, with `--check` |
| `scripts/bootstrap_native_toolchain.sh` | Builds protobuf C++ from source | | `scripts/bootstrap_native_toolchain.sh` | Builds protobuf C++ from source |
| `tests/test_native_shard_protocol.py` | 29 Python tests | | `tests/test_native_shard_protocol.py` | 45 Python tests |
Modified: Modified:
- `packages/node/pyproject.toml` — added `grpcio>=1.60`, `protobuf>=5`; new `proto` - `packages/node/pyproject.toml` — added runtime floors `grpcio>=1.82.1` and
extra pinning `grpcio-tools==1.82.1`. `protobuf>=7.35.0`, matching the committed generated-code requirements; new
`proto` extra pinning `grpcio-tools==1.82.1`.
- `packages/node/meshnet_node/activation_compression.py` — optional bounded zstd
output for untrusted protocol frames; existing callers remain compatible.
- `packages/node/meshnet_node/native_protocol/__init__.py` — exports negotiated
bound constants and whole-session-message validation.
No other working-tree file was touched. `git status` before this story was clean. The canonical PRD marks only DGR-002 passed. `git status` before this story was clean.
## 3. Commands and real results ## 3. Commands and real results
@@ -76,14 +81,45 @@ ctest --test-dir build/native --output-on-failure -> 1/1 Test #1: shard_prot
cmp build/native/cpp_roundtrip.binpb \ cmp build/native/cpp_roundtrip.binpb \
packages/node/native/testdata/session_request_golden.binpb -> identical (exit 0) packages/node/native/testdata/session_request_golden.binpb -> identical (exit 0)
pytest -q tests/test_native_shard_protocol.py -> 29 passed pytest -q tests/test_native_shard_protocol.py -> 45 passed
pytest -q (full suite) -> 712 passed, 12 skipped pytest -q tests/test_native_shard_protocol.py \
tests/test_activation_compression.py -> 51 passed
pytest -q (final full suite) -> 728 passed, 12 skipped
pytest -q tests/test_tracker_routing.py::test_tracker_dashboard_can_cancel_inflight_proxy
(after an earlier flaky full-suite failure) -> 1 passed, 1 passed, 1 passed
clean minimum-runtime import + codec smoke test -> passed
grpcio==1.82.1, protobuf==7.35.0
compileall -q packages tests -> OK (exit 0) compileall -q packages tests -> OK (exit 0)
git diff --check -> clean (exit 0) git diff --check -> clean (exit 0)
``` ```
The C++ lane was rebuilt from scratch (`rm -rf build/native`) using only the The C++ lane was rebuilt from scratch by Ralph (`rm -rf build/native`) using only
documented commands, and reproduced the same result. the documented commands, and reproduced the same result. During controller
review the user explicitly chose not to repeat the destructive build-directory
cleanup, so the independent controller relied on the recorded CMake/CTest run
while reproducing every Python/generation/full-suite gate.
### Controller review corrections
Independent controller review found and fixed two classes of issue before
integration:
1. Generated stubs required gRPC 1.82.1 and Protobuf 7.35.0, while the initial
package metadata allowed much older runtimes that could fail at import time.
2. Flow-control bounds were described but not enforced by the reference decoder.
Tensor declarations, shape rank/dimensions, fragment/tensor counts, fragments,
wire bodies, whole bundles, complete session messages (including envelope
overhead), and zstd window/output expansion are now fail-closed against the
negotiated/default bounds. Unspecified bundle versions, compression and
checksums are rejected rather than interpreted as valid data.
3. Negotiated initial credits could exceed `max_inflight_chunks`; credits are now
capped by the settled in-flight limit.
Controller results: protocol tests `45 passed`; protocol plus shared compression
tests `51 passed`; final full suite `728 passed, 12 skipped`. A clean environment
at the declared minimum gRPC/Protobuf runtime versions imported both generated
stub modules and round-tripped the codec. Generation checks, `compileall`, static
secret scan, and `git diff --check` all passed.
### Full-suite note — a pre-existing flaky test ### Full-suite note — a pre-existing flaky test
@@ -104,10 +140,11 @@ pytest -q tests/test_tracker_routing.py::test_tracker_dashboard_can_cancel_infli
It is a timing race in proxy cancellation (a 3-second in-flight generation raced It is a timing race in proxy cancellation (a 3-second in-flight generation raced
against the cancel assertion), not a deterministic failure, and it touches no code against the cancel assertion), not a deterministic failure, and it touches no code
this story changes. The final full-suite run with DGR-002 applied was green this story changes. One controller full-suite run reported exactly that one failure
(`712 passed, 12 skipped`) — 683 pre-existing tests plus this story's 29. Flagged (`1 failed, 719 passed, 12 skipped`); three immediate isolated retries all passed
for whoever owns the tracker cancel path; **not** fixed here, since silently in 1.11 seconds, and the final exact-code full suite was green (`728 passed,
touching another story's code is out of scope. 12 skipped`). It is flagged for whoever owns the tracker cancel path and is **not**
fixed here, since silently touching another story's code is out of scope.
## 4. Acceptance criteria ## 4. Acceptance criteria
@@ -118,12 +155,12 @@ touching another story's code is out of scope.
| Bounded chunking for prefill; small decode fast path | `ChunkInfo` + `plan_prefill_chunks` (128-token bound, ADR-0008); `DecodeStep`; `test_prefill_is_split_into_bounded_token_aligned_chunks`, `test_decode_fast_path_is_much_smaller_than_a_full_envelope_chunk` | | Bounded chunking for prefill; small decode fast path | `ChunkInfo` + `plan_prefill_chunks` (128-token bound, ADR-0008); `DecodeStep`; `test_prefill_is_split_into_bounded_token_aligned_chunks`, `test_decode_fast_path_is_much_smaller_than_a_full_envelope_chunk` |
| Envelope carries schema version, work id, session id, epoch, fingerprint, range/effective start, phase, position, idempotency step, cache expectation, compression, checksum | `Envelope` + `NamedTensor`; `test_envelope_carries_every_field_the_protocol_promises` asserts against the **descriptor**, so deleting a field from the `.proto` fails the test | | Envelope carries schema version, work id, session id, epoch, fingerprint, range/effective start, phase, position, idempotency step, cache expectation, compression, checksum | `Envelope` + `NamedTensor`; `test_envelope_carries_every_field_the_protocol_promises` asserts against the **descriptor**, so deleting a field from the `.proto` fails the test |
| Versioned named-tensor bundle: name, shape, dtype, byte order, fragments | `TensorBundle`/`NamedTensor`/`TensorFragment`; `test_named_tensor_bundle_is_versioned_and_fully_described`, `test_bundle_round_trips_multiple_named_tensors` | | Versioned named-tensor bundle: name, shape, dtype, byte order, fragments | `TensorBundle`/`NamedTensor`/`TensorFragment`; `test_named_tensor_bundle_is_versioned_and_fully_described`, `test_bundle_round_trips_multiple_named_tensors` |
| Round-trip + compatibility tests in Python and C++ | 29 Python tests; C++ `ctest` 1/1; cross-language byte equality | | Round-trip + compatibility tests in Python and C++ | 45 Python tests; C++ `ctest` 1/1; cross-language byte equality |
| Targeted pytest passes | 29 passed | | Targeted pytest passes | 45 passed |
| `compileall packages tests` | exit 0 | | `compileall packages tests` | exit 0 |
| `git diff --check` | exit 0 | | `git diff --check` | exit 0 |
| Default tests deterministic, download-free, credit-free, GPU-free | Pure in-memory protobuf; no model, no network, no GPU | | Default tests deterministic, download-free, credit-free, GPU-free | Pure in-memory protobuf; no model, no network, no GPU |
| Full deterministic pytest passes, or pre-existing failure recorded | 712 passed, 12 skipped; flaky pre-existing test documented above with clean-tree reproduction | | Full deterministic pytest passes, or pre-existing failure recorded | Final exact-code run: 728 passed, 12 skipped; earlier sole flaky failure documented with clean-tree reproduction and 3/3 passing retries |
## 5. How the cross-language claim is actually earned ## 5. How the cross-language claim is actually earned

View File

@@ -27,3 +27,19 @@ cmp build/native/cpp_roundtrip.binpb packages/node/native/testdata/session_reque
# --- repository gates # --- repository gates
.venv/bin/python -m compileall -q packages tests .venv/bin/python -m compileall -q packages tests
git diff --check git diff --check
# --- independent controller review after Ralph
PYTHONPATH=packages/node .venv/bin/python -m pytest -q tests/test_native_shard_protocol.py
# -> 45 passed
PYTHONPATH=packages/node .venv/bin/python -m pytest -q \
tests/test_native_shard_protocol.py tests/test_activation_compression.py
# -> 51 passed
PYTHONPATH=packages/node .venv/bin/python -m pytest -q
# -> final exact-code run: 728 passed, 12 skipped
for i in 1 2 3; do PYTHONPATH=packages/node .venv/bin/python -m pytest -q \
tests/test_tracker_routing.py::test_tracker_dashboard_can_cancel_inflight_proxy; done
# -> 1 passed, 1 passed, 1 passed
# clean minimum-runtime venv: protobuf==7.35.0 grpcio==1.82.1
# generated pb2 + pb2_grpc imports and one-byte codec round trip -> passed
# The user chose to rely on Ralph's recorded successful C++ CMake/CTest run
# rather than repeat deletion of an isolated generated build directory.

View File

@@ -2,7 +2,7 @@
"schema_version": "SCHEMA_VERSION_1", "schema_version": "SCHEMA_VERSION_1",
"bundle_version": 1, "bundle_version": 1,
"proto_path": "packages/node/native/proto/shard_runtime.proto", "proto_path": "packages/node/native/proto/shard_runtime.proto",
"proto_sha256": "077ee34961ab787c67457b1a373f73ec3d577f67086d6bc6a32435f37395b7d6", "proto_sha256": "9e211660b3fcefc88bcdf3851c3571088c00349aacb5adc5ef45083c83d0cce2",
"protoc": "grpc_tools 1.82.1 (python) / protobuf 33.1 (C++)", "protoc": "grpc_tools 1.82.1 (python) / protobuf 33.1 (C++)",
"service": { "service": {
"GetCapability": { "GetCapability": {
@@ -75,10 +75,21 @@
"max_prefill_chunk_tokens": 128, "max_prefill_chunk_tokens": 128,
"max_chunk_bytes": 4194304, "max_chunk_bytes": 4194304,
"max_fragment_bytes": 1048576, "max_fragment_bytes": 1048576,
"max_inflight_chunks": 8 "max_inflight_chunks": 8,
"max_fragments_per_tensor": 64,
"max_tensors_per_bundle": 64,
"max_tensor_rank": 8,
"max_tensor_dimension": 2147483647,
"whole_session_message_enforced": true
}, },
"golden_vectors": { "golden_vectors": {
"session_request_golden.binpb": "c2c3df8a717ddeae7bd99624d2c7f34c09a518988de990237fe313b75cff0817", "session_request_golden.binpb": "c2c3df8a717ddeae7bd99624d2c7f34c09a518988de990237fe313b75cff0817",
"capability_report_golden.binpb": "71ac5f150775f398515b43a63596a5cbe8d2ad607e7e4de56bd44fbe7987080c" "capability_report_golden.binpb": "71ac5f150775f398515b43a63596a5cbe8d2ad607e7e4de56bd44fbe7987080c"
},
"verification": {
"python_protocol_tests": "45 passed",
"python_protocol_and_compression_tests": "51 passed",
"full_suite": "728 passed, 12 skipped",
"minimum_runtime": "grpcio 1.82.1 / protobuf 7.35.0 passed import and codec smoke"
} }
} }

View File

@@ -54,7 +54,7 @@
"Update only this story issue to Status: done after every acceptance criterion and quality gate passes" "Update only this story issue to Status: done after every acceptance criterion and quality gate passes"
], ],
"priority": 1, "priority": 1,
"passes": false, "passes": true,
"notes": "Source issue: .scratch/distributed-gguf-runtime/issues/02-adopt-the-versioned-grpc-shard-protocol.md", "notes": "Source issue: .scratch/distributed-gguf-runtime/issues/02-adopt-the-versioned-grpc-shard-protocol.md",
"dependsOn": [] "dependsOn": []
}, },

View File

@@ -99,7 +99,12 @@ def compress_activation(body: bytes, policy: CompressionPolicy) -> CompressionRe
return CompressionResult(candidate, "zstd", len(body), len(candidate), time.monotonic() - started, "compressed") return CompressionResult(candidate, "zstd", len(body), len(candidate), time.monotonic() - started, "compressed")
def decompress_activation(body: bytes, encoding: str | None) -> CompressionResult: def decompress_activation(
body: bytes,
encoding: str | None,
*,
max_output_bytes: int | None = None,
) -> CompressionResult:
"""Decode a modern zstd body or preserve a legacy raw body with metrics.""" """Decode a modern zstd body or preserve a legacy raw body with metrics."""
started = time.monotonic() started = time.monotonic()
if not encoding: if not encoding:
@@ -110,8 +115,23 @@ def decompress_activation(body: bytes, encoding: str | None) -> CompressionResul
import zstandard as zstd import zstandard as zstd
except ImportError as exc: except ImportError as exc:
raise ValueError("zstd support is unavailable") from exc raise ValueError("zstd support is unavailable") from exc
if max_output_bytes is not None and max_output_bytes < 0:
raise ValueError("max_output_bytes must be non-negative")
try: try:
if max_output_bytes is None:
raw = zstd.ZstdDecompressor().decompress(body) raw = zstd.ZstdDecompressor().decompress(body)
else:
# Cap both decoder window allocation and bytes read. zstandard's
# max_window_size unit is KiB.
max_window_kib = max(1024, (max_output_bytes + 1023) // 1024)
decompressor = zstd.ZstdDecompressor(max_window_size=max_window_kib)
# `decompress(max_output_size=...)` may trust a frame's advertised
# content size. A bounded stream read enforces the limit regardless
# of frame metadata and detects trailing expansion with one byte.
with decompressor.stream_reader(body) as reader:
raw = reader.read(max_output_bytes + 1)
if len(raw) > max_output_bytes:
raise ValueError("zstd activation body exceeds its output limit")
except zstd.ZstdError as exc: except zstd.ZstdError as exc:
raise ValueError("invalid zstd activation body") from exc raise ValueError("invalid zstd activation body") from exc
return CompressionResult(raw, "zstd", len(body), len(raw), time.monotonic() - started, "decompressed") return CompressionResult(raw, "zstd", len(body), len(raw), time.monotonic() - started, "decompressed")

View File

@@ -17,8 +17,12 @@ from .codec import (
BUNDLE_VERSION, BUNDLE_VERSION,
DEFAULT_MAX_CHUNK_BYTES, DEFAULT_MAX_CHUNK_BYTES,
DEFAULT_MAX_FRAGMENT_BYTES, DEFAULT_MAX_FRAGMENT_BYTES,
DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
DEFAULT_MAX_INFLIGHT_CHUNKS, DEFAULT_MAX_INFLIGHT_CHUNKS,
DEFAULT_MAX_PREFILL_CHUNK_TOKENS, DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
DEFAULT_MAX_TENSORS_PER_BUNDLE,
DEFAULT_MAX_TENSOR_DIMENSION,
DEFAULT_MAX_TENSOR_RANK,
HIDDEN_STATES, HIDDEN_STATES,
SCHEMA_VERSION, SCHEMA_VERSION,
PayloadCorrupt, PayloadCorrupt,
@@ -35,14 +39,19 @@ from .codec import (
itemsize, itemsize,
negotiate_flow_control, negotiate_flow_control,
plan_prefill_chunks, plan_prefill_chunks,
validate_session_message_size,
) )
__all__ = [ __all__ = [
"BUNDLE_VERSION", "BUNDLE_VERSION",
"DEFAULT_MAX_CHUNK_BYTES", "DEFAULT_MAX_CHUNK_BYTES",
"DEFAULT_MAX_FRAGMENT_BYTES", "DEFAULT_MAX_FRAGMENT_BYTES",
"DEFAULT_MAX_FRAGMENTS_PER_TENSOR",
"DEFAULT_MAX_INFLIGHT_CHUNKS", "DEFAULT_MAX_INFLIGHT_CHUNKS",
"DEFAULT_MAX_PREFILL_CHUNK_TOKENS", "DEFAULT_MAX_PREFILL_CHUNK_TOKENS",
"DEFAULT_MAX_TENSORS_PER_BUNDLE",
"DEFAULT_MAX_TENSOR_DIMENSION",
"DEFAULT_MAX_TENSOR_RANK",
"HIDDEN_STATES", "HIDDEN_STATES",
"SCHEMA_VERSION", "SCHEMA_VERSION",
"PayloadCorrupt", "PayloadCorrupt",
@@ -60,4 +69,5 @@ __all__ = [
"negotiate_flow_control", "negotiate_flow_control",
"pb", "pb",
"plan_prefill_chunks", "plan_prefill_chunks",
"validate_session_message_size",
] ]

View File

@@ -48,6 +48,10 @@ DEFAULT_MAX_CHUNK_BYTES = 4 * 1024 * 1024
DEFAULT_MAX_FRAGMENT_BYTES = 1024 * 1024 DEFAULT_MAX_FRAGMENT_BYTES = 1024 * 1024
DEFAULT_MAX_INFLIGHT_CHUNKS = 8 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. # Canonical boundary tensor name for a dense transformer hidden state.
HIDDEN_STATES = "hidden_states" HIDDEN_STATES = "hidden_states"
@@ -79,14 +83,28 @@ def itemsize(dtype: int) -> int:
raise ProtocolError(f"unsupported dtype {dtype}") from None 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.""" """Byte count a tensor of `shape` and `dtype` must occupy."""
if any(dim < 0 for dim in shape): if len(shape) > max_rank:
raise ProtocolError(f"negative dimension in shape {list(shape)}") 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 count = 1
for dim in shape: for dim in shape:
count *= dim 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 ---------------------------------------------------------------- # --- CRC32C ----------------------------------------------------------------
@@ -136,16 +154,18 @@ def encode_tensor(
dtype: int = pb.DTYPE_BFLOAT16, dtype: int = pb.DTYPE_BFLOAT16,
*, *,
policy: CompressionPolicy | None = None, policy: CompressionPolicy | None = None,
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES, max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
) -> pb.NamedTensor: ) -> pb.NamedTensor:
"""Build a NamedTensor, compressing and fragmenting as needed. """Build a NamedTensor, compressing and fragmenting as needed.
`data` is the uncompressed little-endian payload. The checksum is taken over `data` is the uncompressed little-endian payload. The checksum is taken over
it *before* compression so it stays valid whichever framing a hop chooses. it *before* compression so it stays valid whichever framing a hop chooses.
""" """
if max_fragment_bytes <= 0: if max_chunk_bytes <= 0 or max_fragment_bytes <= 0 or max_fragments <= 0:
raise ProtocolError("max_fragment_bytes must be positive") raise ProtocolError("tensor byte/count bounds must be positive")
declared = expected_bytes(shape, dtype) declared = expected_bytes(shape, dtype, max_bytes=max_chunk_bytes)
if len(data) != declared: if len(data) != declared:
raise ProtocolError( raise ProtocolError(
f"tensor {name!r} declares shape {list(shape)} ({declared} bytes) " 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 # A zero-element tensor is legal (e.g. an empty mask) and still needs a
# fragment, so coverage checks have something to verify. # fragment, so coverage checks have something to verify.
slices = [b""] slices = [b""]
if len(slices) > max_fragments:
raise ProtocolError(
f"tensor {name!r} needs {len(slices)} fragments, exceeding limit {max_fragments}"
)
offset = 0 offset = 0
for index, piece in enumerate(slices): for index, piece in enumerate(slices):
@@ -194,20 +218,57 @@ def encode_tensor(
return 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. """Reassemble, decompress and validate a NamedTensor's payload.
Raises PayloadCorrupt rather than returning a payload it cannot fully Raises PayloadCorrupt rather than returning a payload it cannot fully
account for: a hole in the fragments or a bad checksum means the activation 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. 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: if tensor.byte_order == pb.BYTE_ORDER_BIG_ENDIAN:
raise ProtocolError(f"tensor {tensor.name!r} is big-endian; wire order is little-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: if tensor.byte_order != pb.BYTE_ORDER_LITTLE_ENDIAN:
raise ProtocolError(f"tensor {tensor.name!r} declares no byte order") 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: if not tensor.fragments:
raise PayloadCorrupt(f"tensor {tensor.name!r} carries no 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) fragments = sorted(tensor.fragments, key=lambda f: f.byte_offset)
count = fragments[0].fragment_count count = fragments[0].fragment_count
@@ -231,44 +292,90 @@ def decode_tensor(tensor: pb.NamedTensor) -> bytes:
body.extend(fragment.payload) body.extend(fragment.payload)
if tensor.compression == pb.COMPRESSION_ZSTD: if tensor.compression == pb.COMPRESSION_ZSTD:
data = decompress_activation(bytes(body), "zstd").body try:
elif tensor.compression in (pb.COMPRESSION_NONE, pb.COMPRESSION_UNSPECIFIED): 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) data = bytes(body)
else: 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: if len(data) != tensor.total_bytes:
raise PayloadCorrupt( raise PayloadCorrupt(
f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes but " f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes but "
f"reassembled {len(data)}" 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 algorithm = tensor.checksum.algorithm
if algorithm == pb.CHECKSUM_ALGORITHM_CRC32C: if algorithm == pb.CHECKSUM_ALGORITHM_CRC32C:
if tensor.checksum.value != struct.pack(">I", crc32c(data)): if tensor.checksum.value != struct.pack(">I", crc32c(data)):
raise PayloadCorrupt(f"tensor {tensor.name!r} failed its CRC32C check") raise PayloadCorrupt(f"tensor {tensor.name!r} failed its CRC32C check")
elif algorithm not in (pb.CHECKSUM_ALGORITHM_NONE, pb.CHECKSUM_ALGORITHM_UNSPECIFIED): elif algorithm != pb.CHECKSUM_ALGORITHM_NONE:
raise ProtocolError(f"tensor {tensor.name!r} uses unsupported checksum algorithm") raise ProtocolError(
f"tensor {tensor.name!r} uses unspecified or unsupported checksum algorithm"
)
return data return data
def encode_bundle(tensors: Iterable[pb.NamedTensor]) -> pb.TensorBundle: def encode_bundle(
return pb.TensorBundle(bundle_version=BUNDLE_VERSION, tensors=list(tensors)) tensors: Iterable[pb.NamedTensor],
*,
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
def decode_bundle(bundle: pb.TensorBundle) -> dict[str, bytes]: max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
"""Validate every tensor in a bundle and return name -> payload.""" ) -> pb.TensorBundle:
if bundle.bundle_version > BUNDLE_VERSION: 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( raise ProtocolError(
f"bundle version {bundle.bundle_version} is newer than this build " f"bundle carries {len(tensor_list)} tensors, exceeding limit {max_tensors}"
f"understands ({BUNDLE_VERSION})" )
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] = {} payloads: dict[str, bytes] = {}
for tensor in bundle.tensors: 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") raise ProtocolError("bundle carries an unnamed tensor")
if tensor.name in payloads: if tensor.name in payloads:
raise ProtocolError(f"bundle carries duplicate tensor {tensor.name!r}") 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 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 ---------------------------------------------- # --- Bounded prefill chunking ----------------------------------------------
@@ -363,15 +500,22 @@ def negotiate_flow_control(
candidates = [v for v in (a, b) if v > 0] candidates = [v for v in (a, b) if v > 0]
return min(candidates) if candidates else fallback 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( max_inflight_chunks = _min(
proposed.max_inflight_chunks, proposed.max_inflight_chunks,
limits.max_inflight_chunks, limits.max_inflight_chunks,
DEFAULT_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( max_chunk_bytes=_min(
proposed.max_chunk_bytes, limits.max_chunk_bytes, DEFAULT_MAX_CHUNK_BYTES proposed.max_chunk_bytes, limits.max_chunk_bytes, DEFAULT_MAX_CHUNK_BYTES
), ),

View File

@@ -94,9 +94,10 @@ message Checksum {
// One slice of a tensor's wire payload. // One slice of a tensor's wire payload.
// //
// Fragments bound the size of a single gRPC message. `byte_offset` lets a // Fragments bound individual payload pieces and make coverage independently
// receiver verify the fragments tile the body exactly — no hole, no overlap — // verifiable. The negotiated `FlowControl.max_chunk_bytes` bounds the enclosing
// instead of trusting arrival order. // serialized stream message; `byte_offset` proves the pieces tile the tensor
// exactly — no hole, no overlap — instead of trusting arrival order.
// //
// Offsets and payloads describe the *wire* body, which is the compressed frame // Offsets and payloads describe the *wire* body, which is the compressed frame
// when the parent tensor declares a compression. A zstd frame is not decodable // when the parent tensor declares a compression. A zstd frame is not decodable

View File

@@ -13,8 +13,11 @@ dependencies = [
"huggingface-hub>=0.20", "huggingface-hub>=0.20",
"accelerate>=0.28", "accelerate>=0.28",
"bitsandbytes>=0.43", "bitsandbytes>=0.43",
"grpcio>=1.60", # Keep these floors aligned with the committed grpcio-tools 1.82.1 output.
"protobuf>=5", # Older runtimes satisfy broad constraints but fail while importing the
# generated modules before a node can report a useful protocol error.
"grpcio>=1.82.1",
"protobuf>=7.35.0",
"rich>=13", "rich>=13",
"safetensors>=0.4", "safetensors>=0.4",
"torch>=2.1", "torch>=2.1",

View File

@@ -26,7 +26,9 @@ from google.protobuf import descriptor_pb2
from meshnet_node.activation_compression import CompressionPolicy from meshnet_node.activation_compression import CompressionPolicy
from meshnet_node.native_protocol import ( from meshnet_node.native_protocol import (
DEFAULT_MAX_CHUNK_BYTES, DEFAULT_MAX_CHUNK_BYTES,
DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
DEFAULT_MAX_PREFILL_CHUNK_TOKENS, DEFAULT_MAX_PREFILL_CHUNK_TOKENS,
DEFAULT_MAX_TENSORS_PER_BUNDLE,
HIDDEN_STATES, HIDDEN_STATES,
PayloadCorrupt, PayloadCorrupt,
ProtocolError, ProtocolError,
@@ -39,6 +41,7 @@ from meshnet_node.native_protocol import (
negotiate_flow_control, negotiate_flow_control,
pb, pb,
plan_prefill_chunks, plan_prefill_chunks,
validate_session_message_size,
) )
from meshnet_node.native_protocol import conformance from meshnet_node.native_protocol import conformance
@@ -233,10 +236,161 @@ def test_bundle_from_a_newer_layout_is_refused():
bundle = encode_bundle([]) bundle = encode_bundle([])
bundle.bundle_version = 99 bundle.bundle_version = 99
with pytest.raises(ProtocolError, match="newer"): with pytest.raises(ProtocolError, match="not supported"):
decode_bundle(bundle) 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 ---------------------- # --- 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 # A sender cannot talk a worker into unbounded queues by proposing a large
# window: neither peer can raise the other's ceiling. # window: neither peer can raise the other's ceiling.
assert settled.max_inflight_chunks == default_flow_control().max_inflight_chunks 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_chunk_bytes == DEFAULT_MAX_CHUNK_BYTES
assert settled.max_prefill_chunk_tokens == DEFAULT_MAX_PREFILL_CHUNK_TOKENS 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 ---------------------------------------------------------- # --- Compatibility ----------------------------------------------------------