From 703b412b4188fab009be183fd5651767eb5e0752 Mon Sep 17 00:00:00 2001 From: evilgensec Date: Fri, 12 Jun 2026 08:37:15 +0545 Subject: [PATCH] Reject out-of-range GATHER index values at runtime to prevent OOB read The reference GATHER kernel indexes the input tensor with the values in the `positions` tensor and only guards them with TFLITE_DCHECK_GE/LT, which expand to no-ops under -DNDEBUG (the release build). A negative or too-large index supplied as a runtime input therefore makes GatherEval read outside the input tensor (heap-buffer-overflow). GatherPrepare validates axis and batch_dims but never the index values. Validate the index values in GatherEval and return kTfLiteError for any index outside [0, axis_size), mirroring the runtime index check in the full TFLite GATHER kernel. Adds regression tests for an out-of-range and a negative index. --- tensorflow/lite/micro/kernels/gather.cc | 22 +++++++ tensorflow/lite/micro/kernels/gather_test.cc | 62 ++++++++++++++++++++ 2 files changed, 84 insertions(+) diff --git a/tensorflow/lite/micro/kernels/gather.cc b/tensorflow/lite/micro/kernels/gather.cc index a0af4c0edda..2762de3ac3a 100644 --- a/tensorflow/lite/micro/kernels/gather.cc +++ b/tensorflow/lite/micro/kernels/gather.cc @@ -199,6 +199,28 @@ TfLiteStatus GatherEval(TfLiteContext* context, TfLiteNode* node) { tflite::micro::GetEvalOutput(context, node, kOutputTensor); if (coords->type == kTfLiteInt32) { + // The reference Gather() below indexes `input` with the values in `coords` + // and only guards them with TFLITE_DCHECK, which is compiled out in release + // (-DNDEBUG) builds. Validate the index values at runtime so that an + // out-of-range index supplied as a model's runtime input fails closed with + // kTfLiteError instead of reading out of bounds. This mirrors the runtime + // index validation performed by the full TFLite GATHER kernel. + const TfLiteIntArray* input_dims = input->dims; + int axis = params->axis; + if (axis < 0) { + axis += input_dims->size; + } + TF_LITE_ENSURE(context, axis >= 0 && axis < input_dims->size); + const int32_t axis_size = input_dims->data[axis]; + const int32_t* coords_data = tflite::micro::GetTensorData(coords); + int num_coords = 1; + for (int i = 0; i < coords->dims->size; ++i) { + num_coords *= coords->dims->data[i]; + } + for (int i = 0; i < num_coords; ++i) { + TF_LITE_ENSURE(context, coords_data[i] >= 0 && coords_data[i] < axis_size); + } + switch (input->type) { case kTfLiteFloat32: return Gather(params, input, coords, output); diff --git a/tensorflow/lite/micro/kernels/gather_test.cc b/tensorflow/lite/micro/kernels/gather_test.cc index 35ad7dcb17e..3b1a0386e25 100644 --- a/tensorflow/lite/micro/kernels/gather_test.cc +++ b/tensorflow/lite/micro/kernels/gather_test.cc @@ -68,6 +68,40 @@ void TestGather(int* input_dims, const InType* input_data, int* positions_dims, } } +// Builds a GATHER op with the given (out-of-range) index value and checks that +// Invoke() fails closed with kTfLiteError instead of reading the input tensor +// out of bounds. The reference kernel only guards index values with +// TFLITE_DCHECK, which is a no-op in release (-DNDEBUG) builds. +template +void TestGatherFailsForOutOfRangeIndex( + 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); + // Prepare does not inspect index values, so it still succeeds. + EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare()); + // Invoke must reject the out-of-range index rather than read out of bounds. + EXPECT_EQ(kTfLiteError, runner.Invoke()); +} + } // namespace } // namespace testing } // namespace tflite @@ -458,4 +492,32 @@ TEST(GatherTest, GatherOp_BatchDimsEqualIndexDims) { output_data, golden_dims, golden_data, axis, batch_dims); } +TEST(GatherTest, GatherOp_IndexGreaterThanAxisSizeFailsClosed) { + // axis 0 has size 3, so index 3 is out of range and must be rejected at + // runtime rather than reading 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::TestGatherFailsForOutOfRangeIndex( + input_dims, input_data, positions_dims, positions_data, output_dims, + output_data); +} + +TEST(GatherTest, GatherOp_NegativeIndexFailsClosed) { + // A negative index must be rejected at runtime rather than reading 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::TestGatherFailsForOutOfRangeIndex( + input_dims, input_data, positions_dims, positions_data, output_dims, + output_data); +} + TF_LITE_MICRO_TESTS_MAIN