Add validation for negative position_ids in EmbedLayerNorm

onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc

@@ -99,7 +99,7 @@ Status EmbedLayerNorm<T>::Compute(OpKernelContext* context) const {
             return;
           }
           int position_col_index = (position_ids_data == nullptr) ? index % sequence_length : (broadcast_position_ids ? position_ids_data[index % sequence_length] : position_ids_data[index]);
-          if (position_col_index >= position_embedding_length) {
+          if (position_col_index < 0 || position_col_index >= position_embedding_length) {
             failed.store(true, std::memory_order_release);
             return;
           }

onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc

@@ -212,5 +212,64 @@ TEST(EmbedLayerNormTest, EmbedLayerNormBatch_Distill) {
   RunTest(embedlayernorm::EmbedLayerNormBatch_Distill());
 }
 
+// Regression test: negative position_ids must be rejected to not cause OOB read.
+TEST(EmbedLayerNormTest, EmbedLayerNormNegativePositionIds) {
+  int batch_size = 1;
+  int sequence_size = 2;
+  int hidden_size = 4;
+
+  std::vector<int64_t> input_ids_dims = {batch_size, sequence_size};
+  std::vector<int64_t> position_ids_dims = {batch_size, sequence_size};
+
+  // 6 words (rows) x hidden_size=4
+  std::vector<int64_t> word_embedding_dims = {6, hidden_size};
+  // 3 positions x hidden_size=4
+  std::vector<int64_t> position_embedding_dims = {3, hidden_size};
+  // 2 segments x hidden_size=4
+  std::vector<int64_t> segment_embedding_dims = {2, hidden_size};
+
+  std::vector<int64_t> gamma_dims = {hidden_size};
+  std::vector<int64_t> beta_dims = {hidden_size};
+  std::vector<int64_t> output_dims = {batch_size, sequence_size, hidden_size};
+  std::vector<int64_t> mask_index_dims = {batch_size};
+
+  OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain);
+  tester.AddInput<int32_t>("input_ids", input_ids_dims, {1, 3});
+  tester.AddInput<int32_t>("segment_ids", input_ids_dims, {0, 1});
+  tester.AddInput<float>("word_embedding", word_embedding_dims,
+                         {0.2f, 0.1f, 0.4f, -0.6f,
+                          0.3f, 0.2f, 0.5f, 0.6f,
+                          0.6f, 0.7f, 0.0f, -0.1f,
+                          0.8f, 0.6f, 0.9f, 1.2f,
+                          0.1f, 0.3f, 0.5f, 0.9f,
+                          1.0f, -2.0f, 1.1f, 0.8f},
+                         /*is_initializer=*/true);
+  tester.AddInput<float>("position_embedding", position_embedding_dims,
+                         {0.1f, 0.1f, 0.4f, 0.6f,
+                          0.6f, 0.0f, 0.8f, 0.6f,
+                          0.3f, 0.9f, -2.0f, 0.8f},
+                         /*is_initializer=*/true);
+  tester.AddInput<float>("segment_embedding", segment_embedding_dims,
+                         {0.3f, 0.4f, 0.9f, 0.1f,
+                          0.7f, 0.3f, 0.5f, 0.2f},
+                         /*is_initializer=*/true);
+  tester.AddInput<float>("gamma", gamma_dims, {0.25f, 0.15f, 0.45f, -0.66f},
+                         /*is_initializer=*/true);
+  tester.AddInput<float>("beta", beta_dims, {0.6f, 0.2f, 0.5f, -0.6f},
+                         /*is_initializer=*/true);
+  tester.AddAttribute("epsilon", embedlayernorm::kEpsilon);
+  // Skip mask (input 7) - add optional edge
+  tester.AddOptionalInputEdge<int32_t>();
+  // Negative position_ids - this is the malicious input
+  tester.AddInput<int32_t>("position_ids", position_ids_dims, {-5, 1});
+
+  // Dummy outputs (won't be checked since we expect failure)
+  tester.AddOutput<float>("output", output_dims, std::vector<float>(batch_size * sequence_size * hidden_size, 0.0f));
+  tester.AddOutput<int32_t>("mask_index", mask_index_dims, {0});
+
+  // Run CPU only - expect failure due to negative position_ids
+  tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kCudaExecutionProvider, kCudaNHWCExecutionProvider, kDmlExecutionProvider, kOpenVINOExecutionProvider});
+}
+
 }  // namespace test
 }  // namespace onnxruntime