Fix RoiAlign heap out-of-bounds read via unchecked batch_indices

onnxruntime/core/providers/cpu/object_detection/roialign.cc

@@ -294,6 +294,19 @@ Status CheckROIAlignValidInput(const Tensor* X_ptr, const Tensor* rois_ptr, cons
     return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT,
                   "First dimension (num_rois) of batch_indices and rois don't match");
   }
+
+  const int64_t batch_size = X_ptr->Shape()[0];
+  const int64_t* batch_indices_data = batch_indices_ptr->Data<int64_t>();
+  const int64_t num_rois = batch_indices_dims[0];
+  for (int64_t i = 0; i < num_rois; ++i) {
+    if (batch_indices_data[i] < 0 || batch_indices_data[i] >= batch_size) {
+      return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT,
+                    "batch_indices value " + std::to_string(batch_indices_data[i]) +
+                        " at index " + std::to_string(i) +
+                        " is out of range [0, " + std::to_string(batch_size) + ")");
+    }
+  }
+
   return Status::OK();
 }
 

onnxruntime/test/providers/cpu/object_detection/roialign_test.cc

@@ -812,5 +812,43 @@ TEST(RoiAlignTest, MismatchNumRois) {
 
   test.Run(OpTester::ExpectResult::kExpectFailure, "[ShapeInferenceError] Dimension mismatch in unification between 4 and 5");
 }
+
+TEST(RoiAlignTest, BatchIndicesOutOfRange) {
+  OpTester test("RoiAlign", 16);
+  test.AddAttribute<int64_t>("output_height", 2);
+  test.AddAttribute<int64_t>("output_width", 2);
+  test.AddAttribute<int64_t>("sampling_ratio", 2);
+  test.AddAttribute<float>("spatial_scale", 1.0f);
+
+  test.AddInput<float>("X", {1, 1, 4, 4},
+                       {0.f, 1.f, 2.f, 3.f,
+                        4.f, 5.f, 6.f, 7.f,
+                        8.f, 9.f, 10.f, 11.f,
+                        12.f, 13.f, 14.f, 15.f});
+  test.AddInput<float>("rois", {1, 4}, {0.f, 0.f, 3.f, 3.f});
+  test.AddInput<int64_t>("batch_indices", {1}, {1});  // <-- failure condition
+  test.AddOutput<float>("Y", {1, 1, 2, 2}, {0.f, 0.f, 0.f, 0.f});
+
+  test.Run(OpTester::ExpectResult::kExpectFailure, "batch_indices value 1 at index 0 is out of range [0, 1)");
+}
+
+TEST(RoiAlignTest, BatchIndicesNegative) {
+  OpTester test("RoiAlign", 16);
+  test.AddAttribute<int64_t>("output_height", 2);
+  test.AddAttribute<int64_t>("output_width", 2);
+  test.AddAttribute<int64_t>("sampling_ratio", 2);
+  test.AddAttribute<float>("spatial_scale", 1.0f);
+
+  test.AddInput<float>("X", {1, 1, 4, 4},
+                       {0.f, 1.f, 2.f, 3.f,
+                        4.f, 5.f, 6.f, 7.f,
+                        8.f, 9.f, 10.f, 11.f,
+                        12.f, 13.f, 14.f, 15.f});
+  test.AddInput<float>("rois", {1, 4}, {0.f, 0.f, 3.f, 3.f});
+  test.AddInput<int64_t>("batch_indices", {1}, {-1});  // <-- failure condition
+  test.AddOutput<float>("Y", {1, 1, 2, 2}, {0.f, 0.f, 0.f, 0.f});
+
+  test.Run(OpTester::ExpectResult::kExpectFailure, "batch_indices value -1 at index 0 is out of range [0, 1)");
+}
 }  // namespace test
 }  // namespace onnxruntime