CVE-2023-25676: NULL Pointer Dereference (CWE-476)
value: TensorFlow is an open source machine learning platform. When running versions prior to 2.12.0 and 2.11.1 with XLA, `tf.raw_ops.ParallelConcat` segfaults with a nullptr dereference when given a parameter `shape` with rank that is not greater than zero. A fix is available in TensorFlow 2.12.0 and 2.11.1.
CWE description: The product dereferences a pointer that it expects to be valid but is NULL.
CVSS v3 base score: 7.5
Published: 2023-03-25T00:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/inplace_ops.cc, tensorflow/core/ops/array_ops.cc, tensorflow/python/kernel_tests/array_ops/stack_op_test.py, tensorflow/python/ops/array_ops_test.py
Commit date: 2023-01-26 11:51:43 -0800
Commit message: Check for unexpected scalars in the shape argument to ParallelConcat.

PiperOrigin-RevId: 504901518

Diff: diff --git a/tensorflow/core/kernels/inplace_ops.cc b/tensorflow/core/kernels/inplace_ops.cc
index 4aabb0f8df886..5aa38f67a5f52 100644
--- a/tensorflow/core/kernels/inplace_ops.cc
+++ b/tensorflow/core/kernels/inplace_ops.cc
@@ -76,11 +76,11 @@ class ParallelConcatUpdate : public OpKernel {
     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_));
+                                " must be greater than loc_ = ", loc_));
 
     auto update = ctx->input(1);
 
     OP_REQUIRES(
         ctx, value.dims() == update.dims(),
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index ef23f456533fa..856e851aebbff 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -13,10 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
 #include <algorithm>
 #include <ostream>
+#include <vector>
 
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/full_type.pb.h"
 #include "tensorflow/core/framework/kernel_shape_util.h"
 #include "tensorflow/core/framework/op.h"
@@ -307,10 +308,16 @@ REGISTER_OP("ParallelConcat")
       for (int i = 0; i < c->num_inputs(); ++i) {
         if (!c->FullyDefined(c->input(i))) {
           return errors::InvalidArgument(
               "All input shapes must be fully defined.");
         }
+        if (c->Rank(c->input(i)) < 1) {
+          return errors::InvalidArgument(
+              "The rank of all input shapes must be greater than 0, "
+              "but input ",
+              i, " had rank ", c->Rank(c->input(i)), ".");
+        }
         DimensionHandle unused;
         if (!c->WithValue(c->Dim(c->input(i), 0), 1, &unused).ok()) {
           return errors::InvalidArgument("Size of first dimension must be 1.");
         }
         TF_RETURN_WITH_CONTEXT_IF_ERROR(c->Merge(c->input(i), cur, &cur),
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 f46e4f05e030e..a27f3a7a10334 100644
--- a/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
@@ -81,12 +81,13 @@ class StackOpTest(test.TestCase):
     @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"):
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, r"0th dimension .* must be greater than"
+    ):
       f()
 
   def testSimpleParallelGPU(self):
     # tf.parallel_stack is only supported in graph mode.
     with ops.Graph().as_default():
diff --git a/tensorflow/python/ops/array_ops_test.py b/tensorflow/python/ops/array_ops_test.py
index 3f063a24a27cd..0c82f5ac0987d 100644
--- a/tensorflow/python/ops/array_ops_test.py
+++ b/tensorflow/python/ops/array_ops_test.py
@@ -16,10 +16,11 @@
 
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
 from tensorflow.python.platform import test
 
@@ -89,8 +90,22 @@ class ArrayOpTest(test.TestCase):
       return array_ops.zeros([array_ops.identity(array_ops.shape(x)[0])])
 
     conc = g.get_concrete_function(tensor_spec.TensorSpec([10, None]))
     self.assertAllEqual(conc.output_shapes.as_list(), [10])
 
+  @test_util.run_in_graph_and_eager_modes
+  def testParallelConcatFailsWithRankZeroShape(self):
+    op = array_ops.ParallelConcat
+    para = {"shape": 0, "values": [1]}
+
+    def func():
+      y = op(**para)
+      return y
+
+    with self.assertRaisesRegex(
+        Exception, "(rank|dimension) of .* must be greater than .* 0"
+    ):
+      func()
+
 
 if __name__ == "__main__":
   test.main()

Vulnerable code:                                 " is less than loc_=", loc_));
    with self.assertRaisesRegex(errors.InvalidArgumentError,
                                r"0th dimension of value .* is less than"):
Fixed code:                                 " must be greater than loc_ = ", loc_));
#include <vector>
        if (c->Rank(c->input(i)) < 1) {
          return errors::InvalidArgument(
              "The rank of all input shapes must be greater than 0, "
              "but input ",
              i, " had rank ", c->Rank(c->input(i)), ".");
        }
    with self.assertRaisesRegex(
        errors.InvalidArgumentError, r"0th dimension .* must be greater than"
    ):
from tensorflow.python.framework import test_util
  @test_util.run_in_graph_and_eager_modes
  def testParallelConcatFailsWithRankZeroShape(self):
    op = array_ops.ParallelConcat
    para = {"shape": 0, "values": [1]}

    def func():
      y = op(**para)
      return y

    with self.assertRaisesRegex(
        Exception, "(rank|dimension) of .* must be greater than .* 0"
    ):
      func()