Validate input rank/size for BifurcationDetector before indexing

onnxruntime/contrib_ops/cpu/bert/bifurcation_detector.h

@@ -24,6 +24,23 @@ class BifurcationDetector : public OpKernel {
     const Tensor* cur_tokens = context->Input<Tensor>(1);
     const Tensor* prev_suffix_match_idx = context->Input<Tensor>(2);
     const Tensor* pred_tokens = context->Input<Tensor>(3);
+
+    // src_tokens/cur_tokens/pred_tokens are treated as flat 1-D token sequences below (both their
+    // length, via Shape().GetDims()[0], and their backing buffer, via linear DataRaw() indexing).
+    // Requiring rank == 1 guards GetDims()[0] against a 0-D (scalar) shape, whose dims are empty, and
+    // also ensures the reported length always equals the tensor's total element count (a higher-rank
+    // shape could otherwise have a non-empty leading dimension while some other dimension is 0, making
+    // the tensor's data buffer empty even though GetDims()[0] is non-zero).
+    ORT_RETURN_IF_NOT(src_tokens->Shape().NumDimensions() == 1, "src_tokens must be a 1-D tensor");
+    ORT_RETURN_IF_NOT(cur_tokens->Shape().NumDimensions() == 1, "cur_tokens must be a 1-D tensor");
+    // prev_suffix_match_idx (and, by construction below, the suffix_match_idx output that mirrors its
+    // shape) holds a single index value; element [0] is read and/or written unconditionally.
+    ORT_RETURN_IF_NOT(prev_suffix_match_idx->Shape().Size() == 1,
+                      "prev_suffix_match_idx must contain exactly one element");
+    if (pred_tokens != nullptr) {
+      ORT_RETURN_IF_NOT(pred_tokens->Shape().NumDimensions() == 1, "pred_tokens must be a 1-D tensor");
+    }
+
     const auto* src_tokens_data = static_cast<const int64_t*>(src_tokens->DataRaw());
     const auto* cur_tokens_data = static_cast<const int64_t*>(cur_tokens->DataRaw());
     int64_t src_tokens_len = src_tokens->Shape().GetDims()[0];

onnxruntime/test/contrib_ops/bifurcation_detector_op_test.cc

@@ -339,5 +339,95 @@ TEST(BifurcationDetectorTest, PredTokensLengthMismatch) {
              {}, nullptr, &execution_providers);
 }
 
+// A 0-element prev_suffix_match_idx must be rejected: the kernel unconditionally reads and writes
+// element [0] of, respectively, this input and the output tensor that mirrors its shape, so a 0-element
+// shape would otherwise back those accesses with a null/empty buffer.
+TEST(BifurcationDetectorTest, ZeroElementPrevSuffixMatchIdx) {
+  OpTester tester("BifurcationDetector", 1, onnxruntime::kMSDomain);
+
+  tester.AddInput<int64_t>("src_tokens", {4}, {1, 5, 3, 4});
+  tester.AddInput<int64_t>("cur_tokens", {1}, {2});
+  tester.AddInput<int64_t>("prev_suffix_match_idx", {0}, {});
+  tester.AddInput<int64_t>("pred_tokens", {5}, {1, 5, 3, 4, 2});
+  tester.AddOutput<int64_t>("tokens", {1}, {0});
+  tester.AddOutput<int64_t>("suffix_match_idx", {0}, {});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  tester.Run(OpTester::ExpectResult::kExpectFailure,
+             "prev_suffix_match_idx must contain exactly one element",
+             {}, nullptr, &execution_providers);
+}
+
+// Same as above, but with no pred_tokens input: the unconditional output write is still reached.
+TEST(BifurcationDetectorTest, ZeroElementPrevSuffixMatchIdxNoPredTokens) {
+  OpTester tester("BifurcationDetector", 1, onnxruntime::kMSDomain);
+
+  tester.AddInput<int64_t>("src_tokens", {4}, {1, 5, 3, 4});
+  tester.AddInput<int64_t>("cur_tokens", {1}, {2});
+  tester.AddInput<int64_t>("prev_suffix_match_idx", {0}, {});
+  tester.AddOutput<int64_t>("tokens", {1}, {0});
+  tester.AddOutput<int64_t>("suffix_match_idx", {0}, {});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  tester.Run(OpTester::ExpectResult::kExpectFailure,
+             "prev_suffix_match_idx must contain exactly one element",
+             {}, nullptr, &execution_providers);
+}
+
+// A rank-0 (scalar) src_tokens has no dimensions, so Shape().GetDims()[0] would be an out-of-bounds
+// access; a 0-D input must be rejected instead of being read as if it had a length.
+TEST(BifurcationDetectorTest, ScalarSrcTokensRejected) {
+  OpTester tester("BifurcationDetector", 1, onnxruntime::kMSDomain);
+
+  tester.AddInput<int64_t>("src_tokens", {}, {1});
+  tester.AddInput<int64_t>("cur_tokens", {1}, {2});
+  tester.AddInput<int64_t>("prev_suffix_match_idx", {}, {0});
+  tester.AddOutput<int64_t>("tokens", {1}, {0});
+  tester.AddOutput<int64_t>("suffix_match_idx", {}, {0});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  tester.Run(OpTester::ExpectResult::kExpectFailure,
+             "src_tokens must be a 1-D tensor",
+             {}, nullptr, &execution_providers);
+}
+
+// Same rank check for cur_tokens.
+TEST(BifurcationDetectorTest, ScalarCurTokensRejected) {
+  OpTester tester("BifurcationDetector", 1, onnxruntime::kMSDomain);
+
+  tester.AddInput<int64_t>("src_tokens", {4}, {1, 5, 3, 4});
+  tester.AddInput<int64_t>("cur_tokens", {}, {2});
+  tester.AddInput<int64_t>("prev_suffix_match_idx", {}, {0});
+  tester.AddOutput<int64_t>("tokens", {1}, {0});
+  tester.AddOutput<int64_t>("suffix_match_idx", {}, {0});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  tester.Run(OpTester::ExpectResult::kExpectFailure,
+             "cur_tokens must be a 1-D tensor",
+             {}, nullptr, &execution_providers);
+}
+
+// Same rank check for pred_tokens, when present.
+TEST(BifurcationDetectorTest, ScalarPredTokensRejected) {
+  OpTester tester("BifurcationDetector", 1, onnxruntime::kMSDomain);
+
+  tester.AddInput<int64_t>("src_tokens", {4}, {1, 5, 3, 4});
+  tester.AddInput<int64_t>("cur_tokens", {1}, {2});
+  tester.AddInput<int64_t>("prev_suffix_match_idx", {}, {0});
+  tester.AddInput<int64_t>("pred_tokens", {}, {1});
+  tester.AddOutput<int64_t>("tokens", {1}, {0});
+  tester.AddOutput<int64_t>("suffix_match_idx", {}, {0});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  tester.Run(OpTester::ExpectResult::kExpectFailure,
+             "pred_tokens must be a 1-D tensor",
+             {}, nullptr, &execution_providers);
+}
+
 }  // namespace test
 }  // namespace onnxruntime