Fix heap OOB write in Attention PrepareMask via negative mask_index values

onnxruntime/contrib_ops/cpu/bert/attention_base.h

@@ -121,6 +121,20 @@ inline Status AttentionBase::CheckMask(const Tensor* mask_index,
     }
     mask_type = (mask_dims[0] == batch_size ? AttentionMaskType::MASK_1D_KEY_SEQ_LEN : mask_dims[0] == 2 * batch_size ? AttentionMaskType::MASK_1D_END_START
                                                                                                                       : AttentionMaskType::MASK_1D_KEY_SEQ_LEN_START);
+
+    // Validate that end_position values (first batch_size elements) are non-negative.
+    // Negative end_position causes out-of-bounds writes in PrepareMask.
+    // Only validate when mask_index is on CPU; GPU tensors are clamped in the CUDA kernel.
+    if (mask_index->Location().device.Type() == OrtDevice::CPU) {
+      const int32_t* mask_data = mask_index->Data<int32_t>();
+      for (int64_t i = 0; i < batch_size; i++) {
+        if (mask_data[i] < 0) {
+          return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                                 "mask_index value ", mask_data[i], " at index ", i,
+                                 " is negative. mask_index end_position values must be non-negative.");
+        }
+      }
+    }
   } else if (mask_dims.size() == 2) {
     if (mask_dims[0] == batch_size && mask_dims[1] == total_sequence_length) {
       mask_type = AttentionMaskType::MASK_2D_KEY_PADDING;

onnxruntime/contrib_ops/cpu/bert/attention_helper.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <algorithm>
 #include <limits>
 #include "core/util/math.h"
 #include "core/util/math_cpuonly.h"
@@ -95,14 +96,14 @@ void PrepareMask(const int32_t* mask_index,
         // mask_index is 1D: (B) or (2B) => (Bx)T
 
         // Handle right-side padding: mask value at or after the end position will be mask_filter_value
-        int end_position = mask_index[b_i];
+        int end_position = std::max(0, std::min(static_cast<int>(mask_index[b_i]), all_sequence_length));
         for (int m_i = end_position; m_i < all_sequence_length; m_i++) {
           p_mask[m_i] = static_cast<T>(mask_filter_value);
         }
 
         // Handle left-side padding: mask value before the start position will be mask_filter_value
         if (has_mask_start_position) {
-          int start_position = std::min(mask_index[b_i + batch_size], all_sequence_length);
+          int start_position = std::max(0, std::min(static_cast<int>(mask_index[b_i + batch_size]), all_sequence_length));
           for (int m_i = 0; m_i < start_position; m_i++) {
             p_mask[m_i] = static_cast<T>(mask_filter_value);
           }

onnxruntime/contrib_ops/cuda/bert/attention_softmax.cu

@@ -603,7 +603,7 @@
   if (threadIdx.x == 0) {
     const int batch = blockIdx.y;
     start_position = mask_start != nullptr ? max(0, mask_start[batch]) : 0;
-    end_position = min(total_sequence_length, mask_end[batch]);
+    end_position = max(0, min(total_sequence_length, mask_end[batch]));
 
     // Attend to no word has same effect as attend to all words. This is added to get parity with CPU result.
     if (start_position >= end_position) {
@@ -735,7 +735,7 @@
   if (threadIdx.x == 0) {
     const int batch = blockIdx.y;
     start_position = mask_start != nullptr ? max(0, mask_start[batch]) : 0;
-    end_position = min(total_sequence_length, mask_end[batch]);
+    end_position = max(0, min(total_sequence_length, mask_end[batch]));
 
     // Attend to no word has same effect as attend to all words. This is added to get parity with CPU result.
     if (start_position >= end_position) {

onnxruntime/test/contrib_ops/attention_op_test.cc

@@ -2009,6 +2009,56 @@ TEST(ContribOpAttentionTest, AttentionMaskIndexOutOfRange) {
                    AttentionMaskType::MASK_1D_END_START);
 }
 
+TEST(ContribOpAttentionTest, AttentionMaskIndexNegativeEndPosition) {
+  // Test that negative end_position in mask_index is rejected (heap underflow prevention).
+  int batch_size = 2;
+  int sequence_length = 2;
+  int hidden_size = 4;
+  int number_of_heads = 2;
+
+  std::vector<float> input_data = {
+      0.8f, -0.5f, 0.0f, 1.f,
+      0.5f, 0.2f, 0.3f, -0.6f,
+      0.8f, -0.5f, 0.0f, 1.f,
+      0.5f, 0.2f, 0.3f, -0.6f};
+
+  std::vector<float> weight_data = {
+      0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
+      0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
+      0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
+      0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
+
+  std::vector<float> bias_data = {
+      -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
+
+  // Negative end_position values in 1D mask_index
+  std::vector<int32_t> mask_index_data = {-10, -1};
+
+  OpTester tester("Attention", 1, onnxruntime::kMSDomain);
+  tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(number_of_heads));
+  tester.AddAttribute<float>("mask_filter_value", static_cast<float>(-10000.0f));
+
+  std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
+  std::vector<int64_t> weights_dims = {hidden_size, 3 * hidden_size};
+  std::vector<int64_t> bias_dims = {3 * hidden_size};
+  std::vector<int64_t> mask_index_dims = {batch_size};
+  std::vector<int64_t> output_dims = {batch_size, sequence_length, hidden_size};
+
+  tester.AddInput<float>("input", input_dims, input_data);
+  tester.AddInput<float>("weight", weights_dims, weight_data);
+  tester.AddInput<float>("bias", bias_dims, bias_data);
+  tester.AddInput<int32_t>("mask_index", mask_index_dims, mask_index_data);
+  tester.AddOptionalInputEdge<float>();    // past
+  tester.AddOptionalInputEdge<float>();    // attention_bias
+  tester.AddOptionalInputEdge<int32_t>();  // past_sequence_length
+
+  tester.AddOutput<float>("output", output_dims, std::vector<float>(batch_size * sequence_length * hidden_size, 0.0f));
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  tester.Run(OpTester::ExpectResult::kExpectFailure, "mask_index value", {}, nullptr, &execution_providers);
+}
+
 #if !defined(__wasm__)
 // TODO: fix in web assembly
 TEST(ContribOpAttentionTest, AttentionPastState_dynamic) {