Fix GatherND sentinel-value collision enabling OOB memory read

GatherND used err_index==0 as a sentinel for 'no validation error',
but 0 is also a valid tensor index. When a dimension has size 0, index
0 correctly fails the bounds check (0 >= 0), but err_index is set to 0
which is indistinguishable from success. This allowed the operator to
proceed with an out-of-bounds memcpy, leaking process heap memory.

Replace the sentinel pattern with a std::atomic<bool> error flag that
is independent of the index value. The err_index variable is retained
solely for the error message.

Files changed:
- onnxruntime/core/providers/cpu/tensor/gather_nd.cc
- onnxruntime/test/providers/cpu/tensor/gather_nd_op_test.cc

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: bmehta001

onnxruntime/core/providers/cpu/tensor/gather_nd.cc

@@ -1,5 +1,6 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
+#include <atomic>
 #include <core/common/safeint.h>
 #include "gather_nd.h"
 #include "core/platform/threadpool.h"
@@ -85,6 +86,11 @@ Status GatherNDBase::PrepareForCompute(const TensorShape& input_shape, const Ten
     sizes_from_slice_dims[onnxruntime::narrow<size_t>(i)] = input_shape.SizeFromDimension(SafeInt<size_t>(batch_dims_) + i + 1);
   }
 
+  // Use a separate flag to track validation errors instead of relying on the
+  // index value as a sentinel. The value 0 is a valid tensor index, so using
+  // err_index == 0 to mean "no error" creates a collision when a zero-sized
+  // dimension makes index 0 invalid — the error goes undetected.
+  std::atomic<bool> has_invalid_index{false};
   int64_t err_index = 0;
   p.element_bytes = bytes_per_value;
   p.element_count_per_slice = slice_size;
@@ -104,6 +110,7 @@ Status GatherNDBase::PrepareForCompute(const TensorShape& input_shape, const Ten
       const auto upper_limit = input_shape[SafeInt<size_t>(batch_dims_) + dim_idx];
       const auto lower_limit = -upper_limit;
       if (index < lower_limit || index >= upper_limit) {
+        has_invalid_index.store(true, std::memory_order_relaxed);
         err_index = index;
         break;
       }
@@ -123,8 +130,9 @@ Status GatherNDBase::PrepareForCompute(const TensorShape& input_shape, const Ten
         }
       });
 
-  return err_index == 0 ? Status::OK()
-                        : ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "invalid index found, index = ", err_index);
+  return !has_invalid_index.load(std::memory_order_relaxed)
+             ? Status::OK()
+             : ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "invalid index found, index = ", err_index);
 }
 
 template Status GatherNDBase::PrepareForCompute<int32_t>(const TensorShape&,

onnxruntime/test/providers/cpu/tensor/gather_nd_op_test.cc

@@ -372,5 +372,41 @@ TEST(GatherNDOpTest, GatherND_zero_batch_dims_error) {
            &cpu_only_ep);  // force CPU
 }
 
+// Regression test: index 0 into a zero-sized dimension must be rejected.
+// Previously err_index==0 was used as a sentinel for "no error", which collided
+// with the actual index value 0, allowing an OOB memory read.
+TEST(GatherNDOpTest, GatherND_zero_dim_index_zero_rejected) {
+  OpTester test("GatherND", 12, kOnnxDomain);
+  test.AddAttribute<int64_t>("batch_dims", 0);
+
+  // data shape [0]: zero-sized first dimension
+  // indices [0]: index 0 is invalid because dimension size is 0
+  test.AddInput<int32_t>("data", {0}, {});
+  test.AddInput<int64_t>("indices", {1}, {0});
+  test.AddOutput<int32_t>("output", {}, {0});  // dummy, won't be used
+
+  std::vector<std::unique_ptr<onnxruntime::IExecutionProvider>> cpu_only_ep;
+  cpu_only_ep.push_back(DefaultCpuExecutionProvider());
+
+  test.Run(OpTester::ExpectResult::kExpectFailure,
+           "invalid index found",
+           {},
+           nullptr,
+           &cpu_only_ep);
+}
+
+// Verify that index 0 into a non-zero dimension still works correctly.
+TEST(GatherNDOpTest, GatherND_valid_index_zero) {
+  OpTester test("GatherND", 12, kOnnxDomain);
+  test.AddAttribute<int64_t>("batch_dims", 0);
+
+  // data shape [3]: valid dimension, index 0 should return the first element
+  test.AddInput<int32_t>("data", {3}, {10, 20, 30});
+  test.AddInput<int64_t>("indices", {1}, {0});
+  test.AddOutput<int32_t>("output", {}, {10});
+
+  test.Run();
+}
+
 }  // namespace test
 }  // namespace onnxruntime