From f1be338d4fcc58c9ef3e8f26b8b3cb1c84590f5a Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Sun, 11 Oct 2026 14:59:59 +0530 Subject: [PATCH 1/4] security: bound ZIP member decompression --- modelhawk.py | 41 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 37 insertions(+), 4 deletions(-) diff --git a/modelhawk.py b/modelhawk.py index d184d12..a4b2faa 100644 --- a/modelhawk.py +++ b/modelhawk.py @@ -339,15 +339,48 @@ def scan_npy_bytes(data: bytes, label: str = "") -> "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 (OSError, RuntimeError, zipfile.BadZipFile, EOFError): + 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 @@ -357,7 +390,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 "", f.read() From 4b1df0b048e971d9ddb27ab8e16c1320312d207a Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Sun, 11 Oct 2026 15:00:02 +0530 Subject: [PATCH 2/4] test: cover ZIP resource limits --- tests/test_modelhawk.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/test_modelhawk.py b/tests/test_modelhawk.py index 56f16a5..20a8b29 100644 --- a/tests/test_modelhawk.py +++ b/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)) @@ -192,6 +193,40 @@ 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) + 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, @@ -202,6 +237,7 @@ def __reduce__(self): test_npy_object_array_handcrafted, test_npy_float_array_safe, test_npy_real_numpy_object_array, + test_zip_member_resource_limits, ] failed = 0 for fn in tests: From 1f3af080ec8dc2fe198acdd59df2f34212d5b70f Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Sun, 11 Oct 2026 15:00:23 +0530 Subject: [PATCH 3/4] security: handle malformed ZIP member reads --- modelhawk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelhawk.py b/modelhawk.py index a4b2faa..5e4d5ae 100644 --- a/modelhawk.py +++ b/modelhawk.py @@ -374,7 +374,7 @@ def iter_pickle_streams(path: str): try: with z.open(info, "r") as member: raw = member.read(min(_MAX_ARCHIVE_MEMBER_BYTES, remaining) + 1) - except (OSError, RuntimeError, zipfile.BadZipFile, EOFError): + except Exception: yield f"{info.filename}[archive-read-error]", _ARCHIVE_LIMIT_SENTINEL continue if len(raw) > _MAX_ARCHIVE_MEMBER_BYTES or len(raw) > remaining: From b024d0c89d82f64812ffd9917fd5f5e17f0dcfd1 Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Sun, 11 Oct 2026 15:00:25 +0530 Subject: [PATCH 4/4] test: cover ZIP archive member-count limit --- tests/test_modelhawk.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_modelhawk.py b/tests/test_modelhawk.py index 20a8b29..6a8766b 100644 --- a/tests/test_modelhawk.py +++ b/tests/test_modelhawk.py @@ -221,6 +221,18 @@ def test_zip_member_resource_limits(): 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