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..9dc6b1c2b 100644 --- a/bin/pytorch_inference/CModelGraphValidator.cc +++ b/bin/pytorch_inference/CModelGraphValidator.cc @@ -15,12 +15,61 @@ #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 +144,59 @@ 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 + // source. Known stable forms: + // aten::foo -> torch.foo( + // inductor::_reinterpret_tensor -> ops.inductor._reinterpret_tensor( + // prim::CallFunction / prim::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::"}; + constexpr std::string_view INDUCTOR{"inductor::"}; + 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}); + } else if (op.size() > INDUCTOR.size() && op.substr(0, INDUCTOR.size()) == INDUCTOR) { + std::string shortName{op.substr(INDUCTOR.size())}; + signatures.emplace_back("ops.inductor." + 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/CSupportedOperations.cc b/bin/pytorch_inference/CSupportedOperations.cc index 61229a5e0..6eb2e8a3f 100644 --- a/bin/pytorch_inference/CSupportedOperations.cc +++ b/bin/pytorch_inference/CSupportedOperations.cc @@ -22,6 +22,10 @@ const CSupportedOperations::TStringViewSet CSupportedOperations::FORBIDDEN_OPERA "aten::as_strided"sv, "aten::from_file"sv, "aten::save"sv, + // Unchecked storage-offset reinterpret (TorchInductor). Same class of OOB + // heap read/write as as_strided; used to bypass the as_strided forbid + // (HackerOne / elastic/security#12242). + "inductor::_reinterpret_tensor"sv, // After graph inlining, method and function calls should be resolved. // Their presence indicates an opaque call that cannot be validated. "prim::CallFunction"sv, diff --git a/bin/pytorch_inference/Main.cc b/bin/pytorch_inference/Main.cc index 4aaca0049..078153a95 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..400c5df6c 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,125 @@ 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()); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsReinterpretTensorOobRead) { + // Bypass of the as_strided forbid via inductor::_reinterpret_tensor + // (elastic/security#12242). Must be caught before load. + std::string bytes = readFileBytes( + "testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(forbidden, "inductor::_reinterpret_tensor")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsReinterpretTensorOobWrite) { + std::string bytes = readFileBytes( + "testfiles/malicious_models/malicious_reinterpret_tensor_oob_write.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(forbidden, "inductor::_reinterpret_tensor")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanDetectsSetStateReinterpretTensor) { + // Forbidden op only in __setstate__; must be rejected without loading. + std::string bytes = readFileBytes( + "testfiles/malicious_models/malicious_setstate_reinterpret_tensor.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto forbidden = CModelGraphValidator::scanSerialisedCodeForForbiddenOps( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(forbidden, "inductor::_reinterpret_tensor")); +} + +BOOST_AUTO_TEST_CASE(testMaliciousReinterpretTensorRejectedPostLoad) { + // Forward-only fixture: safe to load, then post-load validate must report + // the op as forbidden (not merely unrecognised). + auto module = ::torch::jit::load( + "testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(hasForbiddenOp(result, "inductor::_reinterpret_tensor")); +} + // --- Prepacked model compatibility tests --- // // These load TorchScript models that mirror the ops used by Elasticsearch's diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt new file mode 100644 index 000000000..c7c137d34 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_write.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_write.pt new file mode 100644 index 000000000..d1d66b321 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_write.pt differ 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 000000000..7a60d4f79 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt new file mode 100644 index 000000000..18552abc1 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_reinterpret_tensor.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_reinterpret_tensor.pt new file mode 100644 index 000000000..b1681fda4 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_reinterpret_tensor.pt differ diff --git a/dev-tools/generate_malicious_models.py b/dev-tools/generate_malicious_models.py index fdf6ddbf1..e002789eb 100644 --- a/dev-tools/generate_malicious_models.py +++ b/dev-tools/generate_malicious_models.py @@ -26,7 +26,7 @@ import torch from torch import Tensor -from typing import Optional +from typing import Optional, Tuple # --- Malicious model definitions --- @@ -110,6 +110,134 @@ def forward(self, x: Tensor) -> Tensor: return torch.from_file("/tmp/secret", size=10) +# --- inductor::_reinterpret_tensor OOB models (elastic/security#12242) --- +# +# Bypass of the aten::as_strided forbid: TorchInductor's _reinterpret_tensor +# takes an unchecked storage offset and yields the same class of OOB heap +# read/write. Must be in FORBIDDEN_OPERATIONS and caught by the pre-load scan. + + +class ReinterpretTensorOobRead(torch.nn.Module): + """OOB heap read via inductor::_reinterpret_tensor (reporter PoC).""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + out: list[str] = [""] + for i in range(1, 1000): + target = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], i) + out.append(hex(target.item())) + assert 1 == 2, out + return tmp + + +class ReinterpretTensorOobWrite(torch.nn.Module): + """OOB heap write via inductor::_reinterpret_tensor + fill_ (reporter PoC).""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + acc = torch.zeros(1) + for i in range(1, 100): + target = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], i) + target.fill_(0) + acc = acc + target + return acc + + +# --- 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 + + +class SetStateReinterpretTensor(torch.nn.Module): + """inductor::_reinterpret_tensor in __setstate__ — OOB at load time. + + Forward is benign. Without the pre-load scan this would execute during + torch::jit::load() before the post-load validator runs. + """ + 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: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + leaked = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], 100) + self.w = state[0] + leaked + 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 +350,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 +374,11 @@ 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_reinterpret_tensor_oob_read.pt": ReinterpretTensorOobRead, + "malicious_reinterpret_tensor_oob_write.pt": ReinterpretTensorOobWrite, + "malicious_setstate_file_reader.pt": SetStateFileReaderModel, + "malicious_setstate_file_reader_in_submodule.pt": SetStateFileReaderInSubmodule, + "malicious_setstate_reinterpret_tensor.pt": SetStateReinterpretTensor, "malicious_heap_leak.pt": HeapLeakModel, "malicious_rop_exploit.pt": RopExploitModel, } @@ -254,7 +404,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/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 diff --git a/docs/changelog/3079.yaml b/docs/changelog/3079.yaml new file mode 100644 index 000000000..976249467 --- /dev/null +++ b/docs/changelog/3079.yaml @@ -0,0 +1,5 @@ +area: Machine Learning +issues: [] +pr: 3079 +summary: Forbid inductor::_reinterpret_tensor (as_strided heap-OOB bypass) +type: bug diff --git a/test/test_pytorch_inference_evil_models.py b/test/test_pytorch_inference_evil_models.py index 64b81e2d9..cc257ae92 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,92 @@ 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) + + +class ReinterpretTensorOobRead(torch.nn.Module): + """OOB heap read via inductor::_reinterpret_tensor (security#12242).""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + out: list[str] = [""] + for i in range(1, 1000): + target = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], i) + out.append(hex(target.item())) + assert 1 == 2, out + return tmp + + +class ReinterpretTensorOobWrite(torch.nn.Module): + """OOB heap write via inductor::_reinterpret_tensor + fill_.""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + acc = torch.zeros(1) + for i in range(1, 100): + target = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], i) + target.fill_(0) + acc = acc + target + return acc + + # --------------------------------------------------------------------------- # Binary discovery # --------------------------------------------------------------------------- @@ -242,6 +329,30 @@ 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)", + "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", + }, + "reinterpret_oob_read": { + "class": ReinterpretTensorOobRead, + "expect_rejected": True, + "description": "OOB read via inductor::_reinterpret_tensor", + "expect_stderr_contains": "forbidden operations", + }, + "reinterpret_oob_write": { + "class": ReinterpretTensorOobWrite, + "expect_rejected": True, + "description": "OOB write via inductor::_reinterpret_tensor", + "expect_stderr_contains": "forbidden operations", + }, } diff --git a/test/test_setstate_load_timing.py b/test/test_setstate_load_timing.py new file mode 100644 index 000000000..6148bcc2a --- /dev/null +++ b/test/test_setstate_load_timing.py @@ -0,0 +1,227 @@ +#!/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 + + # 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 = ( + loaded_cleanly + and _LOAD_MARKER in stdout + and stdout.index(_LOAD_MARKER) < stdout.index("LOAD_RETURNED") + ) + + 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 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 + + +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)