From a433457db92752125156760e73dba24841aae0ad Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Thu, 13 Aug 2026 10:39:43 +1200 Subject: [PATCH 1/2] [ML] Fail closed on incomplete TorchScript pre-load state-hook scan Reject unreadable or oversized zip record names so path-length truncation cannot skip the __setstate__/__getstate__ scan before torch::jit::load. Co-authored-by: Cursor --- bin/pytorch_inference/CModelGraphValidator.cc | 27 +++++-- bin/pytorch_inference/CModelGraphValidator.h | 25 +++++- bin/pytorch_inference/Main.cc | 14 +++- .../unittest/CModelGraphValidatorTest.cc | 24 ++++++ .../malicious_setstate_long_path_evasion.pt | Bin 0 -> 4755 bytes dev-tools/generate_malicious_models.py | 71 +++++++++++++++++- 6 files changed, 145 insertions(+), 16 deletions(-) create mode 100644 bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_long_path_evasion.pt diff --git a/bin/pytorch_inference/CModelGraphValidator.cc b/bin/pytorch_inference/CModelGraphValidator.cc index 99a961a24..f7f9be89a 100644 --- a/bin/pytorch_inference/CModelGraphValidator.cc +++ b/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(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; diff --git a/bin/pytorch_inference/CModelGraphValidator.h b/bin/pytorch_inference/CModelGraphValidator.h index ae344f1ab..7e86a7b9d 100644 --- a/bin/pytorch_inference/CModelGraphValidator.h +++ b/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{""}; + + //! 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, diff --git a/bin/pytorch_inference/Main.cc b/bin/pytorch_inference/Main.cc index 5b5fcca9f..cb0e4393a 100644 --- a/bin/pytorch_inference/Main.cc +++ b/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); } } diff --git a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc index 1cee1a084..ecc51c5e9 100644 --- a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc +++ b/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 diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_long_path_evasion.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_long_path_evasion.pt new file mode 100644 index 0000000000000000000000000000000000000000..55a645cf4579ec40e7568f5f295d19bb283480f9 GIT binary patch literal 4755 zcmeHL&2QX96!&hv-If$lRkY$TEfVs=cI{1?qzMXYt2VR@Zo_t~8b#6IwP$zTc>Ot> zY$^^dqID}##DNn?`~z@8g$sylxdGx2;DE#l4jg*n#2ee2bxb3bQ(7rjM(dgPesA79 zzxmCZEaRc!Ljr&9*AJ||^xL^(!|=}AIJBN@>P(jvcomknX+~du_UH!~vyCPchK*%hWUx@6h3Ms<@0s%uhP^o7zbBUY-AF0<@VBmr%7E!!lFI-who z^(oOWNRl=N#zWB5ThuQm5OK+Nb+8)`ivEyR03F32)`sBM@Vfp8F0gm}Q7lOC!rMK= zaBZ6!%yOMjyc~cZ)I4VyO!UV}HCFS-@7?hy8opTBzRLGM(LaA!6a?M`Xg%2g_vNki z$@(++$c<<4k&Rt?1|PfPZ5p;7hG>VRmLRpeHZ3iP2od7xf!?8v1|cD65v{Doqm6E} zxowCAy5(3-3oW3A>)Kq&3Q4nT+oYizYa4oClD6xvf&8+^(^t1wxFR>V0fY5;w?N0F zP1JM)4v<1>Hy6N%bQe~3aJURGgnE4=C3NDA)E-{~x4($Esmey&4Jhjd4oY^dZkkfB zj5bQ?)P%Gs3;9+jxgJk7xEkS36xbX=y)kUhwo4fA*zaa0*q6@~i; zA6wai@Q6zBoTU|I6&AuNv;y%5+L99UWu(AgDS;+m2H-@AO~w8G>S4HVtIAUaq4ewV z`Ni!d_tD*tov(p9c@y_YZWqb%=^(j^2`6!j#7DXIA{=~6muu>oI=ObSet}%6*DhXq zlhn_gefzvGNE(O*zo3=jhi6%yg~~h7i(k}^rVtKGle33j>k|%~>TiL9d^H5Rl&KlY zVPMD<_?KBmNOjm7a>qVlc%E$1^%D#zF_c@>p#i6$j&56q<#t2T;pIjM)D7jq;dI8+ z$=14Vu?p>ZZU9xlh;}*^@)Z!00ky*6404Jf!kqXLDAwve&)r2G7jO{^Tt0+9jr`~E>-cJOTuhXQ@dzHpV|W}-^g;I{ ze-eL%zl(4&Qm2dUN&G$jF~UQUdYX6olC#B(xuvY+0`L5ce~s{Pq~feZ<()t9KM@{@ z)c3Lm+FbG{{wuZlq+95^2D`%IX?-XmzS>&d`pMHj^LGgDF^uFi`Lb&)D{hNc z8c#>I42rFM-sH)9kT*X8 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__//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: From 6687ab46e17969bac2f24f0edd10f7f04a4d40a6 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Thu, 13 Aug 2026 10:41:11 +1200 Subject: [PATCH 2/2] Update docs/changelog/3149.yaml --- docs/changelog/3149.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 docs/changelog/3149.yaml diff --git a/docs/changelog/3149.yaml b/docs/changelog/3149.yaml new file mode 100644 index 000000000..77f8df166 --- /dev/null +++ b/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