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));
}
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<std::vector<int> > edges(nodesMap.size());
std::vector<int> 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;
}
}
}

@ -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<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