[ML] Fail closed on incomplete TorchScript pre-load state-hook scan

bin/pytorch_inference/CModelGraphValidator.cc

@@ -148,6 +148,18 @@ CModelGraphValidator::scanArchiveForCustomStateHooks(const char* data, std::size
     }
 
     for (const auto& name : reader->getAllRecords()) {
+        // Fail closed on names that may have been truncated by miniz
+        // (MZ_ZIP_MAX_ARCHIVE_FILENAME_SIZE == 512, including archive/).
+        // Observed serverless bypasses used ~499-char relative paths so that
+        // getAllRecords() returned a truncated name, getRecord() failed, and a
+        // prior fail-open skip let torch::jit::load run __setstate__.
+        if (name.size() >= MAX_SAFE_ARCHIVE_RECORD_NAME_LENGTH) {
+            LOG_ERROR(<< "Pre-load state-hook scan: refusing archive — record name "
+                      << "length " << name.size() << " exceeds safe limit "
+                      << MAX_SAFE_ARCHIVE_RECORD_NAME_LENGTH << " ('" << name << "')");
+            return {std::string{SCAN_INCOMPLETE_MARKER}};
+        }
+
         try {
             auto[recordData, recordSize] = reader->getRecord(name);
             std::string_view bytes{static_cast<const char*>(recordData.get()), recordSize};
@@ -162,14 +174,13 @@ CModelGraphValidator::scanArchiveForCustomStateHooks(const char* data, std::size
                 hooks.emplace("__getstate__");
             }
         } catch (const std::exception& e) {
-            // A single unreadable record (e.g. a deliberately bad CRC) must not
-            // abort the whole scan: an attacker could otherwise hide a
-            // __setstate__ hook in a later record behind a corrupt earlier one,
-            // slip past the scan, and have torch::jit::load run it at load time.
-            // Warn and keep scanning the remaining records.
-            LOG_WARN(<< "Pre-load state-hook scan: skipping unreadable record '"
-                     << name << "': " << e.what());
-            continue;
+            // Fail closed.  A previous fail-open "skip and continue" allowed
+            // attackers to hide __setstate__ under zip paths that getAllRecords
+            // truncates (so getRecord fails) while torch::jit::load still
+            // resolves the real entry by full logical path.
+            LOG_ERROR(<< "Pre-load state-hook scan: refusing archive — unreadable record '"
+                      << name << "': " << e.what());
+            return {std::string{SCAN_INCOMPLETE_MARKER}};
         }
         if (hooks.size() == 2) {
             break;

bin/pytorch_inference/CModelGraphValidator.h

@@ -83,12 +83,31 @@ class CModelGraphValidator {
     //! only run when methods are invoked (e.g. forward) remain the job of the
     //! post-load allowlist / forbid checks in validate().
     //!
+    //! The scan is fail-closed: if any zip record cannot be read, or any record
+    //! name approaches the PyTorch/miniz filename truncation limit (see
+    //! MAX_SAFE_ARCHIVE_RECORD_NAME_LENGTH), the result contains
+    //! SCAN_INCOMPLETE_MARKER and the caller must refuse to load.  An earlier
+    //! fail-open "skip unreadable record" path was bypassed by archives whose
+    //! hook-bearing paths exceed MZ_ZIP_MAX_ARCHIVE_FILENAME_SIZE (512).
+    //!
     //! \p data / \p size are the raw bytes of the .pt (ZIP) archive.  Returns
-    //! the sorted names of any hooks found, or empty if none are found / the
-    //! archive cannot be parsed (in which case torch::jit::load will surface
-    //! the error).
+    //! the sorted names of any hooks found, a single-element vector containing
+    //! SCAN_INCOMPLETE_MARKER if the scan cannot complete safely, or empty if
+    //! none are found / the archive cannot be opened (in which case
+    //! torch::jit::load will surface the parse error).
     static TStringVec scanArchiveForCustomStateHooks(const char* data, std::size_t size);
 
+    //! Sentinel returned by scanArchiveForCustomStateHooks when the archive
+    //! cannot be scanned completely (unreadable or truncated zip record).
+    static constexpr std::string_view SCAN_INCOMPLETE_MARKER{"<unreadable-or-truncated-record>"};
+
+    //! Relative zip record names at or above this length are rejected.  PyTorch
+    //! getAllRecords() copies names into a 512-byte miniz buffer (including the
+    //! archive/ prefix); names near that limit may be truncated so that
+    //! getRecord() fails while torch::jit::load still resolves the real entry.
+    //! Legitimate ML models use paths well under this threshold.
+    static constexpr std::size_t MAX_SAFE_ARCHIVE_RECORD_NAME_LENGTH{256};
+
 private:
     //! Collect all operation names from a block, recursing into sub-blocks.
     static void collectBlockOps(const ::torch::jit::Block& block,

bin/pytorch_inference/Main.cc

@@ -78,17 +78,23 @@ void verifySafeModel(const torch::jit::script::Module& module_) {
 //! Load executes __setstate__ during deserialization, so post-load graph
 //! validation runs too late.  Matching the recommended remediation for a
 //! privately reported finding, any __setstate__/__getstate__ hooks are refused
-//! outright.
+//! outright.  Incomplete scans (unreadable / truncated zip records) are also
+//! refused — fail closed — so path-length evasions cannot skip the hook check.
 //! Forbidden / unrecognised ops in methods that only run when invoked remain
 //! the job of verifySafeModel() after a successful load.
 //! \p modelData / \p modelSize are the raw bytes of the buffered .pt archive.
 void verifySafeModelBeforeLoad(const char* modelData, std::size_t modelSize) {
     auto hooks = ml::torch::CModelGraphValidator::scanArchiveForCustomStateHooks(
         modelData, modelSize);
-    if (hooks.empty() == false) {
-        std::string names = ml::core::CStringUtils::join(hooks, ", ");
-        HANDLE_FATAL(<< "Model archive contains custom state hooks: " << names);
+    if (hooks.empty()) {
+        return;
     }
+    if (hooks.size() == 1 && hooks[0] == ml::torch::CModelGraphValidator::SCAN_INCOMPLETE_MARKER) {
+        HANDLE_FATAL(<< "Model archive failed pre-load state-hook scan "
+                     << "(unreadable or truncated zip record; possible evasion)");
+    }
+    std::string names = ml::core::CStringUtils::join(hooks, ", ");
+    HANDLE_FATAL(<< "Model archive contains custom state hooks: " << names);
 }
 }
 

bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc

@@ -518,6 +518,30 @@ BOOST_AUTO_TEST_CASE(testPreLoadScanHandlesGarbageInput) {
     BOOST_REQUIRE(hooks.empty());
 }
 
+BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsLongPathEvasion) {
+    // Reproduces a serverless bypass: zip entry names longer than miniz's
+    // MZ_ZIP_MAX_ARCHIVE_FILENAME_SIZE (512) cause getAllRecords() to truncate
+    // and getRecord() to fail.  A prior fail-open skip let torch::jit::load run
+    // __setstate__.  The scan must now fail closed.
+    std::string bytes = readFileBytes(
+        "testfiles/malicious_models/malicious_setstate_long_path_evasion.pt");
+    BOOST_REQUIRE(bytes.empty() == false);
+
+    auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks(
+        bytes.data(), bytes.size());
+
+    BOOST_REQUIRE(hooks.empty() == false);
+    BOOST_REQUIRE_EQUAL(1, hooks.size());
+    BOOST_REQUIRE_EQUAL(std::string{CModelGraphValidator::SCAN_INCOMPLETE_MARKER},
+                        hooks[0]);
+}
+
+BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsOversizedRecordName) {
+    // Even without a getRecord failure, relative names at/above the safe limit
+    // are refused — truncation would make the scan untrustworthy.
+    BOOST_REQUIRE(CModelGraphValidator::MAX_SAFE_ARCHIVE_RECORD_NAME_LENGTH >= 256);
+}
+
 BOOST_AUTO_TEST_CASE(testMaliciousReinterpretTensorRejectedPostLoad) {
     // inductor::_reinterpret_tensor is the as_strided heap-OOB bypass
     // (privately reported finding).  This forward-only fixture carries no custom

dev-tools/generate_malicious_models.py

@@ -22,6 +22,7 @@
 
 import os
 import sys
+import zipfile
 from pathlib import Path
 
 import torch
@@ -359,6 +360,62 @@ def _collect_setstate_ops(scripted: torch.jit.ScriptModule) -> set:
 }
 
 
+def generate_setstate_long_path_evasion(output_dir: Path, source_name: str) -> bool:
+    """Repack a __setstate__ fixture under zip paths that exceed miniz's 512-byte
+    filename buffer so getAllRecords() truncates and getRecord() fails.
+
+    This reproduces a serverless bypass of the pre-load hook scan: the scanner
+    used to skip unreadable truncated names (fail-open) while torch::jit::load
+    still resolved the real entry.  The fixture is for scanArchiveForCustomStateHooks
+    only — loading it via torch.jit.load is not required.
+    """
+    source = output_dir / source_name
+    if not source.is_file():
+        print(f"  malicious_setstate_long_path_evasion.pt... SKIPPED (missing {source_name})")
+        return False
+
+    # Relative path length chosen so archive_prefix + relative exceeds
+    # MZ_ZIP_MAX_ARCHIVE_FILENAME_SIZE (512), matching observed attack paths.
+    long_mid = ("a" * 180) + "/" + ("b" * 180) + "/" + ("c" * 120)
+    dest_name = "malicious_setstate_long_path_evasion.pt"
+    dest = output_dir / dest_name
+    # Short top-level folder maximizes room for the long relative path.
+    archive_root = "x"
+
+    print(f"  {dest_name}...", end=" ")
+    try:
+        with zipfile.ZipFile(source, "r") as zin, zipfile.ZipFile(
+            dest, "w", compression=zipfile.ZIP_STORED
+        ) as zout:
+            for info in zin.infolist():
+                data = zin.read(info.filename)
+                # Strip original archive root; re-home under archive_root with
+                # inflated code/ paths.
+                parts = info.filename.split("/", 1)
+                rel = parts[1] if len(parts) == 2 else parts[0]
+                if rel.startswith("code/"):
+                    # code/__torch__/…/file.py → code/__torch__/<long>/file.py
+                    leaf = rel.rsplit("/", 1)[-1]
+                    new_rel = f"code/__torch__/{long_mid}/{leaf}"
+                else:
+                    new_rel = rel
+                new_name = f"{archive_root}/{new_rel}"
+                zout.writestr(new_name, data)
+
+        # Sanity: at least one full name (with archive root) exceeds 512.
+        with zipfile.ZipFile(dest, "r") as zcheck:
+            max_len = max(len(n) for n in zcheck.namelist())
+        if max_len < 512:
+            raise RuntimeError(
+                f"expected zip entry name length >= 512 for truncation, got {max_len}"
+            )
+        print(f"OK ({dest.stat().st_size} bytes, max entry name {max_len})")
+        return True
+    except Exception as exc:
+        print(f"FAILED: {exc}")
+        return False
+
+
 def generate(output_dir: Path):
     output_dir.mkdir(parents=True, exist_ok=True)
     succeeded = []
@@ -391,7 +448,19 @@ def generate(output_dir: Path):
             print(f"FAILED: {exc}")
             failed.append((filename, str(exc)))
 
-    print(f"\nGenerated {len(succeeded)}/{len(MODELS)} models")
+    if generate_setstate_long_path_evasion(
+        output_dir, "malicious_setstate_file_reader.pt"
+    ):
+        succeeded.append("malicious_setstate_long_path_evasion.pt")
+    else:
+        failed.append(
+            (
+                "malicious_setstate_long_path_evasion.pt",
+                "long-path evasion fixture generation failed",
+            )
+        )
+
+    print(f"\nGenerated {len(succeeded)} models ({len(failed)} failed)")
     if failed:
         print("Failed:")
         for name, err in failed:

docs/changelog/3149.yaml

@@ -0,0 +1,5 @@
+area: Machine Learning
+issues: []
+pr: 3149
+summary: Fail closed on incomplete `TorchScript` pre-load state-hook scan
+type: bug