From a2ad997e972dd6b0535633996221b0dc897d7415 Mon Sep 17 00:00:00 2001 From: rogday Date: Wed, 18 May 2022 23:52:56 +0300 Subject: [PATCH] fix vector access in TF::sortByExecutionOrder --- .../src/tensorflow/tf_graph_simplifier.cpp | 22 +++++++-------- modules/dnn/test/test_tf_importer.cpp | 28 ++++++++++++++----- 2 files changed, 32 insertions(+), 18 deletions(-) diff --git a/modules/dnn/src/tensorflow/tf_graph_simplifier.cpp b/modules/dnn/src/tensorflow/tf_graph_simplifier.cpp index f81ff83ab0..fc0fc94d76 100644 --- a/modules/dnn/src/tensorflow/tf_graph_simplifier.cpp +++ b/modules/dnn/src/tensorflow/tf_graph_simplifier.cpp @@ -990,6 +990,7 @@ void sortByExecutionOrder(tensorflow::GraphDef& net) nodesMap.insert(std::make_pair(node.name(), i)); } + CV_CheckEQ(nodesMap.size(), (size_t)net.node_size(), "Node names must be unique"); // Indices of nodes which use specific node as input. std::vector > edges(nodesMap.size()); std::vector numRefsToAdd(nodesMap.size(), 0); @@ -1007,7 +1008,7 @@ void sortByExecutionOrder(tensorflow::GraphDef& net) nodesMapIt = nodesMap.find(inpName); if (nodesMapIt != nodesMap.end()) { - edges[nodesMapIt->second].push_back(i); + edges.at(nodesMapIt->second).push_back(i); numInputsInGraph += 1; } } @@ -1019,11 +1020,11 @@ void sortByExecutionOrder(tensorflow::GraphDef& net) { int numControlEdges = 0; for (int j = 0; j < numInputsInGraph; ++j) - numControlEdges += node.input(j)[0] == '^'; - numRefsToAdd[i] = numControlEdges + 1; + numControlEdges += node.input(j).at(0) == '^'; + numRefsToAdd.at(i) = numControlEdges + 1; } else - numRefsToAdd[i] = numInputsInGraph; + numRefsToAdd.at(i) = numInputsInGraph; } } @@ -1035,17 +1036,16 @@ 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) + for (int i = 0; i < edges.at(nodeToAdd).size(); ++i) { - int consumerId = edges[nodeToAdd][i]; - if (numRefsToAdd[consumerId] > 0) + int consumerId = edges.at(nodeToAdd).at(i); + if (numRefsToAdd.at(consumerId) > 0) { - if (numRefsToAdd[consumerId] == 1) + if (numRefsToAdd.at(consumerId) == 1) nodesToAdd.push_back(consumerId); else - CV_Assert(numRefsToAdd[consumerId] >= 0); - numRefsToAdd[consumerId] -= 1; + CV_Assert(numRefsToAdd.at(consumerId) >= 0); + numRefsToAdd.at(consumerId) -= 1; } } } diff --git a/modules/dnn/test/test_tf_importer.cpp b/modules/dnn/test/test_tf_importer.cpp index 81880fb66c..d21a5589f9 100644 --- a/modules/dnn/test/test_tf_importer.cpp +++ b/modules/dnn/test/test_tf_importer.cpp @@ -1571,13 +1571,6 @@ 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) @@ -1743,4 +1736,25 @@ TEST_P(Test_TensorFlow_nets, EfficientDet) expectNoFallbacksFromIE(net); } +TEST(Test_TensorFlow_Importer, tf_graph_simplifier_buffer_overflow_21852) +{ + uint8_t payload[] = {0x08, 0x08, 0x0a, 0x00, 0x0a, 0x00}; + EXPECT_ANY_THROW(readNetFromTensorflow(reinterpret_cast(payload), sizeof(payload) / sizeof(payload[0]))); +} + +// can be triggered with -fsanitize=address +TEST(Test_TensorFlow_Importer, tf_graph_simplifier_buffer_overflow_21947) +{ + uint8_t payload[] = {0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, + 0xba, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, + 0x0a, 0xbd, 0x00, 0x1a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0xba, + 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, + 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0xba, 0x0a, 0x00, + 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0xba, + 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, + 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x2a, 0x00, 0xba, 0x0a, 0x00, + 0x0a, 0x00, 0x5d, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x00, 0x0a, 0x40}; + EXPECT_ANY_THROW(readNetFromTensorflow(reinterpret_cast(payload), sizeof(payload) / sizeof(payload[0]))); +} + }