From c59c37e7b2d563967da813fa50fe20b21f4da683 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Wed, 28 Apr 2021 17:50:10 -0700
Subject: [PATCH] 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
---
 tensorflow/lite/kernels/arg_min_max.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/lite/kernels/arg_min_max.cc b/tensorflow/lite/kernels/arg_min_max.cc
index a0ba8cb9f8bbe..291fd61681f2a 100644
--- a/tensorflow/lite/kernels/arg_min_max.cc
+++ b/tensorflow/lite/kernels/arg_min_max.cc
@@ -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;