Fix input validation and null-pointer dereference in STFTDecomposition graph transformer (microsoft#28465)

### Description

Validate that `frame_step` and `dft_size` (derived from `frame_length`)
are positive before they are used in buffer sizing and loop arithmetic.
Use SafeInt for the weight buffer allocation to guard against overflow.
Also fix an unconditional dereference of `window_recipient` which is
nullptr when the STFT node has no window input.

### Motivation and Context

A model with non-positive initializer values for frame_length or
frame_step causes signed-to-unsigned wrapping in size computations,
leading to out-of-bounds writes during graph optimization. The nullptr
dereference is a crash on any windowless STFT node.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: vraspar

onnxruntime/core/optimizer/stft_decomposition.cc

@@ -9,6 +9,7 @@
 #include "core/graph/graph_utils.h"
 #include "core/optimizer/optimizer_execution_frame.h"
 #include "core/optimizer/utils.h"
+#include "core/common/safeint.h"
 #include "core/framework/op_kernel.h"
 #include "core/framework/tensorprotoutils.h"
 #include <numbers>
@@ -210,6 +211,14 @@ Status STFTDecomposition::ApplyImpl(Graph& graph, bool& modified, int graph_leve
       dft_size = window_length_dim.dim_value();
     }
 
+    // Validate model-provided scalar values before using them in size calculations.
+    // These come from untrusted model initializers/shapes and must be positive.
+    if (dft_size <= 0 || frame_step_value <= 0) {
+      LOGS(logger, WARNING) << "STFT decomposition skipped: invalid dft_size (" << dft_size
+                            << ") or frame_step_value (" << frame_step_value << ")";
+      continue;
+    }
+
     bool is_onesided = true;
     auto& attrs = stft.GetAttributes();
     if (attrs.find("onesided") != attrs.end()) {
@@ -227,14 +236,22 @@ Status STFTDecomposition::ApplyImpl(Graph& graph, bool& modified, int graph_leve
     if (is_real) {
       auto output_num_frames = stft.MutableOutputDefs()[0]->Shape()->dim(1).dim_value();
       auto output_frame_length = stft.MutableOutputDefs()[0]->Shape()->dim(2).dim_value();
-      auto weight_size = static_cast<size_t>(dft_unique_bins * dft_size);
+
+      size_t dft_size_sz, dft_unique_bins_sz, weight_size;
+      if (!SafeCast(dft_unique_bins, dft_unique_bins_sz) ||
+          !SafeCast(dft_size, dft_size_sz) ||
+          !SafeMultiply(dft_unique_bins_sz, dft_size_sz, weight_size)) {
+        LOGS(logger, WARNING) << "STFT decomposition skipped: weight size overflow";
+        continue;
+      }
+
       auto real_weights_data = std::vector<float>(weight_size);
       auto imag_weights_data = std::vector<float>(weight_size);
 
       // Populate weights
-      for (size_t k = 0; k < static_cast<size_t>(dft_unique_bins); k++) {
-        for (size_t n = 0; n < static_cast<size_t>(dft_size); n++) {
-          auto index = static_cast<size_t>(k * dft_size + n);
+      for (size_t k = 0; k < dft_unique_bins_sz; k++) {
+        for (size_t n = 0; n < dft_size_sz; n++) {
+          auto index = k * dft_size_sz + n;
           auto theta = -2 * std::numbers::pi_v<float> * k * n / static_cast<float>(dft_size);
           real_weights_data[index] = static_cast<float>(cos(theta));
           imag_weights_data[index] = static_cast<float>(sin(theta));
@@ -356,7 +373,6 @@ Status STFTDecomposition::ApplyImpl(Graph& graph, bool& modified, int graph_leve
 
     // Copy inputs
     auto signal_target_idx = signal_recipient->Index();
-    auto window_target_idx = window_recipient->Index();
     for (auto cur = input_edges.cbegin(), end = input_edges.cend(); cur != end; ++cur) {
       const graph_utils::GraphEdge& edge = *cur;
       NodeIndex target_idx = 0;
@@ -367,8 +383,10 @@ Status STFTDecomposition::ApplyImpl(Graph& graph, bool& modified, int graph_leve
           recipient = signal_recipient;
           break;
         case 2:
-          target_idx = window_target_idx;
-          recipient = window_recipient;
+          if (window_recipient) {
+            target_idx = window_recipient->Index();
+            recipient = window_recipient;
+          }
           break;
       }
 

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -76,6 +76,7 @@
 #include "core/optimizer/rule_based_graph_transformer.h"
 #include "core/optimizer/slice_concat_to_space_to_depth_fusion.h"
 #include "core/optimizer/slice_elimination.h"
+#include "core/optimizer/stft_decomposition.h"
 #include "core/optimizer/unsqueeze_elimination.h"
 #include "core/optimizer/utils.h"
 #include "core/platform/env.h"
@@ -10425,5 +10426,65 @@ TEST_F(GraphTransformationTests, DivMulFusion_MultiElementInitializer) {
 // `dropout_elimination.cc` remains as pure defense-in-depth against future
 // internal callers that may bypass shape inference.
 
+// These tests verify that STFTDecomposition skips malformed models
+// instead of crashing with OOB writes from negative initializer values.
+TEST_F(GraphTransformationTests, STFTDecomposition_NegativeFrameLength) {
+  constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "stft_negative_frame_length.onnx";
+  std::shared_ptr<Model> model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_EQ(op_to_count["STFT"], 1);
+
+  const InlinedHashSet<std::string_view> empty_ep = {};
+  auto stft_transformer = std::make_unique<STFTDecomposition>(empty_ep);
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(stft_transformer), TransformerLevel::Level1));
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+  // STFT node should NOT be decomposed — transformer skips invalid models
+  op_to_count = CountOpsInGraph(graph);
+  ASSERT_EQ(op_to_count["STFT"], 1);
+}
+
+TEST_F(GraphTransformationTests, STFTDecomposition_NegativeFrameStep) {
+  constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "stft_negative_frame_step.onnx";
+  std::shared_ptr<Model> model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_EQ(op_to_count["STFT"], 1);
+
+  const InlinedHashSet<std::string_view> empty_ep = {};
+  auto stft_transformer = std::make_unique<STFTDecomposition>(empty_ep);
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(stft_transformer), TransformerLevel::Level1));
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+  // STFT node should NOT be decomposed — transformer skips invalid models
+  op_to_count = CountOpsInGraph(graph);
+  ASSERT_EQ(op_to_count["STFT"], 1);
+}
+
+TEST_F(GraphTransformationTests, STFTDecomposition_NoWindowInput) {
+  constexpr const ORTCHAR_T* model_uri = MODEL_FOLDER "stft_no_window.onnx";
+  std::shared_ptr<Model> model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_EQ(op_to_count["STFT"], 1);
+
+  const InlinedHashSet<std::string_view> empty_ep = {};
+  auto stft_transformer = std::make_unique<STFTDecomposition>(empty_ep);
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  ASSERT_STATUS_OK(graph_transformation_mgr.Register(std::move(stft_transformer), TransformerLevel::Level1));
+  // Should not crash (previously dereferenced nullptr window_recipient)
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+  // Valid windowless STFT should be successfully decomposed
+  op_to_count = CountOpsInGraph(graph);
+  ASSERT_EQ(op_to_count["STFT"], 0);
+}
+
 }  // namespace test
 }  // namespace onnxruntime

onnxruntime/test/testdata/transform/stft_negative_frame_length.onnx

@@ -0,0 +1,46 @@
+from onnx import TensorProto, helper, save
+
+
+def make_stft_model(frame_length_value, frame_step_value, has_window, filename):
+    batch = 1
+    signal_length = 16
+    dft_size = abs(frame_length_value) if frame_length_value != 0 else 4
+    onesided_bins = dft_size // 2 + 1
+    num_frames = max(1, (signal_length - dft_size) // abs(frame_step_value) + 1) if frame_step_value != 0 else 1
+
+    signal = helper.make_tensor_value_info("signal", TensorProto.FLOAT, [batch, signal_length, 1])
+    output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [batch, num_frames, onesided_bins, 2])
+
+    frame_step_init = helper.make_tensor("frame_step", TensorProto.INT64, [], [frame_step_value])
+    frame_length_init = helper.make_tensor("frame_length", TensorProto.INT64, [], [frame_length_value])
+
+    window_input = "window" if has_window else ""
+    initializers = [frame_step_init, frame_length_init]
+    if has_window:
+        window_init = helper.make_tensor("window", TensorProto.FLOAT, [dft_size], [1.0] * dft_size)
+        initializers.append(window_init)
+
+    stft_node = helper.make_node(
+        "STFT",
+        inputs=["signal", "frame_step", window_input, "frame_length"],
+        outputs=["output"],
+        onesided=1,
+    )
+
+    graph = helper.make_graph(
+        [stft_node],
+        "stft_test",
+        [signal],
+        [output],
+        initializer=initializers,
+    )
+
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    model.ir_version = 8
+    save(model, filename)
+
+
+if __name__ == "__main__":
+    make_stft_model(-2, 4, False, "stft_negative_frame_length.onnx")
+    make_stft_model(4, -2, False, "stft_negative_frame_step.onnx")
+    make_stft_model(4, 4, False, "stft_no_window.onnx")