CVE-2020-15211: 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, saved models in the flatbuffer format use a double indexing scheme: a model has a set of subgraphs, each subgraph has a set of operators and each operator has a set of input/output tensors. The flatbuffer format uses indices for the tensors, indexing into an array of tensors that is owned by the subgraph. This results in a pattern of double array indexing when trying to get the data of each tensor. However, some operators can have some tensors be optional. To handle this scenario, the flatbuffer model uses a negative `-1` value as index for these tensors. This results in special casing during validation at model loading time. Unfortunately, this means that the `-1` index is a valid tensor index for any operator, including those that don't expect optional inputs and including for output tensors. Thus, this allows writing and reading from outside the bounds of heap allocated arrays, although only at a specific offset from the start of these arrays. This results in both read and write gadgets, albeit very limited in scope. The issue is patched in several commits (46d5b0852, 00302787b7, e11f5558, cd31fd0ce, 1970c21, and fff2c83), and is released in TensorFlow versions 1.15.4, 2.0.3, 2.1.2, 2.2.1, or 2.3.1. A potential workaround would be to add a custom `Verifier` to the model loading code to ensure that only operators which accept optional inputs use the `-1` special value and only for the tensors that they expect to be optional. Since this allow-list type approach is erro-prone, we advise upgrading to the patched code.
CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 4.8
Published: 2020-09-25T19:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C
Files changed: tensorflow/lite/kernels/kernel_util.cc, tensorflow/lite/kernels/kernel_util.h
Commit date: 2020-09-18 13:21:55 -0700
Commit message: [tflite] Test for `kTfLiteOptionalTensor` in `GetInput`.

`GetInput`, `GetVariableInput` and `GetOutput` all fail to check for the case where `node->inputs->data[index]` is the special `kTfLiteOptionalTensor` value (-1) which then causes `context->tensors[node->inputs->data[index]]` to read from invalid memory location.

This fix makes `GetInput` and related return `nullptr` in those cases, asking the caller to check for `nullptr`. This is better than having `GetOptionalInputTensor` and `GetOptionalOutputTensor` (does not exist but could be added) as using the patched `GetInput` in error would be caught by a sanitizer test in the default optimized build (due to the `-fsanitize=null` option).

PiperOrigin-RevId: 332512190
Change-Id: Iabca54da2f2de02b6ece3c38b54f76d4277d689e

Diff: diff --git a/tensorflow/lite/kernels/kernel_util.cc b/tensorflow/lite/kernels/kernel_util.cc
index 27d9da84c0a..fab884bc90c 100644
--- a/tensorflow/lite/kernels/kernel_util.cc
+++ b/tensorflow/lite/kernels/kernel_util.cc
@@ -30,15 +30,21 @@ namespace tflite {
 
 namespace {
 
 inline TfLiteTensor* GetMutableInput(const TfLiteContext* context,
                                      const TfLiteNode* node, int index) {
-  if (context->tensors != nullptr) {
-    return &context->tensors[node->inputs->data[index]];
-  } else {
-    return context->GetTensor(context, node->inputs->data[index]);
+  if (index >= 0 && index < node->inputs->size) {
+    const int tensor_index = node->inputs->data[index];
+    if (tensor_index != kTfLiteOptionalTensor) {
+      if (context->tensors != nullptr) {
+        return &context->tensors[tensor_index];
+      } else {
+        return context->GetTensor(context, tensor_index);
+      }
+    }
   }
+  return nullptr;
 }
 
 }  // anonymous namespace.
 
 const TfLiteTensor* GetInput(const TfLiteContext* context,
@@ -52,15 +58,21 @@ TfLiteTensor* GetVariableInput(TfLiteContext* context, const TfLiteNode* node,
   return tensor->is_variable ? tensor : nullptr;
 }
 
 TfLiteTensor* GetOutput(TfLiteContext* context, const TfLiteNode* node,
                         int index) {
-  if (context->tensors != nullptr) {
-    return &context->tensors[node->outputs->data[index]];
-  } else {
-    return context->GetTensor(context, node->outputs->data[index]);
+  if (index >= 0 && index < node->outputs->size) {
+    const int tensor_index = node->outputs->data[index];
+    if (tensor_index != kTfLiteOptionalTensor) {
+      if (context->tensors != nullptr) {
+        return &context->tensors[tensor_index];
+      } else {
+        return context->GetTensor(context, tensor_index);
+      }
+    }
   }
+  return nullptr;
 }
 
 const TfLiteTensor* GetOptionalInputTensor(const TfLiteContext* context,
                                            const TfLiteNode* node, int index) {
   const bool use_tensor = index < node->inputs->size &&
diff --git a/tensorflow/lite/kernels/kernel_util.h b/tensorflow/lite/kernels/kernel_util.h
index 0d6aa8fc790..5950effca0e 100644
--- a/tensorflow/lite/kernels/kernel_util.h
+++ b/tensorflow/lite/kernels/kernel_util.h
@@ -27,39 +27,99 @@ namespace tflite {
 // A fair number of functions in this header have historically been inline.
 // It is ok to change functions to not be inline if the latency with
 // benchmark_model for MobileNet + MobileBERT is unaffected. If such a change is
 // made, move the newly non-inlined function declarations to the top of this
 // header file.
+
+// Note: You must check if result is not null:
+//
+//   TfLiteTensor* my_tensor = GetInput(context, node, kMyTensorIdx);
+//   TF_LITE_ENSURE(context, my_tensor != nullptr);
+//
+// This is because the index might point to the optional tensor constant
+// (kTfLiteOptionalTensor) in which case there is no tensor to return.
 const TfLiteTensor* GetInput(const TfLiteContext* context,
                              const TfLiteNode* node, int index);
 
 // Note: You must check if result is not null:
-// TfLiteTensor* my_tensor = GetVariableInput(context, node, kMyTensorIdx);
-// TF_LITE_ENSURE(context, my_tensor != nullptr);
+//
+//   TfLiteTensor* my_tensor = GetVariableInput(context, node, kMyTensorIdx);
+//   TF_LITE_ENSURE(context, my_tensor != nullptr);
+//
+// This is because the index might point to the optional tensor constant
+// (kTfLiteOptionalTensor) in which case there is no tensor to return.
 TfLiteTensor* GetVariableInput(TfLiteContext* context, const TfLiteNode* node,
                                int index);
 
+// Note: You must check if result is not null:
+//
+//   TfLiteTensor* my_tensor = GetOutput(context, node, kMyTensorIdx);
+//   TF_LITE_ENSURE(context, my_tensor != nullptr);
+//
+// This is because the index might point to the optional tensor constant
+// (kTfLiteOptionalTensor) in which case there is no tensor to return.
 TfLiteTensor* GetOutput(TfLiteContext* context, const TfLiteNode* node,
                         int index);
 
+// Note: You must check if result is not null:
+//
+//   TfLiteTensor* my_tensor = GetOptionalInputTensor(context, node, kIdx);
+//   TF_LITE_ENSURE(context, my_tensor != nullptr);
+//
+// This is because the index might point to the optional tensor constant
+// (kTfLiteOptionalTensor) in which case there is no tensor to return.
+//
+// Deprecated. GetInput has the same functionality.
 const TfLiteTensor* GetOptionalInputTensor(const TfLiteContext* context,
                                            const TfLiteNode* node, int index);
 
 inline int NumDimensions(const TfLiteTensor* t) { return t->dims->size; }
 inline int SizeOfDimension(const TfLiteTensor* t, int dim) {
   return t->dims->data[dim];
 }
 
 #ifndef TF_LITE_STATIC_MEMORY
+// Note: You must check if result is not null:
+//
+//   TfLiteTensor* my_tensor = GetTemporary(context, node, kMyTensorIdx);
+//   TF_LITE_ENSURE(context, my_tensor != nullptr);
+//
+// This is because the index might point to the optional tensor constant
+// (kTfLiteOptionalTensor) in which case there is no tensor to return.
 inline TfLiteTensor* GetTemporary(TfLiteContext* context,
                                   const TfLiteNode* node, int index) {
-  return &context->tensors[node->temporaries->data[index]];
+  if (index >= 0 && index < node->temporaries->size) {
+    const int tensor_index = node->temporaries->data[index];
+    if (tensor_index != kTfLiteOptionalTensor) {
+      if (context->tensors != nullptr) {
+        return &context->tensors[tensor_index];
+      }
+    }
+  }
+  return nullptr;
 }
+
+// Note: You must check if result is not null:
+//
+//   TfLiteTensor* my_tensor = GetIntermediates(context, node, kMyTensorIdx);
+//   TF_LITE_ENSURE(context, my_tensor != nullptr);
+//
+// This is because the index might point to the optional tensor constant
+// (kTfLiteOptionalTensor) in which case there is no tensor to return.
 inline const TfLiteTensor* GetIntermediates(TfLiteContext* context,
                                             const TfLiteNode* node, int index) {
-  return &context->tensors[node->intermediates->data[index]];
+  if (index >= 0 && index < node->intermediates->size) {
+    const int tensor_index = node->intermediates->data[index];
+    if (tensor_index != kTfLiteOptionalTensor) {
+      if (context->tensors != nullptr) {
+        return &context->tensors[tensor_index];
+      }
+    }
+  }
+  return nullptr;
 }
+
 inline int NumIntermediates(const TfLiteNode* node) {
   return node->intermediates->size;
 }
 #endif  // TF_LITE_STATIC_MEMORY
 inline int NumInputs(const TfLiteNode* node) { return node->inputs->size; }

Vulnerable code:   if (context->tensors != nullptr) {
    return &context->tensors[node->inputs->data[index]];
  } else {
    return context->GetTensor(context, node->inputs->data[index]);
  if (context->tensors != nullptr) {
    return &context->tensors[node->outputs->data[index]];
  } else {
    return context->GetTensor(context, node->outputs->data[index]);
// TfLiteTensor* my_tensor = GetVariableInput(context, node, kMyTensorIdx);
// TF_LITE_ENSURE(context, my_tensor != nullptr);
  return &context->tensors[node->temporaries->data[index]];
  return &context->tensors[node->intermediates->data[index]];
Fixed code:   if (index >= 0 && index < node->inputs->size) {
    const int tensor_index = node->inputs->data[index];
    if (tensor_index != kTfLiteOptionalTensor) {
      if (context->tensors != nullptr) {
        return &context->tensors[tensor_index];
      } else {
        return context->GetTensor(context, tensor_index);
      }
    }
  return nullptr;
  if (index >= 0 && index < node->outputs->size) {
    const int tensor_index = node->outputs->data[index];
    if (tensor_index != kTfLiteOptionalTensor) {
      if (context->tensors != nullptr) {
        return &context->tensors[tensor_index];
      } else {
        return context->GetTensor(context, tensor_index);
      }
    }
  return nullptr;

// Note: You must check if result is not null:
//
//   TfLiteTensor* my_tensor = GetInput(context, node, kMyTensorIdx);
//   TF_LITE_ENSURE(context, my_tensor != nullptr);
//
// This is because the index might point to the optional tensor constant
// (kTfLiteOptionalTensor) in which case there is no tensor to return.
//
//   TfLiteTensor* my_tensor = GetVariableInput(context, node, kMyTensorIdx);
//   TF_LITE_ENSURE(context, my_tensor != nullptr);
//
// This is because the index might point to the optional tensor constant
// (kTfLiteOptionalTensor) in which case there is no tensor to return.
// Note: You must check if result is not null:
//
//   TfLiteTensor* my_tensor = GetOutput(context, node, kMyTensorIdx);
//   TF_LITE_ENSURE(context, my_tensor != nullptr);
//
// This is because the index might point to the optional tensor constant
// (kTfLiteOptionalTensor) in which case there is no tensor to return.
// Note: You must check if result is not null:
//
//   TfLiteTensor* my_tensor = GetOptionalInputTensor(context, node, kIdx);
//   TF_LITE_ENSURE(context, my_tensor != nullptr);
//
// This is because the index might point to the optional tensor constant
// (kTfLiteOptionalTensor) in which case there is no tensor to return.
//
// Deprecated. GetInput has the same functionality.
// Note: You must check if result is not null:
//
//   TfLiteTensor* my_tensor = GetTemporary(context, node, kMyTensorIdx);
//   TF_LITE_ENSURE(context, my_tensor != nullptr);
//
// This is because the index might point to the optional tensor constant
// (kTfLiteOptionalTensor) in which case there is no tensor to return.
  if (index >= 0 && index < node->temporaries->size) {
    const int tensor_index = node->temporaries->data[index];
    if (tensor_index != kTfLiteOptionalTensor) {
      if (context->tensors != nullptr) {
        return &context->tensors[tensor_index];
      }
    }
  }
  return nullptr;

// Note: You must check if result is not null:
//
//   TfLiteTensor* my_tensor = GetIntermediates(context, node, kMyTensorIdx);
//   TF_LITE_ENSURE(context, my_tensor != nullptr);
//
// This is because the index might point to the optional tensor constant
// (kTfLiteOptionalTensor) in which case there is no tensor to return.
  if (index >= 0 && index < node->intermediates->size) {
    const int tensor_index = node->intermediates->data[index];
    if (tensor_index != kTfLiteOptionalTensor) {
      if (context->tensors != nullptr) {
        return &context->tensors[tensor_index];
      }
    }
  }
  return nullptr;