Skip to content

Commit

Permalink
Prevent an OOB read in expand_dims.cc
Browse files Browse the repository at this point in the history
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
  • Loading branch information
mihaimaruseac authored and geetachavan1 committed Jul 28, 2021
1 parent 8d5f87f commit 6fadd98
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions tensorflow/lite/kernels/expand_dims.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down

0 comments on commit 6fadd98

Please sign in to comment.