From 13c2101ecfefe898ed8efc801e57aad3f9818a10 Mon Sep 17 00:00:00 2001
From: danielsongmicrosoft <danielsongmicrosoft@users.noreply.github.com>
Date: Sun, 27 Sep 2026 22:25:01 -0700
Subject: [PATCH] Reject recursive function graph attributes

---
 onnxruntime/core/graph/function_utils.cc    | 105 ++++++++++--
 onnxruntime/core/graph/function_utils.h     |   6 +-
 onnxruntime/core/graph/graph.cc             |   2 +-
 onnxruntime/test/framework/function_test.cc | 181 ++++++++++++++++++++
 4 files changed, 274 insertions(+), 20 deletions(-)

diff --git a/onnxruntime/core/graph/function_utils.cc b/onnxruntime/core/graph/function_utils.cc
index 7a851799454b0..eb9bc9c4440ef 100644
--- a/onnxruntime/core/graph/function_utils.cc
+++ b/onnxruntime/core/graph/function_utils.cc
@@ -1,5 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+#include <algorithm>
+#include <sstream>
 #include <unordered_map>
 
 #include "core/graph/function_utils.h"
@@ -15,6 +17,33 @@ namespace function_utils {
 using string = std::string;
 using namespace ONNX_NAMESPACE;
 
+namespace {
+
+bool AttributeContainsSubgraphs(const AttributeProto& attr) {
+  return attr.has_g() || attr.graphs_size() > 0;
+}
+
+std::string FormatAttributeReferenceChain(const std::vector<std::string>& active_attr_refs, const std::string& next_attr_ref) {
+  std::ostringstream stream;
+
+  for (size_t i = 0; i < active_attr_refs.size(); ++i) {
+    if (i != 0) {
+      stream << " -> ";
+    }
+
+    stream << active_attr_refs[i];
+  }
+
+  if (!active_attr_refs.empty()) {
+    stream << " -> ";
+  }
+
+  stream << next_attr_ref;
+  return stream.str();
+}
+
+}  // namespace
+
 // Utilify function to get the imported version of domain from opset imports
 // Returns -1 if requested domain is not found in the opset_imports
 template <template <typename, typename> class M>
@@ -365,6 +394,7 @@ class Inliner {
   std::string prefix_;
   const onnxruntime::NodeAttributes& attr_map_;
   std::vector<InlinedHashMap<std::string, std::string>> rename_scopes_;
+  std::vector<std::string> active_attr_refs_;
 
   Inliner(const std::string& prefix, const onnxruntime::NodeAttributes& attr_map) : prefix_(prefix),
                                                                                     attr_map_(attr_map) {
@@ -432,7 +462,7 @@ class Inliner {
   }
 
   // Process a node:
-  void transform(NodeProto& n) {
+  Status transform(NodeProto& n) {
     if (!n.name().empty())
       n.set_name(prefix_ + "_" + n.name());
 
@@ -446,11 +476,22 @@ class Inliner {
     const auto attr_map_end = attr_map_.cend();
     for (auto attr_iter = attributes.begin(); attr_iter != attributes.end();) {
       auto& attr = *attr_iter;
+      bool tracked_attr_ref = false;
       if (!attr.ref_attr_name().empty()) {
         // Attribute-references must be replaced by the corresponding attribute-value in the call-node
         // if the call-node contains the attribute. Otherwise, this attribute must be removed.
         auto entry = attr_map_.find(attr.ref_attr_name());
         if (entry != attr_map_end) {
+          if (AttributeContainsSubgraphs(entry->second)) {
+            const auto active_attr_ref = std::find(active_attr_refs_.cbegin(), active_attr_refs_.cend(), attr.ref_attr_name());
+            ORT_RETURN_IF(active_attr_ref != active_attr_refs_.cend(),
+                          "Function attribute graph expansion is recursive: ",
+                          FormatAttributeReferenceChain(active_attr_refs_, attr.ref_attr_name()));
+
+            active_attr_refs_.push_back(attr.ref_attr_name());
+            tracked_attr_ref = true;
+          }
+
           // Copy value of attribute, but retain original name:
           std::string name = attr.name();
           attr = entry->second;
@@ -462,16 +503,38 @@ class Inliner {
       }
       // Subgraphs must be recursively processed.
       if (attr.has_g()) {
-        transform(*attr.mutable_g());
+        auto status = transform(*attr.mutable_g());
+        if (!status.IsOK()) {
+          if (tracked_attr_ref) {
+            active_attr_refs_.pop_back();
+          }
+
+          return status;
+        }
+      }
+      for (auto& graph : *attr.mutable_graphs()) {
+        auto status = transform(graph);
+        if (!status.IsOK()) {
+          if (tracked_attr_ref) {
+            active_attr_refs_.pop_back();
+          }
+
+          return status;
+        }
+      }
+
+      if (tracked_attr_ref) {
+        active_attr_refs_.pop_back();
       }
-      for (auto& graph : *attr.mutable_graphs())
-        transform(graph);
+
       ++attr_iter;
     }
+
+    return Status::OK();
   }
 
   // Process a sub-graph, contained as an attribute in a control-flow op node.
-  void transform(GraphProto& graph) {
+  Status transform(GraphProto& graph) {
     rename_scopes_.emplace_back();
     for (auto& x : *graph.mutable_input())
       make_unique(*x.mutable_name());
@@ -479,31 +542,41 @@ class Inliner {
       make_unique(*init.mutable_name());
     for (auto& y : *graph.mutable_output())
       make_unique(*y.mutable_name());
-    for (auto& n : *graph.mutable_node())
-      transform(n);
+    for (auto& n : *graph.mutable_node()) {
+      auto status = transform(n);
+      if (!status.IsOK()) {
+        rename_scopes_.pop_back();
+        return status;
+      }
+    }
     rename_scopes_.pop_back();
+
+    return Status::OK();
   }
 
  public:
   // The main specialization method: specialize a FunctionProto for a particular call-site.
-  static void specialize(const NodeProto& callnode, FunctionProto& callee, const onnxruntime::NodeAttributes& attr_map,
-                         const std::string& unique_prefix) {
+  static Status specialize(const NodeProto& callnode, FunctionProto& callee, const onnxruntime::NodeAttributes& attr_map,
+                           const std::string& unique_prefix) {
     Inliner inliner(unique_prefix, attr_map);
 
     inliner.bind<false>(*callee.mutable_input(), callnode.input());
     inliner.bind<true>(*callee.mutable_output(), callnode.output());
 
-    for (auto& n : *callee.mutable_node())
-      inliner.transform(n);
+    for (auto& n : *callee.mutable_node()) {
+      ORT_RETURN_IF_ERROR(inliner.transform(n));
+    }
+
+    return Status::OK();
   }
 };
 
-void Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const ONNX_NAMESPACE::NodeProto& calling_node,
-                const onnxruntime::NodeAttributes& attr_map, const std::string& unique_prefix) {
-  Inliner::specialize(calling_node, called_function, attr_map, unique_prefix);
+Status Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const ONNX_NAMESPACE::NodeProto& calling_node,
+                  const onnxruntime::NodeAttributes& attr_map, const std::string& unique_prefix) {
+  return Inliner::specialize(calling_node, called_function, attr_map, unique_prefix);
 }
 
-void Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const Node& calling_node, const std::string& unique_prefix) {
+Status Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const Node& calling_node, const std::string& unique_prefix) {
   ONNX_NAMESPACE::NodeProto calling_node_proto;
   calling_node.ToProto(calling_node_proto);
 
@@ -511,7 +584,7 @@ void Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const Node& call
   for (auto& attribute_proto : called_function.attribute_proto()) {
     ORT_IGNORE_RETURN_VALUE(attr_map.emplace(attribute_proto.name(), attribute_proto));
   }
-  Specialize(called_function, calling_node_proto, attr_map, unique_prefix);
+  return Specialize(called_function, calling_node_proto, attr_map, unique_prefix);
 }
 
 }  // namespace function_utils
diff --git a/onnxruntime/core/graph/function_utils.h b/onnxruntime/core/graph/function_utils.h
index 6b92b9a8955db..f1944473b12bb 100644
--- a/onnxruntime/core/graph/function_utils.h
+++ b/onnxruntime/core/graph/function_utils.h
@@ -63,10 +63,10 @@ inline std::string GetFunctionIdentifier(std::string_view function_domain, std::
   return id;
 }
 
-void Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const ONNX_NAMESPACE::NodeProto& calling_node,
-                const onnxruntime::NodeAttributes& attr_map, const std::string& unique_prefix);
+Status Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const ONNX_NAMESPACE::NodeProto& calling_node,
+                  const onnxruntime::NodeAttributes& attr_map, const std::string& unique_prefix);
 
-void Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const Node& calling_node, const std::string& unique_prefix);
+Status Specialize(ONNX_NAMESPACE::FunctionProto& called_function, const Node& calling_node, const std::string& unique_prefix);
 
 }  // namespace function_utils
 
diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc
index 572a34a408cc1..cea41fb74bc72 100644
--- a/onnxruntime/core/graph/graph.cc
+++ b/onnxruntime/core/graph/graph.cc
@@ -6584,7 +6584,7 @@ Status Graph::InlineFunction(Node& callnode) {
     ORT_ENFORCE(callnode.TryGetFunctionProto(inlined_fp), "Node has no function body and cannot be inlined.");
 
     // Make all the names unique and resolve nested graphs inputs to the outer scope.
-    function_utils::Specialize(inlined_fp, callnode, uniq_identifier);
+    ORT_RETURN_IF_ERROR(function_utils::Specialize(inlined_fp, callnode, uniq_identifier));
 
     // In this case, global Resolve() will take care of everything.
     ORT_RETURN_IF_ERROR(InlineFunctionProto(inlined_fp, parent_annotation));
diff --git a/onnxruntime/test/framework/function_test.cc b/onnxruntime/test/framework/function_test.cc
index d212071f361b5..efb9b3b48ade6 100644
--- a/onnxruntime/test/framework/function_test.cc
+++ b/onnxruntime/test/framework/function_test.cc
@@ -13,6 +13,7 @@
 #include "core/common/span_utils.h"
 #include "core/framework/customregistry.h"
 #include "core/framework/op_kernel.h"
+#include "core/graph/function_utils.h"
 #include "core/graph/model.h"
 #include "core/graph/model_helpers.h"
 #include "core/providers/cpu/cpu_execution_provider.h"
@@ -1333,6 +1334,186 @@ static ONNX_NAMESPACE::FunctionProto MakeIdentityFunction(
   return function;
 }
 
+static ONNX_NAMESPACE::NodeProto MakeFunctionCallNodeProto() {
+  ONNX_NAMESPACE::NodeProto call_node;
+  call_node.add_input("actual_x");
+  call_node.add_output("actual_y");
+  return call_node;
+}
+
+static NodeAttributes CreateDefaultAttributeMap(const ONNX_NAMESPACE::FunctionProto& function) {
+  NodeAttributes attr_map;
+  for (const auto& attribute_proto : function.attribute_proto()) {
+    ORT_IGNORE_RETURN_VALUE(attr_map.emplace(attribute_proto.name(), attribute_proto));
+  }
+
+  return attr_map;
+}
+
+static Status SpecializeWithDefaultAttributes(ONNX_NAMESPACE::FunctionProto& function) {
+  return function_utils::Specialize(function, MakeFunctionCallNodeProto(), CreateDefaultAttributeMap(function),
+                                    "test_inliner");
+}
+
+static ONNX_NAMESPACE::AttributeProto MakeGraphRefAttribute(const std::string& name, const std::string& ref_attr_name,
+                                                            ONNX_NAMESPACE::AttributeProto_AttributeType type) {
+  ONNX_NAMESPACE::AttributeProto attr;
+  attr.set_name(name);
+  attr.set_ref_attr_name(ref_attr_name);
+  attr.set_type(type);
+  return attr;
+}
+
+static ONNX_NAMESPACE::GraphProto MakeRecursiveDefaultGraph(const std::string& ref_attr_name) {
+  ONNX_NAMESPACE::GraphProto graph;
+  graph.set_name("default_graph");
+
+  auto* node = graph.add_node();
+  node->set_name("body_node");
+  node->set_op_type("Identity");
+  node->add_input("x");
+  node->add_output("y");
+  *node->add_attribute() = MakeGraphRefAttribute("nested", ref_attr_name, ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH);
+
+  return graph;
+}
+
+static ONNX_NAMESPACE::GraphProto MakeNonRecursiveDefaultGraph() {
+  ONNX_NAMESPACE::GraphProto graph;
+  graph.set_name("default_graph");
+
+  auto* node = graph.add_node();
+  node->set_name("body_node");
+  node->set_op_type("Identity");
+  node->add_input("x");
+  node->add_output("y");
+  node->add_attribute()->mutable_g()->set_name("leaf_graph");
+
+  return graph;
+}
+
+static ONNX_NAMESPACE::FunctionProto MakeFunctionWithDefaultGraphAttributes(
+    const std::vector<ONNX_NAMESPACE::AttributeProto>& default_attrs,
+    const std::vector<ONNX_NAMESPACE::AttributeProto>& body_attrs) {
+  ONNX_NAMESPACE::FunctionProto function;
+  function.set_domain("local");
+  function.set_name("myfun");
+  function.add_input("x");
+  function.add_output("y");
+
+  for (const auto& default_attr : default_attrs) {
+    function.add_attribute(default_attr.name());
+    *function.add_attribute_proto() = default_attr;
+  }
+
+  auto* node = function.add_node();
+  node->set_name("body_root");
+  node->set_op_type("Identity");
+  node->add_input("x");
+  node->add_output("y");
+  for (const auto& body_attr : body_attrs) {
+    *node->add_attribute() = body_attr;
+  }
+
+  return function;
+}
+
+static ONNX_NAMESPACE::FunctionProto MakeFunctionWithDefaultGraphAttribute(const ONNX_NAMESPACE::AttributeProto& default_attr) {
+  return MakeFunctionWithDefaultGraphAttributes(
+      {default_attr},
+      {MakeGraphRefAttribute("body_attr", default_attr.name(), default_attr.type())});
+}
+
+TEST(FunctionTest, SpecializeRejectsRecursiveDefaultGraphAttributeExpansion) {
+  ONNX_NAMESPACE::AttributeProto default_attr;
+  default_attr.set_name("body");
+  default_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH);
+  *default_attr.mutable_g() = MakeRecursiveDefaultGraph(default_attr.name());
+
+  auto function = MakeFunctionWithDefaultGraphAttribute(default_attr);
+  const auto status = SpecializeWithDefaultAttributes(function);
+
+  ASSERT_FALSE(status.IsOK());
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("Function attribute graph expansion is recursive"));
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("body -> body"));
+}
+
+TEST(FunctionTest, SpecializeRejectsRecursiveDefaultGraphsAttributeExpansion) {
+  ONNX_NAMESPACE::AttributeProto default_attr;
+  default_attr.set_name("body_list");
+  default_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPHS);
+  *default_attr.add_graphs() = MakeRecursiveDefaultGraph(default_attr.name());
+
+  auto function = MakeFunctionWithDefaultGraphAttribute(default_attr);
+  const auto status = SpecializeWithDefaultAttributes(function);
+
+  ASSERT_FALSE(status.IsOK());
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("Function attribute graph expansion is recursive"));
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("body_list -> body_list"));
+}
+
+TEST(FunctionTest, SpecializeRejectsRecursiveDefaultGraphAttributeExpansionAcrossDefaults) {
+  ONNX_NAMESPACE::AttributeProto first_default_attr;
+  first_default_attr.set_name("first");
+  first_default_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH);
+  *first_default_attr.mutable_g() = MakeRecursiveDefaultGraph("second");
+
+  ONNX_NAMESPACE::AttributeProto second_default_attr;
+  second_default_attr.set_name("second");
+  second_default_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH);
+  *second_default_attr.mutable_g() = MakeRecursiveDefaultGraph("first");
+
+  auto function = MakeFunctionWithDefaultGraphAttributes(
+      {first_default_attr, second_default_attr},
+      {MakeGraphRefAttribute("body_attr", first_default_attr.name(), first_default_attr.type())});
+  const auto status = SpecializeWithDefaultAttributes(function);
+
+  ASSERT_FALSE(status.IsOK());
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("Function attribute graph expansion is recursive"));
+  EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("first -> second -> first"));
+}
+
+TEST(FunctionTest, SpecializeAllowsNonRecursiveDefaultGraphAttributeExpansion) {
+  ONNX_NAMESPACE::AttributeProto default_attr;
+  default_attr.set_name("body");
+  default_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH);
+  *default_attr.mutable_g() = MakeNonRecursiveDefaultGraph();
+
+  auto function = MakeFunctionWithDefaultGraphAttribute(default_attr);
+  auto status = SpecializeWithDefaultAttributes(function);
+
+  ASSERT_TRUE(status.IsOK()) << status.ErrorMessage();
+
+  const auto& attr = function.node(0).attribute(0);
+  EXPECT_EQ(attr.name(), "body_attr");
+  EXPECT_TRUE(attr.ref_attr_name().empty());
+  EXPECT_TRUE(attr.has_g());
+}
+
+TEST(FunctionTest, SpecializeAllowsSequentialReuseOfNonRecursiveDefaultGraphAttribute) {
+  ONNX_NAMESPACE::AttributeProto default_attr;
+  default_attr.set_name("body");
+  default_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_GRAPH);
+  *default_attr.mutable_g() = MakeNonRecursiveDefaultGraph();
+
+  auto function = MakeFunctionWithDefaultGraphAttributes(
+      {default_attr},
+      {
+          MakeGraphRefAttribute("body_attr_0", default_attr.name(), default_attr.type()),
+          MakeGraphRefAttribute("body_attr_1", default_attr.name(), default_attr.type()),
+      });
+  const auto status = SpecializeWithDefaultAttributes(function);
+
+  ASSERT_TRUE(status.IsOK()) << status.ErrorMessage();
+
+  const auto& attrs = function.node(0).attribute();
+  ASSERT_EQ(attrs.size(), 2);
+  for (const auto& attr : attrs) {
+    EXPECT_TRUE(attr.ref_attr_name().empty());
+    EXPECT_TRUE(attr.has_g());
+  }
+}
+
 TEST(FunctionTest, RejectDuplicateFunctionIdentifiersFromModelProto) {
   ONNX_NAMESPACE::ModelProto model_proto;
   model_proto.set_ir_version(10);