CVE-2020-15209: NULL Pointer Dereference (CWE-476)
value: In tensorflow-lite before versions 1.15.4, 2.0.3, 2.1.2, 2.2.1 and 2.3.1, a crafted TFLite model can force a node to have as input a tensor backed by a `nullptr` buffer. This can be achieved by changing a buffer index in the flatbuffer serialization to convert a read-only tensor to a read-write one. The runtime assumes that these buffers are written to before a possible read, hence they are initialized with `nullptr`. However, by changing the buffer index for a tensor and implicitly converting that tensor to be a read-write one, as there is nothing in the model that writes to it, we get a null pointer dereference. The issue is patched in commit 0b5662bc, and is released in TensorFlow versions 1.15.4, 2.0.3, 2.1.2, 2.2.1, or 2.3.1.
CWE description: The product dereferences a pointer that it expects to be valid but is NULL.
CVSS v3 base score: 5.9
Published: 2020-09-25T19:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/lite/BUILD, tensorflow/lite/core/subgraph.cc, tensorflow/lite/model_test.cc
Commit date: 2020-09-18 14:34:47 -0700
Commit message: [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

Diff: diff --git a/tensorflow/lite/BUILD b/tensorflow/lite/BUILD
index e50cd4549b2..d08a2e2fbfa 100644
--- a/tensorflow/lite/BUILD
+++ b/tensorflow/lite/BUILD
@@ -240,10 +240,11 @@ cc_library(
     deps = [
         ":allocation",
         ":arena_planner",
         ":external_cpu_backend_context",
         ":graph_info",
+        ":kernel_api",
         ":memory_planner",
         ":minimal_logging",
         ":shared_library",
         ":simple_memory_arena",
         ":string",
@@ -467,10 +468,11 @@ cc_test(
         "testdata/0_subgraphs.bin",
         "testdata/2_subgraphs.bin",
         "testdata/add_shared_tensors.bin",
         "testdata/empty_model.bin",
         "testdata/multi_add_flex.bin",
+        "testdata/segment_sum_invalid_buffer.bin",
         "testdata/sparse_tensor.bin",
         "testdata/test_min_runtime.bin",
         "testdata/test_model.bin",
         "testdata/test_model_broken.bin",
     ],
diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc
index 47d0fddb12b..235a776c56f 100644
--- a/tensorflow/lite/core/subgraph.cc
+++ b/tensorflow/lite/core/subgraph.cc
@@ -17,10 +17,11 @@ limitations under the License.
 
 #include <algorithm>
 #include <cstdint>
 
 #include "tensorflow/lite/arena_planner.h"
+#include "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"
 #include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
 #include "tensorflow/lite/graph_info.h"
@@ -1028,10 +1029,23 @@ TfLiteStatus Subgraph::Invoke() {
       TfLiteTensor* tensor = &tensors_[tensor_index];
       if (tensor->delegate && tensor->delegate != node.delegate &&
           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 &&
         check_cancelled_func_(cancellation_data_)) {
       ReportError("Client requested cancel during Invoke()");
diff --git a/tensorflow/lite/model_test.cc b/tensorflow/lite/model_test.cc
index 6993e350a48..110c54aa571 100644
--- a/tensorflow/lite/model_test.cc
+++ b/tensorflow/lite/model_test.cc
@@ -436,28 +436,52 @@ TEST(BasicFlatBufferModel, TestParseModelWithSparseTensor) {
   ASSERT_EQ(t1->sparsity->dim_metadata[3].array_segments, nullptr);
   ASSERT_EQ(t1->sparsity->dim_metadata[3].array_indices, nullptr);
 }
 
 // TODO(b/150072943): Add malformed model with sparse tensor tests.
-TEST(BasicFlatBufferModel, TestHandleMalformedModel) {
-  const auto model_paths = {
-      // These models use the same tensor as both input and ouput of a node
-      "tensorflow/lite/testdata/add_shared_tensors.bin",
-  };
-
-  for (const auto& model_path : model_paths) {
-    std::unique_ptr<tflite::FlatBufferModel> model =
-        FlatBufferModel::BuildFromFile(model_path);
-    ASSERT_NE(model, nullptr);
-
-    tflite::ops::builtin::BuiltinOpResolver resolver;
-    InterpreterBuilder builder(*model, resolver);
-    std::unique_ptr<Interpreter> interpreter;
-    ASSERT_EQ(builder(&interpreter), kTfLiteOk);
-    ASSERT_NE(interpreter, nullptr);
-    ASSERT_NE(interpreter->AllocateTensors(), kTfLiteOk);
-  }
+
+// The models here have at least a node that uses the same tensor as input and
+// output. This causes segfaults when trying to eval the operator, hence we try
+// to prevent this scenario. The earliest place we can check this is in
+// `AllocateTensors`, hence the test checks that `interpreter->AllocateTensors`
+// detects these bad models.
+TEST(BasicFlatBufferModel, TestHandleMalformedModelReuseTensor) {
+  const auto model_path =
+      "tensorflow/lite/testdata/add_shared_tensors.bin";
+
+  std::unique_ptr<tflite::FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(model_path);
+  ASSERT_NE(model, nullptr);
+
+  tflite::ops::builtin::BuiltinOpResolver resolver;
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> interpreter;
+  ASSERT_EQ(builder(&interpreter), kTfLiteOk);
+  ASSERT_NE(interpreter, nullptr);
+  ASSERT_NE(interpreter->AllocateTensors(), kTfLiteOk);
+}
+
+// The models here have a buffer index for a tensor pointing to a null buffer.
+// This results in the tensor being interpreted as read-write, but the model
+// assumes the tensor is read-only. As such, `interpreter->Invoke()` would
+// segfault if no precondition check is added. The test checks that the
+// precondition check exists.
+TEST(BasicFlatBufferModel, TestHandleMalformedModelInvalidBuffer) {
+  const auto model_path =
+      "tensorflow/lite/testdata/segment_sum_invalid_buffer.bin";
+
+  std::unique_ptr<tflite::FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(model_path);
+  ASSERT_NE(model, nullptr);
+
+  tflite::ops::builtin::BuiltinOpResolver resolver;
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> interpreter;
+  ASSERT_EQ(builder(&interpreter), kTfLiteOk);
+  ASSERT_NE(interpreter, nullptr);
+  ASSERT_EQ(interpreter->AllocateTensors(), kTfLiteOk);
+  ASSERT_NE(interpreter->Invoke(), kTfLiteOk);
 }
 
 // TODO(aselle): Add tests for serialization of builtin op data types.
 // These tests will occur with the evaluation tests of individual operators,
 // not here.
diff --git a/tensorflow/lite/testdata/segment_sum_invalid_buffer.bin b/tensorflow/lite/testdata/segment_sum_invalid_buffer.bin
new file mode 100644
index 00000000000..b7b81cde6e7
Binary files /dev/null and b/tensorflow/lite/testdata/segment_sum_invalid_buffer.bin differ

Vulnerable code: TEST(BasicFlatBufferModel, TestHandleMalformedModel) {
  const auto model_paths = {
      // These models use the same tensor as both input and ouput of a node
      "tensorflow/lite/testdata/add_shared_tensors.bin",
  };

  for (const auto& model_path : model_paths) {
    std::unique_ptr<tflite::FlatBufferModel> model =
        FlatBufferModel::BuildFromFile(model_path);
    ASSERT_NE(model, nullptr);

    tflite::ops::builtin::BuiltinOpResolver resolver;
    InterpreterBuilder builder(*model, resolver);
    std::unique_ptr<Interpreter> interpreter;
    ASSERT_EQ(builder(&interpreter), kTfLiteOk);
    ASSERT_NE(interpreter, nullptr);
    ASSERT_NE(interpreter->AllocateTensors(), kTfLiteOk);
  }
Fixed code:         ":kernel_api",
        "testdata/segment_sum_invalid_buffer.bin",
#include "tensorflow/lite/builtin_ops.h"
      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;
        }
      }

// The models here have at least a node that uses the same tensor as input and
// output. This causes segfaults when trying to eval the operator, hence we try
// to prevent this scenario. The earliest place we can check this is in
// `AllocateTensors`, hence the test checks that `interpreter->AllocateTensors`
// detects these bad models.
TEST(BasicFlatBufferModel, TestHandleMalformedModelReuseTensor) {
  const auto model_path =
      "tensorflow/lite/testdata/add_shared_tensors.bin";

  std::unique_ptr<tflite::FlatBufferModel> model =
      FlatBufferModel::BuildFromFile(model_path);
  ASSERT_NE(model, nullptr);

  tflite::ops::builtin::BuiltinOpResolver resolver;
  InterpreterBuilder builder(*model, resolver);
  std::unique_ptr<Interpreter> interpreter;
  ASSERT_EQ(builder(&interpreter), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);
  ASSERT_NE(interpreter->AllocateTensors(), kTfLiteOk);
}

// The models here have a buffer index for a tensor pointing to a null buffer.
// This results in the tensor being interpreted as read-write, but the model
// assumes the tensor is read-only. As such, `interpreter->Invoke()` would
// segfault if no precondition check is added. The test checks that the
// precondition check exists.
TEST(BasicFlatBufferModel, TestHandleMalformedModelInvalidBuffer) {
  const auto model_path =
      "tensorflow/lite/testdata/segment_sum_invalid_buffer.bin";

  std::unique_ptr<tflite::FlatBufferModel> model =
      FlatBufferModel::BuildFromFile(model_path);
  ASSERT_NE(model, nullptr);

  tflite::ops::builtin::BuiltinOpResolver resolver;
  InterpreterBuilder builder(*model, resolver);
  std::unique_ptr<Interpreter> interpreter;
  ASSERT_EQ(builder(&interpreter), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);
  ASSERT_EQ(interpreter->AllocateTensors(), kTfLiteOk);
  ASSERT_NE(interpreter->Invoke(), kTfLiteOk);