Commit message: TFLite: Error out when the graph has a recurion.

Recursion is currently unsupported.

PiperOrigin-RevId: 371708957
Change-Id: I8dfad0d85cbfe08e39ae8ea7bad21254ddee5003
CVE: CVE-2021-29591
CWE: CWE-703
Vulnerability description: 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.

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestHandleModelWithWhileOpContainsForwardingInput) {
  const auto model_path =
      "tensorflow/lite/testdata/while_op_with_forwarding_input.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);

  int32_t* tensor_data = interpreter->typed_tensor<int32_t>(0);
  tensor_data[0] = 20;

  auto tensor = interpreter->tensor(1);
  DynamicBuffer buf;
  buf.AddString("a", 1);
  buf.WriteToTensor(tensor, /*new_shape=*/nullptr);

  ASSERT_EQ(interpreter->Invoke(), kTfLiteOk);
}

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestBuildFromInvalidAllocation) {
  TestErrorReporter reporter;
  std::unique_ptr<Allocation> model_allocation(
      new MemoryAllocation(nullptr, 0, nullptr));

  auto model =
      FlatBufferModel::BuildFromAllocation(std::move(model_allocation));
  ASSERT_FALSE(model);
}

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestSetNumThreads) {
  TestErrorReporter reporter;
  auto model = FlatBufferModel::BuildFromFile(
      "tensorflow/lite/testdata/test_model.bin", &reporter);
  ASSERT_TRUE(model);
  std::unique_ptr<Interpreter> interpreter;
  TrivialResolver resolver(&dummy_reg);
  InterpreterBuilder builder(*model, resolver);

  ASSERT_EQ(builder.SetNumThreads(42), kTfLiteOk);
  interpreter.reset();
  ASSERT_EQ(builder(&interpreter), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);

  ASSERT_EQ(builder.SetNumThreads(0), kTfLiteOk);
  interpreter.reset();
  ASSERT_EQ(builder(&interpreter), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);

  ASSERT_EQ(builder.SetNumThreads(-1), kTfLiteOk);
  interpreter.reset();
  ASSERT_EQ(builder(&interpreter), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);

  ASSERT_EQ(reporter.num_calls(), 0);
  ASSERT_EQ(builder.SetNumThreads(-2), kTfLiteError);
  interpreter.reset();
  ASSERT_EQ(builder(&interpreter), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);
  ASSERT_EQ(reporter.num_calls(), 1);
  ASSERT_PRED_FORMAT2(testing::IsSubstring,
                      "num_threads should be >= 0 or just -1",
                      reporter.error_messages());
}

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestNullDestination) {
  auto model = FlatBufferModel::BuildFromFile(
      "tensorflow/lite/testdata/empty_model.bin");
  ASSERT_TRUE(model);
  // Test that building with null destination fails.
  ASSERT_NE(InterpreterBuilder(*model, TrivialResolver())(nullptr), kTfLiteOk);
}

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestSubgraphName) {
  auto m = FlatBufferModel::BuildFromFile(
      "tensorflow/lite/testdata/"
      "2_subgraphs_dont_delegate_name.bin");
  ASSERT_TRUE(m);
  std::unique_ptr<Interpreter> interpreter;
  ASSERT_EQ(InterpreterBuilder(*m, TrivialResolver())(&interpreter), kTfLiteOk);
  EXPECT_EQ(interpreter->subgraphs_size(), 2);
  EXPECT_EQ(interpreter->subgraph(0)->GetName(), "");
  EXPECT_EQ(interpreter->subgraph(1)->GetName(), "VALIDATION:main");
}

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestBuildFromAllocation) {
  TestErrorReporter reporter;
  std::unique_ptr<Allocation> model_allocation(new FileCopyAllocation(
      "tensorflow/lite/testdata/test_model.bin", &reporter));
  ASSERT_TRUE(model_allocation->valid());

  auto model =
      FlatBufferModel::BuildFromAllocation(std::move(model_allocation));
  ASSERT_TRUE(model);

  std::unique_ptr<Interpreter> interpreter;
  ASSERT_EQ(
      InterpreterBuilder(*model, TrivialResolver(&dummy_reg))(&interpreter),
      kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);
}

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestBuildFromNullAllocation) {
  TestErrorReporter reporter;
  std::unique_ptr<Allocation> model_allocation;

  auto model =
      FlatBufferModel::BuildFromAllocation(std::move(model_allocation));
  ASSERT_FALSE(model);
}

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestEmptyModels) {
  auto model = FlatBufferModel::BuildFromFile(
      "tensorflow/lite/testdata/empty_model.bin");
  ASSERT_TRUE(model);
  // Now try to build it into a model.
  std::unique_ptr<Interpreter> interpreter;
  ASSERT_EQ(InterpreterBuilder(*model, TrivialResolver())(&interpreter),
            kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);
}

Project: tensorflow
File: subgraph.cc

  ~SubgraphGuard() {
    // If tht original status was OK, recover the boolean flag.
    if (status_ == kTfLiteOk) {
      *is_subgraph_in_use_ = false;
    }
  }

Project: tensorflow
File: subgraph.cc

TfLiteStatus Subgraph::AllocateTensors() {
  TFLITE_SCOPED_TAGGED_DEFAULT_PROFILE(profiler_.get(), "AllocateTensors");

  if (!consistent_) {
    ReportError("AllocateTensors() called on inconsistent model.");
    return kTfLiteError;
  }

  // Restore delegation state if applicable.
  TF_LITE_ENSURE_STATUS(RedoAllDelegates());

  // Explicit (re)allocation is necessary if nodes have been changed or tensors
  // have been resized. For inputs marked as dynamic, we can't short-circuit the
  // allocation as the client may have done the resize manually.
  if (state_ != kStateUninvokable &&
      !HasDynamicTensorImpl(context_, inputs())) {
    if (memory_planner_ && !memory_planner_->HasNonPersistentMemory()) {
      // If the only change was the release of non-persistent memory via
      // ReleaseNonPersistentMemory(), just re-allocate it. For any other type
      // of memory-planning change (for eg, ResizeInputTensor), the state would
      // be kStateUninvokable.
      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());
  }

  TF_LITE_ENSURE_STATUS(PrepareOpsAndTensors());

  state_ = kStateInvokable;

  // Reset the variable tensors to zero after (re)allocating the tensors.
  // Developers shouldn't rely on the side effect of this function to reset
  // variable tensors. They should call `ResetVariableTensors` directly
  // instead.
  ResetVariableTensors();

  return kTfLiteOk;
}

Project: tensorflow
File: while.cc

TfLiteStatus CopyTensorsData(TfLiteContext* context, Subgraph* src_subgraph,
                             const SrcVector& src_tensor_indices,
                             Subgraph* dst_subgraph,
                             const DstVector& dst_tensor_indices) {
  TF_LITE_ENSURE_EQ(context, src_tensor_indices.size(),
                    dst_tensor_indices.size());
  for (int i = 0; i < src_tensor_indices.size(); ++i) {
    const TfLiteTensor* src_tensor =
        src_subgraph->tensor(src_tensor_indices[i]);
    TfLiteTensor* dst_tensor = dst_subgraph->tensor(dst_tensor_indices[i]);
    if (IsDynamicTensor(dst_tensor)) {
      TfLiteTensorRealloc(src_tensor->bytes, dst_tensor);
    }
    TF_LITE_ENSURE_EQ(context, src_tensor->bytes, dst_tensor->bytes);
    memcpy(dst_tensor->data.raw, src_tensor->data.raw, src_tensor->bytes);
  }
  return kTfLiteOk;
}

Project: tensorflow
File: model_test.cc

TEST(TestAddDelegateOwnership, AddDelegateDoesNotTakeOwnership) {
  class TestDelegate : public TfLiteDelegate {
   public:
    TestDelegate(bool* destroyed, bool* prepared)
        : TfLiteDelegate(TfLiteDelegateCreate()),
          destroyed_(destroyed),
          prepared_(prepared) {
      flags = kTfLiteDelegateFlagsNone;
      Prepare = [](TfLiteContext*, TfLiteDelegate* delegate) -> TfLiteStatus {
        *(static_cast<TestDelegate*>(delegate)->prepared_) = true;
        return kTfLiteOk;
      };
    }
    ~TestDelegate() { *destroyed_ = true; }

   private:
    bool* destroyed_;
    bool* prepared_;
  };

  // Construct a delegate with flags for indicating preparation/destruction.
  bool destroyed = false;
  bool prepared = false;
  {
    std::unique_ptr<TestDelegate> delegate(
        new TestDelegate(&destroyed, &prepared));
    {
      // Load a model.
      auto model = FlatBufferModel::BuildFromFile(
          "tensorflow/lite/testdata/empty_model.bin");
      ASSERT_TRUE(model);
      // Now try to build it into an interpreter.
      std::unique_ptr<Interpreter> interpreter;
      InterpreterBuilder builder(*model, TrivialResolver());
      builder.AddDelegate(delegate.get());  // Does not transfer ownership.
      // Loop to check we can construct multiple interpreters from one builder.
      for (int i = 0; i < 3; i++) {
        prepared = false;
        ASSERT_EQ(builder(&interpreter), kTfLiteOk);
        ASSERT_NE(interpreter, nullptr);

        // The delegate should be prepared as normal, and should be preserved.
        EXPECT_TRUE(prepared);
        EXPECT_FALSE(destroyed);

        // Interpreter interaction should not impact the delegate's validity.
        interpreter->AllocateTensors();
        interpreter->Invoke();
        EXPECT_FALSE(destroyed);
      }
    }
    EXPECT_NE(delegate, nullptr);
    EXPECT_FALSE(destroyed);
  }
  // Only after the delegate itself goes out of scope should the delegate be
  // destroyed.
  EXPECT_TRUE(destroyed);
}

Project: tensorflow
File: subgraph.cc

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

Project: tensorflow
File: model_test.cc

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

Project: tensorflow
File: subgraph.cc

  TfLiteStatus status() const { return status_; }

Project: tensorflow
File: model_test.cc

TEST(BasicFlatBufferModel, TestWithNumThreads) {
  TestErrorReporter reporter;
  auto model = FlatBufferModel::BuildFromFile(
      "tensorflow/lite/testdata/test_model.bin", &reporter);
  ASSERT_TRUE(model);
  TrivialResolver resolver(&dummy_reg);
  InterpreterBuilder builder(*model, resolver);

  std::unique_ptr<Interpreter> interpreter;
  ASSERT_EQ(builder(&interpreter, /*num_threads=*/42), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);
  ASSERT_EQ(interpreter->subgraph(0)->context()->recommended_num_threads, 42);

  interpreter.reset();
  ASSERT_EQ(builder(&interpreter, 0), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);
  ASSERT_EQ(interpreter->subgraph(0)->context()->recommended_num_threads, 1);

  interpreter.reset();
  ASSERT_EQ(builder(&interpreter, -1), kTfLiteOk);
  ASSERT_NE(interpreter, nullptr);
  ASSERT_EQ(interpreter->subgraph(0)->context()->recommended_num_threads, -1);

  ASSERT_EQ(reporter.num_calls(), 0);
  interpreter.reset(new Interpreter);
  ASSERT_EQ(builder(&interpreter, -2), kTfLiteError);
  ASSERT_EQ(interpreter, nullptr);
  ASSERT_EQ(reporter.num_calls(), 1);
  ASSERT_PRED_FORMAT2(testing::IsSubstring,
                      "num_threads should be >= 0 or just -1",
                      reporter.error_messages());
}

Project: tensorflow
File: subgraph.cc

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

  TfLiteStatus status = kTfLiteOk;
  if (state_ == kStateUninvokable) {
    ReportError("Invoke called on model that is not ready.");
    return kTfLiteError;
  } else if (memory_planner_ && !memory_planner_->HasNonPersistentMemory()) {
    ReportError("Non-persistent memory is not available.");
    return kTfLiteError;
  }
  TFLITE_SCOPED_TAGGED_DEFAULT_PROFILE(profiler_.get(), "Invoke");

  // Invocations are always done in node order.
  // Note that calling Invoke repeatedly will cause the original memory plan to
  // be reused, unless either ResizeInputTensor() or AllocateTensors() has been
  // called.
  for (int execution_plan_index = 0;
       execution_plan_index < execution_plan_.size(); execution_plan_index++) {
    if (execution_plan_index == next_execution_plan_index_to_prepare_) {
      TF_LITE_ENSURE_STATUS(PrepareOpsAndTensors());
      TF_LITE_ENSURE(&context_, next_execution_plan_index_to_prepare_ >=
                                    execution_plan_index);
    }
    int node_index = execution_plan_[execution_plan_index];
    TfLiteNode& node = nodes_and_registration_[node_index].first;
    const TfLiteRegistration& registration =
        nodes_and_registration_[node_index].second;

    const char* op_name = nullptr;
    if (profiler_) op_name = GetTFLiteOpName(registration);
    TFLITE_SCOPED_TAGGED_OPERATOR_PROFILE(profiler_.get(), op_name, node_index);

    for (int i = 0; i < node.inputs->size; ++i) {
      int tensor_index = node.inputs->data[i];
      if (tensor_index == kTfLiteOptionalTensor) {
        continue;
      }
      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 &&
            tensor->dims->size != 1) {
          // In general, having a tensor here with no buffer will be an error.
          // However, for the reshape operator, the second input tensor is
          // sometimes only used for the shape, not for the data. Thus, null
          // buffer is ok in this situation.
          // The situation where null buffer is not ok for reshape operator is
          // only when there are 2 inputs given to the node and the one
          // corresponding to the shape (i == 1) is a vector that contains all
          // dimensions. See `GetOutputShape()` function in
          // `tensorflow/lite/kernels/reshape.cc`
          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()");
      return kTfLiteError;
    }

    EnsureTensorsVectorCapacity();
    tensor_resized_since_op_invoke_ = false;
    if (OpInvoke(registration, &node) != kTfLiteOk) {
      return ReportOpError(&context_, node, registration, node_index,
                           "failed to invoke");
    }

    // Force execution prep for downstream ops if the latest op triggered the
    // resize of a dynamic tensor.
    if (tensor_resized_since_op_invoke_ &&
        HasDynamicTensor(context_, node.outputs)) {
      next_execution_plan_index_to_prepare_ = execution_plan_index + 1;

      // This happens when an intermediate dynamic tensor is resized.
      // We don't have to prepare all the ops, but we need to recompute
      // the allocation plan.
      if (next_execution_plan_index_to_plan_allocation_ >
          next_execution_plan_index_to_prepare_) {
        next_execution_plan_index_to_plan_allocation_ =
            next_execution_plan_index_to_prepare_;
        if (memory_planner_) {
          TF_LITE_ENSURE_STATUS(memory_planner_->ResetAllocationsAfter(
              next_execution_plan_index_to_plan_allocation_ - 1));
        }
      }
    }
  }

  return status;
}