CVE-2020-15210: Out-of-bounds Write (CWE-787)
value: In tensorflow-lite before versions 1.15.4, 2.0.3, 2.1.2, 2.2.1 and 2.3.1, if a TFLite saved model uses the same tensor as both input and output of an operator, then, depending on the operator, we can observe a segmentation fault or just memory corruption. We have patched the issue in d58c96946b and will release patch releases for all versions between 1.15 and 2.3. We recommend users to upgrade to TensorFlow 1.15.4, 2.0.3, 2.1.2, 2.2.1, or 2.3.1.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 6.5
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/core/subgraph.h, tensorflow/lite/model_test.cc
Commit date: 2020-09-18 14:21:55 -0700
Commit message: [tflite] Ensure inputs and outputs don't overlap.

If a model uses the same tensor for both an input and an output then this can result in data loss and memory corruption. This should not happen.

PiperOrigin-RevId: 332522916
Change-Id: If0905b142415a9dfceaf2d181872f2a8fb88f48a

Diff: diff --git a/tensorflow/lite/BUILD b/tensorflow/lite/BUILD
index 76e9f5de7c1..e50cd4549b2 100644
--- a/tensorflow/lite/BUILD
+++ b/tensorflow/lite/BUILD
@@ -464,10 +464,11 @@ cc_test(
     size = "small",
     srcs = ["model_test.cc"],
     data = [
         "testdata/0_subgraphs.bin",
         "testdata/2_subgraphs.bin",
+        "testdata/add_shared_tensors.bin",
         "testdata/empty_model.bin",
         "testdata/multi_add_flex.bin",
         "testdata/sparse_tensor.bin",
         "testdata/test_min_runtime.bin",
         "testdata/test_model.bin",
diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc
index 2fe8099d372..47d0fddb12b 100644
--- a/tensorflow/lite/core/subgraph.cc
+++ b/tensorflow/lite/core/subgraph.cc
@@ -579,10 +579,37 @@ TfLiteStatus Subgraph::CheckTensorIndices(const char* label, const int* indices,
     }
   }
   return kTfLiteOk;
 }
 
+// We have two arrays and we need to check that elements from one array don't
+// show up in the other. We could sort both arrays and then iterate with two
+// pointers from start to finish always increasing the smaller one but since
+// these arrays are usually short (<25 elements for inputs, usually <3 for
+// outputs), this might be slower than the naive approach (if arrays have size n
+// and m, with n >> m ~ O(1), first approach is O(nlogn) whereas the other is
+// O(n)). Plus, sorting the input and output arrays might not be something we
+// want as it destroys ordering of elements.
+//
+// If it turns out that this is an issue, we can switch to the other algorithm.
+TfLiteStatus Subgraph::CheckInputAndOutputForOverlap(const int* input_indices,
+                                                     int num_inputs,
+                                                     const int* output_indices,
+                                                     int num_outputs) {
+  for (int i = 0; i < num_inputs; i++) {
+    for (int j = 0; j < num_outputs; j++) {
+      if (input_indices[i] == output_indices[j]) {
+        ReportError("Tensor %d is both input %d and output %d\n",
+                    input_indices[i], i, j);
+        consistent_ = false;
+        return kTfLiteError;
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+
 namespace {
 // Multiply two sizes and return true if overflow occurred;
 // This is based off tensorflow/overflow.h but is simpler as we already
 // have unsigned numbers. It is also generalized to work where sizeof(size_t)
 // is not 8.
@@ -705,10 +732,20 @@ TfLiteStatus Subgraph::AddNodeWithParameters(
                                                   inputs.size()));
   TF_LITE_ENSURE_OK(
       &context_,
       CheckTensorIndices("node outputs", outputs.data(), outputs.size()));
 
+  // For builtin ops, inputs and outputs must not overlap. Custom ops must do
+  // this check by themselves if they don't support overlapping tensors. This
+  // distinction is to allow custom ops to just forward a tensor, reusing it as
+  // both input and output.
+  if (builtin_data != nullptr) {
+    TF_LITE_ENSURE_OK(&context_, CheckInputAndOutputForOverlap(
+                                     inputs.data(), inputs.size(),
+                                     outputs.data(), outputs.size()));
+  }
+
   int new_node_index = nodes_and_registration_.size();
   if (node_index) *node_index = new_node_index;
   nodes_and_registration_.resize(nodes_and_registration_.size() + 1);
   auto& node_and_reg = nodes_and_registration_.back();
   TfLiteNode& node = node_and_reg.first;
diff --git a/tensorflow/lite/core/subgraph.h b/tensorflow/lite/core/subgraph.h
index c74611a6ca4..b94d1a0b2bc 100644
--- a/tensorflow/lite/core/subgraph.h
+++ b/tensorflow/lite/core/subgraph.h
@@ -449,10 +449,19 @@ class Subgraph {
   // array.
   // NOTE: this changes consistent_ to be false if indices are out of bounds.
   TfLiteStatus CheckTensorIndices(const char* label, const int* indices,
                                   int length);
 
+  // Check that the input indices and the output indices don't overlap.
+  // This is needed because same tensor must not be used both as input and
+  // output for an operator.
+  // NOTE: this changes consistent_ to be false if indices are out of bounds.
+  TfLiteStatus CheckInputAndOutputForOverlap(const int* input_indices,
+                                             int num_inputs,
+                                             const int* output_indices,
+                                             int num_outputs);
+
   // Compute the number of bytes required to represent a tensor with dimensions
   // specified by the array dims (of length dims_size). Returns the status code
   // and bytes.
   TfLiteStatus BytesRequired(TfLiteType type, const int* dims, size_t dims_size,
                              size_t* bytes);
diff --git a/tensorflow/lite/model_test.cc b/tensorflow/lite/model_test.cc
index ba96494225c..6993e350a48 100644
--- a/tensorflow/lite/model_test.cc
+++ b/tensorflow/lite/model_test.cc
@@ -436,10 +436,29 @@ 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);
+  }
+}
 
 // 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/add_shared_tensors.bin b/tensorflow/lite/testdata/add_shared_tensors.bin
new file mode 100644
index 00000000000..cf9f9775818
Binary files /dev/null and b/tensorflow/lite/testdata/add_shared_tensors.bin differ
diff --git a/tensorflow/lite/testdata/sparse_tensor.bin b/tensorflow/lite/testdata/sparse_tensor.bin
index ef023280887..ec4c8dcf1fc 100644
Binary files a/tensorflow/lite/testdata/sparse_tensor.bin and b/tensorflow/lite/testdata/sparse_tensor.bin differ

Fixed code:         "testdata/add_shared_tensors.bin",
// We have two arrays and we need to check that elements from one array don't
// show up in the other. We could sort both arrays and then iterate with two
// pointers from start to finish always increasing the smaller one but since
// these arrays are usually short (<25 elements for inputs, usually <3 for
// outputs), this might be slower than the naive approach (if arrays have size n
// and m, with n >> m ~ O(1), first approach is O(nlogn) whereas the other is
// O(n)). Plus, sorting the input and output arrays might not be something we
// want as it destroys ordering of elements.
//
// If it turns out that this is an issue, we can switch to the other algorithm.
TfLiteStatus Subgraph::CheckInputAndOutputForOverlap(const int* input_indices,
                                                     int num_inputs,
                                                     const int* output_indices,
                                                     int num_outputs) {
  for (int i = 0; i < num_inputs; i++) {
    for (int j = 0; j < num_outputs; j++) {
      if (input_indices[i] == output_indices[j]) {
        ReportError("Tensor %d is both input %d and output %d\n",
                    input_indices[i], i, j);
        consistent_ = false;
        return kTfLiteError;
      }
    }
  }
  return kTfLiteOk;
}

  // For builtin ops, inputs and outputs must not overlap. Custom ops must do
  // this check by themselves if they don't support overlapping tensors. This
  // distinction is to allow custom ops to just forward a tensor, reusing it as
  // both input and output.
  if (builtin_data != nullptr) {
    TF_LITE_ENSURE_OK(&context_, CheckInputAndOutputForOverlap(
                                     inputs.data(), inputs.size(),
                                     outputs.data(), outputs.size()));
  }

  // Check that the input indices and the output indices don't overlap.
  // This is needed because same tensor must not be used both as input and
  // output for an operator.
  // NOTE: this changes consistent_ to be false if indices are out of bounds.
  TfLiteStatus CheckInputAndOutputForOverlap(const int* input_indices,
                                             int num_inputs,
                                             const int* output_indices,
                                             int num_outputs);

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);
  }
}