Skip to content

Commit 448a161

Browse files
Prevent stack overflow when FunctionLib in GraphDef has a self-recursive function.
It is likely that no recursivity is supported, but we should handle this separately. PiperOrigin-RevId: 414860329 Change-Id: I02a2270e86282b37362ddd485eeef16fb986a9e0
1 parent 7b1eba4 commit 448a161

File tree

1 file changed

+18
-0
lines changed

1 file changed

+18
-0
lines changed

tensorflow/cc/saved_model/loader.cc

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ limitations under the License.
2525
#include "tensorflow/core/framework/attr_value.pb.h"
2626
#include "tensorflow/core/framework/function.pb.h"
2727
#include "tensorflow/core/framework/node_def.pb.h"
28+
#include "tensorflow/core/framework/op_def.pb.h"
2829
#include "tensorflow/core/framework/tensor.pb.h"
2930
#include "tensorflow/core/lib/io/path.h"
3031
#include "tensorflow/core/lib/monitoring/counter.h"
@@ -99,6 +100,19 @@ static Status ValidateNode(const NodeDef& node) {
99100
return Status::OK();
100101
}
101102

103+
static Status ValidateFunctionNotRecursive(const FunctionDef& function) {
104+
const auto& function_name = function.signature().name();
105+
for (const auto& node : function.node_def()) {
106+
if (node.op() == function_name) {
107+
return errors::FailedPrecondition(
108+
"Function ", function_name,
109+
" is self recursive and TensorFlow does not support this scenario.");
110+
}
111+
}
112+
113+
return Status::OK();
114+
}
115+
102116
static Status ValidateSavedTensors(const GraphDef& graph_def) {
103117
for (const auto& node : graph_def.node()) {
104118
TF_RETURN_IF_ERROR(ValidateNode(node));
@@ -110,6 +124,10 @@ static Status ValidateSavedTensors(const GraphDef& graph_def) {
110124
for (const auto& node : function.node_def()) {
111125
TF_RETURN_IF_ERROR(ValidateNode(node));
112126
}
127+
128+
// Also check that there is no recursivity in the library
129+
// TODO(mihaimaruseac): Do more than self-recursivity
130+
TF_RETURN_IF_ERROR(ValidateFunctionNotRecursive(function));
113131
}
114132
}
115133

0 commit comments

Comments
 (0)