diff --git a/modelhawk.py b/modelhawk.py index d184d12..1615dbe 100644 --- a/modelhawk.py +++ b/modelhawk.py @@ -283,20 +283,50 @@ def scan_yaml_bytes(data: bytes, label: str = "") -> "StreamReport": # Numpy .npy scanning # --------------------------------------------------------------------------- # +# Reject hostile or corrupt headers before decoding/parsing them. NumPy's +# format uses a small Python-literal dictionary; a 1 MiB ceiling is generous +# for real headers while bounding work on untrusted artifacts. +_MAX_NPY_HEADER_BYTES = 1 << 20 + + def _parse_npy_header(data: bytes) -> tuple[str, int]: - """Return (dtype_str, data_start_offset). Raises ValueError on bad files.""" - if not data.startswith(_NPY_MAGIC): - raise ValueError("missing numpy .npy magic bytes") - major = data[6] + """Return (dtype_str, data_start_offset); reject malformed/oversized headers.""" + if len(data) < 8 or not data.startswith(_NPY_MAGIC): + raise ValueError("missing or truncated numpy .npy magic/version bytes") + + major, minor = data[6], data[7] if major == 1: - header_len = int.from_bytes(data[8:10], "little") - hdr_start, data_start = 10, 10 + header_len - else: # v2+: 4-byte header length field - header_len = int.from_bytes(data[8:12], "little") - hdr_start, data_start = 12, 12 + header_len + length_start, length_size, hdr_start = 8, 2, 10 + encoding = "latin-1" + elif major in (2, 3): + length_start, length_size, hdr_start = 8, 4, 12 + encoding = "utf-8" if major == 3 else "latin-1" + else: + raise ValueError(f"unsupported numpy .npy format version {major}.{minor}") + + if len(data) < hdr_start: + raise ValueError("truncated numpy .npy header-length field") + + header_len = int.from_bytes(data[length_start:length_start + length_size], "little") + if header_len <= 0: + raise ValueError("empty numpy .npy header") + if header_len > _MAX_NPY_HEADER_BYTES: + raise ValueError( + f"numpy .npy header exceeds {_MAX_NPY_HEADER_BYTES} byte safety limit" + ) + + data_start = hdr_start + header_len + if data_start > len(data): + raise ValueError("truncated numpy .npy header body") + raw_hdr = data[hdr_start:data_start] - header = ast.literal_eval(raw_hdr.decode("latin-1").strip().rstrip("\x00").rstrip()) - return header.get("descr", ""), data_start + try: + header = ast.literal_eval(raw_hdr.decode(encoding).strip().rstrip("\x00").rstrip()) + except (UnicodeDecodeError, SyntaxError, ValueError) as exc: + raise ValueError(f"invalid numpy .npy header: {exc}") from exc + if not isinstance(header, dict) or not isinstance(header.get("descr"), str): + raise ValueError("numpy .npy header must contain a string 'descr' field") + return header["descr"], data_start def _npy_extract_pickle(data: bytes): diff --git a/tests/test_modelhawk.py b/tests/test_modelhawk.py index 56f16a5..46cbebc 100644 --- a/tests/test_modelhawk.py +++ b/tests/test_modelhawk.py @@ -168,6 +168,40 @@ def test_npy_float_array_safe(): assert stream.severity == "SAFE", f"expected SAFE, got {stream.severity}" + +def test_npy_header_bounds_and_versions(): + """Reject truncated/oversized headers and parse supported v3 UTF-8 headers.""" + import struct + + # Truncated version and length fields must fail closed. + for data in (b"\x93NUMPY", b"\x93NUMPY\x01\x00\x05"): + try: + modelhawk._parse_npy_header(data) + except ValueError: + pass + else: + raise AssertionError(f"truncated header accepted: {data!r}") + + # An attacker-controlled length must be rejected before header parsing. + oversized = ( + b"\x93NUMPY\x02\x00" + + struct.pack("