security: validate msgpack ext-type decoding

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<size>'), 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<size>' 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"]

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", ["<f8", "<i4", "bool", "int8", "S5", "<U4"])
+def test_ndarray_round_trips_across_dtype_byte_orders(dtype: str) -> 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"<f8|3"))  # no NUL
+    with pytest.raises(ExtDecodeError):
+        unpack(blob)
+
+
+def test_ndarray_non_ascii_header_refused() -> 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"<f8|a,b", b"\x00" * 8))
+
+
+def test_ndarray_negative_shape_refused() -> None:
+    with pytest.raises(ExtDecodeError):
+        unpack(ndarray_ext(b"<f8|-1", b"\x00" * 8))
+
+
+def test_ndarray_shape_buffer_mismatch_refused() -> None:
+    with pytest.raises(ExtDecodeError):
+        unpack(ndarray_ext(b"<f8|3", b"\x00" * 8))  # 3 float64 needs 24 bytes
+
+
+def test_ndarray_excessive_dimensions_refused() -> 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"<f8|" + b",".join([b"1"] * 100)
+    with pytest.raises(ExtDecodeError):
+        unpack(ndarray_ext(header, b"\x00" * 8))
+
+
+def test_ndarray_subarray_dtype_refused() -> 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)