fix(security): add SafeInt overflow protection in Expand and constant folding output size limit (microsoft#28055)

tianleiwu · web-flow · commit 04f744032cf2 · 2026-05-14T14:08:25.000-07:00
### Description

Harden the constant folding optimizer and the `Expand` CPU kernel
against integer overflow attacks from crafted ONNX models.

**Problem:** The `Expand::Compute()` kernel performs cumulative
dimension multiplications (`input_count *= input_dim`, `output_count *=
output_dim`) using raw `int64_t` arithmetic. When triggered during
constant folding at `CreateSession()` time via a crafted model with
extreme shape values, signed integer overflow can produce corrupted
values used for buffer offset calculations and `memcpy` lengths,
creating a potential out-of-bounds write. The downstream SafeInt check
in the allocator catches overflow only when the total byte count wraps,
but carefully chosen dimensions can make the overflowed value appear
valid.

Additionally, the constant folding optimizer has no output size budget —
any deterministic node with constant inputs is eligible for constant
folding regardless of output size, enabling memory exhaustion attacks at
model load time.

### Key Changes

**1. SafeInt-protected arithmetic in `expand.cc`**

Wraps all dimension accumulation and offset/length calculations with
`SafeInt&lt;int64_t&gt;` or `SafeInt&lt;size_t&gt;` to catch overflow before it can
corrupt buffer arithmetic:

| Location | Before | After |
|---|---|---|
| Accumulator loop (L97-98) | `input_count *= input_dim` |
`SafeInt&lt;int64_t&gt;(input_count) * input_dim` |
| Accumulator loop (L109) | `last_dim_size *= expand_dim_size[...]` |
`SafeInt&lt;int64_t&gt;(last_dim_size) * ...` |
| copy_byte (L116) | `copy_len * sizeof(T)` | `SafeInt&lt;size_t&gt;(copy_len)
* sizeof(T)` |
| input_offset (L122) | `i * copy_len` | `SafeInt&lt;int64_t&gt;(i) *
copy_len` |
| output_offset (L126) | `output_offset += current_count * ...` |
`SafeInt&lt;int64_t&gt;(output_offset) + SafeInt&lt;int64_t&gt;(current_count) *
...` |

**2. Constant folding output size limit in `constant_folding.cc`**

- **Pre-execution check**: `EstimateNodeOutputSizeInBytes()` uses shape
inference results with SafeInt-protected arithmetic to estimate total
output bytes. Nodes exceeding the limit are skipped.
- **Post-execution check**: After `kernel-&gt;Compute()`, actual output
`SizeInBytes()` is verified against the limit (catches cases where shape
inference couldn't determine output size).
- **Exception isolation**: `kernel-&gt;Compute()` is wrapped in `try/catch`
so that SafeInt overflow exceptions from individual nodes skip the node
rather than aborting the entire optimization pass.
- **Configurable limit**: New session option
`optimization.constant_folding_max_output_size_in_bytes` (default: 1 GB,
`"0"` to disable).

**3. Session option**

New key `kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes` in
`onnxruntime_session_options_config_keys.h`.

### Motivation and Context

This addresses a security vulnerability where a malicious ONNX model can
cause signed integer overflow in the Expand kernel during constant
folding at model load time (`CreateSession()`), potentially leading to
out-of-bounds memory writes. The constant folding size limit provides
defense-in-depth against memory exhaustion attacks from untrusted
models.

### Testing

- `ConstantFoldingOutputSizeLimit` — Verifies 4 MB Expand is blocked at
1 MB limit, allowed at 8 MB limit.
- `ConstantFoldingDefaultLimitBlocksLargeExpand` — Verifies 1 GB
ConstantOfShape is blocked at 512 MB limit.
- `ConstantFoldingSmallOutputAllowed` — Verifies small Expand (64 bytes)
is still folded normally.
- `ConstantFoldingExpandOverflowDimsSkipped` — Verifies Expand with
`[2^32, 2^32]` dimensions (int64 overflow) is gracefully skipped during
constant folding.
diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h
@@ -124,6 +124,18 @@ static const char* const kOrtSessionOptionsMemoryOptimizerProbeConfig = "optimiz
 // Default is an empty string which means no optimizers are disabled.
 static const char* const kOrtSessionOptionsDisableSpecifiedOptimizers = "optimization.disable_specified_optimizers";
 
+// Maximum total output size in bytes that the constant folding optimizer is allowed to produce per node.
+// Prevents malicious models from causing excessive memory allocation during optimization.
+// If the estimated or actual output size of a constant-foldable node exceeds this limit, the node will
+// not be constant folded and will instead be executed at runtime.
+//
+// Option values:
+// - A positive integer (as string): Maximum allowed output size in bytes per constant-folded node.
+//   Default is "1073741824" (1 GB).
+// - "0": Disable the size limit (not recommended for untrusted models).
+static const char* const kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes =
+    "optimization.constant_folding_max_output_size_in_bytes";
+
 // It controls whether to run graph optimizations in loop or not.
 //
 // "0": disable. Graph Optimization Loop is disabled.
diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc
@@ -2,15 +2,18 @@
 // Licensed under the MIT License.
 
 #include <limits>
+#include <string>
 
 #include "core/optimizer/constant_folding.h"
 #include "core/optimizer/initializer.h"
 #include "core/optimizer/utils.h"
 #include "core/graph/graph_utils.h"
 #include "core/optimizer/optimizer_execution_frame.h"
-#include "core/optimizer/utils.h"
 #include "core/framework/op_kernel.h"
 #include "core/framework/tensorprotoutils.h"
+#include "core/session/onnxruntime_session_options_config_keys.h"
+#include "core/common/safeint.h"
+#include "core/common/parse_string.h"
 
 using namespace onnxruntime::common;
 
@@ -140,6 +143,154 @@ static Status ConstantFoldIfNode(Graph& graph, Node& if_node, const logging::Log
   return status;
 }
 
+// Default maximum output size per constant-folded node: 1 GB.
+// This prevents malicious models from causing excessive memory allocation during optimization.
+static constexpr int64_t kDefaultConstantFoldingMaxOutputSizeInBytes = 1024 * 1024 * 1024;
+
+static size_t GetElementSizeForConstantFolding(ONNX_NAMESPACE::TensorProto_DataType elem_type) {
+  const size_t element_size = utils::GetElementSizeOfTensor(elem_type);
+  if (element_size != 0) {
+    return element_size;
+  }
+
+  // String tensors allocate storage for std::string slots even though the payload size is variable.
+  return elem_type == ONNX_NAMESPACE::TensorProto_DataType_STRING ? sizeof(std::string) : 0;
+}
+
+static int64_t EstimateTensorElementCount(const ONNX_NAMESPACE::TensorShapeProto& shape) {
+  SafeInt<int64_t> num_elements = 1;
+  for (int i = 0; i < shape.dim_size(); ++i) {
+    const auto& dim = shape.dim(i);
+    if (!utils::HasDimValue(dim)) {
+      return -1;  // Symbolic dimension
+    }
+    int64_t dim_value = dim.dim_value();
+    if (dim_value < 0) {
+      return -1;  // Invalid dimension
+    }
+    num_elements *= dim_value;
+  }
+
+  return num_elements;
+}
+
+static int64_t EstimateTensorSizeInBytes(const NodeArg& node_arg) {
+  const auto* type_proto = node_arg.TypeAsProto();
+  if (type_proto == nullptr || !utils::HasTensorType(*type_proto)) {
+    return -1;  // Cannot estimate non-tensor or unknown types
+  }
+
+  const auto* shape = node_arg.Shape();
+  if (shape == nullptr) {
+    return -1;  // Unknown shape
+  }
+
+  auto elem_type = static_cast<ONNX_NAMESPACE::TensorProto_DataType>(
+      type_proto->tensor_type().elem_type());
+  size_t element_size = GetElementSizeForConstantFolding(elem_type);
+  if (element_size == 0) {
+    return -1;  // Unknown element type
+  }
+
+  int64_t num_elements = EstimateTensorElementCount(*shape);
+  if (num_elements < 0) {
+    return -1;
+  }
+
+  return SafeInt<int64_t>(num_elements) * static_cast<int64_t>(element_size);
+}
+
+static int64_t EstimateUniqueOutputSizeInBytes(const Node& node) {
+  const auto& input_defs = node.InputDefs();
+  if (input_defs.empty() || input_defs[0] == nullptr) {
+    return -1;
+  }
+
+  const int64_t input_num_elements = input_defs[0]->Shape() != nullptr
+                                         ? EstimateTensorElementCount(*input_defs[0]->Shape())
+                                         : -1;
+  if (input_num_elements < 0) {
+    return -1;
+  }
+
+  const auto* input_type_proto = input_defs[0]->TypeAsProto();
+  if (input_type_proto == nullptr || !utils::HasTensorType(*input_type_proto)) {
+    return -1;
+  }
+
+  auto input_elem_type = static_cast<ONNX_NAMESPACE::TensorProto_DataType>(
+      input_type_proto->tensor_type().elem_type());
+  const size_t input_element_size = GetElementSizeForConstantFolding(input_elem_type);
+  if (input_element_size == 0) {
+    return -1;
+  }
+
+  SafeInt<int64_t> total_size = 0;
+  const auto& output_defs = node.OutputDefs();
+  for (size_t output_idx = 0; output_idx < output_defs.size(); ++output_idx) {
+    const auto* output_def = output_defs[output_idx];
+    if (!output_def->Exists()) {
+      continue;
+    }
+
+    const size_t element_size = output_idx == 0 ? input_element_size : sizeof(int64_t);
+    total_size += SafeInt<int64_t>(input_num_elements) * static_cast<int64_t>(element_size);
+  }
+
+  return total_size;
+}
+
+static int64_t EstimateIdentityOutputSizeInBytes(const Node& node) {
+  const auto& input_defs = node.InputDefs();
+  if (input_defs.empty() || input_defs[0] == nullptr) {
+    return -1;
+  }
+
+  return EstimateTensorSizeInBytes(*input_defs[0]);
+}
+
+// Estimate the total output size in bytes for a node using shape inference results.
+// Returns -1 if the output size cannot be estimated (e.g., unknown shapes or types).
+static int64_t EstimateNodeOutputSizeInBytes(const Node& node) {
+  if (node.OpType() == "Identity" && node.Domain().empty()) {
+    return EstimateIdentityOutputSizeInBytes(node);
+  }
+
+  if (node.OpType() == "Unique" && node.Domain().empty()) {
+    return EstimateUniqueOutputSizeInBytes(node);
+  }
+
+  SafeInt<int64_t> total_size = 0;
+  for (const auto* output_def : node.OutputDefs()) {
+    if (!output_def->Exists()) {
+      continue;
+    }
+
+    const int64_t output_size = EstimateTensorSizeInBytes(*output_def);
+    if (output_size < 0) {
+      return -1;
+    }
+
+    total_size += output_size;
+  }
+
+  return total_size;
+}
+
+// Get the configured max output size from session options, or use the default.
+static int64_t GetConstantFoldingMaxOutputSize(const ConfigOptions& config_options) {
+  std::string max_size_str = config_options.GetConfigOrDefault(
+      kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes,
+      std::to_string(kDefaultConstantFoldingMaxOutputSizeInBytes));
+
+  int64_t max_size = 0;
+  if (!TryParseStringWithClassicLocale(max_size_str, max_size) || max_size < 0) {
+    max_size = kDefaultConstantFoldingMaxOutputSizeInBytes;
+  }
+
+  return max_size;
+}
+
 Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const {
   bool have_updated_nodes = false;
   GraphViewer graph_viewer(graph);
@@ -151,6 +302,8 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level,
   };
 #endif
 
+  const int64_t max_output_size = GetConstantFoldingMaxOutputSize(config_options_);
+
   for (NodeIndex i : order) {
     auto* node = graph.GetNode(i);
     if (!node || !AllowConstantFolding(*node)) {
@@ -233,6 +386,35 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level,
         }
       }
 
+      // Check if the estimated output size exceeds the configured limit.
+      // This prevents malicious models from causing excessive memory allocation during constant folding.
+      if (max_output_size > 0) {
+        int64_t estimated_size = -1;
+        try {
+          estimated_size = EstimateNodeOutputSizeInBytes(*node);
+        } catch (const std::exception&) {
+          // SafeInt overflow means the size is astronomically large - definitely skip
+          LOGS(logger, WARNING) << "Integer overflow while estimating output size of "
+                                << node->OpType() << " node '" << node->Name()
+                                << "'. Skipping constant folding for this node.";
+          continue;
+        }
+
+        if (estimated_size > max_output_size) {
+          LOGS(logger, WARNING) << "Skipping constant folding for " << node->OpType()
+                                << " node '" << node->Name()
+                                << "' because estimated output size (" << estimated_size
+                                << " bytes) exceeds the limit (" << max_output_size << " bytes).";
+          continue;
+        }
+        if (estimated_size < 0) {
+          LOGS(logger, INFO) << "Skipping constant folding for " << node->OpType()
+                             << " node '" << node->Name()
+                             << "' because output size could not be estimated before execution.";
+          continue;
+        }
+      }
+
 #if !defined(DISABLE_SPARSE_TENSORS)
       // Create execution frame for executing constant nodes.
       OptimizerExecutionFrame::Info info({node}, constant_inputs, graph.ModelPath(), execution_provider_,
@@ -312,14 +494,59 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level,
 #pragma warning(disable : 6387)
 #endif
       OpKernelContext op_kernel_context(&frame, kernel.get(), /*stream*/ nullptr, nullptr, logger);
-      ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context));
+
+      // Skip the current node if Compute fails so one bad constant-fold candidate does not abort
+      // the entire constant folding pass.
+      Status compute_status = Status::OK();
+      try {
+        compute_status = kernel->Compute(&op_kernel_context);
+      } catch (const std::exception& ex) {
+        LOGS(logger, WARNING) << "Exception during constant folding of " << node->OpType()
+                              << " node '" << node->Name() << "': " << ex.what()
+                              << ". Skipping constant folding for this node.";
+        continue;
+      }
+
+      if (!compute_status.IsOK()) {
+        LOGS(logger, WARNING) << "Failure during constant folding of " << node->OpType()
+                              << " node '" << node->Name() << "': " << compute_status.ErrorMessage()
+                              << ". Skipping constant folding for this node.";
+        continue;
+      }
 #ifdef _WIN32
 #pragma warning(pop)
 #endif
 
       std::vector<OrtValue> fetches;
       ORT_RETURN_IF_ERROR(frame.GetOutputs(fetches));
 
+      // Post-execution size check: verify actual output sizes don't exceed the limit.
+      // This catches cases where pre-execution shape inference couldn't determine the output size.
+      if (max_output_size > 0) {
+        SafeInt<int64_t> actual_total_size = 0;
+        bool size_exceeded = false;
+        try {
+          for (size_t fetch_idx = 0; fetch_idx < fetches.size(); ++fetch_idx) {
+            if (fetches[fetch_idx].IsAllocated() && fetches[fetch_idx].IsTensor()) {
+              const auto& tensor = fetches[fetch_idx].Get<Tensor>();
+              actual_total_size += tensor.SizeInBytes();
+            }
+          }
+          size_exceeded = actual_total_size > max_output_size;
+        } catch (const std::exception&) {
+          // SafeInt overflow means total size is astronomically large
+          size_exceeded = true;
+        }
+
+        if (size_exceeded) {
+          LOGS(logger, WARNING) << "Skipping constant folding for " << node->OpType()
+                                << " node '" << node->Name()
+                                << "' because actual output size exceeds the limit ("
+                                << max_output_size << " bytes).";
+          continue;
+        }
+      }
+
       // Go over all output node args and substitute them with the newly computed tensors, which will be
       // added to the graph as initializers.
       ORT_ENFORCE(fetches.size() == fetch_to_output_idx.size());
diff --git a/onnxruntime/core/providers/cpu/tensor/expand.cc b/onnxruntime/core/providers/cpu/tensor/expand.cc
@@ -94,8 +94,8 @@ Status Expand<T>::Compute(OpKernelContext* context) const {
     auto input_dim = input_dims_iter > -1 ? input_dims[input_dims_iter] : 1;
     auto output_dim = output_dims[output_dims_iter];
 
-    input_count *= input_dim;
-    output_count *= output_dim;
+    input_count = SafeInt<int64_t>(input_count) * input_dim;
+    output_count = SafeInt<int64_t>(output_count) * output_dim;
 
     if (0 == input_count || 0 == output_count) {
       return Status::OK();
@@ -106,26 +106,26 @@ Status Expand<T>::Compute(OpKernelContext* context) const {
       input_dim_group[onnxruntime::narrow<size_t>(dim_group_start)] = input_count;
       output_dim_group[onnxruntime::narrow<size_t>(dim_group_start)] = output_count;
       expand_dim_size[onnxruntime::narrow<size_t>(dim_group_start)] = output_count / input_count / last_dim_size;
-      last_dim_size *= expand_dim_size[onnxruntime::narrow<size_t>(dim_group_start)];
+      last_dim_size = SafeInt<int64_t>(last_dim_size) * expand_dim_size[onnxruntime::narrow<size_t>(dim_group_start)];
     }
   }
 
   auto distribute_count = input_dim_group[onnxruntime::narrow<size_t>(dim_group_start)] / input_dim_group[SafeInt<size_t>(max_dims_size) - 1];
   std::vector<int64_t> output_offsets(onnxruntime::narrow<size_t>(distribute_count), 0);
   int64_t copy_len = input_dim_group[SafeInt<size_t>(max_dims_size) - 1];
-  auto copy_byte = copy_len * sizeof(T);
+  size_t copy_byte = SafeInt<size_t>(copy_len) * sizeof(T);
 
   auto distribute_fn =
       [&](ptrdiff_t i_start, ptrdiff_t i_end) {
         for (auto i = i_start; i < i_end; i++) {
-          auto input_offset = i * copy_len;
+          int64_t input_offset = SafeInt<int64_t>(i) * copy_len;
           int64_t output_offset = 0;
           for (auto j = dim_group_start + 1, remains = input_offset; j < max_dims_size; ++j) {
             auto current_count = remains / input_dim_group[onnxruntime::narrow<size_t>(j)];
-            output_offset += current_count * output_dim_group[onnxruntime::narrow<size_t>(j)];
+            output_offset = SafeInt<int64_t>(output_offset) + SafeInt<int64_t>(current_count) * output_dim_group[onnxruntime::narrow<size_t>(j)];
             remains = remains % input_dim_group[onnxruntime::narrow<size_t>(j)];
           }  // for j
-          memcpy(output_data + output_offset, input_data + input_offset, onnxruntime::narrow<size_t>(copy_byte));
+          memcpy(output_data + output_offset, input_data + input_offset, copy_byte);
           output_offsets[onnxruntime::narrow<size_t>(i)] = output_offset;
         }  // for i
       };  // distribute_fn
diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc
