OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 12:20:11 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Merge pull request #49758 from geetachavan1/cherrypicks_X9FKE

Browse Source 
Prevent a null pointer dereference in TFLite.
This commit is contained in:
Mihai Maruseac 2021-05-26 11:05:43 -07:00 committed by GitHub
commit d4039e5779
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
tensorflow/lite/core/subgraph.cc
View File

@ -1033,10 +1033,17 @@ TfLiteStatus Subgraph::Invoke() {
TF_LITE_ENSURE_STATUS(EnsureTensorDataIsReadable(tensor_index));
}
if (tensor->data.raw == nullptr && tensor->bytes > 0) {
if (registration.builtin_code == kTfLiteBuiltinReshape && i == 1) {
if (registration.builtin_code == kTfLiteBuiltinReshape && i == 1 &&
tensor->dims->size != 1) {
// In general, having a tensor here with no buffer will be an error.
// However, for the reshape operator, the second input tensor is only
// used for the shape, not for the data. Thus, null buffer is ok.
// However, for the reshape operator, the second input tensor is
// sometimes only used for the shape, not for the data. Thus, null
// buffer is ok in this situation.
// The situation where null buffer is not ok for reshape operator is
// only when there are 2 inputs given to the node and the one
// corresponding to the shape (i == 1) is a vector that contains all
// dimensions. See `GetOutputShape()` function in
// `tensorflow/lite/kernels/reshape.cc`
continue;
} else {
// In all other cases, we need to return an error as otherwise we will