popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: harden DGR-002 protocol bounds

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

24
packages/node/meshnet_node/activation_compression.py
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")
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."""
started = time.monotonic()
if not encoding:
@@ -110,8 +115,23 @@ def decompress_activation(body: bytes, encoding: str | None) -> CompressionResul
import zstandard as zstd
except ImportError as 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:
raw = zstd.ZstdDecompressor().decompress(body)
if max_output_bytes is None:
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:
raise ValueError("invalid zstd activation body") from exc
return CompressionResult(raw, "zstd", len(body), len(raw), time.monotonic() - started, "decompressed")

10
packages/node/meshnet_node/native_protocol/__init__.py
View File

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

216
packages/node/meshnet_node/native_protocol/codec.py
View File

@@ -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
),

7
packages/node/native/proto/shard_runtime.proto
View File

@@ -94,9 +94,10 @@ message Checksum {
// One slice of a tensor's wire payload.
//
// Fragments bound the size of a single gRPC message. `byte_offset` lets a
// receiver verify the fragments tile the body exactly — no hole, no overlap —
// instead of trusting arrival order.
// Fragments bound individual payload pieces and make coverage independently
// verifiable. The negotiated `FlowControl.max_chunk_bytes` bounds the enclosing
// 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
// when the parent tensor declares a compression. A zstd frame is not decodable

7
packages/node/pyproject.toml
View File

@@ -13,8 +13,11 @@ dependencies = [
"huggingface-hub>=0.20",
"accelerate>=0.28",
"bitsandbytes>=0.43",
"grpcio>=1.60",
"protobuf>=5",
# Keep these floors aligned with the committed grpcio-tools 1.82.1 output.
# 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",
"safetensors>=0.4",
"torch>=2.1",