From 9c1dc920d8ffb4893d6c9d27d1f039607b326743 Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Tue, 27 Apr 2021 17:47:46 -0700 Subject: [PATCH] Prevent infinite loop/stack overflow in TFLite `while` op. PiperOrigin-RevId: 370800333 Change-Id: I6a2e4ff849da339545c449db2af7e11ce6ff02c3 --- tensorflow/lite/kernels/while.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorflow/lite/kernels/while.cc b/tensorflow/lite/kernels/while.cc index e05959fe2a682..74ab81c2a9599 100644 --- a/tensorflow/lite/kernels/while.cc +++ b/tensorflow/lite/kernels/while.cc @@ -138,6 +138,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { 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();