From 2d5be13e7f7699d406a65cf3f88f5608dfa648d9 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Tue, 21 Jul 2026 09:26:43 +1200 Subject: [PATCH 1/8] Reproduce __setstate__ load-time execution gap in model validation TorchScript runs a module's __setstate__ during torch::jit::load(), before the loaded module reaches CModelGraphValidator (which only walks the inlined forward graph). A forbidden op hidden in __setstate__ is therefore invisible to the validator and has already executed by the time validation would run. Add fixtures and a self-contained repro that demonstrate the gap: - generate_malicious_models.py: add SetStateFileReaderModel and a submodule variant (aten::from_file in __setstate__, benign forward); surface __setstate__ ops in the generator output. - test/test_setstate_load_timing.py: torch-only, build-free repro proving load-time execution (benign print probe) and op-hiding (static inspection). - test_pytorch_inference_evil_models.py: register the fixtures as regression guards for the forthcoming pre-load scan. No fix yet; this only reproduces and documents the vulnerability. Co-authored-by: Cursor --- .../malicious_setstate_file_reader.pt | Bin 0 -> 2991 bytes ...cious_setstate_file_reader_in_submodule.pt | Bin 0 -> 3287 bytes dev-tools/generate_malicious_models.py | 100 +++++++- test/test_pytorch_inference_evil_models.py | 73 ++++++ test/test_setstate_load_timing.py | 215 ++++++++++++++++++ 5 files changed, 386 insertions(+), 2 deletions(-) create mode 100644 bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt create mode 100644 bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt create mode 100644 test/test_setstate_load_timing.py diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt new file mode 100644 index 0000000000000000000000000000000000000000..7a60d4f7990e125015f69eb722ed5b9c78c29bc8 GIT binary patch literal 2991 zcmWIWW@cev;NW1u0ICeO47rIpnaP>?rN!~ZsU^iFi6yD=X_+~x@kObLDXB&JDTyVC z`UY-JmJCtk0+6)<-s~JQ_gxf$4q*V{0Gw`8XHX}`O?m~{Ic`qM#JH+~DLFpAB)=#* zBR-xhIJG1gWwE7FTEEy>I&<^qO5QfX#RN_8XWG5G`D8Ir)h|bHRdKh0GC*Ks5%1ED_8=I=P^3}MCYNv(vTFvF1Ql|0wiR+F6>{l; zLau%FtgrikMu9LWb`sZyr3x8KQ4#by}x=aKCiJi zYg^3Z>}z}L%@;>(I+3(Nv&-hvh8=-UizC0BS-mJt_?2+*<4r2(7kr4Hw_>HV|M5q5 z=c9II<-V$PJy`#d7Ziat6B1|J0~5t@CVUZS#UM#?1nQ*#Q)zlUI5&u5kILA9v;J2c zMDkC63bi*dR6Lb=34!R<1;X|Wb71CMu`|L5ti zx@3P~ap(VLb=!)sPfE%sc24$8i>s4=|My<8`i`{7YtI(%+JB8>exCWZL&tl6eLMGe z!^Um@3YsIYIvf_;++6urpjIenQrfxQg8r+RETqCLgrT5hvCn#7D|n){z4HT0zpan5Vy2G4qr%hK%ee~!D>&u!PuTdk zmN!U#@<_VT7bJCgkMg0ic?nusM>$usTt3{l+UWYHTZ&Un--udH*OTTA>t#vQ2;hpe zYCSCR*USB-h~_)Kpfi(e@0PBgpRImAV7U~xSx<<$=ibSjLY}#+k^^k7WU6G&6ZE{j z;^~U;R+UqJCm8#?#4NA%EIGeoXO2yT*`A=DJ_AMV6ocx?&);tF3|@1tN+WrFN6QPf z%~gI?UE1E=-NkydIr1dgn)hZk*G;;6)OPRd#+Cne=NRqyXdn5hBhAG6$MVQaHsPDN zLpMD7EZWA}+aFZ&-P*+0Y1YXpTzfBeU)`&9$|@p3vux%67q7$LgmpIX@ay2!Y~J#m5S3^(lxXLiVI#aZIETybKssyRr}q|XHGeS0i}q;LZ`CP2zuKeeNBMs&P6+vV-lsXN^G{1W*Rd)7zE3V(S(5s3x6ZMt z{|{Z+V=^&Z?b9undo%L#&)UhGK8oM|S2Frw%O`uu87bi*Z}mZGxu>Hq&l{MIHvtPb z+{L&KgEZl^oCmCL@=A&!Ww|)^)EgDmz@!N-1#|T|!wJ_PY782LEG|neDh8IjZcZw= ztVF5mjJZJGzBn&E6zFRP5XS9o4Z^l3RhFdYg9?0N?a#?9DalC%xuMneu_PPNeh|j( z22%!cLLs0BuGWiFi!y<=e^nx=7K+bIadQ&J79a*DhK43aW=57~2IdCFhDPS*rWOWf z#-^r5rj`~!IU_SmLrX(*bC3%oWAy|zfUX1K0B=T6V8L4+$o0P}NCE{A)eb>761iej zM=@>+kVk}(pri*e6j7%l7uE(ShMof!;RFoDsA17fN6tQ0D5m!VGZd+&BZ`0Yph3

8oTpgvH;k-?W6q6O4K4De?rN!~ZsU^iFi6yD=X_+~x@kObLDXB&AnR)TW zrAfK@DWy57`YDMeiTVa^PC5)xBmj`L0p9E!GWT5+fev8+;Q*X&GGUM+*-d%{**R`b zVmMsYz?2*xUy@&xoDm<-6`WcU4DzHK(2qf2FM8$$!#v1U$e`I+s}TVTs*<9_%)HFJ z^g_nY2r(c(KC?JJtu!YmJ}EIdyF9TdB|alRKf930FM<^)o1BrElTyfxWKle(P0k=i zT!k!}jmco==*5?oWaboe0mC|}6ety6lA2eXUldCKn4hRCMW0Tr@XG^vnZ4-`-9R4V7|0L=wqP&|nM(++0h zv1QOCDV~z^Q&RQe@t{{w>E@)0Bf`+q*4|0}+070DN8kSxarW1{aU`boOOVEw3u3eG zI7ygUsZF{%^{(jU{QG|wcP{n39LLr4?#!Ob{Q4Q3>r8b{>hPL>3E`?zpEtQ`-SoHH zCRJ){RDIN+=Eoo)y{^KwZHik9>l*WYyK=i9d7spXyOz4QS!7!fmw8B*PLyK7>CMi^ zL;I>V`M3SQHY4a-%$mCkE~ottda|uCW67y+9l?L(K3wFz`EU-G#*!j#W#2^-r|b5F z2q%>Gn zT4wmLKj~7*GuP1C;=y?6uq&6N#HtfTZ@qI~-LSjAl+(w9$6P_|AETYhiU4-82ua3` zhxnL|+&^e3|0JSnQ_1%)zboE<|6`CnZ-UCpPrb`F-rJOC9u>ZJ_S&@M$6Z&W_LuDK zQ$71+LyK`*YrK}d-m4XJH6(9zt`-z2(Rx+jn!9j<%gIGidkrerX;@^>o8@t7vcfHY zr_vSk0)s=(e*AObYFjtgw6ocBmfu&>P_!!Q^b>M-Z}eU9ULZxiZ~6D8Rl>)2b(gI& z>0(+`9@rr?TWNJV`@YBSwMk!<#95b!-9K*m`H5XGJJ-3{$JJX*r3(MwvYmhJfrdvxDsS%uWq=j+>_wwrw5j6dYKcCF$q7EYBm z{$9O*lIg~AT@GxU6qfGb(U+Xp{F#Br=*)`rJ(fE-c-bP;?s#@<7MVXjoAc()6*cp% zmFGVs`mHmbFjFv{wJ2V0^5=jCr6~(+rtV3tklVH4u)^`4<(tg;nJbPcobO4`G<1zS z(R5Q=@TPC~@eeJIjgt8;Df+te4qLOirqu}W+ATYgeIqp>>6p6UPo1oNOOD1xO_}*I z_R1fDiSv~{dArng=hYvvb1#)R9TXm52{n}*q@VGQj(Jjazm@_V@Wok{UD6n4NfEm zyB@e%Qk+_p32e$#C4wrt_{3$G7?~Mani-fI7#kXyo10n~7#kTI8JU_G zSX!8ySQ?v|n;M&gT Tensor: return torch.from_file("/tmp/secret", size=10) +# --- Load-time execution models (__setstate__) --- +# +# TorchScript serialises a module's __setstate__ method and runs it *during* +# torch::jit::load(), before the loaded Module is ever handed to +# CModelGraphValidator. Because the validator only walks the (inlined) forward +# graph, a forbidden op hidden in __setstate__ is invisible to it AND has +# already executed by the time validation would run. These fixtures reproduce +# that gap: their forward graphs are benign, so the current validator accepts +# them, yet loading them triggers aten::from_file at deserialization time. + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__ — executes during load. + + The forbidden op never appears in forward(), so the post-load graph + validator (which inspects only the forward graph) does not see it. Worse, + torch::jit::load() runs __setstate__ before validation happens at all, so + the file read has already fired. This is the load-time execution gap + identified in security review. + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + # Runs at deserialization time, before any graph validation. + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderInSubmodule(torch.nn.Module): + """Same load-time attack, but hidden in a submodule's __setstate__. + + Confirms that a pre-load scan must recurse into every serialised + __setstate__ in the archive, not just the top-level module's. + """ + def __init__(self) -> None: + super().__init__() + self.child = _SetStateFileReaderChild() + + def forward(self, x: Tensor) -> Tensor: + return x + self.child(x) + + +class _SetStateFileReaderChild(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/tmp/secret", size=10) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + # --- Sandbox2 attack models (PR #2873) --- # # These reproduce real-world attack vectors that exploit torch.as_strided @@ -222,6 +294,23 @@ def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: # --- Generation logic --- +def _collect_setstate_ops(scripted: torch.jit.ScriptModule) -> set: + """Collect ops from every __setstate__ in the module tree. + + __setstate__ runs during torch::jit::load(), so ops here execute before + (and are invisible to) the forward-graph validator. Surfacing them makes + the load-time attack fixtures self-documenting. + """ + ops: set = set() + for submodule in scripted.modules(): + try: + graph = submodule._c._get_method("__setstate__").graph + except (RuntimeError, AttributeError): + continue + ops.update(node.kind() for node in graph.nodes()) + return ops + + MODELS = { "malicious_file_reader.pt": FileReaderModel, "malicious_mixed_file_reader.pt": MixedFileReaderModel, @@ -229,6 +318,8 @@ def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: "malicious_conditional.pt": ConditionalMalicious, "malicious_many_unrecognised.pt": ManyUnrecognisedOps, "malicious_file_reader_in_submodule.pt": FileReaderInSubmodule, + "malicious_setstate_file_reader.pt": SetStateFileReaderModel, + "malicious_setstate_file_reader_in_submodule.pt": SetStateFileReaderInSubmodule, "malicious_heap_leak.pt": HeapLeakModel, "malicious_rop_exploit.pt": RopExploitModel, } @@ -254,7 +345,12 @@ def generate(output_dir: Path): graph = scripted.forward.graph.copy() torch._C._jit_pass_inline(graph) ops = sorted(set(n.kind() for n in graph.nodes())) - print(f" ops: {ops}") + print(f" forward ops: {ops}") + + # Surface ops hidden in __setstate__ (executed at load time). + setstate_ops = _collect_setstate_ops(scripted) + if setstate_ops: + print(f" __setstate__ ops (run at load): {sorted(setstate_ops)}") succeeded.append(filename) except Exception as exc: diff --git a/test/test_pytorch_inference_evil_models.py b/test/test_pytorch_inference_evil_models.py index 64b81e2d9..82c04b53f 100644 --- a/test/test_pytorch_inference_evil_models.py +++ b/test/test_pytorch_inference_evil_models.py @@ -37,6 +37,7 @@ import torch from torch import Tensor +from typing import Tuple # --------------------------------------------------------------------------- @@ -156,6 +157,68 @@ def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: return torch.zeros(0) +# --------------------------------------------------------------------------- +# Load-time execution models (__setstate__) +# --------------------------------------------------------------------------- +# +# TorchScript runs a module's __setstate__ during torch::jit::load(), before the +# loaded Module reaches CModelGraphValidator (which only walks the forward +# graph). A forbidden op in __setstate__ is therefore invisible to the current +# validator AND has already executed by load time. Pre-fix, feeding these models +# to the binary crashes/errors at load with no validator rejection message +# (INCONCLUSIVE below). Once a pre-load static scan is added, they must be +# cleanly REJECTED like the forward-graph attacks. + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__; forward() is benign.""" + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class _SetStateFileReaderChild(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/tmp/secret", size=10) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderInSubmodule(torch.nn.Module): + """Load-time aten::from_file hidden in a submodule's __setstate__.""" + def __init__(self) -> None: + super().__init__() + self.child = _SetStateFileReaderChild() + + def forward(self, x: Tensor) -> Tensor: + return x + self.child(x) + + # --------------------------------------------------------------------------- # Binary discovery # --------------------------------------------------------------------------- @@ -242,6 +305,16 @@ def find_pytorch_inference() -> str: "description": "ROP-chain file-write via aten::as_strided", "expect_stderr_contains": "Unrecognised operations", }, + "setstate_file_reader": { + "class": SetStateFileReaderModel, + "expect_rejected": True, + "description": "aten::from_file in __setstate__ (runs at load time)", + }, + "setstate_file_reader_submodule": { + "class": SetStateFileReaderInSubmodule, + "expect_rejected": True, + "description": "aten::from_file in a submodule's __setstate__", + }, } diff --git a/test/test_setstate_load_timing.py b/test/test_setstate_load_timing.py new file mode 100644 index 000000000..e2ba89f02 --- /dev/null +++ b/test/test_setstate_load_timing.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +# +# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +# or more contributor license agreements. Licensed under the Elastic License +# 2.0 and the following additional limitation. Functionality enabled by the +# files subject to the Elastic License 2.0 may only be used in production when +# invoked by an Elasticsearch process with a license key installed that permits +# use of machine learning features. You may not use this file except in +# compliance with the Elastic License 2.0 and the foregoing additional +# limitation. +# +"""Reproduce the __setstate__ load-time execution gap in model validation. + +TorchScript serialises a module's __setstate__ method and executes it *during* +torch::jit::load(). CModelGraphValidator, however, only inspects the (inlined) +forward graph of an already-loaded module. A forbidden operation hidden in +__setstate__ is therefore: + + 1. invisible to the validator (it never appears in the forward graph), and + 2. already executed by the time the validator would run (load happens first). + +This is the timing gap raised in security review: a malicious model can read +files / crash the process at load time, before any graph check. This test +demonstrates both facts *without needing a built binary* and only depends on +torch, so it can run in CI as a guard. + +It is deliberately safe: the load-time execution proof uses a benign probe +whose __setstate__ merely print()s a marker. The forbidden-op proofs use +static graph inspection on the saved archive and never load an attack model. + +Usage: + python3 test_setstate_load_timing.py +""" + +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Tuple + +import torch +from torch import Tensor + + +# Mirror of CSupportedOperations::FORBIDDEN_OPERATIONS (subset relevant here). +FORBIDDEN_OPERATIONS = { + "aten::from_file", + "aten::save", + "aten::as_strided", +} + +_LOAD_MARKER = "SETSTATE_EXECUTED_DURING_LOAD" + + +class BenignSetStateProbe(torch.nn.Module): + """Harmless probe: __setstate__ only print()s a marker (TorchScript-safe). + + Used to prove — without running any dangerous op — that __setstate__ code + executes during torch.jit.load(). + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + self.w = state[0] + self.training = state[1] + # Literal must be inlined: TorchScript cannot close over globals. + print("SETSTATE_EXECUTED_DURING_LOAD") + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__; forward() is benign. + + This is the attack shape: the forbidden op runs at load time and never + appears in the forward graph the validator inspects. We only ever inspect + this model statically (never load it), so no file is read here. + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +def _forward_ops(scripted: torch.jit.ScriptModule) -> set: + graph = scripted.forward.graph.copy() + torch._C._jit_pass_inline(graph) + return {node.kind() for node in graph.nodes()} + + +def _setstate_ops(scripted: torch.jit.ScriptModule) -> set: + ops: set = set() + for submodule in scripted.modules(): + try: + graph = submodule._c._get_method("__setstate__").graph + except (RuntimeError, AttributeError): + continue + ops.update(node.kind() for node in graph.nodes()) + return ops + + +def test_setstate_executes_during_load() -> bool: + """Prove that __setstate__ code runs during torch.jit.load().""" + print("--- 1. __setstate__ executes during load (benign probe) ---") + ok = True + with tempfile.TemporaryDirectory(prefix="setstate_timing_") as tmp: + path = Path(tmp) / "probe.pt" + scripted = torch.jit.script(BenignSetStateProbe()) + torch.jit.save(scripted, str(path)) + + # Load in a fresh subprocess; capture what is printed *during* load. + loader = ( + "import torch;" + f"torch.jit.load(r'{path}');" + "print('LOAD_RETURNED')" + ) + proc = subprocess.run( + [sys.executable, "-c", loader], + capture_output=True, text=True, timeout=120, + ) + stdout = proc.stdout + marker_before_return = ( + _LOAD_MARKER in stdout + and stdout.index(_LOAD_MARKER) < stdout.index("LOAD_RETURNED") + ) if "LOAD_RETURNED" in stdout else (_LOAD_MARKER in stdout) + + if marker_before_return: + print(f" OK: '{_LOAD_MARKER}' printed during load — " + "code ran before load() returned (and before validation).") + else: + print(" FAIL: marker not observed during load.") + print(f" stdout: {stdout!r}") + ok = False + print() + return ok + + +def test_forbidden_op_hidden_from_forward_graph() -> bool: + """The validator inspects only forward; the forbidden op lives in setstate.""" + print("--- 2. forbidden op is invisible to the forward-graph validator ---") + ok = True + scripted = torch.jit.script(SetStateFileReaderModel().eval()) + + forward = _forward_ops(scripted) + setstate = _setstate_ops(scripted) + + forward_forbidden = forward & FORBIDDEN_OPERATIONS + setstate_forbidden = setstate & FORBIDDEN_OPERATIONS + + print(f" forward ops : {sorted(forward)}") + print(f" __setstate__ ops: {sorted(setstate)}") + + if forward_forbidden: + print(f" FAIL: forbidden op unexpectedly in forward: {sorted(forward_forbidden)}") + ok = False + else: + print(" OK: forward graph contains no forbidden op — " + "the current validator would ACCEPT this model.") + + if "aten::from_file" in setstate_forbidden: + print(" OK: aten::from_file present in __setstate__ — " + "a correct pre-load scan must catch this.") + else: + print(" FAIL: aten::from_file not found in __setstate__ (fixture broken).") + ok = False + print() + return ok + + +def run_tests() -> bool: + print("=" * 72) + print("Repro: __setstate__ load-time execution gap (security review)") + print("=" * 72) + print(f"torch version: {torch.__version__}") + print() + + results = [ + test_setstate_executes_during_load(), + test_forbidden_op_hidden_from_forward_graph(), + ] + + all_passed = all(results) + print("=" * 72) + if all_passed: + print("ALL CHECKS PASSED — the load-time gap is reproduced and documented.") + print("A fix must statically scan every serialised __setstate__ for") + print("forbidden ops BEFORE torch::jit::load() is called.") + else: + print("SOME CHECKS FAILED — see above.") + print("=" * 72) + return all_passed + + +if __name__ == "__main__": + sys.exit(0 if run_tests() else 1) From 1700b1458bd658dbf594e91a36b3542f7316ad4b Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Tue, 21 Jul 2026 09:40:49 +1200 Subject: [PATCH 2/8] Reject forbidden ops in serialised code before load (__setstate__ gap) Close the load-time execution gap: torch::jit::load runs a module's __setstate__ during deserialization, before CModelGraphValidator (which walks an already-loaded module's graph) can run. A forbidden op hidden in __setstate__ therefore executes at load time and never appears in the forward graph the validator inspects. Add CModelGraphValidator::scanSerialisedCodeForForbiddenOps, which uses a PyTorchStreamReader over the buffered archive bytes to textually scan the serialised TorchScript code (code/**/*.py, including every submodule's __setstate__) for the forbidden aten ops, WITHOUT loading the model. Main.cc runs this scan before torch::jit::load and rejects via HANDLE_FATAL, so no model code executes. Expose the buffered bytes via CBufferedIStreamAdapter::buffer(). This is a defense-in-depth textual check targeted at the curated forbidden set (from_file, as_strided, save); the post-load allowlist validation is unchanged, and seccomp remains the syscall backstop. Unit tests cover both setstate fixtures, forward-graph attacks, a benign model (no false positive), and malformed input; the binary integration test now expects the setstate models to be rejected with the forbidden-operations message. Co-authored-by: Cursor --- .../CBufferedIStreamAdapter.h | 5 + bin/pytorch_inference/CModelGraphValidator.cc | 95 +++++++++++++++++++ bin/pytorch_inference/CModelGraphValidator.h | 21 ++++ bin/pytorch_inference/Main.cc | 21 ++++ .../unittest/CModelGraphValidatorTest.cc | 86 +++++++++++++++++ test/test_pytorch_inference_evil_models.py | 2 + 6 files changed, 230 insertions(+) diff --git a/bin/pytorch_inference/CBufferedIStreamAdapter.h b/bin/pytorch_inference/CBufferedIStreamAdapter.h index d1f845e75..2f9bf9b3b 100644 --- a/bin/pytorch_inference/CBufferedIStreamAdapter.h +++ b/bin/pytorch_inference/CBufferedIStreamAdapter.h @@ -51,6 +51,11 @@ class CBufferedIStreamAdapter : public caffe2::serialize::ReadAdapterInterface { std::size_t size() const override; std::size_t read(std::uint64_t pos, void* buf, std::size_t n, const char* what = "") const override; + //! Read-only view of the buffered model bytes. Valid until this object is + //! destroyed. Used to statically scan the archive before it is handed to + //! torch::jit::load (which would otherwise execute __setstate__ code). + const char* buffer() const { return m_Buffer.get(); } + CBufferedIStreamAdapter(const CBufferedIStreamAdapter&) = delete; CBufferedIStreamAdapter& operator=(const CBufferedIStreamAdapter&) = delete; diff --git a/bin/pytorch_inference/CModelGraphValidator.cc b/bin/pytorch_inference/CModelGraphValidator.cc index 01658b440..b292919d7 100644 --- a/bin/pytorch_inference/CModelGraphValidator.cc +++ b/bin/pytorch_inference/CModelGraphValidator.cc @@ -15,12 +15,60 @@ #include +#include +#include #include #include +#include +#include +#include +#include +#include +#include namespace ml { namespace torch { +namespace { + +//! Minimal in-memory ReadAdapterInterface so a PyTorchStreamReader can parse an +//! archive that is already fully buffered in memory, without taking ownership. +class CMemoryReadAdapter final : public caffe2::serialize::ReadAdapterInterface { +public: + CMemoryReadAdapter(const char* data, std::size_t size) + : m_Data{data}, m_Size{size} {} + + std::size_t size() const override { return m_Size; } + + std::size_t read(std::uint64_t pos, void* buf, std::size_t n, const char* /*what*/ = "") const override { + if (pos >= m_Size) { + return 0; + } + n = std::min(n, m_Size - static_cast(pos)); + std::memcpy(buf, m_Data + pos, n); + return n; + } + +private: + const char* m_Data; + std::size_t m_Size; +}; + +//! True for the serialised TorchScript source records (code/**/*.py) that hold +//! the executable code of the module and its submodules — including +//! __setstate__. Other records (data.pkl, constants.pkl, version, tensors) are +//! not executable source and are skipped. +bool isSerialisedCodeRecord(const std::string& name) { + constexpr std::string_view PY_EXT{".py"}; + if (name.size() < PY_EXT.size() || + name.compare(name.size() - PY_EXT.size(), PY_EXT.size(), PY_EXT) != 0) { + return false; + } + // Archive entries are typically "/code/....py"; accept a leading + // "code/" too in case the root prefix is absent. + return name.find("/code/") != std::string::npos || name.rfind("code/", 0) == 0; +} +} CModelGraphValidator::SResult CModelGraphValidator::validate(const ::torch::jit::Module& module) { @@ -95,6 +143,53 @@ void CModelGraphValidator::collectBlockOps(const ::torch::jit::Block& block, } } +CModelGraphValidator::TStringVec +CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::size_t size) { + TStringSet found; + + // Build the textual signatures we are looking for. In serialised TorchScript + // code an aten op "aten::foo" is printed as a call "torch.foo(". We only + // scan for the aten:: entries of the forbidden set: the prim:: entries + // (CallFunction/CallMethod) describe unresolved calls in an inlined graph and + // have no stable textual form in the serialised source. + std::vector> signatures; // (needle, qualified op) + constexpr std::string_view ATEN{"aten::"}; + for (const auto& op : CSupportedOperations::FORBIDDEN_OPERATIONS) { + if (op.size() > ATEN.size() && op.substr(0, ATEN.size()) == ATEN) { + std::string shortName{op.substr(ATEN.size())}; + signatures.emplace_back("torch." + shortName + "(", std::string{op}); + } + } + + try { + auto adapter = std::make_shared(data, size); + caffe2::serialize::PyTorchStreamReader reader{adapter}; + + for (const auto& name : reader.getAllRecords()) { + if (isSerialisedCodeRecord(name) == false) { + continue; + } + auto [recordData, recordSize] = reader.getRecord(name); + std::string_view code{static_cast(recordData.get()), recordSize}; + for (const auto& [needle, qualifiedOp] : signatures) { + if (code.find(needle) != std::string_view::npos) { + found.emplace(qualifiedOp); + } + } + } + } catch (const std::exception& e) { + // If the archive cannot be parsed as a stream we do not treat this as a + // rejection here: torch::jit::load will attempt the same parse and + // surface a clear error. The post-load graph validator still applies. + LOG_WARN(<< "Pre-load model code scan skipped (could not parse archive): " << e.what()); + return {}; + } + + TStringVec result{found.begin(), found.end()}; + std::sort(result.begin(), result.end()); + return result; +} + void CModelGraphValidator::collectModuleOps(const ::torch::jit::Module& module, TStringSet& ops, std::size_t& nodeCount) { diff --git a/bin/pytorch_inference/CModelGraphValidator.h b/bin/pytorch_inference/CModelGraphValidator.h index 2c589dab5..e24ddd121 100644 --- a/bin/pytorch_inference/CModelGraphValidator.h +++ b/bin/pytorch_inference/CModelGraphValidator.h @@ -14,6 +14,7 @@ #include +#include #include #include #include @@ -71,6 +72,26 @@ class CModelGraphValidator { const std::unordered_set& allowedOps, const std::unordered_set& forbiddenOps); + //! Statically scan the *serialised* TorchScript code inside a model archive + //! for forbidden operations, WITHOUT loading the model. + //! + //! This closes a load-time execution gap: torch::jit::load runs a module's + //! __setstate__ during deserialization, before validate() (which inspects an + //! already-loaded module) could ever run. A forbidden op hidden in + //! __setstate__ would therefore execute at load time and never appear in the + //! forward graph. Scanning the serialised code first lets us reject such a + //! model before any of its code runs. + //! + //! \p data / \p size are the raw bytes of the .pt (ZIP) archive. Returns the + //! sorted, de-duplicated qualified names of any forbidden operations found in + //! the serialised code. Returns empty if none are found, or if the archive + //! cannot be parsed (in which case torch::jit::load will surface the error). + //! + //! This is a textual, defense-in-depth check targeted at the curated + //! forbidden set; the post-load graph validation in validate() remains the + //! authoritative allowlist check. + static TStringVec scanSerialisedCodeForForbiddenOps(const char* data, std::size_t size); + 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 4aaca0049..f4371ccdb 100644 --- a/bin/pytorch_inference/Main.cc +++ b/bin/pytorch_inference/Main.cc @@ -73,6 +73,21 @@ void verifySafeModel(const torch::jit::script::Module& module_) { HANDLE_FATAL(<< "Model graph validation failed: " << e.what()); } } + +//! Reject models that carry a forbidden operation in their serialised code +//! BEFORE the model is loaded. torch::jit::load executes a module's +//! __setstate__ during deserialization, so a purely post-load check would run +//! too late for an op hidden there. \p modelData / \p modelSize are the raw +//! bytes of the buffered .pt archive. +void verifySafeModelBeforeLoad(const char* modelData, std::size_t modelSize) { + auto forbiddenOps = + ml::torch::CModelGraphValidator::scanSerialisedCodeForForbiddenOps(modelData, modelSize); + if (forbiddenOps.empty() == false) { + std::string ops = ml::core::CStringUtils::join(forbiddenOps, ", "); + HANDLE_FATAL(<< "Model contains forbidden operations: " << ops + << " (detected in serialised code before load)"); + } +} } torch::Tensor infer(torch::jit::script::Module& module_, @@ -315,6 +330,12 @@ int main(int argc, char** argv) { if (readAdapter->init() == false) { return EXIT_FAILURE; } + if (skipModelValidation == false) { + // Scan the serialised code before loading so that a forbidden op + // hidden in __setstate__ (which torch::jit::load would execute during + // deserialization) is rejected before any of the model's code runs. + verifySafeModelBeforeLoad(readAdapter->buffer(), readAdapter->size()); + } module_ = torch::jit::load(std::move(readAdapter)); if (skipModelValidation) { LOG_WARN(<< "Model graph validation SKIPPED — --skipModelValidation flag is set. " diff --git a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc index 9c4f7d6d5..dba78ade3 100644 --- a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc +++ b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc @@ -338,6 +338,17 @@ bool hasUnrecognisedOp(const CModelGraphValidator::SResult& result, const std::s return std::find(result.s_UnrecognisedOps.begin(), result.s_UnrecognisedOps.end(), op) != result.s_UnrecognisedOps.end(); } + +std::string readFileBytes(const std::string& path) { + std::ifstream file(path, std::ios::binary); + std::ostringstream buffer; + buffer << file.rdbuf(); + return buffer.str(); +} + +bool scanContains(const CModelGraphValidator::TStringVec& ops, const std::string& op) { + return std::find(ops.begin(), ops.end(), op) != ops.end(); +} } BOOST_AUTO_TEST_CASE(testMaliciousFileReader) { @@ -439,6 +450,81 @@ BOOST_AUTO_TEST_CASE(testMaliciousRopExploit) { BOOST_REQUIRE(hasForbiddenOp(result, "aten::as_strided")); } +// --- Pre-load __setstate__ scan tests --- +// +// torch::jit::load executes a module's __setstate__ during deserialization, +// before CModelGraphValidator::validate (which inspects an already-loaded +// module) could run. A forbidden op hidden in __setstate__ therefore never +// appears in the forward graph AND has already executed by load time. +// scanSerialisedCodeForForbiddenOps closes this gap by inspecting the +// serialised code textually, WITHOUT loading the model — so these tests never +// call torch::jit::load on the attack fixtures. + +BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReader) { + // aten::from_file lives only in __setstate__; the forward graph is benign, + // so the post-load validator would miss it. The pre-load scan must catch it + // without loading (and thus without executing) the model. + std::string bytes = + readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(forbidden, "aten::from_file")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReaderInSubmodule) { + // The forbidden op is hidden in a submodule's __setstate__; the scan must + // inspect every serialised code record, not just the top-level module's. + std::string bytes = readFileBytes( + "testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(forbidden, "aten::from_file")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsForwardForbiddenOps) { + // The scan is textual, so it also flags forbidden ops that appear in the + // forward graph (a superset of the load-time surface). + std::string fileReader = + readFileBytes("testfiles/malicious_models/malicious_file_reader.pt"); + BOOST_REQUIRE(fileReader.empty() == false); + auto forbiddenFileReader = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + fileReader.data(), fileReader.size()); + BOOST_REQUIRE(scanContains(forbiddenFileReader, "aten::from_file")); + + std::string heapLeak = + readFileBytes("testfiles/malicious_models/malicious_heap_leak.pt"); + BOOST_REQUIRE(heapLeak.empty() == false); + auto forbiddenHeapLeak = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + heapLeak.data(), heapLeak.size()); + BOOST_REQUIRE(scanContains(forbiddenHeapLeak, "aten::as_strided")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsBenignModel) { + // A legitimate model must not trip the scan (no false positives). + std::string bytes = readFileBytes("testfiles/e5_with_norm.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(forbidden.empty()); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanHandlesGarbageInput) { + // Non-archive input must not throw or crash; torch::jit::load will surface + // the real error later. The scan returns no forbidden ops. + std::string garbage = "this is not a valid .pt archive"; + auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + garbage.data(), garbage.size()); + BOOST_REQUIRE(forbidden.empty()); +} + // --- Prepacked model compatibility tests --- // // These load TorchScript models that mirror the ops used by Elasticsearch's diff --git a/test/test_pytorch_inference_evil_models.py b/test/test_pytorch_inference_evil_models.py index 82c04b53f..442d70a75 100644 --- a/test/test_pytorch_inference_evil_models.py +++ b/test/test_pytorch_inference_evil_models.py @@ -309,11 +309,13 @@ def find_pytorch_inference() -> str: "class": SetStateFileReaderModel, "expect_rejected": True, "description": "aten::from_file in __setstate__ (runs at load time)", + "expect_stderr_contains": "forbidden operations", }, "setstate_file_reader_submodule": { "class": SetStateFileReaderInSubmodule, "expect_rejected": True, "description": "aten::from_file in a submodule's __setstate__", + "expect_stderr_contains": "forbidden operations", }, } From 4438b8a79414c9520fe86ee26e1da04204c76247 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Tue, 21 Jul 2026 09:55:42 +1200 Subject: [PATCH 3/8] Update docs/changelog/3078.yaml --- docs/changelog/3078.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 docs/changelog/3078.yaml diff --git a/docs/changelog/3078.yaml b/docs/changelog/3078.yaml new file mode 100644 index 000000000..b9d7cc21b --- /dev/null +++ b/docs/changelog/3078.yaml @@ -0,0 +1,5 @@ +area: Machine Learning +issues: [] +pr: 3078 +summary: Reproduce and fix __setstate__ load-time execution gap in model validation +type: bug From f58df058e40084f6dfec67e058f178ad56242a53 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Tue, 21 Jul 2026 10:21:40 +1200 Subject: [PATCH 4/8] Formatting --- bin/pytorch_inference/CModelGraphValidator.cc | 10 ++++++---- bin/pytorch_inference/Main.cc | 4 ++-- .../unittest/CModelGraphValidatorTest.cc | 9 +++------ 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/bin/pytorch_inference/CModelGraphValidator.cc b/bin/pytorch_inference/CModelGraphValidator.cc index b292919d7..160364d9c 100644 --- a/bin/pytorch_inference/CModelGraphValidator.cc +++ b/bin/pytorch_inference/CModelGraphValidator.cc @@ -40,7 +40,8 @@ class CMemoryReadAdapter final : public caffe2::serialize::ReadAdapterInterface std::size_t size() const override { return m_Size; } - std::size_t read(std::uint64_t pos, void* buf, std::size_t n, const char* /*what*/ = "") const override { + std::size_t + read(std::uint64_t pos, void* buf, std::size_t n, const char* /*what*/ = "") const override { if (pos >= m_Size) { return 0; } @@ -169,9 +170,9 @@ CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::s if (isSerialisedCodeRecord(name) == false) { continue; } - auto [recordData, recordSize] = reader.getRecord(name); + auto[recordData, recordSize] = reader.getRecord(name); std::string_view code{static_cast(recordData.get()), recordSize}; - for (const auto& [needle, qualifiedOp] : signatures) { + for (const auto & [ needle, qualifiedOp ] : signatures) { if (code.find(needle) != std::string_view::npos) { found.emplace(qualifiedOp); } @@ -181,7 +182,8 @@ CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::s // If the archive cannot be parsed as a stream we do not treat this as a // rejection here: torch::jit::load will attempt the same parse and // surface a clear error. The post-load graph validator still applies. - LOG_WARN(<< "Pre-load model code scan skipped (could not parse archive): " << e.what()); + LOG_WARN(<< "Pre-load model code scan skipped (could not parse archive): " + << e.what()); return {}; } diff --git a/bin/pytorch_inference/Main.cc b/bin/pytorch_inference/Main.cc index f4371ccdb..078153a95 100644 --- a/bin/pytorch_inference/Main.cc +++ b/bin/pytorch_inference/Main.cc @@ -80,8 +80,8 @@ void verifySafeModel(const torch::jit::script::Module& module_) { //! too late for an op hidden there. \p modelData / \p modelSize are the raw //! bytes of the buffered .pt archive. void verifySafeModelBeforeLoad(const char* modelData, std::size_t modelSize) { - auto forbiddenOps = - ml::torch::CModelGraphValidator::scanSerialisedCodeForForbiddenOps(modelData, modelSize); + auto forbiddenOps = ml::torch::CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + modelData, modelSize); if (forbiddenOps.empty() == false) { std::string ops = ml::core::CStringUtils::join(forbiddenOps, ", "); HANDLE_FATAL(<< "Model contains forbidden operations: " << ops diff --git a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc index dba78ade3..583205efa 100644 --- a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc +++ b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc @@ -464,8 +464,7 @@ BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReader) { // aten::from_file lives only in __setstate__; the forward graph is benign, // so the post-load validator would miss it. The pre-load scan must catch it // without loading (and thus without executing) the model. - std::string bytes = - readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt"); + std::string bytes = readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt"); BOOST_REQUIRE(bytes.empty() == false); auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( @@ -490,15 +489,13 @@ BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReaderInSubmodule) { BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsForwardForbiddenOps) { // The scan is textual, so it also flags forbidden ops that appear in the // forward graph (a superset of the load-time surface). - std::string fileReader = - readFileBytes("testfiles/malicious_models/malicious_file_reader.pt"); + std::string fileReader = readFileBytes("testfiles/malicious_models/malicious_file_reader.pt"); BOOST_REQUIRE(fileReader.empty() == false); auto forbiddenFileReader = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( fileReader.data(), fileReader.size()); BOOST_REQUIRE(scanContains(forbiddenFileReader, "aten::from_file")); - std::string heapLeak = - readFileBytes("testfiles/malicious_models/malicious_heap_leak.pt"); + std::string heapLeak = readFileBytes("testfiles/malicious_models/malicious_heap_leak.pt"); BOOST_REQUIRE(heapLeak.empty() == false); auto forbiddenHeapLeak = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( heapLeak.data(), heapLeak.size()); From 5957acdb53755eabac7a79b463d5331eb6ed2f17 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Tue, 21 Jul 2026 10:27:50 +1200 Subject: [PATCH 5/8] Tighten setstate load-timing probe to require a clean load Address Copilot review: the benign probe must load cleanly for the marker to prove load-time execution. Require returncode == 0 and that LOAD_RETURNED was printed, and print returncode + stderr tail on failure, so a crash that happens to print the marker first is no longer a false positive. Co-authored-by: Cursor --- test/test_setstate_load_timing.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/test/test_setstate_load_timing.py b/test/test_setstate_load_timing.py index e2ba89f02..6148bcc2a 100644 --- a/test/test_setstate_load_timing.py +++ b/test/test_setstate_load_timing.py @@ -139,17 +139,29 @@ def test_setstate_executes_during_load() -> bool: capture_output=True, text=True, timeout=120, ) stdout = proc.stdout + + # The benign probe must load *cleanly*: require a zero exit and that + # load() actually returned (LOAD_RETURNED). Otherwise a crash that + # happened to print the marker first would be a false positive — so we + # only accept the marker as proof when the whole load completed and the + # marker was printed strictly before load() returned. + loaded_cleanly = proc.returncode == 0 and "LOAD_RETURNED" in stdout marker_before_return = ( - _LOAD_MARKER in stdout + loaded_cleanly + and _LOAD_MARKER in stdout and stdout.index(_LOAD_MARKER) < stdout.index("LOAD_RETURNED") - ) if "LOAD_RETURNED" in stdout else (_LOAD_MARKER in stdout) + ) if marker_before_return: print(f" OK: '{_LOAD_MARKER}' printed during load — " "code ran before load() returned (and before validation).") else: - print(" FAIL: marker not observed during load.") + print(" FAIL: marker not observed during a clean load.") + print(f" returncode: {proc.returncode}") print(f" stdout: {stdout!r}") + if proc.stderr.strip(): + tail = "\n".join(proc.stderr.strip().splitlines()[-5:]) + print(f" stderr (tail):\n{tail}") ok = False print() return ok From e88f958faabd499f21ab41c8e77c0034a86b0c37 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Tue, 21 Jul 2026 15:38:35 +1200 Subject: [PATCH 6/8] Reject TorchScript custom state hooks before load Match the #12621 reporter remediation by refusing archives that embed __setstate__/__getstate__, and keep the forbidden-op pre-scan as defense in depth. Allowlisted-op RCE gadgets in setstate are otherwise missed. Co-authored-by: Cursor --- bin/pytorch_inference/CModelGraphValidator.cc | 58 +++++++++++------ bin/pytorch_inference/CModelGraphValidator.h | 49 +++++++++----- bin/pytorch_inference/Main.cc | 28 ++++---- .../unittest/CModelGraphValidatorTest.cc | 64 +++++++++++++------ docs/changelog/3078.yaml | 2 +- 5 files changed, 135 insertions(+), 66 deletions(-) diff --git a/bin/pytorch_inference/CModelGraphValidator.cc b/bin/pytorch_inference/CModelGraphValidator.cc index 160364d9c..138fc27a6 100644 --- a/bin/pytorch_inference/CModelGraphValidator.cc +++ b/bin/pytorch_inference/CModelGraphValidator.cc @@ -144,15 +144,15 @@ void CModelGraphValidator::collectBlockOps(const ::torch::jit::Block& block, } } -CModelGraphValidator::TStringVec -CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::size_t size) { - TStringSet found; - - // Build the textual signatures we are looking for. In serialised TorchScript - // code an aten op "aten::foo" is printed as a call "torch.foo(". We only - // scan for the aten:: entries of the forbidden set: the prim:: entries - // (CallFunction/CallMethod) describe unresolved calls in an inlined graph and - // have no stable textual form in the serialised source. +CModelGraphValidator::SPreLoadScanResult +CModelGraphValidator::scanArchiveBeforeLoad(const char* data, std::size_t size) { + SPreLoadScanResult result; + TStringSet hooks; + TStringSet forbidden; + + // In serialised TorchScript code an aten op "aten::foo" is printed as + // "torch.foo(". Only aten:: entries of the forbidden set have a stable + // textual form; prim::CallFunction/CallMethod do not. std::vector> signatures; // (needle, qualified op) constexpr std::string_view ATEN{"aten::"}; for (const auto& op : CSupportedOperations::FORBIDDEN_OPERATIONS) { @@ -162,19 +162,32 @@ CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::s } } + constexpr std::string_view SETSTATE{"__setstate__"}; + constexpr std::string_view GETSTATE{"__getstate__"}; + try { auto adapter = std::make_shared(data, size); caffe2::serialize::PyTorchStreamReader reader{adapter}; for (const auto& name : reader.getAllRecords()) { - if (isSerialisedCodeRecord(name) == false) { - continue; - } auto[recordData, recordSize] = reader.getRecord(name); - std::string_view code{static_cast(recordData.get()), recordSize}; - for (const auto & [ needle, qualifiedOp ] : signatures) { - if (code.find(needle) != std::string_view::npos) { - found.emplace(qualifiedOp); + std::string_view bytes{static_cast(recordData.get()), recordSize}; + + // Reporter remediation (elastic/security#12621): reject any archive + // that embeds custom state hooks. Scan every record — not just + // code/*.py — so hooks cannot hide in debug_pkl / adjacent entries. + if (bytes.find(SETSTATE) != std::string_view::npos) { + hooks.emplace("__setstate__"); + } + if (bytes.find(GETSTATE) != std::string_view::npos) { + hooks.emplace("__getstate__"); + } + + if (isSerialisedCodeRecord(name)) { + for (const auto & [ needle, qualifiedOp ] : signatures) { + if (bytes.find(needle) != std::string_view::npos) { + forbidden.emplace(qualifiedOp); + } } } } @@ -182,16 +195,23 @@ CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::s // If the archive cannot be parsed as a stream we do not treat this as a // rejection here: torch::jit::load will attempt the same parse and // surface a clear error. The post-load graph validator still applies. - LOG_WARN(<< "Pre-load model code scan skipped (could not parse archive): " + LOG_WARN(<< "Pre-load model archive scan skipped (could not parse archive): " << e.what()); return {}; } - TStringVec result{found.begin(), found.end()}; - std::sort(result.begin(), result.end()); + result.s_CustomStateHooks.assign(hooks.begin(), hooks.end()); + result.s_ForbiddenOps.assign(forbidden.begin(), forbidden.end()); + std::sort(result.s_CustomStateHooks.begin(), result.s_CustomStateHooks.end()); + std::sort(result.s_ForbiddenOps.begin(), result.s_ForbiddenOps.end()); return result; } +CModelGraphValidator::TStringVec +CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::size_t size) { + return scanArchiveBeforeLoad(data, size).s_ForbiddenOps; +} + void CModelGraphValidator::collectModuleOps(const ::torch::jit::Module& module, TStringSet& ops, std::size_t& nodeCount) { diff --git a/bin/pytorch_inference/CModelGraphValidator.h b/bin/pytorch_inference/CModelGraphValidator.h index e24ddd121..e7d878052 100644 --- a/bin/pytorch_inference/CModelGraphValidator.h +++ b/bin/pytorch_inference/CModelGraphValidator.h @@ -72,24 +72,41 @@ class CModelGraphValidator { const std::unordered_set& allowedOps, const std::unordered_set& forbiddenOps); - //! Statically scan the *serialised* TorchScript code inside a model archive - //! for forbidden operations, WITHOUT loading the model. - //! - //! This closes a load-time execution gap: torch::jit::load runs a module's - //! __setstate__ during deserialization, before validate() (which inspects an - //! already-loaded module) could ever run. A forbidden op hidden in - //! __setstate__ would therefore execute at load time and never appear in the - //! forward graph. Scanning the serialised code first lets us reject such a - //! model before any of its code runs. + //! Result of a pre-load archive scan. Both fields must be empty for the + //! model to be considered safe to hand to torch::jit::load. + struct SPreLoadScanResult { + //! Custom TorchScript state hooks found anywhere in the archive + //! ("__setstate__", "__getstate__"). torch::jit::load executes + //! __setstate__ during deserialization — before post-load graph + //! validation — so any custom hooks are rejected outright. + TStringVec s_CustomStateHooks; + //! Forbidden operations found in serialised code/**/*.py records + //! (textual "torch.(" needles derived from FORBIDDEN_OPERATIONS). + TStringVec s_ForbiddenOps; + }; + + //! Statically scan a model archive BEFORE torch::jit::load, without + //! executing any of the model's code. //! - //! \p data / \p size are the raw bytes of the .pt (ZIP) archive. Returns the - //! sorted, de-duplicated qualified names of any forbidden operations found in - //! the serialised code. Returns empty if none are found, or if the archive - //! cannot be parsed (in which case torch::jit::load will surface the error). + //! Closes the load-time execution gap exploited in elastic/security#12621: + //! torch::jit::load runs a module's __setstate__ during deserialization, + //! and an attacker can hide allowlisted-but-dangerous logic (or memory- + //! corruption gadgets) there that never appears in the forward graph. + //! Matching the reporter's remediation, any archive containing the + //! literal substrings "__setstate__" or "__getstate__" is rejected. + //! As defense in depth, serialised code is also scanned for the curated + //! forbidden-op set (so a forbidden op in forward is still caught here + //! even if no custom state hooks are present). //! - //! This is a textual, defense-in-depth check targeted at the curated - //! forbidden set; the post-load graph validation in validate() remains the - //! authoritative allowlist check. + //! \p data / \p size are the raw bytes of the .pt (ZIP) archive. Returns + //! empty vectors if nothing is found, or if the archive cannot be parsed + //! (in which case torch::jit::load will surface the error). The post-load + //! graph validation in validate() remains the authoritative allowlist check. + static SPreLoadScanResult scanArchiveBeforeLoad(const char* data, std::size_t size); + + //! Convenience wrapper: returns only the forbidden-op half of + //! scanArchiveBeforeLoad(). Kept for unit tests that target that check + //! in isolation. static TStringVec scanSerialisedCodeForForbiddenOps(const char* data, std::size_t size); private: diff --git a/bin/pytorch_inference/Main.cc b/bin/pytorch_inference/Main.cc index 078153a95..3df788741 100644 --- a/bin/pytorch_inference/Main.cc +++ b/bin/pytorch_inference/Main.cc @@ -74,16 +74,20 @@ void verifySafeModel(const torch::jit::script::Module& module_) { } } -//! Reject models that carry a forbidden operation in their serialised code -//! BEFORE the model is loaded. torch::jit::load executes a module's -//! __setstate__ during deserialization, so a purely post-load check would run -//! too late for an op hidden there. \p modelData / \p modelSize are the raw -//! bytes of the buffered .pt archive. +//! Reject unsafe models BEFORE torch::jit::load. Load executes a module's +//! __setstate__ during deserialization, so a purely post-load check runs too +//! late. Matching elastic/security#12621 / the reporter's remediation, any +//! custom __setstate__/__getstate__ hooks are refused outright; forbidden ops +//! in serialised code are also refused as defense in depth. +//! \p modelData / \p modelSize are the raw bytes of the buffered .pt archive. void verifySafeModelBeforeLoad(const char* modelData, std::size_t modelSize) { - auto forbiddenOps = ml::torch::CModelGraphValidator::scanSerialisedCodeForForbiddenOps( - modelData, modelSize); - if (forbiddenOps.empty() == false) { - std::string ops = ml::core::CStringUtils::join(forbiddenOps, ", "); + auto scan = ml::torch::CModelGraphValidator::scanArchiveBeforeLoad(modelData, modelSize); + if (scan.s_CustomStateHooks.empty() == false) { + std::string hooks = ml::core::CStringUtils::join(scan.s_CustomStateHooks, ", "); + HANDLE_FATAL(<< "Model archive contains custom state hooks: " << hooks); + } + if (scan.s_ForbiddenOps.empty() == false) { + std::string ops = ml::core::CStringUtils::join(scan.s_ForbiddenOps, ", "); HANDLE_FATAL(<< "Model contains forbidden operations: " << ops << " (detected in serialised code before load)"); } @@ -331,9 +335,9 @@ int main(int argc, char** argv) { return EXIT_FAILURE; } if (skipModelValidation == false) { - // Scan the serialised code before loading so that a forbidden op - // hidden in __setstate__ (which torch::jit::load would execute during - // deserialization) is rejected before any of the model's code runs. + // Reject custom state hooks and forbidden ops before loading so + // that __setstate__ (which torch::jit::load would execute during + // deserialization) never runs. verifySafeModelBeforeLoad(readAdapter->buffer(), readAdapter->size()); } module_ = torch::jit::load(std::move(readAdapter)); diff --git a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc index 583205efa..33893e5c1 100644 --- a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc +++ b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc @@ -450,15 +450,39 @@ BOOST_AUTO_TEST_CASE(testMaliciousRopExploit) { BOOST_REQUIRE(hasForbiddenOp(result, "aten::as_strided")); } -// --- Pre-load __setstate__ scan tests --- +// --- Pre-load archive scan tests --- // // torch::jit::load executes a module's __setstate__ during deserialization, // before CModelGraphValidator::validate (which inspects an already-loaded -// module) could run. A forbidden op hidden in __setstate__ therefore never -// appears in the forward graph AND has already executed by load time. -// scanSerialisedCodeForForbiddenOps closes this gap by inspecting the -// serialised code textually, WITHOUT loading the model — so these tests never -// call torch::jit::load on the attack fixtures. +// module) could run. Matching elastic/security#12621, custom state hooks are +// rejected outright; forbidden ops in serialised code are also flagged as +// defense in depth. These tests never call torch::jit::load on attack fixtures. + +BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooks) { + // The HackerOne #12621 RCE uses allowlisted ops inside __setstate__; a + // forbidden-op-only scan would miss it. Rejecting the hooks themselves + // is the reporter's remediation and must fire here. + std::string bytes = readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto scan = CModelGraphValidator::scanArchiveBeforeLoad(bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(scan.s_CustomStateHooks, "__setstate__")); + BOOST_REQUIRE(scanContains(scan.s_CustomStateHooks, "__getstate__")); + // Defense in depth: the fixture also embeds aten::from_file in setstate. + BOOST_REQUIRE(scanContains(scan.s_ForbiddenOps, "aten::from_file")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooksInSubmodule) { + std::string bytes = readFileBytes( + "testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto scan = CModelGraphValidator::scanArchiveBeforeLoad(bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(scan.s_CustomStateHooks, "__setstate__")); + BOOST_REQUIRE(scanContains(scan.s_ForbiddenOps, "aten::from_file")); +} BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReader) { // aten::from_file lives only in __setstate__; the forward graph is benign, @@ -491,35 +515,39 @@ BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsForwardForbiddenOps) { // forward graph (a superset of the load-time surface). std::string fileReader = readFileBytes("testfiles/malicious_models/malicious_file_reader.pt"); BOOST_REQUIRE(fileReader.empty() == false); - auto forbiddenFileReader = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + auto scanFileReader = CModelGraphValidator::scanArchiveBeforeLoad( fileReader.data(), fileReader.size()); - BOOST_REQUIRE(scanContains(forbiddenFileReader, "aten::from_file")); + BOOST_REQUIRE(scanFileReader.s_CustomStateHooks.empty()); + BOOST_REQUIRE(scanContains(scanFileReader.s_ForbiddenOps, "aten::from_file")); std::string heapLeak = readFileBytes("testfiles/malicious_models/malicious_heap_leak.pt"); BOOST_REQUIRE(heapLeak.empty() == false); - auto forbiddenHeapLeak = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + auto scanHeapLeak = CModelGraphValidator::scanArchiveBeforeLoad( heapLeak.data(), heapLeak.size()); - BOOST_REQUIRE(scanContains(forbiddenHeapLeak, "aten::as_strided")); + BOOST_REQUIRE(scanHeapLeak.s_CustomStateHooks.empty()); + BOOST_REQUIRE(scanContains(scanHeapLeak.s_ForbiddenOps, "aten::as_strided")); } BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsBenignModel) { - // A legitimate model must not trip the scan (no false positives). + // A legitimate model must not trip the scan (no false positives on hooks + // or forbidden ops — typical scripted transformers do not embed custom + // __setstate__/__getstate__). std::string bytes = readFileBytes("testfiles/e5_with_norm.pt"); BOOST_REQUIRE(bytes.empty() == false); - auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( - bytes.data(), bytes.size()); + auto scan = CModelGraphValidator::scanArchiveBeforeLoad(bytes.data(), bytes.size()); - BOOST_REQUIRE(forbidden.empty()); + BOOST_REQUIRE(scan.s_CustomStateHooks.empty()); + BOOST_REQUIRE(scan.s_ForbiddenOps.empty()); } BOOST_AUTO_TEST_CASE(testPreLoadScanHandlesGarbageInput) { // Non-archive input must not throw or crash; torch::jit::load will surface - // the real error later. The scan returns no forbidden ops. + // the real error later. The scan returns empty results. std::string garbage = "this is not a valid .pt archive"; - auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( - garbage.data(), garbage.size()); - BOOST_REQUIRE(forbidden.empty()); + auto scan = CModelGraphValidator::scanArchiveBeforeLoad(garbage.data(), garbage.size()); + BOOST_REQUIRE(scan.s_CustomStateHooks.empty()); + BOOST_REQUIRE(scan.s_ForbiddenOps.empty()); } // --- Prepacked model compatibility tests --- diff --git a/docs/changelog/3078.yaml b/docs/changelog/3078.yaml index b9d7cc21b..eb77de67e 100644 --- a/docs/changelog/3078.yaml +++ b/docs/changelog/3078.yaml @@ -1,5 +1,5 @@ area: Machine Learning issues: [] pr: 3078 -summary: Reproduce and fix __setstate__ load-time execution gap in model validation +summary: Reject TorchScript custom state hooks before load and scan for forbidden ops type: bug From 2770954a32445452c4a3c8eb84645294d9beae72 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Tue, 21 Jul 2026 15:45:03 +1200 Subject: [PATCH 7/8] Drop redundant pre-load forbidden-op scan Custom state-hook rejection is enough for the load-time execution gap; forward-graph ops remain covered by post-load validate(). Co-authored-by: Cursor --- bin/pytorch_inference/CModelGraphValidator.cc | 57 ++---------- bin/pytorch_inference/CModelGraphValidator.h | 40 +++------ bin/pytorch_inference/Main.cc | 31 +++---- .../unittest/CModelGraphValidatorTest.cc | 89 ++++++------------- docs/changelog/3078.yaml | 2 +- 5 files changed, 59 insertions(+), 160 deletions(-) diff --git a/bin/pytorch_inference/CModelGraphValidator.cc b/bin/pytorch_inference/CModelGraphValidator.cc index 138fc27a6..c9b2bd21e 100644 --- a/bin/pytorch_inference/CModelGraphValidator.cc +++ b/bin/pytorch_inference/CModelGraphValidator.cc @@ -23,8 +23,6 @@ #include #include #include -#include -#include #include namespace ml { @@ -54,21 +52,6 @@ class CMemoryReadAdapter final : public caffe2::serialize::ReadAdapterInterface const char* m_Data; std::size_t m_Size; }; - -//! True for the serialised TorchScript source records (code/**/*.py) that hold -//! the executable code of the module and its submodules — including -//! __setstate__. Other records (data.pkl, constants.pkl, version, tensors) are -//! not executable source and are skipped. -bool isSerialisedCodeRecord(const std::string& name) { - constexpr std::string_view PY_EXT{".py"}; - if (name.size() < PY_EXT.size() || - name.compare(name.size() - PY_EXT.size(), PY_EXT.size(), PY_EXT) != 0) { - return false; - } - // Archive entries are typically "/code/....py"; accept a leading - // "code/" too in case the root prefix is absent. - return name.find("/code/") != std::string::npos || name.rfind("code/", 0) == 0; -} } CModelGraphValidator::SResult CModelGraphValidator::validate(const ::torch::jit::Module& module) { @@ -144,23 +127,9 @@ void CModelGraphValidator::collectBlockOps(const ::torch::jit::Block& block, } } -CModelGraphValidator::SPreLoadScanResult -CModelGraphValidator::scanArchiveBeforeLoad(const char* data, std::size_t size) { - SPreLoadScanResult result; +CModelGraphValidator::TStringVec +CModelGraphValidator::scanArchiveForCustomStateHooks(const char* data, std::size_t size) { TStringSet hooks; - TStringSet forbidden; - - // In serialised TorchScript code an aten op "aten::foo" is printed as - // "torch.foo(". Only aten:: entries of the forbidden set have a stable - // textual form; prim::CallFunction/CallMethod do not. - std::vector> signatures; // (needle, qualified op) - constexpr std::string_view ATEN{"aten::"}; - for (const auto& op : CSupportedOperations::FORBIDDEN_OPERATIONS) { - if (op.size() > ATEN.size() && op.substr(0, ATEN.size()) == ATEN) { - std::string shortName{op.substr(ATEN.size())}; - signatures.emplace_back("torch." + shortName + "(", std::string{op}); - } - } constexpr std::string_view SETSTATE{"__setstate__"}; constexpr std::string_view GETSTATE{"__getstate__"}; @@ -182,36 +151,24 @@ CModelGraphValidator::scanArchiveBeforeLoad(const char* data, std::size_t size) if (bytes.find(GETSTATE) != std::string_view::npos) { hooks.emplace("__getstate__"); } - - if (isSerialisedCodeRecord(name)) { - for (const auto & [ needle, qualifiedOp ] : signatures) { - if (bytes.find(needle) != std::string_view::npos) { - forbidden.emplace(qualifiedOp); - } - } + if (hooks.size() == 2) { + break; } } } catch (const std::exception& e) { // If the archive cannot be parsed as a stream we do not treat this as a // rejection here: torch::jit::load will attempt the same parse and // surface a clear error. The post-load graph validator still applies. - LOG_WARN(<< "Pre-load model archive scan skipped (could not parse archive): " + LOG_WARN(<< "Pre-load state-hook scan skipped (could not parse archive): " << e.what()); return {}; } - result.s_CustomStateHooks.assign(hooks.begin(), hooks.end()); - result.s_ForbiddenOps.assign(forbidden.begin(), forbidden.end()); - std::sort(result.s_CustomStateHooks.begin(), result.s_CustomStateHooks.end()); - std::sort(result.s_ForbiddenOps.begin(), result.s_ForbiddenOps.end()); + TStringVec result{hooks.begin(), hooks.end()}; + std::sort(result.begin(), result.end()); return result; } -CModelGraphValidator::TStringVec -CModelGraphValidator::scanSerialisedCodeForForbiddenOps(const char* data, std::size_t size) { - return scanArchiveBeforeLoad(data, size).s_ForbiddenOps; -} - void CModelGraphValidator::collectModuleOps(const ::torch::jit::Module& module, TStringSet& ops, std::size_t& nodeCount) { diff --git a/bin/pytorch_inference/CModelGraphValidator.h b/bin/pytorch_inference/CModelGraphValidator.h index e7d878052..198800af5 100644 --- a/bin/pytorch_inference/CModelGraphValidator.h +++ b/bin/pytorch_inference/CModelGraphValidator.h @@ -72,42 +72,22 @@ class CModelGraphValidator { const std::unordered_set& allowedOps, const std::unordered_set& forbiddenOps); - //! Result of a pre-load archive scan. Both fields must be empty for the - //! model to be considered safe to hand to torch::jit::load. - struct SPreLoadScanResult { - //! Custom TorchScript state hooks found anywhere in the archive - //! ("__setstate__", "__getstate__"). torch::jit::load executes - //! __setstate__ during deserialization — before post-load graph - //! validation — so any custom hooks are rejected outright. - TStringVec s_CustomStateHooks; - //! Forbidden operations found in serialised code/**/*.py records - //! (textual "torch.(" needles derived from FORBIDDEN_OPERATIONS). - TStringVec s_ForbiddenOps; - }; - - //! Statically scan a model archive BEFORE torch::jit::load, without + //! Statically scan a model archive BEFORE torch::jit::load for custom + //! TorchScript state hooks ("__setstate__", "__getstate__"), without //! executing any of the model's code. //! //! Closes the load-time execution gap exploited in elastic/security#12621: //! torch::jit::load runs a module's __setstate__ during deserialization, - //! and an attacker can hide allowlisted-but-dangerous logic (or memory- - //! corruption gadgets) there that never appears in the forward graph. - //! Matching the reporter's remediation, any archive containing the - //! literal substrings "__setstate__" or "__getstate__" is rejected. - //! As defense in depth, serialised code is also scanned for the curated - //! forbidden-op set (so a forbidden op in forward is still caught here - //! even if no custom state hooks are present). + //! before post-load graph validation. Matching the reporter's remediation, + //! any archive containing those literal substrings is rejected. Ops that + //! only run when methods are invoked (e.g. forward) remain the job of the + //! post-load allowlist / forbid checks in validate(). //! //! \p data / \p size are the raw bytes of the .pt (ZIP) archive. Returns - //! empty vectors if nothing is found, or if the archive cannot be parsed - //! (in which case torch::jit::load will surface the error). The post-load - //! graph validation in validate() remains the authoritative allowlist check. - static SPreLoadScanResult scanArchiveBeforeLoad(const char* data, std::size_t size); - - //! Convenience wrapper: returns only the forbidden-op half of - //! scanArchiveBeforeLoad(). Kept for unit tests that target that check - //! in isolation. - static TStringVec scanSerialisedCodeForForbiddenOps(const char* data, std::size_t size); + //! 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). + static TStringVec scanArchiveForCustomStateHooks(const char* data, std::size_t size); private: //! Collect all operation names from a block, recursing into sub-blocks. diff --git a/bin/pytorch_inference/Main.cc b/bin/pytorch_inference/Main.cc index 3df788741..76c57d610 100644 --- a/bin/pytorch_inference/Main.cc +++ b/bin/pytorch_inference/Main.cc @@ -74,22 +74,19 @@ void verifySafeModel(const torch::jit::script::Module& module_) { } } -//! Reject unsafe models BEFORE torch::jit::load. Load executes a module's -//! __setstate__ during deserialization, so a purely post-load check runs too -//! late. Matching elastic/security#12621 / the reporter's remediation, any -//! custom __setstate__/__getstate__ hooks are refused outright; forbidden ops -//! in serialised code are also refused as defense in depth. +//! Reject models with custom TorchScript state hooks BEFORE torch::jit::load. +//! Load executes __setstate__ during deserialization, so post-load graph +//! validation runs too late. Matching elastic/security#12621 / the reporter's +//! remediation, any __setstate__/__getstate__ hooks are refused outright. +//! 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 scan = ml::torch::CModelGraphValidator::scanArchiveBeforeLoad(modelData, modelSize); - if (scan.s_CustomStateHooks.empty() == false) { - std::string hooks = ml::core::CStringUtils::join(scan.s_CustomStateHooks, ", "); - HANDLE_FATAL(<< "Model archive contains custom state hooks: " << hooks); - } - if (scan.s_ForbiddenOps.empty() == false) { - std::string ops = ml::core::CStringUtils::join(scan.s_ForbiddenOps, ", "); - HANDLE_FATAL(<< "Model contains forbidden operations: " << ops - << " (detected in serialised code before load)"); + 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); } } } @@ -335,9 +332,9 @@ int main(int argc, char** argv) { return EXIT_FAILURE; } if (skipModelValidation == false) { - // Reject custom state hooks and forbidden ops before loading so - // that __setstate__ (which torch::jit::load would execute during - // deserialization) never runs. + // Reject custom state hooks before loading so that __setstate__ + // (which torch::jit::load would execute during deserialization) + // never runs. Op allowlisting remains a post-load check. verifySafeModelBeforeLoad(readAdapter->buffer(), readAdapter->size()); } module_ = torch::jit::load(std::move(readAdapter)); diff --git a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc index 33893e5c1..0204123d3 100644 --- a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc +++ b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc @@ -450,27 +450,26 @@ BOOST_AUTO_TEST_CASE(testMaliciousRopExploit) { BOOST_REQUIRE(hasForbiddenOp(result, "aten::as_strided")); } -// --- Pre-load archive scan tests --- +// --- Pre-load custom state-hook scan tests --- // // torch::jit::load executes a module's __setstate__ during deserialization, // before CModelGraphValidator::validate (which inspects an already-loaded // module) could run. Matching elastic/security#12621, custom state hooks are -// rejected outright; forbidden ops in serialised code are also flagged as -// defense in depth. These tests never call torch::jit::load on attack fixtures. +// rejected outright before load. Ops that only run when methods are invoked +// remain covered by the post-load allowlist / forbid checks. These tests +// never call torch::jit::load on attack fixtures. BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooks) { - // The HackerOne #12621 RCE uses allowlisted ops inside __setstate__; a - // forbidden-op-only scan would miss it. Rejecting the hooks themselves - // is the reporter's remediation and must fire here. + // The HackerOne #12621 RCE uses allowlisted ops inside __setstate__; + // rejecting the hooks themselves is the reporter's remediation. std::string bytes = readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt"); BOOST_REQUIRE(bytes.empty() == false); - auto scan = CModelGraphValidator::scanArchiveBeforeLoad(bytes.data(), bytes.size()); + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); - BOOST_REQUIRE(scanContains(scan.s_CustomStateHooks, "__setstate__")); - BOOST_REQUIRE(scanContains(scan.s_CustomStateHooks, "__getstate__")); - // Defense in depth: the fixture also embeds aten::from_file in setstate. - BOOST_REQUIRE(scanContains(scan.s_ForbiddenOps, "aten::from_file")); + BOOST_REQUIRE(scanContains(hooks, "__setstate__")); + BOOST_REQUIRE(scanContains(hooks, "__getstate__")); } BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooksInSubmodule) { @@ -478,76 +477,42 @@ BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooksInSubmodule) { "testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt"); BOOST_REQUIRE(bytes.empty() == false); - auto scan = CModelGraphValidator::scanArchiveBeforeLoad(bytes.data(), bytes.size()); - - BOOST_REQUIRE(scanContains(scan.s_CustomStateHooks, "__setstate__")); - BOOST_REQUIRE(scanContains(scan.s_ForbiddenOps, "aten::from_file")); -} - -BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReader) { - // aten::from_file lives only in __setstate__; the forward graph is benign, - // so the post-load validator would miss it. The pre-load scan must catch it - // without loading (and thus without executing) the model. - std::string bytes = readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt"); - BOOST_REQUIRE(bytes.empty() == false); - - auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( bytes.data(), bytes.size()); - BOOST_REQUIRE(scanContains(forbidden, "aten::from_file")); + BOOST_REQUIRE(scanContains(hooks, "__setstate__")); } -BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateFileReaderInSubmodule) { - // The forbidden op is hidden in a submodule's __setstate__; the scan must - // inspect every serialised code record, not just the top-level module's. - std::string bytes = readFileBytes( - "testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt"); +BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsBenignModel) { + // Typical scripted transformers do not embed custom __setstate__/__getstate__. + std::string bytes = readFileBytes("testfiles/e5_with_norm.pt"); BOOST_REQUIRE(bytes.empty() == false); - auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( bytes.data(), bytes.size()); - BOOST_REQUIRE(scanContains(forbidden, "aten::from_file")); -} - -BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsForwardForbiddenOps) { - // The scan is textual, so it also flags forbidden ops that appear in the - // forward graph (a superset of the load-time surface). - std::string fileReader = readFileBytes("testfiles/malicious_models/malicious_file_reader.pt"); - BOOST_REQUIRE(fileReader.empty() == false); - auto scanFileReader = CModelGraphValidator::scanArchiveBeforeLoad( - fileReader.data(), fileReader.size()); - BOOST_REQUIRE(scanFileReader.s_CustomStateHooks.empty()); - BOOST_REQUIRE(scanContains(scanFileReader.s_ForbiddenOps, "aten::from_file")); - - std::string heapLeak = readFileBytes("testfiles/malicious_models/malicious_heap_leak.pt"); - BOOST_REQUIRE(heapLeak.empty() == false); - auto scanHeapLeak = CModelGraphValidator::scanArchiveBeforeLoad( - heapLeak.data(), heapLeak.size()); - BOOST_REQUIRE(scanHeapLeak.s_CustomStateHooks.empty()); - BOOST_REQUIRE(scanContains(scanHeapLeak.s_ForbiddenOps, "aten::as_strided")); + BOOST_REQUIRE(hooks.empty()); } -BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsBenignModel) { - // A legitimate model must not trip the scan (no false positives on hooks - // or forbidden ops — typical scripted transformers do not embed custom - // __setstate__/__getstate__). - std::string bytes = readFileBytes("testfiles/e5_with_norm.pt"); +BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsForwardOnlyMaliciousModel) { + // Forbidden ops in forward (no custom state hooks) are not a pre-load + // concern — they are caught by post-load validate() after jit::load. + std::string bytes = readFileBytes("testfiles/malicious_models/malicious_file_reader.pt"); BOOST_REQUIRE(bytes.empty() == false); - auto scan = CModelGraphValidator::scanArchiveBeforeLoad(bytes.data(), bytes.size()); + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); - BOOST_REQUIRE(scan.s_CustomStateHooks.empty()); - BOOST_REQUIRE(scan.s_ForbiddenOps.empty()); + BOOST_REQUIRE(hooks.empty()); } BOOST_AUTO_TEST_CASE(testPreLoadScanHandlesGarbageInput) { // Non-archive input must not throw or crash; torch::jit::load will surface // the real error later. The scan returns empty results. std::string garbage = "this is not a valid .pt archive"; - auto scan = CModelGraphValidator::scanArchiveBeforeLoad(garbage.data(), garbage.size()); - BOOST_REQUIRE(scan.s_CustomStateHooks.empty()); - BOOST_REQUIRE(scan.s_ForbiddenOps.empty()); + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + garbage.data(), garbage.size()); + BOOST_REQUIRE(hooks.empty()); } // --- Prepacked model compatibility tests --- diff --git a/docs/changelog/3078.yaml b/docs/changelog/3078.yaml index eb77de67e..f2ef914ac 100644 --- a/docs/changelog/3078.yaml +++ b/docs/changelog/3078.yaml @@ -1,5 +1,5 @@ area: Machine Learning issues: [] pr: 3078 -summary: Reject TorchScript custom state hooks before load and scan for forbidden ops +summary: Reject TorchScript custom state hooks before load type: bug From 4b828ab920dfca37999f2aa7d66fdecaefd80b15 Mon Sep 17 00:00:00 2001 From: Ed Savage Date: Tue, 21 Jul 2026 15:45:35 +1200 Subject: [PATCH 8/8] Align setstate timing repro message with hook-ban fix Co-authored-by: Cursor --- test/test_setstate_load_timing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_setstate_load_timing.py b/test/test_setstate_load_timing.py index 6148bcc2a..a2e6daeca 100644 --- a/test/test_setstate_load_timing.py +++ b/test/test_setstate_load_timing.py @@ -215,8 +215,8 @@ def run_tests() -> bool: print("=" * 72) if all_passed: print("ALL CHECKS PASSED — the load-time gap is reproduced and documented.") - print("A fix must statically scan every serialised __setstate__ for") - print("forbidden ops BEFORE torch::jit::load() is called.") + print("A fix must refuse archives containing __setstate__/__getstate__") + print("BEFORE torch::jit::load() is called.") else: print("SOME CHECKS FAILED — see above.") print("=" * 72)