From 62934d0824244e52e09835082b36ca7c4a9f5388 Mon Sep 17 00:00:00 2001 From: evilgensec Date: Sat, 1 Aug 2026 23:05:38 +0545 Subject: [PATCH] gather: bounds-check runtime position values in Eval The GATHER kernel copies each output slice using a position value taken from the positions input tensor: input_data + (((batch * outer_size) + outer) * axis_size + coords_data[batch * coord_size + coord]) * inner_size Those values were guarded only by TFLITE_DCHECK, which compiles out under NDEBUG, so a release build used an out-of-range position directly as an offset and read outside the input tensor. The DCHECKs also read coords_data[coord] while the memcpy indexes coords_data[batch * coord_size + coord], so once batch_size is greater than 1 they did not guard the value actually used. Read the position once, reject it with a raw branch when it falls outside [0, axis_size), and use that same value in the memcpy. The error handling guide calls for a raw branch rather than TF_LITE_ENSURE for runtime index data in Eval, so no error string is added to .rodata. Cost in release is one compare and branch per copied slice. Adds three tests covering a position above the axis size, a negative position, and an out-of-range position in the second batch, which is the case the previous DCHECKs never read. --- tensorflow/lite/micro/kernels/gather.cc | 15 +++- tensorflow/lite/micro/kernels/gather_test.cc | 78 ++++++++++++++++++++ 2 files changed, 90 insertions(+), 3 deletions(-) diff --git a/tensorflow/lite/micro/kernels/gather.cc b/tensorflow/lite/micro/kernels/gather.cc index a0af4c0edda..b2d2a9be967 100644 --- a/tensorflow/lite/micro/kernels/gather.cc +++ b/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); } diff --git a/tensorflow/lite/micro/kernels/gather_test.cc b/tensorflow/lite/micro/kernels/gather_test.cc index 35ad7dcb17e..a856cd6f6d9 100644 --- a/tensorflow/lite/micro/kernels/gather_test.cc +++ b/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 +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, ¶ms); + 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( + 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( + 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( + input_dims, input_data, positions_dims, positions_data, output_dims, + output_data, axis, batch_dims); +} + TF_LITE_MICRO_TESTS_MAIN