Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent array write out-of-bounds.
If user passes an invalid axis, then we copy one too many dimensions to the output in the loop below these checks. Even if we didn't do that, there will be further issues with an invalid axis, so we check for that right now.

PiperOrigin-RevId: 371023299
Change-Id: I9eca37ffc2b29e8e48710f500701270ef0790224
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Apr 29, 2021
1 parent 102b211 commit c59c37e
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tensorflow/lite/kernels/arg_min_max.cc
Expand Up @@ -48,6 +48,9 @@ TfLiteStatus ResizeOutput(TfLiteContext* context, const TfLiteTensor* input,
axis_value += NumDimensions(input);
}

TF_LITE_ENSURE(context, axis_value >= 0);
TF_LITE_ENSURE(context, axis_value < NumDimensions(input));

// Copy the input dimensions to output except the axis dimension.
TfLiteIntArray* output_dims = TfLiteIntArrayCreate(NumDimensions(input) - 1);
int j = 0;
Expand Down

0 comments on commit c59c37e

Please sign in to comment.