From ae2daeb45abfe2c6dda539cf8d0d6f653d3ef412 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Wed, 28 Apr 2021 18:12:15 -0700
Subject: [PATCH] Prevent array OOB read/write

PiperOrigin-RevId: 371026165
Change-Id: I26ac6372c87246e03c7eb8c94e84c84d86054b36
---
 tensorflow/lite/kernels/split_v.cc | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/lite/kernels/split_v.cc b/tensorflow/lite/kernels/split_v.cc
index 054e00572f5a6..ed8a4851c1e97 100644
--- a/tensorflow/lite/kernels/split_v.cc
+++ b/tensorflow/lite/kernels/split_v.cc
@@ -96,6 +96,8 @@ TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node,
     }
   }
 
+  TF_LITE_ENSURE(context, axis_value >= 0);
+  TF_LITE_ENSURE(context, axis_value < NumDimensions(input));
   const int input_size = SizeOfDimension(input, axis_value);
 
   if (minus_one_index != -1) {