Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent null dereference read in GetInitOp.
We have a map of maps. We test that the key exists in the first map but then we don't have any validation that this also means the second map has the needed key. In the scenarios where this is not the case, we'll dereference a nullptr, if we don't have this check

PiperOrigin-RevId: 408739325
Change-Id: If9bb7ed759aba1f3b56a34913f209508dbaf65ce
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Nov 10, 2021
1 parent 7b84b7c commit 4f38b1a
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions tensorflow/cc/saved_model/loader_util.cc
Expand Up @@ -34,9 +34,14 @@ Status GetInitOp(const string& export_dir, const MetaGraphDef& meta_graph_def,
const auto& init_op_sig_it =
meta_graph_def.signature_def().find(kSavedModelInitOpSignatureKey);
if (init_op_sig_it != sig_def_map.end()) {
*init_op_name = init_op_sig_it->second.outputs()
.find(kSavedModelInitOpSignatureKey)
->second.name();
const auto& sig_def_outputs = init_op_sig_it->second.outputs();
const auto& sig_def_outputs_it =
sig_def_outputs.find(kSavedModelInitOpSignatureKey);
if (sig_def_outputs_it == sig_def_outputs.end()) {
return errors::FailedPrecondition("Could not find output ",
kSavedModelInitOpSignatureKey);
}
*init_op_name = sig_def_outputs_it->second.name();
return Status::OK();
}

Expand Down

0 comments on commit 4f38b1a

Please sign in to comment.