EnigmaConsultant/pytorch-pt2-archive-streamreader-leak-poc

PyTorch PyTorchStreamReader (.pt2 / TorchScript zip archive) unbounded native heap leak

Status: gated PoC, private disclosure preparation. Do not use against systems you do not own or have authorization to test.

Summary

caffe2::serialize::PyTorchStreamReader (caffe2/serialize/inline_container.cc), the C++ class that backs torch._C.PyTorchFileReader and therefore torch.export.pt2_archive.PT2ArchiveReader (the .pt2 package zip-container reader), torch.jit.load, and legacy zip-based torch.load() checkpoints, leaks native heap memory every single time its constructor throws an exception after the internal mz_zip_reader_init() call has already succeeded.

This happens because:

• PyTorchStreamReader::init() (called from the constructor body, after the ar_ member is already fully constructed) calls mz_zip_reader_init(ar_.get(), size, 0), which makes miniz allocate and populate internal central-directory state (mz_zip_internal_state, its sorted-offsets array, its filename buffer, etc.) inside *ar_.
• init() then performs several more validation checks — is the archive non-empty, is every entry nested under a subdirectory, does a version record exist, does it parse as an integer, is it within [kMinSupportedFileFormatVersion, kMaxSupportedFileFormatVersion] — each of which throws a C++ exception (c10::Error / CAFFE_THROW) on failure.
• The only code that frees the miniz-internal state allocated by mz_zip_reader_init() is PyTorchStreamReader::~PyTorchStreamReader(), which calls mz_zip_reader_end(ar_.get()).
• Per C++ object-construction semantics, if a constructor throws, the destructor of the object under construction is never invoked — only already-fully-constructed subobjects get destroyed. ar_ is a std::unique_ptr<mz_zip_archive>; its own destructor merely free()s the fixed-size mz_zip_archive struct itself, it does not call mz_zip_reader_end().
• Result: every attacker-crafted archive that is well-formed enough to pass mz_zip_reader_init() but trips any of the later checks in init() leaks all of miniz's internal per-archive state, permanently, for the life of the process.

No pickle/graph-module deserialization is ever reached — this is purely in the zip-container layer that runs before any tensor/graph payload is touched, so it fires even against a strict "reject anything that isn't a clean, minimal, empty-model archive" validator.

Attacker-controlled input → sink

1. Attacker supplies a byte string that is opened via torch._C.PyTorchFileReader(...) — reachable via torch.export.load() / torch.export.pt2_archive._package.load_pt2() → PT2ArchiveReader.__init__, or via torch.jit.load() / legacy torch.load() on the zip-based container format, or directly if an application exposes format validation/scanning of uploaded "model files" (a very common pattern for AI model-hosting/scanning services, which is exactly the threat model bug-bounty programs like huntr's MFF care about).
2. The bytes form a syntactically valid ZIP (so mz_zip_reader_init() succeeds and allocates central-directory state) but violate one of PyTorchStreamReader::init()'s post-init checks, e.g.:
   • the top-level entry name has no / (not nested in a subdirectory), or
   • there is no version / .data/version record, or
   • the version record isn't parseable as an integer (std::stoull throws), or
   • the version is outside [kMinSupportedFileFormatVersion, kMaxSupportedFileFormatVersion].
3. PyTorchStreamReader's constructor throws. The exception is caught by pybind11 and surfaces to Python as a RuntimeError (e.g. application code wraps the load call in try/except to validate/reject bad uploads — completely standard and expected here).
4. All miniz-internal heap allocations for that archive attempt (proportional to the number of zip central-directory entries and their name lengths) are never freed.

Because the leaked amount is proportional to the number of entries in the archive (not their contents — every entry can be 0 bytes), the leak is cheaply amplifiable: a 1.52 MB crafted file containing 5000 zero-byte entries with padded names leaks ~840 KB of unreachable native heap per rejected load (confirmed both in an isolated harness and against the real, compiled, official PyTorch build — see below). Any service that repeatedly opens/validates/rejects untrusted model files (a scanning/validation gate, a retry loop, a multi-tenant inference host processing uploads) will have its memory usage grow without bound until it is OOM-killed — Denial of Service.

Reproduction

1. Against the real, official, compiled PyTorch (no custom harness needed)

verify_against_real_torch.py calls torch._C.PyTorchFileReader directly — the actual production binding, unmodified — from a normal pip install torch (tested against torch==2.12.1+cpu). It measures process RSS:

• Control: a well-formed .pt2-shaped archive opened+destroyed 200 times → flat, ~0.6 KB/iteration of noise (well-formed opens do not leak).

• Test: the same archive but with an unparseable .data/version value (528 KB file, 2000 padded zero-byte entries) opened 60 times, sampled every 10 iterations, with gc.collect() forced before each sample (proving it's not a Python-level reference-retention artifact — the leaked memory is entirely below the Python object layer, inside libc's heap, invisible to Python's GC):

  after  10 failed opens: RSS=230732KB
  after  20 failed opens: RSS=233420KB
  after  30 failed opens: RSS=236364KB
  after  40 failed opens: RSS=239180KB
  after  50 failed opens: RSS=242124KB
  after  60 failed opens: RSS=244940KB
  per-10-iteration deltas (KB): [2688, 2944, 2816, 2944, 2816]
  avg leak per failed open (KB): 284.16
  

  Clean, monotonic, unbounded growth (~284 KB/attempt for this archive shape) that gc.collect() cannot reclaim. The reported exception is:

  RuntimeError [enforce fail at inline_container.cc:215] . Couldn't parse the version NOTANUMBER
   as Long Long.
  

  confirming the exact inline_container.cc throw site.

Run it yourself:

pip install torch --index-url https://download.pytorch.org/whl/cpu
python3 verify_against_real_torch.py

2. Root-cause isolation harness (ASan + AFL++, unmodified target source)

fuzz_harness_source.cc links the real, unmodified caffe2/serialize/inline_container.cc (+ istream_adapter.cc, file_adapter.cc, read_adapter_interface.cc) and the real, unmodified vendored miniz-3.0.2 from the pytorch/pytorch tree, against minimal shims for the surrounding c10/ATen framework glue only (allocator, exception macros, logging — not the parsing logic itself). Built with -fsanitize=address,undefined (plain) and with afl-clang-fast++ (AFL_USE_ASAN=1) for the fuzzing campaign.

A 18-minute, 3-instance AFL++ campaign (1.37M total executions) starting from two valid .pt2-shaped zip seeds surfaced this leak repeatedly (LeakSanitizer flags it as a fatal error on process exit) — every triaged non-timeout "hang" bucket entry that actually terminated with an error reported the identical stack:

Direct leak of 152 byte(s) in 1 object(s) allocated from:
    #0 malloc
    #1 mz_zip_reader_init_internal miniz.c:3580
Indirect leak of NNN byte(s) in 3 object(s) allocated from:
    #0 realloc
    #1 mz_zip_array_ensure_capacity miniz.c:3446

minimal_repro_bad_version.pt2 (151 bytes, single-entry archive with an unparseable version string) and repro_no_subdirectory.pt2 (a top-level, non-nested entry) both independently trigger the identical leak signature, confirming this is a structural exception-safety defect in PyTorchStreamReader's constructor, not tied to one specific validation check.

amplified_leak_5000entries.pt2 demonstrates the amplification: 1.52 MB in → ~840 KB leaked per rejected load (standalone ASan/LeakSanitizer build).

Dedup / prior-art check

• pytorch/pytorch#74798 — a different bug class: a heap-buffer-overflow (OOB read) inside mz_zip_reader_read_central_dir in the older bundled miniz-2.0.8; fixed upstream by miniz version bump. PyTorch now bundles miniz-3.0.2 (the version this report targets) and does not reproduce that OOB read. This report is a distinct bug (an unbounded leak, not an OOB read) in the reader-lifetime/exception-safety glue PyTorch itself wrote around miniz, not in miniz's parsing logic.
• pytorch/pytorch#102334 ("memory leak in torch.load") and #10348 ("possible memory leak when exceptions are raised") were checked and have unrelated root causes (retained tensor/storage references and Python traceback reference cycles, respectively) — neither touches PyTorchStreamReader's C++-level exception-safety during construction.
• No CVE/GHSA/issue found describing this specific PyTorchStreamReader constructor leak.

Suggested fix

Give PyTorchStreamReader::init() real RAII/exception-safety: either wrap the mz_zip_reader_init() call in a scope guard that calls mz_zip_reader_end() on the unwind path before re-throwing, or move the whole body of init() into a constructor-delegate/factory pattern so a partially-initialized mz_zip_archive is always torn down via mz_zip_reader_end() regardless of which check fails.