CVE-2021-41212: Out-of-bounds Read (CWE-125)
value: TensorFlow is an open source platform for machine learning. In affected versions the shape inference code for `tf.ragged.cross` can trigger a read outside of bounds of heap allocated array. 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 reads data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.1
Published: 2021-11-05T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/ops/ragged_array_ops.cc, tensorflow/python/ops/ragged/ragged_cross_op_test.py
Commit date: 2021-09-30 16:10:47 -0700
Commit message: Fix null pointer exception in shape inference function when tf.ragged.cross() is called with invalid inputs.

PiperOrigin-RevId: 400045848
Change-Id: Ia65501583b85cf1ec14a252d83fbdd716817a516

Diff: diff --git a/tensorflow/core/ops/ragged_array_ops.cc b/tensorflow/core/ops/ragged_array_ops.cc
index 4e9b820eec2a8..b86f3968153da 100644
--- a/tensorflow/core/ops/ragged_array_ops.cc
+++ b/tensorflow/core/ops/ragged_array_ops.cc
@@ -97,10 +97,17 @@ REGISTER_OP("RaggedCross")
 
       // Merge the batch size of each dense input into out_splits.
       int dense_start = num_ragged * 2 + num_sparse * 3;
       for (int i = 0; i < dense_types.size(); ++i) {
         ShapeHandle dense_input = c->input(i + dense_start);
+        int32 rank = c->Rank(dense_input);
+        if (rank == InferenceContext::kUnknownRank) {
+          continue;
+        } else if (rank != 2) {
+          return errors::InvalidArgument(
+              "tf.ragged.cross only supports inputs with rank=2");
+        }
         int64_t batch_size = c->Value(c->Dim(dense_input, 0));
         if (batch_size != InferenceContext::kUnknownDim) {
           ShapeHandle row_splits = c->Vector(batch_size + 1);
           if (!c->Merge(out_splits, row_splits, &out_splits).ok()) {
             return errors::InvalidArgument(
diff --git a/tensorflow/python/ops/ragged/ragged_cross_op_test.py b/tensorflow/python/ops/ragged/ragged_cross_op_test.py
index 99f88d06a811a..71085257f6872 100644
--- a/tensorflow/python/ops/ragged/ragged_cross_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_cross_op_test.py
@@ -16,14 +16,16 @@
 
 from absl.testing import parameterized
 
 import numpy as np
 
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import sparse_ops
 from tensorflow.python.ops.ragged import ragged_array_ops
 from tensorflow.python.ops.ragged import ragged_factory_ops
 from tensorflow.python.ops.ragged import ragged_tensor
@@ -356,31 +358,60 @@ class RaggedCrossOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
           testcase_name='StaticBatchSizeMismatch2',
           inputs=[ragged_const([[1]]),
                   dense_const([[2], [3]])],
           exception=(ValueError, errors.InvalidArgumentError),
           message='inputs must all have the same batch dimension size'),
+      dict(
+          testcase_name='3DDenseTensor',
+          inputs=[dense_const([[[1]]])],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2'),
+      dict(
+          testcase_name='0DDenseTensor',
+          inputs=[dense_const(1)],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2'),
   ])
   def testStaticError(self, inputs, exception=ValueError, message=None):
     with self.assertRaisesRegex(exception, message):
       ragged_array_ops.cross(inputs)
 
   @parameterized.named_parameters([
       dict(
           testcase_name='3DRaggedTensor',
           inputs=[ragged_const([[[1]]], ragged_rank=1)],
           message='tf.ragged.cross only supports inputs with rank=2'),
+      dict(
+          testcase_name='0DDenseTensor',
+          inputs=[dense_const(1)],
+          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2'),
+      dict(
+          testcase_name='1DDenseTensor',
+          inputs=[dense_const([1])],
+          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2'),
       dict(
           testcase_name='3DDenseTensor',
           inputs=[dense_const([[[1]]])],
+          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
+          exception=(ValueError, errors.InvalidArgumentError),
           message='tf.ragged.cross only supports inputs with rank=2'),
   ])
   def testRuntimeError(self,
                        inputs,
                        exception=errors.InvalidArgumentError,
-                       message=None):
+                       message=None,
+                       signature=None):
+    @def_function.function(input_signature=signature)
+    def fn(x):
+      return ragged_array_ops.cross(x)
+
     with self.assertRaisesRegex(exception, message):
-      self.evaluate(ragged_array_ops.cross(inputs))
+      self.evaluate(fn(inputs))
 
   def _ragged_to_sparse(self, t):
     if ragged_tensor.is_ragged(t):
       return ragged_tensor.convert_to_tensor_or_ragged_tensor(t).to_sparse()
     elif sparse_tensor.is_sparse(t):

Vulnerable code:                        message=None):
      self.evaluate(ragged_array_ops.cross(inputs))
Fixed code:         int32 rank = c->Rank(dense_input);
        if (rank == InferenceContext::kUnknownRank) {
          continue;
        } else if (rank != 2) {
          return errors::InvalidArgument(
              "tf.ragged.cross only supports inputs with rank=2");
        }
from tensorflow.python.eager import def_function
from tensorflow.python.framework import tensor_spec
      dict(
          testcase_name='3DDenseTensor',
          inputs=[dense_const([[[1]]])],
          exception=(ValueError, errors.InvalidArgumentError),
          message='tf.ragged.cross only supports inputs with rank=2'),
      dict(
          testcase_name='0DDenseTensor',
          inputs=[dense_const(1)],
          exception=(ValueError, errors.InvalidArgumentError),
          message='tf.ragged.cross only supports inputs with rank=2'),
      dict(
          testcase_name='0DDenseTensor',
          inputs=[dense_const(1)],
          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
          exception=(ValueError, errors.InvalidArgumentError),
          message='tf.ragged.cross only supports inputs with rank=2'),
      dict(
          testcase_name='1DDenseTensor',
          inputs=[dense_const([1])],
          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
          exception=(ValueError, errors.InvalidArgumentError),
          message='tf.ragged.cross only supports inputs with rank=2'),
          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
          exception=(ValueError, errors.InvalidArgumentError),
                       message=None,
                       signature=None):
    @def_function.function(input_signature=signature)
    def fn(x):
      return ragged_array_ops.cross(x)

      self.evaluate(fn(inputs))