CVE-2020-15206: Improper Input Validation (CWE-20)
value: In Tensorflow before versions 1.15.4, 2.0.3, 2.1.2, 2.2.1 and 2.3.1, changing the TensorFlow's `SavedModel` protocol buffer and altering the name of required keys results in segfaults and data corruption while loading the model. This can cause a denial of service in products using `tensorflow-serving` or other inference-as-a-service installments. Fixed were added in commits f760f88b4267d981e13f4b302c437ae800445968 and fcfef195637c6e365577829c4d67681695956e7d (both going into TensorFlow 2.2.0 and 2.3.0 but not yet backported to earlier versions). However, this was not enough, as #41097 reports a different failure mode. The issue is patched in commit adf095206f25471e864a8e63a0f1caef53a0e3a6, and is released in TensorFlow versions 1.15.4, 2.0.3, 2.1.2, 2.2.1, or 2.3.1.
CWE description: The product receives input or data, but it does
        not validate or incorrectly validates that the input has the
        properties that are required to process the data safely and
        correctly.
CVSS v3 base score: 7.5
Published: 2020-09-25T19:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/cc/saved_model/loader.cc, tensorflow/cc/saved_model/saved_model_bundle_test.cc
Commit date: 2020-09-18 15:17:30 -0700
Commit message: Validate `NodeDef`s from `FunctionDefLibrary` of a `GraphDef`.

We already validated `NodeDef`s from a `GraphDef` but missed validating those from the `FunctionDefLibrary`. Thus, some maliciously crafted models could evade detection and cause denial of service due to a `CHECK`-fail.

PiperOrigin-RevId: 332536309
Change-Id: I052efe919ff1fe2f90815e286a1aa4c54c7b94ff

Diff: diff --git a/tensorflow/cc/saved_model/loader.cc b/tensorflow/cc/saved_model/loader.cc
index ecefe7d0406..70d080a682f 100644
--- a/tensorflow/cc/saved_model/loader.cc
+++ b/tensorflow/cc/saved_model/loader.cc
@@ -19,10 +19,11 @@ limitations under the License.
 
 #include "tensorflow/cc/saved_model/constants.h"
 #include "tensorflow/cc/saved_model/loader_util.h"
 #include "tensorflow/cc/saved_model/reader.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"
 #include "tensorflow/core/framework/tensor.pb.h"
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/lib/monitoring/counter.h"
 #include "tensorflow/core/lib/monitoring/sampler.h"
@@ -71,30 +72,45 @@ uint64 GetLatencyMicroseconds(const uint64 start_microseconds) {
 }
 
 // Ensure that constant tensors loaded from the saved model have valid shape.
 // Also ensure that constant nodes have a value assigned to them.
 // TODO(b/154763635): this is temporary and will be replaced with a better audit
+static Status ValidateNode(const NodeDef& node) {
+  const auto node_iterator = node.attr().find("value");
+  if (node_iterator != node.attr().end()) {
+    AttrValue node_value = node_iterator->second;
+    if (node_value.has_tensor()) {
+      const PartialTensorShape node_shape(node_value.tensor().tensor_shape());
+      if (node_shape.num_elements() < 0) {
+        return errors::FailedPrecondition(
+            "Saved model contains node \"", node.name(), "\" (op \"", node.op(),
+            "\") which initializes from a tensor with ",
+            node_shape.num_elements(), " elements");
+      }
+    }
+  } else if (node.op() == "Const") {
+    return errors::FailedPrecondition(
+        "Saved model contains node \"", node.name(),
+        "\" which is a constant tensor but no value has been provided");
+  }
+  return Status::OK();
+}
+
 static Status ValidateSavedTensors(const GraphDef& graph_def) {
   for (const auto& node : graph_def.node()) {
-    const auto node_iterator = node.attr().find("value");
-    if (node_iterator != node.attr().end()) {
-      AttrValue node_value = node_iterator->second;
-      if (node_value.has_tensor()) {
-        const PartialTensorShape node_shape(node_value.tensor().tensor_shape());
-        if (node_shape.num_elements() < 0) {
-          return errors::FailedPrecondition(
-              "Saved model contains node \"", node.name(), "\" (op \"",
-              node.op(), "\") which initializes from a tensor with ",
-              node_shape.num_elements(), " elements");
-        }
+    TF_RETURN_IF_ERROR(ValidateNode(node));
+  }
+
+  if (graph_def.has_library()) {
+    const FunctionDefLibrary& library = graph_def.library();
+    for (const auto& function : library.function()) {
+      for (const auto& node : function.node_def()) {
+        TF_RETURN_IF_ERROR(ValidateNode(node));
       }
-    } else if (node.op() == "Const") {
-      return errors::FailedPrecondition(
-          "Saved model contains node \"", node.name(),
-          "\" which is a constant tensor but no value has been provided");
     }
   }
+
   return Status::OK();
 }
 
 Tensor CreateStringTensor(const string& value) {
   Tensor tensor(DT_STRING, TensorShape({}));
diff --git a/tensorflow/cc/saved_model/saved_model_bundle_test.cc b/tensorflow/cc/saved_model/saved_model_bundle_test.cc
index 3f258745fa4..127176002b9 100644
--- a/tensorflow/cc/saved_model/saved_model_bundle_test.cc
+++ b/tensorflow/cc/saved_model/saved_model_bundle_test.cc
@@ -43,10 +43,12 @@ constexpr char kTestDataV2DebugInfo[] =
     "cc/saved_model/testdata/x_plus_y_v2_debuginfo";
 constexpr char kTestFuzzGeneratedNegativeShape[] =
     "cc/saved_model/testdata/fuzz_generated/negative_shape";
 constexpr char kTestFuzzGeneratedConstWithNoValue[] =
     "cc/saved_model/testdata/fuzz_generated/const_with_no_value";
+constexpr char kTestFuzzGeneratedBadNodeAttr[] =
+    "cc/saved_model/testdata/fuzz_generated/bad_node_attr";
 
 class LoaderTest : public ::testing::Test {
  protected:
   LoaderTest() {}
 
@@ -326,7 +328,22 @@ TEST_F(LoaderTest, ConstNoValue) {
   EXPECT_NE(
       st.error_message().find("constant tensor but no value has been provided"),
       std::string::npos);
 }
 
+TEST_F(LoaderTest, BadNodeAttr) {
+  SavedModelBundle bundle;
+  RunOptions run_options;
+  SessionOptions session_options;
+
+  const string export_dir =
+      io::JoinPath(testing::TensorFlowSrcRoot(), kTestFuzzGeneratedBadNodeAttr);
+  Status st = LoadSavedModel(session_options, run_options, export_dir,
+                             {kSavedModelTagServe}, &bundle);
+  EXPECT_FALSE(st.ok());
+  EXPECT_NE(
+      st.error_message().find("constant tensor but no value has been provided"),
+      std::string::npos);
+}
+
 }  // namespace
 }  // namespace tensorflow
diff --git a/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/assets/empty b/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/assets/empty
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/saved_model.pb b/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/saved_model.pb
new file mode 100644
index 00000000000..0b33dbe7352
Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/saved_model.pb differ
diff --git a/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/variables/variables.data-00000-of-00001 b/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/variables/variables.data-00000-of-00001
new file mode 100644
index 00000000000..3fd3ba2223d
Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/variables/variables.data-00000-of-00001 differ
diff --git a/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/variables/variables.index b/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/variables/variables.index
new file mode 100644
index 00000000000..7357e8d57ed
Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/variables/variables.index differ

Vulnerable code:     const auto node_iterator = node.attr().find("value");
    if (node_iterator != node.attr().end()) {
      AttrValue node_value = node_iterator->second;
      if (node_value.has_tensor()) {
        const PartialTensorShape node_shape(node_value.tensor().tensor_shape());
        if (node_shape.num_elements() < 0) {
          return errors::FailedPrecondition(
              "Saved model contains node \"", node.name(), "\" (op \"",
              node.op(), "\") which initializes from a tensor with ",
              node_shape.num_elements(), " elements");
        }
    } else if (node.op() == "Const") {
      return errors::FailedPrecondition(
          "Saved model contains node \"", node.name(),
          "\" which is a constant tensor but no value has been provided");
Fixed code: #include "tensorflow/core/framework/function.pb.h"
static Status ValidateNode(const NodeDef& node) {
  const auto node_iterator = node.attr().find("value");
  if (node_iterator != node.attr().end()) {
    AttrValue node_value = node_iterator->second;
    if (node_value.has_tensor()) {
      const PartialTensorShape node_shape(node_value.tensor().tensor_shape());
      if (node_shape.num_elements() < 0) {
        return errors::FailedPrecondition(
            "Saved model contains node \"", node.name(), "\" (op \"", node.op(),
            "\") which initializes from a tensor with ",
            node_shape.num_elements(), " elements");
      }
    }
  } else if (node.op() == "Const") {
    return errors::FailedPrecondition(
        "Saved model contains node \"", node.name(),
        "\" which is a constant tensor but no value has been provided");
  }
  return Status::OK();
}

    TF_RETURN_IF_ERROR(ValidateNode(node));
  }

  if (graph_def.has_library()) {
    const FunctionDefLibrary& library = graph_def.library();
    for (const auto& function : library.function()) {
      for (const auto& node : function.node_def()) {
        TF_RETURN_IF_ERROR(ValidateNode(node));

constexpr char kTestFuzzGeneratedBadNodeAttr[] =
    "cc/saved_model/testdata/fuzz_generated/bad_node_attr";
TEST_F(LoaderTest, BadNodeAttr) {
  SavedModelBundle bundle;
  RunOptions run_options;
  SessionOptions session_options;

  const string export_dir =
      io::JoinPath(testing::TensorFlowSrcRoot(), kTestFuzzGeneratedBadNodeAttr);
  Status st = LoadSavedModel(session_options, run_options, export_dir,
                             {kSavedModelTagServe}, &bundle);
  EXPECT_FALSE(st.ok());
  EXPECT_NE(
      st.error_message().find("constant tensor but no value has been provided"),
      std::string::npos);
}