Limit nested model graph depth

onnxruntime/core/graph/model.cc

@@ -136,6 +136,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());
@@ -270,6 +271,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

@@ -64,6 +64,38 @@ void CollectLocalFunctionCalls(
 
 }  // namespace
 
+Status ValidateModelSubgraphDepth(const ONNX_NAMESPACE::ModelProto& model_proto) {
+  using NodeRange = const google::protobuf::RepeatedPtrField<ONNX_NAMESPACE::NodeProto>*;
+  InlinedVector<std::pair<NodeRange, size_t>> pending{{&model_proto.graph().node(), 0}};
+  for (const auto& function : model_proto.functions()) {
+    pending.push_back({&function.node(), 0});
+  }
+
+  while (!pending.empty()) {
+    const auto [nodes, depth] = pending.back();
+    pending.pop_back();
+    for (const auto& node : *nodes) {
+      for (const auto& attr : node.attribute()) {
+        const size_t subgraph_depth = depth + 1;
+        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 BuildLocalFunctionCallGraph(
     const std::unordered_map<std::string, const ONNX_NAMESPACE::FunctionProto*>& model_local_functions,
     LocalFunctionCallGraph& call_graph) {

onnxruntime/core/graph/model_helpers.h

@@ -14,6 +14,7 @@
 
 namespace ONNX_NAMESPACE {
 class FunctionProto;
+class ModelProto;
 }
 
 namespace onnxruntime {
@@ -22,6 +23,10 @@ namespace onnxruntime {
 /// 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);
+
 /// 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,35 @@ 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;
+}
+
+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, 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,31 @@ 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"));
+}
+
 }  // namespace test
 }  // namespace onnxruntime