From 61bfb9f3768bf932268df5d1f8bf823e350dfff3 Mon Sep 17 00:00:00 2001 From: TristanInSec Date: Sun, 26 Jul 2026 09:50:52 -0400 Subject: [PATCH] Fix OOB read in Gather kernel via runtime bounds check on coordinate indices The Gather kernel uses TFLITE_DCHECK_GE/TFLITE_DCHECK_LT to validate coordinate indices, but these compile to no-ops in release builds (NDEBUG). This allows attacker-controlled indices from an input tensor to read arbitrary heap memory past the input buffer. Per the Error Handling Guide, control data from input tensors evaluated at runtime should use raw if/return validation in Eval to prevent memory corruption. This matches the pattern in gather_nd.cc which already validates from_pos at runtime. Replace the TFLITE_DCHECK pair with a runtime bounds check that returns kTfLiteError for out-of-range indices, consistent with Section 4 (Preventing Buffer Overflows in Eval) of the guide. --- tensorflow/lite/micro/kernels/gather.cc | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tensorflow/lite/micro/kernels/gather.cc b/tensorflow/lite/micro/kernels/gather.cc index a0af4c0edda..0a5f10b9fde 100644 --- a/tensorflow/lite/micro/kernels/gather.cc +++ b/tensorflow/lite/micro/kernels/gather.cc @@ -82,14 +82,16 @@ 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); + const CoordsT idx = coords_data[batch * coord_size + coord]; + if (idx < 0 || idx >= 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]) * - inner_size, + input_data + + (((batch * outer_size) + outer) * axis_size + idx) * + inner_size, sizeof(InputT) * inner_size); } }