security: bound ZIP/NPZ member decompression (#6)

* security: bound ZIP member decompression

* test: cover ZIP resource limits

* security: handle malformed ZIP member reads

* test: cover ZIP archive member-count limit

Author: Pyhroff

modelhawk.py

@@ -369,15 +369,48 @@ def scan_npy_bytes(data: bytes, label: str = "<npy>") -> "StreamReport":
     return stream
 
 
+# Resource limits for untrusted ZIP/NPZ containers. Oversized entries are
+# surfaced as parse errors (MEDIUM) instead of silently skipped.
+_MAX_ARCHIVE_MEMBERS = 10_000
+_MAX_ARCHIVE_MEMBER_BYTES = 64 * 1024 * 1024
+_MAX_ARCHIVE_TOTAL_BYTES = 256 * 1024 * 1024
+_ARCHIVE_LIMIT_SENTINEL = b""
+
 def iter_pickle_streams(path: str):
-    """Yield (label, bytes) for each pickle stream in a file or zip container."""
+    """Yield (label, bytes) for each pickle stream in a file or zip container.
+
+    ZIP metadata is untrusted. Bound entry count, each decompressed member, and
+    total bytes read; never call ZipFile.read(), which materializes an entire
+    member before the scanner can enforce a limit. Limit violations produce a
+    malformed-stream sentinel so scan_file reports MEDIUM rather than SAFE.
+    """
     if zipfile.is_zipfile(path):
         try:
             with zipfile.ZipFile(path) as z:
-                for info in z.infolist():
+                total_read = 0
+                for index, info in enumerate(z.infolist()):
                     if info.is_dir():
                         continue
-                    raw = z.read(info)
+                    if index >= _MAX_ARCHIVE_MEMBERS:
+                        yield "[archive-limit:member-count]", _ARCHIVE_LIMIT_SENTINEL
+                        break
+                    if info.file_size > _MAX_ARCHIVE_MEMBER_BYTES:
+                        yield f"{info.filename}[archive-limit:member-size]", _ARCHIVE_LIMIT_SENTINEL
+                        continue
+                    remaining = _MAX_ARCHIVE_TOTAL_BYTES - total_read
+                    if info.file_size > remaining:
+                        yield f"{info.filename}[archive-limit:total-size]", _ARCHIVE_LIMIT_SENTINEL
+                        break
+                    try:
+                        with z.open(info, "r") as member:
+                            raw = member.read(min(_MAX_ARCHIVE_MEMBER_BYTES, remaining) + 1)
+                    except Exception:
+                        yield f"{info.filename}[archive-read-error]", _ARCHIVE_LIMIT_SENTINEL
+                        continue
+                    if len(raw) > _MAX_ARCHIVE_MEMBER_BYTES or len(raw) > remaining:
+                        yield f"{info.filename}[archive-limit:decompressed-size]", _ARCHIVE_LIMIT_SENTINEL
+                        break
+                    total_read += len(raw)
                     name_lower = info.filename.lower()
                     if name_lower.endswith(_PICKLEISH_NAMES) or _looks_like_pickle(raw):
                         yield info.filename, raw
@@ -387,7 +420,7 @@ def iter_pickle_streams(path: str):
                         if pkl is not None:
                             yield f"{info.filename}[object-array]", pkl
         except zipfile.BadZipFile:
-            return
+            yield "[archive-read-error]", _ARCHIVE_LIMIT_SENTINEL
         return
     with open(path, "rb") as f:
         yield "<pickle>", f.read()

tests/test_modelhawk.py

@@ -22,6 +22,7 @@
 import pickle
 import sys
 import tempfile
+import zipfile
 
 sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
 
@@ -226,6 +227,52 @@ def __reduce__(self):
         )
 
 
+def test_zip_member_resource_limits():
+    """ZIP scanning must bound decompression and report limit hits as suspicious."""
+    old_member_limit = modelhawk._MAX_ARCHIVE_MEMBER_BYTES
+    old_total_limit = modelhawk._MAX_ARCHIVE_TOTAL_BYTES
+    old_count_limit = modelhawk._MAX_ARCHIVE_MEMBERS
+    try:
+        modelhawk._MAX_ARCHIVE_MEMBER_BYTES = 32
+        modelhawk._MAX_ARCHIVE_TOTAL_BYTES = 48
+        modelhawk._MAX_ARCHIVE_MEMBERS = 100
+        with tempfile.TemporaryDirectory() as d:
+            oversized = os.path.join(d, "oversized.zip")
+            with zipfile.ZipFile(oversized, "w", compression=zipfile.ZIP_DEFLATED) as z:
+                z.writestr("huge.pt", b"A" * 128)
+            report = modelhawk.scan_file(oversized)
+            assert report.verdict == "MEDIUM", (
+                f"oversized ZIP member must be suspicious, got {report.verdict}"
+            )
+            assert any("archive-limit:member-size" in s.label for s in report.streams)
+
+            aggregate = os.path.join(d, "aggregate.zip")
+            with zipfile.ZipFile(aggregate, "w", compression=zipfile.ZIP_STORED) as z:
+                z.writestr("first.bin", b"A" * 24)
+                z.writestr("second.bin", b"B" * 32)
+            report = modelhawk.scan_file(aggregate)
+            assert report.verdict == "MEDIUM", (
+                f"aggregate ZIP limit must be suspicious, got {report.verdict}"
+            )
+            assert any("archive-limit:total-size" in s.label for s in report.streams)
+
+            modelhawk._MAX_ARCHIVE_TOTAL_BYTES = 1024
+            modelhawk._MAX_ARCHIVE_MEMBERS = 1
+            count_limited = os.path.join(d, "count-limited.zip")
+            with zipfile.ZipFile(count_limited, "w", compression=zipfile.ZIP_STORED) as z:
+                z.writestr("one.bin", b"A")
+                z.writestr("two.bin", b"B")
+            report = modelhawk.scan_file(count_limited)
+            assert report.verdict == "MEDIUM", (
+                f"archive member-count limit must be suspicious, got {report.verdict}"
+            )
+            assert any("archive-limit:member-count" in s.label for s in report.streams)
+    finally:
+        modelhawk._MAX_ARCHIVE_MEMBER_BYTES = old_member_limit
+        modelhawk._MAX_ARCHIVE_TOTAL_BYTES = old_total_limit
+        modelhawk._MAX_ARCHIVE_MEMBERS = old_count_limit
+
+
 if __name__ == "__main__":
     tests = [
         test_detection_and_no_false_positives,
@@ -237,6 +284,7 @@ def __reduce__(self):
         test_npy_float_array_safe,
         test_npy_header_bounds_and_versions,
         test_npy_real_numpy_object_array,
+        test_zip_member_resource_limits,
     ]
     failed = 0
     for fn in tests: