tensorflow: Unexpected crash when loading a modified saved_model.pb file

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04

• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: No

• TensorFlow installed from (source or binary): source

• TensorFlow version (use command below): v2.0.2-0-g2c2fdd3205 2.0.2

• Python version:

• Bazel version (if compiling from source): 0.26.1

• GCC/Compiler version (if compiling from source): clang++ 10.0

• CUDA/cuDNN version: No

• GPU model and memory: No

You can collect some of this information using our environment capture script You can also obtain the TensorFlow version with:

1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"
2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior I used the LibFuzzer to mutate an intact saved_model.pb file and uses LoadSavedModel C API to load it. Instead of returning a Status whose .ok() is false, the program directly crashes with following errors:

2020-06-20 15:29:05.816403: I tensorflow/cc/saved_model/reader.cc:31] Reading SavedModel from: /home/xxx/Playground/tensorflow/saved_model/crashes
2020-06-20 15:29:05.817167: I tensorflow/cc/saved_model/reader.cc:54] Reading meta graph with tags { serve }
2020-06-20 15:29:05.829727: I tensorflow/cc/saved_model/loader.cc:202] Restoring SavedModel bundle.
[libprotobuf FATAL external/com_google_protobuf/src/google/protobuf/map.h:1059] CHECK failed: it != end(): key not found: value
terminate called after throwing an instance of 'google::protobuf::FatalException'
  what():  CHECK failed: it != end(): key not found: value
[1]    26722 abort (core dumped)  ./loader_test

Describe the expected behavior Since LoadSavedModel returns an Status object, the Status.ok() should return false, rather than a direct crash.

Standalone code to reproduce the issue Provide a reproducible test case that is the bare minimum necessary to generate the problem. If possible, please share a link to Colab/Jupyter/any notebook.

#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/tag_constants.h"

#include <fstream>
using namespace tensorflow;
int main(){
        const string export_dir = "/home/xxx/Playground/saved_model/intact"; // original pb file --> Status.ok() == true
        const string export_dir2 = "/home/xxx/Playground/saved_model/test"; // modified pb file --> Status.ok() == false
        const string export_dir3 = "/home/xxx/Playground/saved_model/crashes"; // modified pb file --> crash
        SavedModelBundle bundle;
        SessionOptions session_options;
        RunOptions run_options;
	    std::cout<<"hello"<<std::endl;
        if(LoadSavedModel(session_options, run_options, export_dir, {kSavedModelTagServe}, &bundle).ok()){
		std::cout<<"Load Successful!"<<std::endl;
	}
	std::cout<<"-------------------------"<<std::endl;
	if(LoadSavedModel(session_options, run_options, export_dir2, {kSavedModelTagServe}, &bundle).ok()){
		std::cout<<"Load Successful!"<<std::endl;
	}
	std::cout<<"-------------------------"<<std::endl;
	if(LoadSavedModel(session_options, run_options, export_dir3, {kSavedModelTagServe}, &bundle).ok()){
		std::cout<<"Load Successful!"<<std::endl;
	}
	std::cout<<"-------------------------"<<std::endl;

    }

saved_model.zip libfuzzer output.txt

Other info / logs Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

I attached three saved_model.pb files and the log of libfuzzer as well.

Personal understanding I think the following codes in tensorflow/core/grappler/costs/graph_properties.cc lead to the error

  Status MaybeUpdateNodeContextOutput(const NodeDef& node, const bool is_fed,
                                      NodeContext* c) {
    // Propagate tensors and shape tensors unless the node is fed.
    // TODO(bsteiner) We should still propagate the shapes to the ports that
    // aren't fed in the case of a ShapeN node.

    InferenceContext* ic = c->inference_context.get();
    if (!is_fed) {
      if (IsConstant(node)) {
        c->output_tensor_protos.resize(1);
        const TensorProto& tensor_proto = node.attr().at("value").tensor(); // at("xx") if "xx" does not exist will bring a crash

Once the value attribute cannot be obtained, an exception will be thrown.