CVE-2021-41207: Divide By Zero (CWE-369)
value: TensorFlow is an open source platform for machine learning. In affected versions the implementation of `ParallelConcat` misses some input validation and can produce a division by 0. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.
CWE description: The product divides a value by zero.
CVSS v3 base score: 5.5
Published: 2021-11-05T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/inplace_ops.cc, tensorflow/python/kernel_tests/array_ops/stack_op_test.py
Commit date: 2021-10-12 08:36:20 -0700
Commit message: Adding more validation checks to _ParallelConcatUpdate to avoid NPE.

PiperOrigin-RevId: 402569467
Change-Id: I2db122dab68be2a5e4e8dd3375f5a70c4d2307ec

Diff: diff --git a/tensorflow/core/kernels/inplace_ops.cc b/tensorflow/core/kernels/inplace_ops.cc
index 6cf2337eeacbd..b53ebf942232b 100644
--- a/tensorflow/core/kernels/inplace_ops.cc
+++ b/tensorflow/core/kernels/inplace_ops.cc
@@ -69,10 +69,19 @@ class ParallelConcatUpdate : public OpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("loc", &loc_));
   }
 
   void Compute(OpKernelContext* ctx) override {
     auto value = ctx->input(0);
+    // Value should be at least rank 1. Also the 0th dimension should be
+    // at least loc_.
+    OP_REQUIRES(ctx, value.dims() >= 1,
+                errors::InvalidArgument("value should be at least rank 1."));
+    OP_REQUIRES(
+        ctx, value.dim_size(0) > loc_,
+        errors::InvalidArgument("0th dimension of value = ", value.dim_size(0),
+                                " is less than loc_=", loc_));
+
     auto update = ctx->input(1);
 
     OP_REQUIRES(
         ctx, value.dims() == update.dims(),
         errors::InvalidArgument("value and update shape doesn't match: ",
diff --git a/tensorflow/python/kernel_tests/array_ops/stack_op_test.py b/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
index 1a5cc5f02938f..4ee468cd091eb 100644
--- a/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
@@ -14,16 +14,20 @@
 # ==============================================================================
 """Functional tests for Stack and ParallelStack Ops."""
 
 import numpy as np
 
+from tensorflow.python import tf2
 from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.ops import gradient_checker_v2
 from tensorflow.python.platform import test
 
 
 def np_split_squeeze(array, axis):
@@ -67,10 +71,23 @@ class StackOpTest(test.TestCase):
             data = self.randn(shape, np.float32)
             xs = list(map(constant_op.constant, data))
             c = array_ops.parallel_stack(xs)
             self.assertAllEqual(c, data)
 
+  def testParallelConcatShapeZero(self):
+    if not tf2.enabled():
+      self.skipTest("only fails in TF2")
+
+    @def_function.function
+    def f():
+      y = gen_array_ops.parallel_concat(values=[["tf"]], shape=0)
+      return y
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r"0th dimension of value .* is less than"):
+      f()
+
   def testSimpleParallelGPU(self):
     # tf.parallel_stack is only supported in graph mode.
     with ops.Graph().as_default():
       with test_util.device(use_gpu=True):
         for shape in (2,), (3,), (2, 3), (3, 2), (4, 3, 2), (100, 24, 24, 3):

Fixed code:     // Value should be at least rank 1. Also the 0th dimension should be
    // at least loc_.
    OP_REQUIRES(ctx, value.dims() >= 1,
                errors::InvalidArgument("value should be at least rank 1."));
    OP_REQUIRES(
        ctx, value.dim_size(0) > loc_,
        errors::InvalidArgument("0th dimension of value = ", value.dim_size(0),
                                " is less than loc_=", loc_));

from tensorflow.python import tf2
from tensorflow.python.eager import def_function
from tensorflow.python.framework import errors
from tensorflow.python.ops import gen_array_ops
  def testParallelConcatShapeZero(self):
    if not tf2.enabled():
      self.skipTest("only fails in TF2")

    @def_function.function
    def f():
      y = gen_array_ops.parallel_concat(values=[["tf"]], shape=0)
      return y

    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                r"0th dimension of value .* is less than"):
      f()