Prevent an OOB read in expand_dims.cc

mihaimaruseac · tensorflower-gardener · commit d94ffe08a654 · 2021-07-27T14:52:03.000-07:00
The for loop that follows this check assumes that `axis` is between `0` and `input_dims.size`. If user supplied `axis` is negative, the if code before this check is supposed to bring it back to positive (similar to how in Python one can do `l[-3]` to mean `l[-3 + len(l)]`).

PiperOrigin-RevId: 387200206
Change-Id: I162f4feba12d547c3a4340833ae682016a2ebfab
diff --git a/tensorflow/lite/kernels/expand_dims.cc b/tensorflow/lite/kernels/expand_dims.cc
@@ -37,6 +37,7 @@ TfLiteStatus ExpandTensorDim(TfLiteContext* context, const TfLiteTensor& input,
     axis = input_dims.size + 1 + axis;
   }
   TF_LITE_ENSURE(context, axis <= input_dims.size);
+  TF_LITE_ENSURE(context, axis >= 0);
 
   TfLiteIntArray* output_dims = TfLiteIntArrayCreate(input_dims.size + 1);
   for (int i = 0; i < output_dims->size; ++i) {

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,7 @@ TfLiteStatus ExpandTensorDim(TfLiteContext* context, const TfLiteTensor& input,`
`37`	`37`	`axis = input_dims.size + 1 + axis;`
`38`	`38`	`}`
`39`	`39`	`TF_LITE_ENSURE(context, axis <= input_dims.size);`
	`40`	`+ TF_LITE_ENSURE(context, axis >= 0);`
`40`	`41`
`41`	`42`	`TfLiteIntArray* output_dims = TfLiteIntArrayCreate(input_dims.size + 1);`
`42`	`43`	`for (int i = 0; i < output_dims->size; ++i) {`