Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent heap OOB read in TFLite's gather_nd.cc.
Passing negative indices is illegal but there was a missing check so that resulted in OOB accesses.

PiperOrigin-RevId: 387208551
Change-Id: I6b7a8a62d3e7c13a16d81619e5bc23ae2cdbc7fd
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Jul 27, 2021
1 parent 247e60c commit bb6a038
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions tensorflow/lite/kernels/gather_nd.cc
Expand Up @@ -123,6 +123,17 @@ TfLiteStatus GatherNdString(const TfLiteTensor* params,
template <typename IndicesT>
TfLiteStatus EvalGatherNd(TfLiteContext* context, const TfLiteTensor* params,
const TfLiteTensor* indices, TfLiteTensor* output) {
bool indices_has_only_positive_elements = true;
const auto* indices_values = GetTensorData<IndicesT>(indices);
const size_t num_indices = indices->bytes / sizeof(IndicesT);
for (size_t i = 0; i < num_indices; i++) {
if (indices_values[i] < 0) {
indices_has_only_positive_elements = false;
break;
}
}
TF_LITE_ENSURE(context, indices_has_only_positive_elements);

switch (params->type) {
case kTfLiteFloat32:
return GatherNd<float, IndicesT>(params, indices, output);
Expand Down

0 comments on commit bb6a038

Please sign in to comment.