gather: bounds-check runtime position values in Eval

tensorflow/lite/micro/kernels/gather.cc

@@ -82,13 +82,22 @@ TfLiteStatus Gather(const TfLiteGatherParams* params,
   for (int batch = 0; batch < batch_size; ++batch) {
     for (int outer = 0; outer < outer_size; ++outer) {
       for (int coord = 0; coord < coord_size; ++coord) {
-        TFLITE_DCHECK_GE(coords_data[coord], 0);
-        TFLITE_DCHECK_LT(coords_data[coord], axis_size);
+        // The positions tensor is a runtime input, so its values are control
+        // data that has to be bounds-checked here rather than in Prepare. A
+        // raw branch is used instead of TF_LITE_ENSURE so no error string is
+        // placed in .rodata. The previous TFLITE_DCHECKs compiled out in
+        // release builds, and they also read coords_data[coord] while the
+        // memcpy below indexes coords_data[batch * coord_size + coord], so
+        // they did not guard the value actually used once batch_size > 1.
+        const CoordsT coord_value = coords_data[batch * coord_size + coord];
+        if (coord_value < 0 || coord_value >= axis_size) {
+          return kTfLiteError;
+        }
         std::memcpy(output_data +
                         (((batch * outer_size) + outer) * coord_size + coord) *
                             inner_size,
                     input_data + (((batch * outer_size) + outer) * axis_size +
-                                  coords_data[batch * coord_size + coord]) *
+                                  coord_value) *
                                      inner_size,
                     sizeof(InputT) * inner_size);
       }

tensorflow/lite/micro/kernels/gather_test.cc

@@ -68,6 +68,39 @@ void TestGather(int* input_dims, const InType* input_data, int* positions_dims,
   }
 }
 
+// Runs GATHER with an out-of-range value in the positions tensor and checks
+// that Invoke() returns an error instead of indexing the input tensor out of
+// bounds. Prepare() does not inspect position values, so it still succeeds.
+template <typename InType, typename PosType>
+void TestGatherOutOfRangePosition(int* input_dims, const InType* input_data,
+                                  int* positions_dims,
+                                  const PosType* positions_data,
+                                  int* output_dims, InType* output_data,
+                                  const int axis = 0,
+                                  const int batch_dims = 0) {
+  TfLiteIntArray* in_dims = IntArrayFromInts(input_dims);
+  TfLiteIntArray* pos_dims = IntArrayFromInts(positions_dims);
+  TfLiteIntArray* out_dims = IntArrayFromInts(output_dims);
+  TfLiteGatherParams params = {axis, batch_dims};
+
+  constexpr int tensors_size = 3;
+  TfLiteTensor tensors[tensors_size] = {
+      CreateTensor(input_data, in_dims),
+      CreateTensor(positions_data, pos_dims),
+      CreateTensor(output_data, out_dims, true),
+  };
+  int inputs_array_data[] = {2, 0, 1};
+  TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
+  int outputs_array_data[] = {1, 2};
+  TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
+
+  const TFLMRegistration registration = Register_GATHER();
+  micro::KernelRunner runner(registration, tensors, tensors_size, inputs_array,
+                             outputs_array, &params);
+  EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare());
+  EXPECT_EQ(kTfLiteError, runner.Invoke());
+}
+
 }  // namespace
 }  // namespace testing
 }  // namespace tflite
@@ -458,4 +491,49 @@ TEST(GatherTest, GatherOp_BatchDimsEqualIndexDims) {
       output_data, golden_dims, golden_data, axis, batch_dims);
 }
 
+TEST(GatherTest, GatherOp_PositionAboveAxisSizeReturnsError) {
+  // Axis 0 has size 3, so position 3 is out of range and must be rejected
+  // rather than read past the end of the input tensor.
+  int input_dims[] = {2, 3, 4};
+  const float input_data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11};
+  int positions_dims[] = {1, 1};
+  const int32_t positions_data[] = {3};
+  float output_data[4];
+  int output_dims[] = {2, 0, 0};
+  tflite::testing::TestGatherOutOfRangePosition<float, int32_t>(
+      input_dims, input_data, positions_dims, positions_data, output_dims,
+      output_data);
+}
+
+TEST(GatherTest, GatherOp_NegativePositionReturnsError) {
+  // A negative position must be rejected rather than read before the start of
+  // the input tensor.
+  int input_dims[] = {2, 3, 4};
+  const float input_data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11};
+  int positions_dims[] = {1, 1};
+  const int32_t positions_data[] = {-1};
+  float output_data[4];
+  int output_dims[] = {2, 0, 0};
+  tflite::testing::TestGatherOutOfRangePosition<float, int32_t>(
+      input_dims, input_data, positions_dims, positions_data, output_dims,
+      output_data);
+}
+
+TEST(GatherTest, GatherOp_OutOfRangePositionInSecondBatchReturnsError) {
+  // With batch_dims = 1 the offending value sits in the second batch, at
+  // coords_data[batch * coord_size + coord]. The previous DCHECKs only read
+  // coords_data[coord], so this case was unguarded even in debug builds.
+  const int axis = 1;
+  const int batch_dims = 1;
+  int input_dims[] = {3, 2, 2, 2};
+  const float input_data[] = {0, 1, 2, 3, 4, 5, 6, 7};
+  int positions_dims[] = {2, 2, 1};
+  const int32_t positions_data[] = {0, 2};
+  float output_data[4];
+  int output_dims[] = {3, 0, 0, 0};
+  tflite::testing::TestGatherOutOfRangePosition<float, int32_t>(
+      input_dims, input_data, positions_dims, positions_data, output_dims,
+      output_data, axis, batch_dims);
+}
+
 TF_LITE_MICRO_TESTS_MAIN