From 1a506aef22640364ca629ebefa14605a33d2efae Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Sat, 19 Sep 2020 18:00:09 -0700 Subject: [PATCH] [tflite] Ensure input tensors don't have `nullptr` buffers. A crafted TFLite model can force a node to have as input a tensor backed by a `nullptr` buffer. That is, by carefully changing the buffer index in the flatbuffer serialization, we can force the TFLite interpreter to consider a read-only tensor to be a read-write one and assume that there is an operator that has this tensor as output, writing to it and allocating memory before the tensor is used as input. If this does not happen, we get memory corruption. PiperOrigin-RevId: 332524692 Change-Id: I57ef175152a29020af9ab041dc959e5631dce40f --- tensorflow/lite/core/subgraph.cc | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc index 5ad8b51d433..0646adebf8c 100644 --- a/tensorflow/lite/core/subgraph.cc +++ b/tensorflow/lite/core/subgraph.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "tensorflow/lite/arena_planner.h" +#include "third_party/tensorflow/lite/builtin_ops.h" #include "tensorflow/lite/c/common.h" #include "tensorflow/lite/context_util.h" #include "tensorflow/lite/core/api/tensor_utils.h" @@ -971,6 +972,19 @@ TfLiteStatus Subgraph::Invoke() { tensor->data_is_stale) { TF_LITE_ENSURE_STATUS(EnsureTensorDataIsReadable(tensor_index)); } + if (tensor->data.raw == nullptr && tensor->bytes > 0) { + if (registration.builtin_code == kTfLiteBuiltinReshape && i == 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. + continue; + } else { + // In all other cases, we need to return an error as otherwise we will + // trigger a null pointer dereference (likely). + ReportError("Input tensor %d lacks data", tensor_index); + return kTfLiteError; + } + } } if (check_cancelled_func_ != nullptr &&