fix vector access in TF::sortByExecutionOrder

pull/22006/head
rogday 3 years ago
parent 90820cd044
commit a2ad997e97
  1. 22
      modules/dnn/src/tensorflow/tf_graph_simplifier.cpp
  2. 28
      modules/dnn/test/test_tf_importer.cpp

@ -990,6 +990,7 @@ void sortByExecutionOrder(tensorflow::GraphDef& net)
nodesMap.insert(std::make_pair(node.name(), i)); 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. // Indices of nodes which use specific node as input.
std::vector<std::vector<int> > edges(nodesMap.size()); std::vector<std::vector<int> > edges(nodesMap.size());
std::vector<int> numRefsToAdd(nodesMap.size(), 0); std::vector<int> numRefsToAdd(nodesMap.size(), 0);
@ -1007,7 +1008,7 @@ void sortByExecutionOrder(tensorflow::GraphDef& net)
nodesMapIt = nodesMap.find(inpName); nodesMapIt = nodesMap.find(inpName);
if (nodesMapIt != nodesMap.end()) if (nodesMapIt != nodesMap.end())
{ {
edges[nodesMapIt->second].push_back(i); edges.at(nodesMapIt->second).push_back(i);
numInputsInGraph += 1; numInputsInGraph += 1;
} }
} }
@ -1019,11 +1020,11 @@ void sortByExecutionOrder(tensorflow::GraphDef& net)
{ {
int numControlEdges = 0; int numControlEdges = 0;
for (int j = 0; j < numInputsInGraph; ++j) for (int j = 0; j < numInputsInGraph; ++j)
numControlEdges += node.input(j)[0] == '^'; numControlEdges += node.input(j).at(0) == '^';
numRefsToAdd[i] = numControlEdges + 1; numRefsToAdd.at(i) = numControlEdges + 1;
} }
else else
numRefsToAdd[i] = numInputsInGraph; numRefsToAdd.at(i) = numInputsInGraph;
} }
} }
@ -1035,17 +1036,16 @@ void sortByExecutionOrder(tensorflow::GraphDef& net)
nodesToAdd.pop_back(); nodesToAdd.pop_back();
permIds.push_back(nodeToAdd); permIds.push_back(nodeToAdd);
CV_Assert(nodeToAdd < edges.size()); for (int i = 0; i < edges.at(nodeToAdd).size(); ++i)
for (int i = 0; i < edges[nodeToAdd].size(); ++i)
{ {
int consumerId = edges[nodeToAdd][i]; int consumerId = edges.at(nodeToAdd).at(i);
if (numRefsToAdd[consumerId] > 0) if (numRefsToAdd.at(consumerId) > 0)
{ {
if (numRefsToAdd[consumerId] == 1) if (numRefsToAdd.at(consumerId) == 1)
nodesToAdd.push_back(consumerId); nodesToAdd.push_back(consumerId);
else else
CV_Assert(numRefsToAdd[consumerId] >= 0); CV_Assert(numRefsToAdd.at(consumerId) >= 0);
numRefsToAdd[consumerId] -= 1; numRefsToAdd.at(consumerId) -= 1;
} }
} }
} }

@ -1571,13 +1571,6 @@ TEST_P(Test_TensorFlow_layers, tf2_permute_nhwc_ncwh)
runTensorFlowNet("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) TEST_P(Test_TensorFlow_layers, squeeze)
{ {
#if defined(INF_ENGINE_RELEASE) #if defined(INF_ENGINE_RELEASE)
@ -1743,4 +1736,25 @@ TEST_P(Test_TensorFlow_nets, EfficientDet)
expectNoFallbacksFromIE(net); 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<const char*>(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<const char*>(payload), sizeof(payload) / sizeof(payload[0])));
}
} }

Loading…
Cancel
Save