Limit nested model graph depth

onnxruntime/core/graph/model.cc

@@ -80,6 +80,11 @@ void Model::RemoveLocalFunctionsProtos(const InlinedHashSet<std::string>& retain
 
 static constexpr int DEFAULT_PROTOBUF_BLOCK_SIZE = 4 * 1024 * 1024;
 
+static ModelProto ValidateAndCopyModelProto(const ModelProto& model_proto) {
+  ORT_THROW_IF_ERROR(ValidateModelSubgraphDepth(model_proto));
+  return model_proto;
+}
+
 Model::Model(const std::string& graph_name,
              bool is_onnx_domain_only,
              const ModelMetaData& model_metadata,
@@ -128,6 +133,10 @@ Model::Model(const std::string& graph_name,
     opset_id_proto->set_version(version);
   }
 
+  for (const auto& func : model_local_functions) {
+    ORT_THROW_IF_ERROR(ValidateFunctionSubgraphDepth(func));
+  }
+
   model_local_functions_.reserve(model_local_functions.size());
   for (auto& func : model_local_functions) {
     auto func_ptr = model_proto_.add_functions();
@@ -136,6 +145,7 @@ Model::Model(const std::string& graph_name,
                                             func_ptr);
   }
 
+  ORT_THROW_IF_ERROR(ValidateModelSubgraphDepth(model_proto_));
   ORT_THROW_IF_ERROR(ValidateModelLocalFunctionAcyclic(model_local_functions_));
 
   model_local_function_templates_maps_.reserve(model_proto_.functions().size());
@@ -166,7 +176,7 @@ Model::Model(const std::string& graph_name,
 Model::Model(const ModelProto& model_proto, const PathString& model_path,
              const IOnnxRuntimeOpSchemaRegistryList* local_registries, const logging::Logger& logger,
              const ModelOptions& options)
-    : Model(ModelProto(model_proto), model_path, local_registries, logger, options) {
+    : Model(ValidateAndCopyModelProto(model_proto), model_path, local_registries, logger, options) {
 }
 
 Model::Model(ModelProto&& model_proto, const PathString& model_path,
@@ -270,6 +280,7 @@ Model::Model(ModelProto&& model_proto, const PathString& model_path,
     model_local_functions_.insert_or_assign(function_utils::GetFunctionIdentifier(func.domain(), func.name(), func.overload()), &func);
   }
 
+  ORT_THROW_IF_ERROR(ValidateModelSubgraphDepth(model_proto_));
   ORT_THROW_IF_ERROR(ValidateModelLocalFunctionAcyclic(model_local_functions_));
 
   model_local_function_templates_maps_.reserve(model_proto_.functions().size());

onnxruntime/core/graph/model_helpers.cc

@@ -18,6 +18,52 @@ namespace onnxruntime {
 
 namespace {
 
+using NodeRange = const google::protobuf::RepeatedPtrField<ONNX_NAMESPACE::NodeProto>*;
+using PendingNodeRanges = InlinedVector<std::pair<NodeRange, size_t>>;
+
+Status AddAttributeSubgraphs(const ONNX_NAMESPACE::AttributeProto& attr,
+                             size_t subgraph_depth,
+                             PendingNodeRanges& pending) {
+  if ((attr.has_g() || !attr.graphs().empty()) && subgraph_depth > kMaxModelSubgraphDepth) {
+    return ORT_MAKE_STATUS(
+        ONNXRUNTIME, NOT_IMPLEMENTED,
+        "Model subgraph depth ", subgraph_depth,
+        " exceeds the maximum supported depth of ", kMaxModelSubgraphDepth, ".");
+  }
+
+  if (attr.has_g()) {
+    pending.push_back({&attr.g().node(), subgraph_depth});
+  }
+  for (const auto& graph : attr.graphs()) {
+    pending.push_back({&graph.node(), subgraph_depth});
+  }
+
+  return Status::OK();
+}
+
+Status ValidateSubgraphDepth(
+    const google::protobuf::RepeatedPtrField<ONNX_NAMESPACE::NodeProto>& root_nodes,
+    const google::protobuf::RepeatedPtrField<ONNX_NAMESPACE::AttributeProto>* root_attributes = nullptr) {
+  PendingNodeRanges pending{{&root_nodes, 0}};
+  if (root_attributes != nullptr) {
+    for (const auto& attr : *root_attributes) {
+      ORT_RETURN_IF_ERROR(AddAttributeSubgraphs(attr, 1, pending));
+    }
+  }
+
+  while (!pending.empty()) {
+    const auto [nodes, depth] = pending.back();
+    pending.pop_back();
+    for (const auto& node : *nodes) {
+      for (const auto& attr : node.attribute()) {
+        ORT_RETURN_IF_ERROR(AddAttributeSubgraphs(attr, depth + 1, pending));
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
 // Iterative collection of local function calls from a sequence of nodes,
 // including nodes inside nested subgraph attributes. Avoids recursion to
 // prevent stack overflow from maliciously deep subgraph nesting.
@@ -64,6 +110,19 @@ void CollectLocalFunctionCalls(
 
 }  // namespace
 
+Status ValidateModelSubgraphDepth(const ONNX_NAMESPACE::ModelProto& model_proto) {
+  ORT_RETURN_IF_ERROR(ValidateSubgraphDepth(model_proto.graph().node()));
+  for (const auto& function : model_proto.functions()) {
+    ORT_RETURN_IF_ERROR(ValidateFunctionSubgraphDepth(function));
+  }
+
+  return Status::OK();
+}
+
+Status ValidateFunctionSubgraphDepth(const ONNX_NAMESPACE::FunctionProto& function_proto) {
+  return ValidateSubgraphDepth(function_proto.node(), &function_proto.attribute_proto());
+}
+
 Status BuildLocalFunctionCallGraph(
     const std::unordered_map<std::string, const ONNX_NAMESPACE::FunctionProto*>& model_local_functions,
     LocalFunctionCallGraph& call_graph) {

onnxruntime/core/graph/model_helpers.h

@@ -14,14 +14,20 @@
 
 namespace ONNX_NAMESPACE {
 class FunctionProto;
-}
+class ModelProto;
+}  // namespace ONNX_NAMESPACE
 
 namespace onnxruntime {
 
 /// Adjacency list representation of a local function call graph.
 /// Keys and values are string_views into stable storage (e.g. map keys that outlive this structure).
 using LocalFunctionCallGraph = InlinedHashMap<std::string_view, InlinedVector<std::string_view>>;
 
+constexpr size_t kMaxModelSubgraphDepth = 32;
+
+Status ValidateModelSubgraphDepth(const ONNX_NAMESPACE::ModelProto& model_proto);
+Status ValidateFunctionSubgraphDepth(const ONNX_NAMESPACE::FunctionProto& function_proto);
+
 /// Build a call graph adjacency list from model local functions.
 /// String views in the returned graph point into the keys of @p model_local_functions.
 Status BuildLocalFunctionCallGraph(

onnxruntime/test/framework/function_test.cc

@@ -444,6 +444,63 @@ TEST(FunctionTest, RejectsRecursionThroughSubgraph) {
 // --- Synthetic adjacency-list tests for ValidateCallGraphAcyclic ---
 // These test the cycle detection algorithm directly without constructing ONNX models.
 
+static ONNX_NAMESPACE::ModelProto CreateNestedLocalFunctionModel(size_t depth, bool use_graphs_attribute) {
+  ONNX_NAMESPACE::ModelProto model_proto;
+  auto* nodes = model_proto.add_functions()->mutable_node();
+  for (size_t i = 0; i < depth; ++i) {
+    auto* node = nodes->Add();
+    auto* attr = node->add_attribute();
+    if (use_graphs_attribute) {
+      nodes = attr->add_graphs()->mutable_node();
+    } else {
+      nodes = attr->mutable_g()->mutable_node();
+    }
+  }
+
+  return model_proto;
+}
+
+static ONNX_NAMESPACE::ModelProto CreateNestedLocalFunctionDefaultAttributeModel(
+    size_t depth, bool use_graphs_attribute) {
+  ONNX_NAMESPACE::ModelProto model_proto;
+  auto* function = model_proto.add_functions();
+  auto* attr = function->add_attribute_proto();
+  auto* graph = use_graphs_attribute ? attr->add_graphs() : attr->mutable_g();
+  for (size_t i = 1; i < depth; ++i) {
+    auto* node = graph->add_node();
+    attr = node->add_attribute();
+    graph = attr->mutable_g();
+  }
+
+  return model_proto;
+}
+
+TEST(FunctionTest, LocalFunctionSubgraphDepthValidated) {
+  EXPECT_STATUS_OK(ValidateModelSubgraphDepth(
+      CreateNestedLocalFunctionModel(kMaxModelSubgraphDepth, false)));
+  EXPECT_EQ(ValidateModelSubgraphDepth(
+                CreateNestedLocalFunctionModel(kMaxModelSubgraphDepth + 1, false))
+                .Code(),
+            common::NOT_IMPLEMENTED);
+  EXPECT_EQ(ValidateModelSubgraphDepth(
+                CreateNestedLocalFunctionModel(kMaxModelSubgraphDepth + 1, true))
+                .Code(),
+            common::NOT_IMPLEMENTED);
+}
+
+TEST(FunctionTest, LocalFunctionDefaultAttributeSubgraphDepthValidated) {
+  EXPECT_STATUS_OK(ValidateModelSubgraphDepth(
+      CreateNestedLocalFunctionDefaultAttributeModel(kMaxModelSubgraphDepth, false)));
+  EXPECT_EQ(ValidateModelSubgraphDepth(
+                CreateNestedLocalFunctionDefaultAttributeModel(kMaxModelSubgraphDepth + 1, false))
+                .Code(),
+            common::NOT_IMPLEMENTED);
+  EXPECT_EQ(ValidateModelSubgraphDepth(
+                CreateNestedLocalFunctionDefaultAttributeModel(kMaxModelSubgraphDepth + 1, true))
+                .Code(),
+            common::NOT_IMPLEMENTED);
+}
+
 TEST(FunctionTest, CallGraphAcyclic_EmptyGraph) {
   onnxruntime::LocalFunctionCallGraph call_graph;
   ASSERT_STATUS_OK(onnxruntime::ValidateCallGraphAcyclic(call_graph));

onnxruntime/test/ir/graph_test.cc

@@ -13,6 +13,7 @@
 #include "core/graph/graph_viewer.h"
 #include "core/graph/graph_utils.h"
 #include "core/graph/model.h"
+#include "core/graph/model_helpers.h"
 #include "core/graph/op.h"
 #include "core/graph/ort_format_load_options.h"
 #include "core/session/inference_session.h"
@@ -3876,5 +3877,40 @@ TEST_F(GraphTest, DeeplyNestedLoopSubgraphsResolveInReasonableTime) {
                                     "regression has returned.";
 }
 
+static ModelProto CreateNestedSubgraphModel(size_t depth) {
+  ModelProto model_proto;
+  model_proto.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION);
+  auto* opset = model_proto.add_opset_import();
+  opset->set_domain(kOnnxDomain);
+  opset->set_version(21);
+
+  auto* graph = model_proto.mutable_graph();
+  for (size_t i = 0; i < depth; ++i) {
+    auto* node = graph->add_node();
+    auto* attr = node->add_attribute();
+    graph = attr->mutable_g();
+  }
+
+  return model_proto;
+}
+
+TEST_F(GraphTest, ExcessiveSubgraphDepthRejected) {
+  auto model_proto = CreateNestedSubgraphModel(kMaxModelSubgraphDepth + 1);
+  std::shared_ptr<Model> model;
+  const auto status = Model::Load(std::move(model_proto), model, nullptr, *logger_);
+  ASSERT_FALSE(status.IsOK());
+  EXPECT_EQ(status.Code(), common::NOT_IMPLEMENTED);
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("exceeds the maximum supported depth"));
+}
+
+TEST_F(GraphTest, ExcessiveSubgraphDepthRejectedFromLvalueProto) {
+  const auto model_proto = CreateNestedSubgraphModel(kMaxModelSubgraphDepth + 1);
+  std::shared_ptr<Model> model;
+  const auto status = Model::Load(model_proto, model, nullptr, *logger_);
+  ASSERT_FALSE(status.IsOK());
+  EXPECT_EQ(status.Code(), common::NOT_IMPLEMENTED);
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("exceeds the maximum supported depth"));
+}
+
 }  // namespace test
 }  // namespace onnxruntime