CVE-2021-29616: NULL Pointer Dereference (CWE-476)
value: TensorFlow is an end-to-end open source platform for machine learning. The implementation of TrySimplify(https://github.com/tensorflow/tensorflow/blob/c22d88d6ff33031aa113e48aa3fc9aa74ed79595/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc#L390-L401) has undefined behavior due to dereferencing a null pointer in corner cases that result in optimizing a node with no inputs. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.
CWE description: The product dereferences a pointer that it expects to be valid but is NULL.
CVSS v3 base score: 7.8
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc, tensorflow/core/grappler/optimizers/dependency_optimizer.cc
Commit date: 2021-04-19 10:06:46 -0700
Commit message: Handle a special grappler case resulting in crash.

It might happen that a malformed input could be used to trick Grappler into trying to optimize a node with no inputs. This, in turn, would produce a null pointer dereference and a segfault.

PiperOrigin-RevId: 369242852
Change-Id: I2e5cbe7aec243d34a6d60220ac8ac9b16f136f6b

Diff: diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index 484983a7a549c..774fdf3484788 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -2045,10 +2045,16 @@ class ReorderCastLikeAndValuePreserving : public ArithmeticOptimizerStage {
            !IsControlFlow(*node) && !IsInPreserveSet(*node);
   }
 
   Status TrySimplify(NodeDef* consumer, string* simplified_node_name) override {
     NodeDef* producer;
+
+    if (consumer->input_size() < 1) {
+      return errors::FailedPrecondition("Node ", simplified_node_name,
+                                        " lacks inputs");
+    }
+
     TF_RETURN_IF_ERROR(GetInputNode(consumer->input(0), &producer));
     const bool producer_is_cast = IsCastLike(*producer);
     const bool can_optimize =
         !IsCheckNumerics(*producer) &&
         ((producer_is_cast && IsValuePreserving(*consumer)) ||
@@ -2536,10 +2542,15 @@ class ReplaceMulWithSquare : public ArithmeticOptimizerStage {
                                 const ArithmeticOptimizerContext& ctx_ext)
       : ArithmeticOptimizerStage("ReplaceMulWithSquare", ctx, ctx_ext) {}
   ~ReplaceMulWithSquare() override = default;
 
   bool IsSupported(const NodeDef* node) const override {
+    if (!node || node->input_size() < 2) {
+      // Invalid node
+      return false;
+    }
+
     return IsAnyMul(*node) && node->input(0) == node->input(1);
   }
 
   Status TrySimplify(NodeDef* node, string* simplified_node_name) override {
     const NodeScopeAndName mul = ParseNodeScopeAndName(node->name());
diff --git a/tensorflow/core/grappler/optimizers/dependency_optimizer.cc b/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
index 332d2d4e6d966..1be7f2692e0f7 100644
--- a/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/dependency_optimizer.cc
@@ -66,10 +66,16 @@ bool DependencyOptimizer::SafeToRemoveIdentity(const NodeDef& node) const {
   }
   if (!fetch_nodes_known_) {
     // The output values of this node may be needed.
     return false;
   }
+
+  if (node.input_size() < 1) {
+    // Node lacks input, is invalid
+    return false;
+  }
+
   const NodeDef* input = node_map_->GetNode(NodeName(node.input(0)));
   CHECK(input != nullptr) << "node = " << node.name()
                           << " input = " << node.input(0);
   // Don't remove Identity nodes corresponding to Variable reads or following
   // Recv.

Fixed code: 
    if (consumer->input_size() < 1) {
      return errors::FailedPrecondition("Node ", simplified_node_name,
                                        " lacks inputs");
    }

    if (!node || node->input_size() < 2) {
      // Invalid node
      return false;
    }


  if (node.input_size() < 1) {
    // Node lacks input, is invalid
    return false;
  }