diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc index 72adfa025da57..db0e76f839293 100644 --- a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc @@ -99,7 +99,7 @@ Status EmbedLayerNorm::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; } diff --git a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc index 6cd84fc55ea86..de6fdc3238a22 100644 --- a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc +++ b/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 input_ids_dims = {batch_size, sequence_size}; + std::vector position_ids_dims = {batch_size, sequence_size}; + + // 6 words (rows) x hidden_size=4 + std::vector word_embedding_dims = {6, hidden_size}; + // 3 positions x hidden_size=4 + std::vector position_embedding_dims = {3, hidden_size}; + // 2 segments x hidden_size=4 + std::vector segment_embedding_dims = {2, hidden_size}; + + std::vector gamma_dims = {hidden_size}; + std::vector beta_dims = {hidden_size}; + std::vector output_dims = {batch_size, sequence_size, hidden_size}; + std::vector mask_index_dims = {batch_size}; + + OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain); + tester.AddInput("input_ids", input_ids_dims, {1, 3}); + tester.AddInput("segment_ids", input_ids_dims, {0, 1}); + tester.AddInput("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("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("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("gamma", gamma_dims, {0.25f, 0.15f, 0.45f, -0.66f}, + /*is_initializer=*/true); + tester.AddInput("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(); + // Negative position_ids - this is the malicious input + tester.AddInput("position_ids", position_ids_dims, {-5, 1}); + + // Dummy outputs (won't be checked since we expect failure) + tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_size * hidden_size, 0.0f)); + tester.AddOutput("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