CVE-2021-29591: Uncontrolled Recursion (CWE-674)
value: TensorFlow is an end-to-end open source platform for machine learning. TFlite graphs must not have loops between nodes. However, this condition was not checked and an attacker could craft models that would result in infinite loop during evaluation. In certain cases, the infinite loop would be replaced by stack overflow due to too many recursive calls. For example, the `While` implementation(https://github.com/tensorflow/tensorflow/blob/106d8f4fb89335a2c52d7c895b7a7485465ca8d9/tensorflow/lite/kernels/while.cc) could be tricked into a scneario where both the body and the loop subgraphs are the same. Evaluating one of the subgraphs means calling the `Eval` function for the other and this quickly exhaust all stack space. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range. Please consult our security guide(https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for more information regarding the security model and how to contact us with issues and questions.
CWE description: The product does not properly control the amount of recursion that takes place,  consuming excessive resources, such as allocated memory or the program stack.
CVSS v3 base score: 7.8
Published: 2021-05-14T20: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/kernels/while.cc, tensorflow/lite/model_test.cc
Commit date: 2021-05-03 09:16:54 -0700
Commit message: TFLite: Error out when the graph has a recurion.

Recursion is currently unsupported.

PiperOrigin-RevId: 371708957
Change-Id: I8dfad0d85cbfe08e39ae8ea7bad21254ddee5003

Diff: diff --git a/tensorflow/lite/BUILD b/tensorflow/lite/BUILD
index e6974846796a9..00fb51f9a08c7 100644
--- a/tensorflow/lite/BUILD
+++ b/tensorflow/lite/BUILD
@@ -641,10 +641,11 @@ cc_test(
         "testdata/segment_sum_invalid_buffer.bin",
         "testdata/sparse_tensor.bin",
         "testdata/test_min_runtime.bin",
         "testdata/test_model.bin",
         "testdata/test_model_broken.bin",
+        "testdata/unsupported_recursion.bin",
         "testdata/while_op_with_forwarding_input.bin",
     ],
     tags = [
         "tflite_not_portable",
         "tflite_smoke_test",
diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc
index 0273018b3bf03..123d8cfd0343f 100644
--- a/tensorflow/lite/core/subgraph.cc
+++ b/tensorflow/lite/core/subgraph.cc
@@ -154,10 +154,46 @@ const char* GetTFLiteOpName(const TfLiteRegistration& op_reg) {
     return op_reg.custom_name;
   }
   return tflite::EnumNamesBuiltinOperator()[op_reg.builtin_code];
 }
 
+// An utility test to detect if the subgraph is abused:
+// 1. Detects if recursion exists in the graph (recursion is not currently
+//    supported.
+// 2. Detects if the interpreter / subgraph is used in multiple subgraphs.
+//    Note: It's clearly documented that the interpreter / subgraph are not
+//    thread-safe. This serves as a check with possible false negatives
+//    unless we switch to atomic boolean flags.
+class SubgraphGuard {
+ public:
+  SubgraphGuard(TfLiteContext* context, bool* is_subgraph_in_use)
+      : is_subgraph_in_use_(is_subgraph_in_use) {
+    if (*is_subgraph_in_use_) {
+      TF_LITE_KERNEL_LOG(
+          context,
+          "Subgraph is already in use. Using an interpreter or a subgraph in "
+          "multiple threads is not supported. Recursion in the graph is not "
+          "supported.");
+      status_ = kTfLiteError;
+    } else {
+      *is_subgraph_in_use_ = true;
+    }
+  }
+  ~SubgraphGuard() {
+    // If tht original status was OK, recover the boolean flag.
+    if (status_ == kTfLiteOk) {
+      *is_subgraph_in_use_ = false;
+    }
+  }
+
+  TfLiteStatus status() const { return status_; }
+
+ private:
+  TfLiteStatus status_ = kTfLiteOk;
+  bool* is_subgraph_in_use_;
+};
+
 }  // namespace
 
 // A trivial implementation of GraphInfo around the Interpreter.
 // NOTE: this interpreter info represents the subset of the
 // graph that is executed according to execution plan. Thus,
@@ -653,10 +689,11 @@ TfLiteStatus Subgraph::BytesRequired(TfLiteType type, const int* dims,
   return kTfLiteOk;
 }
 
 TfLiteStatus Subgraph::AllocateTensors() {
   TFLITE_SCOPED_TAGGED_DEFAULT_PROFILE(profiler_.get(), "AllocateTensors");
+
   if (!consistent_) {
     ReportError("AllocateTensors() called on inconsistent model.");
     return kTfLiteError;
   }
 
@@ -676,10 +713,16 @@ TfLiteStatus Subgraph::AllocateTensors() {
       memory_planner_->AcquireNonPersistentMemory();
     }
     return kTfLiteOk;
   }
 
+  // Note `AllocateTensors` sometimes calls itself recursively above
+  // for delegates. Therefore only the logic below need to be guarded
+  // by `SubgraphGuard`.
+  SubgraphGuard guard(&context_, &is_subgraph_in_use_);
+  TF_LITE_ENSURE_OK(&context_, guard.status());
+
   next_execution_plan_index_to_prepare_ = 0;
   next_execution_plan_index_to_plan_allocation_ = 0;
   next_original_execution_plan_index_to_prepare_ = 0;
   if (memory_planner_) {
     TF_LITE_ENSURE_STATUS(memory_planner_->ResetAllocations());
@@ -1012,10 +1055,13 @@ TfLiteStatus Subgraph::PrepareOpsAndTensors() {
 
   return kTfLiteOk;
 }
 
 TfLiteStatus Subgraph::Invoke() {
+  SubgraphGuard guard(&context_, &is_subgraph_in_use_);
+  TF_LITE_ENSURE_OK(&context_, guard.status());
+
   if (!consistent_) {
     ReportError("Invoke called on model that is not consistent.");
     return kTfLiteError;
   }
 
diff --git a/tensorflow/lite/core/subgraph.h b/tensorflow/lite/core/subgraph.h
index 793d1de792ab2..d6912d7bed04f 100644
--- a/tensorflow/lite/core/subgraph.h
+++ b/tensorflow/lite/core/subgraph.h
@@ -757,9 +757,13 @@ class Subgraph {
   std::string name_;
 
   // Whether memory planner should be instantiated to retain intermediates for
   // debugging.
   bool preserve_all_tensors_ = false;
+
+  // Whether the subgraph is currently in use (e.g. running the `Invoke`
+  // or `AllocateTensors` functions).
+  bool is_subgraph_in_use_ = false;
 };
 
 }  // namespace tflite
 #endif  // TENSORFLOW_LITE_CORE_SUBGRAPH_H_
diff --git a/tensorflow/lite/kernels/while.cc b/tensorflow/lite/kernels/while.cc
index 74ab81c2a9599..e05959fe2a682 100644
--- a/tensorflow/lite/kernels/while.cc
+++ b/tensorflow/lite/kernels/while.cc
@@ -136,12 +136,10 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   // Check subgraph indices and get subgraphs.
   Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_);
   auto* subgraphs = this_subgraph->GetSubgraphs();
   TF_LITE_ENSURE(context, op_data->cond_subgraph_index < subgraphs->size());
   TF_LITE_ENSURE(context, op_data->body_subgraph_index < subgraphs->size());
-  TF_LITE_ENSURE(context,
-                 op_data->cond_subgraph_index != op_data->body_subgraph_index);
 
   Subgraph* cond_subgraph = (*subgraphs)[op_data->cond_subgraph_index].get();
   Subgraph* body_subgraph = (*subgraphs)[op_data->body_subgraph_index].get();
 
   // Check input & output count of the condition subgraph.
diff --git a/tensorflow/lite/model_test.cc b/tensorflow/lite/model_test.cc
index 59ba00652e4b8..1dbe5c12abc5b 100644
--- a/tensorflow/lite/model_test.cc
+++ b/tensorflow/lite/model_test.cc
@@ -598,10 +598,29 @@ TEST(BasicFlatBufferModel, TestHandleMalformedModelReuseTensor) {
   ASSERT_EQ(builder(&interpreter), kTfLiteOk);
   ASSERT_NE(interpreter, nullptr);
   ASSERT_NE(interpreter->AllocateTensors(), kTfLiteOk);
 }
 
+// Recursion & reentrant are not supported in TFLite.
+// The test ensures it fails gracefullly instead of crashing with
+// a stack overflow.
+TEST(BasicFlatBufferModel, TestUnsupportedRecursion) {
+  const auto model_path =
+      "tensorflow/lite/testdata/unsupported_recursion.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.
diff --git a/tensorflow/lite/testdata/unsupported_recursion.bin b/tensorflow/lite/testdata/unsupported_recursion.bin
new file mode 100644
index 0000000000000..525c5383ab4ef
Binary files /dev/null and b/tensorflow/lite/testdata/unsupported_recursion.bin differ

Vulnerable code:   TF_LITE_ENSURE(context,
                 op_data->cond_subgraph_index != op_data->body_subgraph_index);
Fixed code:         "testdata/unsupported_recursion.bin",
// An utility test to detect if the subgraph is abused:
// 1. Detects if recursion exists in the graph (recursion is not currently
//    supported.
// 2. Detects if the interpreter / subgraph is used in multiple subgraphs.
//    Note: It's clearly documented that the interpreter / subgraph are not
//    thread-safe. This serves as a check with possible false negatives
//    unless we switch to atomic boolean flags.
class SubgraphGuard {
 public:
  SubgraphGuard(TfLiteContext* context, bool* is_subgraph_in_use)
      : is_subgraph_in_use_(is_subgraph_in_use) {
    if (*is_subgraph_in_use_) {
      TF_LITE_KERNEL_LOG(
          context,
          "Subgraph is already in use. Using an interpreter or a subgraph in "
          "multiple threads is not supported. Recursion in the graph is not "
          "supported.");
      status_ = kTfLiteError;
    } else {
      *is_subgraph_in_use_ = true;
    }
  }
  ~SubgraphGuard() {
    // If tht original status was OK, recover the boolean flag.
    if (status_ == kTfLiteOk) {
      *is_subgraph_in_use_ = false;
    }
  }

  TfLiteStatus status() const { return status_; }

 private:
  TfLiteStatus status_ = kTfLiteOk;
  bool* is_subgraph_in_use_;
};


  // Note `AllocateTensors` sometimes calls itself recursively above
  // for delegates. Therefore only the logic below need to be guarded
  // by `SubgraphGuard`.
  SubgraphGuard guard(&context_, &is_subgraph_in_use_);
  TF_LITE_ENSURE_OK(&context_, guard.status());

  SubgraphGuard guard(&context_, &is_subgraph_in_use_);
  TF_LITE_ENSURE_OK(&context_, guard.status());


  // Whether the subgraph is currently in use (e.g. running the `Invoke`
  // or `AllocateTensors` functions).
  bool is_subgraph_in_use_ = false;
// Recursion & reentrant are not supported in TFLite.
// The test ensures it fails gracefullly instead of crashing with
// a stack overflow.
TEST(BasicFlatBufferModel, TestUnsupportedRecursion) {
  const auto model_path =
      "tensorflow/lite/testdata/unsupported_recursion.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);
}