diff --git a/agentix/runtime/shared/codec.py b/agentix/runtime/shared/codec.py index 4b5101a..8688c45 100644 --- a/agentix/runtime/shared/codec.py +++ b/agentix/runtime/shared/codec.py @@ -18,11 +18,20 @@ Numpy is optional — if it's not installed, the ndarray hook is just skipped (the type never appears on the wire). pydantic is a hard dep because the rest of the framework uses it. + +Decode validation (#140): `unpack` runs on payloads the peer shaped — +including host-side unpacks of sandbox-emitted side-channel events — so +ext decoding is validated and bounded. A malformed ext payload raises +`ExtDecodeError` (never an arbitrary numpy/msgpack error from mid-decode), +the ndarray header is checked before the buffer is interpreted (dtype must +parse, carry no objects, and agree with the buffer size), and pydantic ext +nesting is depth-bounded so a payload cannot recurse the unpacker. """ from __future__ import annotations import importlib.util +import math from typing import Any import msgpack @@ -39,6 +48,16 @@ _EXT_NDARRAY = 1 _EXT_PYDANTIC = 2 +# Legitimate ext nesting is shallow: a pydantic payload carrying an ndarray +# or another dumped model. The bound exists so a wire payload cannot drive +# unpacker recursion arbitrarily deep. +_MAX_EXT_DEPTH = 16 +_ext_depth = 0 # single-threaded loop assumption — mirrors `_PACKER` below + + +class ExtDecodeError(ValueError): + """A msgpack extension payload failed validation during decode.""" + def _numpy() -> Any: """Lazy numpy import. Cached on the module.""" @@ -54,6 +73,11 @@ def _encode_ext(obj: Any) -> msgpack.ExtType: if _HAS_NUMPY: np = _numpy() if isinstance(obj, np.ndarray): + if obj.dtype.hasobject or obj.dtype.names is not None: + # `dtype.str` drops field names/types (a structured dtype + # collapses to '|V'), and object arrays are pointers — + # neither round-trips. Refuse loudly at the source. + raise TypeError(f"agentix.codec: cannot encode ndarray of dtype {obj.dtype}") header = f"{obj.dtype.str}|{','.join(map(str, obj.shape))}".encode() return msgpack.ExtType(_EXT_NDARRAY, header + b"\x00" + obj.tobytes()) if isinstance(obj, BaseModel): @@ -66,18 +90,67 @@ def _encode_ext(obj: Any) -> msgpack.ExtType: raise TypeError(f"agentix.codec: cannot encode {type(obj).__name__}") +def _decode_ndarray(data: bytes) -> Any: + np = _numpy() + header, sep, raw = data.partition(b"\x00") + if not sep: + raise ExtDecodeError("ndarray ext: missing header terminator") + try: + text = header.decode("ascii") + except UnicodeDecodeError as exc: + raise ExtDecodeError("ndarray ext: header is not ASCII") from exc + # `dtype.str` may itself lead with the separator ("|b1", "|S5"), so the + # shape is everything after the LAST separator. + dtype_str, sep2, shape_str = text.rpartition("|") + if not sep2: + raise ExtDecodeError("ndarray ext: malformed header (expected 'dtype|shape')") + try: + dtype = np.dtype(dtype_str) + except Exception as exc: + raise ExtDecodeError(f"ndarray ext: unknown dtype {dtype_str!r}") from exc + if dtype.hasobject or dtype.itemsize == 0 or dtype.subdtype is not None or dtype.kind == "V": + # object: pointers; V/void: a structured dtype's `.str` collapses to + # '|V' with fields stripped; subarray: frombuffer expands extra + # elements. No real encode produces any of them (see `_encode_ext`). + raise ExtDecodeError(f"ndarray ext: refusing dtype {dtype_str!r}") + try: + shape = tuple(int(s) for s in shape_str.split(",") if s) + except ValueError as exc: + raise ExtDecodeError("ndarray ext: non-integer shape entry") from exc + if any(n < 0 for n in shape): + raise ExtDecodeError("ndarray ext: negative shape entry") + if math.prod(shape) * dtype.itemsize != len(raw): + raise ExtDecodeError("ndarray ext: shape does not match buffer size") + try: + return np.frombuffer(raw, dtype=dtype).reshape(shape) + except Exception as exc: + # Belt-and-braces for numpy refusals the checks above don't model + # (e.g. ndim above numpy's cap) — the decode guarantee is that a bad + # payload surfaces as ExtDecodeError, never a raw numpy error. + raise ExtDecodeError(f"ndarray ext: {exc!r}") from exc + + def _decode_ext(code: int, data: bytes) -> Any: + global _ext_depth if code == _EXT_NDARRAY: if not _HAS_NUMPY: - raise RuntimeError("ndarray ext received but numpy not installed") - np = _numpy() - header, raw = data.split(b"\x00", 1) - dtype_str, shape_str = header.decode().split("|") - shape = tuple(int(s) for s in shape_str.split(",") if s) - return np.frombuffer(raw, dtype=np.dtype(dtype_str)).reshape(shape) + raise ExtDecodeError("ndarray ext received but numpy not installed") + return _decode_ndarray(data) if code == _EXT_PYDANTIC: # Decoded as a plain dict for callers to interpret. - return msgpack.unpackb(data, ext_hook=_decode_ext, raw=False) + if _ext_depth >= _MAX_EXT_DEPTH: + raise ExtDecodeError("pydantic ext: nesting exceeds the decode depth bound") + _ext_depth += 1 + try: + return msgpack.unpackb(data, ext_hook=_decode_ext, raw=False) + except ExtDecodeError: + raise + except Exception as exc: + # `{exc!r}`: msgpack's FormatError stringifies to "" — keep the + # cause class visible in log lines that print only str(error). + raise ExtDecodeError(f"pydantic ext: malformed payload: {exc!r}") from exc + finally: + _ext_depth -= 1 return msgpack.ExtType(code, data) @@ -104,4 +177,4 @@ def unpack(blob: bytes | bytearray | memoryview) -> Any: return msgpack.unpackb(blob, ext_hook=_decode_ext, raw=False) -__all__ = ["pack", "unpack"] +__all__ = ["ExtDecodeError", "pack", "unpack"] diff --git a/tests/runtime/shared/test_codec.py b/tests/runtime/shared/test_codec.py new file mode 100644 index 0000000..ae5e23e --- /dev/null +++ b/tests/runtime/shared/test_codec.py @@ -0,0 +1,180 @@ +"""Codec ext-type decoding — round-trip fidelity and decode validation. + +The wire codec decodes msgpack extension types on every unpack, including +host-side unpacks of sandbox-emitted side-channel payloads (`/trace`, `/log`, +plugin namespaces). A malformed ext payload must surface as one typed error +(`ExtDecodeError`) instead of whatever numpy or msgpack happens to raise +mid-decode, and decoding must stay bounded (#140). +""" + +from __future__ import annotations + +import msgpack +import numpy as np +import pytest +from pydantic import BaseModel + +from agentix.runtime.shared.codec import ExtDecodeError, pack, unpack + +_EXT_NDARRAY = 1 +_EXT_PYDANTIC = 2 + + +def ndarray_ext(header: bytes, raw: bytes = b"") -> bytes: + """Wire bytes for an ndarray ext with an arbitrary (possibly bogus) header.""" + return msgpack.packb(msgpack.ExtType(_EXT_NDARRAY, header + b"\x00" + raw)) + + +# ---------------------------------------------------------------- round-trip + + +@pytest.mark.parametrize("dtype", [" None: + # bool/int8/S5 have `dtype.str` forms that LEAD with the header + # separator ("|b1", "|i1", "|S5") — the decoder must split the header + # on the LAST separator, not the first. + array = np.zeros((2, 3), dtype=dtype) + decoded = unpack(pack(array)) + assert np.array_equal(decoded, array) and decoded.dtype == array.dtype + + +def test_ndarray_scalar_round_trips() -> None: + array = np.array(3.5) + decoded = unpack(pack(array)) + assert decoded.shape == () and decoded == array + + +def test_pydantic_model_decodes_to_plain_dict() -> None: + class Point(BaseModel): + x: int + y: int + + assert unpack(pack(Point(x=1, y=2))) == {"x": 1, "y": 2} + + +def test_ndarray_inside_model_payload_decodes() -> None: + from typing import Any + + class Sample(BaseModel): + model_config = {"arbitrary_types_allowed": True} + name: str + data: Any + + decoded = unpack(pack(Sample(name="s", data=np.arange(4)))) + assert decoded["name"] == "s" and np.array_equal(decoded["data"], np.arange(4)) + + +def test_unknown_ext_code_passes_through_inert() -> None: + blob = msgpack.packb(msgpack.ExtType(42, b"opaque")) + assert unpack(blob) == msgpack.ExtType(42, b"opaque") + + +# ---------------------------------------------------------------- validation + + +def test_ndarray_missing_header_terminator_refused() -> None: + blob = msgpack.packb(msgpack.ExtType(_EXT_NDARRAY, b" None: + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b"\xff\xfe|1", b"\x00" * 8)) + + +def test_ndarray_header_without_separator_refused() -> None: + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b"f8", b"\x00" * 8)) + + +def test_ndarray_unknown_dtype_refused() -> None: + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b"notadtype|1", b"\x00" * 8)) + + +def test_ndarray_object_dtype_refused() -> None: + # `np.dtype("O")` parses fine; the decoder must refuse it before the + # buffer is ever interpreted (frombuffer on object dtype is unsound). + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b"O|1", b"\x00" * 8)) + + +def test_ndarray_zero_itemsize_dtype_refused() -> None: + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b"U0|0")) + + +def test_ndarray_non_integer_shape_refused() -> None: + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b" None: + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b" None: + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b" None: + # prod((1,)*100) == 1 passes the size check, but numpy caps ndim at 64 — + # the refusal must be an ExtDecodeError, not numpy's ValueError. + header = b" None: + # '(2,2)f8' parses to a subarray dtype (itemsize 32) that frombuffer + # expands to extra elements — no real encode produces it. + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b"(2,2)f8|2", b"\x00" * 64)) + + +def test_ndarray_void_dtype_refused() -> None: + # '|V12' is what a structured dtype's `.str` collapses to — decoding it + # would return silently field-stripped data. + with pytest.raises(ExtDecodeError): + unpack(ndarray_ext(b"V12|3", b"\x00" * 36)) + + +def test_structured_array_refused_at_encode() -> None: + # `dtype.str` drops field names/types, so a structured array cannot + # round-trip — refuse loudly at the source instead of corrupting. + with pytest.raises(TypeError): + pack(np.zeros(3, dtype="i4,f8")) + + +def test_object_array_refused_at_encode() -> None: + with pytest.raises(TypeError): + pack(np.array([object()], dtype=object)) + + +def test_pydantic_ext_nesting_is_bounded() -> None: + # Each ext-2 level re-enters the unpacker; unbounded nesting would + # recurse toward a RecursionError. Legitimate nesting is shallow + # (a model payload carrying an array or another dumped model). + payload = msgpack.packb("x") + for _ in range(64): + payload = msgpack.packb(msgpack.ExtType(_EXT_PYDANTIC, payload)) + with pytest.raises(ExtDecodeError): + unpack(payload) + + +def test_pydantic_ext_shallow_nesting_decodes() -> None: + payload = msgpack.packb("x") + for _ in range(8): + payload = msgpack.packb(msgpack.ExtType(_EXT_PYDANTIC, payload)) + assert unpack(payload) == "x" + + +def test_pydantic_ext_malformed_payload_refused() -> None: + blob = msgpack.packb(msgpack.ExtType(_EXT_PYDANTIC, b"\xc1")) # 0xc1: never valid + with pytest.raises(ExtDecodeError, match="FormatError"): + # msgpack's FormatError stringifies to "" — the wrapped message must + # still name the cause class. + unpack(blob)