[ML] Port TorchScript graph validation and __setstate__ hardening to 8.19

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;
 

bin/pytorch_inference/CCmdLineParser.cc

@@ -41,7 +41,8 @@ bool CCmdLineParser::parse(int argc,
                            std::size_t& cacheMemorylimitBytes,
                            bool& validElasticLicenseKeyConfirmed,
                            bool& lowPriority,
-                           bool& useImmediateExecutor) {
+                           bool& useImmediateExecutor,
+                           bool& skipModelValidation) {
     try {
         boost::program_options::options_description desc(DESCRIPTION);
         // clang-format off
@@ -75,6 +76,7 @@ bool CCmdLineParser::parse(int argc,
             ("lowPriority", "Execute process in low priority")
             ("useImmediateExecutor", "Execute requests on the main thread. This mode should only used for "
             "benchmarking purposes to ensure requests are processed in order)")
+            ("skipModelValidation", "Skip TorchScript model graph validation. WARNING: disables security checks on model operations.")
         ;
         // clang-format on
 
@@ -148,6 +150,9 @@ bool CCmdLineParser::parse(int argc,
                 return false;
             }
         }
+        if (vm.count("skipModelValidation") > 0) {
+            skipModelValidation = true;
+        }
     } catch (std::exception& e) {
         std::cerr << "Error processing command line: " << e.what() << std::endl;
         return false;

bin/pytorch_inference/CCmdLineParser.h

@@ -52,7 +52,8 @@ class CCmdLineParser {
                       std::size_t& cacheMemorylimitBytes,
                       bool& validElasticLicenseKeyConfirmed,
                       bool& lowPriority,
-                      bool& useImmediateExecutor);
+                      bool& useImmediateExecutor,
+                      bool& skipModelValidation);
 
 private:
     static const std::string DESCRIPTION;

bin/pytorch_inference/CMakeLists.txt

@@ -35,7 +35,9 @@ ml_add_executable(pytorch_inference
   CBufferedIStreamAdapter.cc
   CCmdLineParser.cc
   CCommandParser.cc
+  CModelGraphValidator.cc
   CResultWriter.cc
+  CSupportedOperations.cc
   CThreadSettings.cc
   )
 

bin/pytorch_inference/CModelGraphValidator.cc

@@ -0,0 +1,202 @@
+/*
+ * 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.
+ */
+
+#include "CModelGraphValidator.h"
+
+#include "CSupportedOperations.h"
+
+#include <core/CLogger.h>
+
+#include <caffe2/serialize/inline_container.h>
+#include <caffe2/serialize/read_adapter_interface.h>
+#include <torch/csrc/jit/passes/inliner.h>
+
+#include <algorithm>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <vector>
+
+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<std::size_t>(pos));
+        std::memcpy(buf, m_Data + pos, n);
+        return n;
+    }
+
+private:
+    const char* m_Data;
+    std::size_t m_Size;
+};
+}
+
+CModelGraphValidator::SResult CModelGraphValidator::validate(const ::torch::jit::Module& module) {
+
+    TStringSet observedOps;
+    std::size_t nodeCount{0};
+    collectModuleOps(module, observedOps, nodeCount);
+
+    if (nodeCount > MAX_NODE_COUNT) {
+        LOG_ERROR(<< "Model graph is too large: " << nodeCount
+                  << " nodes exceeds limit of " << MAX_NODE_COUNT);
+        return {false, {}, {}, nodeCount};
+    }
+
+    LOG_DEBUG(<< "Model graph contains " << observedOps.size()
+              << " distinct operations across " << nodeCount << " nodes");
+    for (const auto& op : observedOps) {
+        LOG_DEBUG(<< "  observed op: " << op);
+    }
+
+    auto result = validate(observedOps, CSupportedOperations::ALLOWED_OPERATIONS,
+                           CSupportedOperations::FORBIDDEN_OPERATIONS);
+    result.s_NodeCount = nodeCount;
+    return result;
+}
+
+CModelGraphValidator::SResult
+CModelGraphValidator::validate(const TStringSet& observedOps,
+                               const std::unordered_set<std::string_view>& allowedOps,
+                               const std::unordered_set<std::string_view>& forbiddenOps) {
+
+    SResult result;
+
+    // Two-pass check: forbidden ops first, then unrecognised.  This lets us
+    // fail fast when a known-dangerous operation is present and avoids the
+    // cost of scanning for unrecognised ops on a model we will reject anyway.
+    // Use find() rather than contains() — 8.19 builds as C++17.
+    for (const auto& op : observedOps) {
+        if (forbiddenOps.find(op) != forbiddenOps.end()) {
+            result.s_IsValid = false;
+            result.s_ForbiddenOps.push_back(op);
+        }
+    }
+
+    if (result.s_ForbiddenOps.empty()) {
+        for (const auto& op : observedOps) {
+            if (allowedOps.find(op) == allowedOps.end()) {
+                result.s_IsValid = false;
+                result.s_UnrecognisedOps.push_back(op);
+            }
+        }
+    }
+
+    std::sort(result.s_ForbiddenOps.begin(), result.s_ForbiddenOps.end());
+    std::sort(result.s_UnrecognisedOps.begin(), result.s_UnrecognisedOps.end());
+
+    return result;
+}
+
+void CModelGraphValidator::collectBlockOps(const ::torch::jit::Block& block,
+                                           TStringSet& ops,
+                                           std::size_t& nodeCount) {
+    for (const auto* node : block.nodes()) {
+        if (++nodeCount > MAX_NODE_COUNT) {
+            return;
+        }
+        ops.emplace(node->kind().toQualString());
+        for (const auto* subBlock : node->blocks()) {
+            collectBlockOps(*subBlock, ops, nodeCount);
+            if (nodeCount > MAX_NODE_COUNT) {
+                return;
+            }
+        }
+    }
+}
+
+CModelGraphValidator::TStringVec
+CModelGraphValidator::scanArchiveForCustomStateHooks(const char* data, std::size_t size) {
+    TStringSet hooks;
+
+    constexpr std::string_view SETSTATE{"__setstate__"};
+    constexpr std::string_view GETSTATE{"__getstate__"};
+
+    std::unique_ptr<caffe2::serialize::PyTorchStreamReader> reader;
+    try {
+        auto adapter = std::make_shared<CMemoryReadAdapter>(data, size);
+        reader = std::make_unique<caffe2::serialize::PyTorchStreamReader>(std::move(adapter));
+    } 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 state-hook scan skipped (could not parse archive): "
+                 << e.what());
+        return {};
+    }
+
+    for (const auto& name : reader->getAllRecords()) {
+        try {
+            auto[recordData, recordSize] = reader->getRecord(name);
+            std::string_view bytes{static_cast<const char*>(recordData.get()), recordSize};
+
+            // Remediation for a privately reported finding: 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__");
+            }
+        } catch (const std::exception& e) {
+            // A single unreadable record (e.g. a deliberately bad CRC) must not
+            // abort the whole scan: an attacker could otherwise hide a
+            // __setstate__ hook in a later record behind a corrupt earlier one,
+            // slip past the scan, and have torch::jit::load run it at load time.
+            // Warn and keep scanning the remaining records.
+            LOG_WARN(<< "Pre-load state-hook scan: skipping unreadable record '"
+                     << name << "': " << e.what());
+            continue;
+        }
+        if (hooks.size() == 2) {
+            break;
+        }
+    }
+
+    TStringVec result{hooks.begin(), hooks.end()};
+    std::sort(result.begin(), result.end());
+    return result;
+}
+
+void CModelGraphValidator::collectModuleOps(const ::torch::jit::Module& module,
+                                            TStringSet& ops,
+                                            std::size_t& nodeCount) {
+    for (const auto& method : module.get_methods()) {
+        // Inline all method calls so that operations hidden behind
+        // prim::CallMethod are surfaced.  After inlining, any remaining
+        // prim::CallMethod indicates a call that could not be resolved
+        // statically and will be flagged as unrecognised.
+        auto graph = method.graph()->copy();
+        ::torch::jit::Inline(*graph);
+        collectBlockOps(*graph->block(), ops, nodeCount);
+        if (nodeCount > MAX_NODE_COUNT) {
+            return;
+        }
+    }
+}
+}
+}

bin/pytorch_inference/CModelGraphValidator.h

@@ -0,0 +1,109 @@
+/*
+ * 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.
+ */
+
+#ifndef INCLUDED_ml_torch_CModelGraphValidator_h
+#define INCLUDED_ml_torch_CModelGraphValidator_h
+
+#include <torch/script.h>
+
+#include <cstddef>
+#include <string>
+#include <string_view>
+#include <unordered_set>
+#include <vector>
+
+namespace ml {
+namespace torch {
+
+//! \brief
+//! Validates TorchScript model computation graphs against a set of
+//! allowed operations.
+//!
+//! DESCRIPTION:\n
+//! Provides defense-in-depth by statically inspecting the TorchScript
+//! graph of a loaded model and rejecting any model that contains
+//! operations not present in the allowlist derived from supported
+//! transformer architectures.
+//!
+//! IMPLEMENTATION DECISIONS:\n
+//! The validation walks all methods of the module and its submodules
+//! recursively, collecting every distinct operation.  Any operation
+//! that appears in the forbidden set causes immediate rejection.
+//! Any operation not in the allowed set is collected and reported.
+//! This ensures that even operations buried in helper methods or
+//! nested submodules are inspected.
+//!
+class CModelGraphValidator {
+public:
+    using TStringSet = std::unordered_set<std::string>;
+    using TStringVec = std::vector<std::string>;
+
+    //! Upper bound on the number of graph nodes we are willing to inspect.
+    //! Transformer models typically have O(10k) nodes after inlining; a
+    //! limit of 1M provides generous headroom while preventing a
+    //! pathologically large graph from consuming unbounded memory or CPU.
+    static constexpr std::size_t MAX_NODE_COUNT{1000000};
+
+    //! Result of validating a model graph.
+    struct SResult {
+        bool s_IsValid{true};
+        TStringVec s_ForbiddenOps;
+        TStringVec s_UnrecognisedOps;
+        std::size_t s_NodeCount{0};
+    };
+
+public:
+    //! Validate the computation graph of the given module against the
+    //! supported operation allowlist.  Recursively inspects all methods
+    //! across all submodules.
+    static SResult validate(const ::torch::jit::Module& module);
+
+    //! Validate a pre-collected set of operation names.  Useful for
+    //! unit testing the matching logic without requiring a real model.
+    static SResult validate(const TStringSet& observedOps,
+                            const std::unordered_set<std::string_view>& allowedOps,
+                            const std::unordered_set<std::string_view>& forbiddenOps);
+
+    //! 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 a load-time execution gap from a privately reported finding:
+    //! torch::jit::load runs a module's __setstate__ during deserialization,
+    //! before post-load graph validation.  Matching the recommended 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
+    //! 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.
+    static void collectBlockOps(const ::torch::jit::Block& block,
+                                TStringSet& ops,
+                                std::size_t& nodeCount);
+
+    //! Inline all method calls and collect ops from the flattened graph.
+    //! After inlining, prim::CallMethod should not appear; if it does,
+    //! the call could not be resolved statically and is treated as
+    //! unrecognised.
+    static void collectModuleOps(const ::torch::jit::Module& module,
+                                 TStringSet& ops,
+                                 std::size_t& nodeCount);
+};
+}
+}
+
+#endif // INCLUDED_ml_torch_CModelGraphValidator_h