From 7535f6beb7ba95bf54e1513b0c2c51b844a7a49f Mon Sep 17 00:00:00 2001 From: Thomas Bastiani Date: Fri, 6 Apr 2018 17:56:07 +0100 Subject: [PATCH] Bounds-check node ID before getting it's name (#18090) When the edge is either a frame enter or exit edge then DescribeCycle() would segfault. --- tensorflow/compiler/jit/mark_for_compilation_pass.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index f651768..3b631d6 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -35,6 +35,7 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/control_flow.h" +#include "tensorflow/core/kernels/bounds_check.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/public/version.h" @@ -432,6 +433,9 @@ string DescribeCycle(const GraphCycles& cycles, const Graph& graph, int src, } auto node_name = [&cycles, &graph](int node_id) { + if (!FastBoundsCheck(node_id, graph.num_node_ids())) { + return string("(null)"); + } auto* node = graph.FindNodeId(node_id); if (node == nullptr) { return string("(null)"); -- 2.7.4