From a2b84e989759aabbc260e8db1b56e75e4e9d97fd Mon Sep 17 00:00:00 2001 From: rogday Date: Wed, 13 Apr 2022 15:55:36 +0300 Subject: [PATCH] add assert to tf graph simplifier to address security concerns --- modules/dnn/src/tensorflow/tf_graph_simplifier.cpp | 2 +- modules/dnn/test/test_tf_importer.cpp | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/modules/dnn/src/tensorflow/tf_graph_simplifier.cpp b/modules/dnn/src/tensorflow/tf_graph_simplifier.cpp index 54395504c7..f81ff83ab0 100644 --- a/modules/dnn/src/tensorflow/tf_graph_simplifier.cpp +++ b/modules/dnn/src/tensorflow/tf_graph_simplifier.cpp @@ -1035,7 +1035,7 @@ void sortByExecutionOrder(tensorflow::GraphDef& net) nodesToAdd.pop_back(); permIds.push_back(nodeToAdd); - + CV_Assert(nodeToAdd < edges.size()); for (int i = 0; i < edges[nodeToAdd].size(); ++i) { int consumerId = edges[nodeToAdd][i]; diff --git a/modules/dnn/test/test_tf_importer.cpp b/modules/dnn/test/test_tf_importer.cpp index d3ee5d3300..81880fb66c 100644 --- a/modules/dnn/test/test_tf_importer.cpp +++ b/modules/dnn/test/test_tf_importer.cpp @@ -1571,6 +1571,13 @@ TEST_P(Test_TensorFlow_layers, tf2_permute_nhwc_ncwh) runTensorFlowNet("tf2_permute_nhwc_ncwh"); } +// issue #21852 +TEST_P(Test_TensorFlow_layers, tf_graph_simplifier_buffer_overflow) +{ + // This just shouldn't segfault, otherwise it's fine + EXPECT_ANY_THROW(readNetFromTensorflow(path("tf_graph_simplifier_buffer_overflow_net.pb"))); +} + TEST_P(Test_TensorFlow_layers, squeeze) { #if defined(INF_ENGINE_RELEASE)