Prevent array OOB read/write

mihaimaruseac · tensorflower-gardener · commit ae2daeb45abf · 2021-04-28T18:15:39.000-07:00
PiperOrigin-RevId: 371026165
Change-Id: I26ac6372c87246e03c7eb8c94e84c84d86054b36
diff --git 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) {

Original file line number	Diff line number	Diff line change
`@@ -96,6 +96,8 @@ TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node,`
`96`	`96`	`}`
`97`	`97`	`}`
`98`	`98`
	`99`	`+ TF_LITE_ENSURE(context, axis_value >= 0);`
	`100`	`+ TF_LITE_ENSURE(context, axis_value < NumDimensions(input));`
`99`	`101`	`const int input_size = SizeOfDimension(input, axis_value);`
`100`	`102`
`101`	`103`	`if (minus_one_index != -1) {`