CVE-2022-21736: NULL Pointer Dereference (CWE-476)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of `SparseTensorSliceDataset` has an undefined behavior: under certain condition it can be made to dereference a `nullptr` value. The 3 input arguments to `SparseTensorSliceDataset` represent a sparse tensor. However, there are some preconditions that these arguments must satisfy but these are not validated in the implementation. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, 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: 6.5
Published: 2022-02-03T12:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc, tensorflow/python/data/kernel_tests/from_sparse_tensor_slices_test.py
Commit date: 2021-12-09 15:55:37 -0800
Commit message: Properly validate sparse tensor in `SparseTensorSliceDataset`

Existing validation was incomplete.

PiperOrigin-RevId: 415375048
Change-Id: I14cd18f29ede73286f3ffac35171bd15828997e9

Diff: diff --git a/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc b/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc
index c8f3db30c37f4..c6901de5c3426 100644
--- a/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc
+++ b/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc
@@ -238,32 +238,33 @@ class SparseTensorSliceDatasetOp : public DatasetOpKernel {
     OP_REQUIRES_OK(ctx, ctx->input("values", &values));
     const Tensor* dense_shape;
     OP_REQUIRES_OK(ctx, ctx->input("dense_shape", &dense_shape));
 
     OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(indices->shape()),
-                errors::InvalidArgument(
-                    "Input indices should be a matrix but received shape ",
-                    indices->shape().DebugString()));
-
-    const auto num_indices = indices->NumElements();
-    const auto num_values = values->NumElements();
-    if (num_indices == 0 || num_values == 0) {
-      OP_REQUIRES(ctx, num_indices == num_values,
-                  errors::InvalidArgument(
-                      "If indices or values are empty, the other one must also "
-                      "be. Got indices of shape ",
-                      indices->shape().DebugString(), " and values of shape ",
-                      values->shape().DebugString()));
-    }
+                errors::InvalidArgument("Input indices must be a matrix. Got: ",
+                                        indices->shape().DebugString()));
     OP_REQUIRES(ctx, TensorShapeUtils::IsVector(values->shape()),
-                errors::InvalidArgument(
-                    "Input values should be a vector but received shape ",
-                    indices->shape().DebugString()));
+                errors::InvalidArgument("Input values must be a vector. Got: ",
+                                        values->shape().DebugString()));
     OP_REQUIRES(ctx, TensorShapeUtils::IsVector(dense_shape->shape()),
+                errors::InvalidArgument("Input shape must be a vector. Got: ",
+                                        dense_shape->shape().DebugString()));
+    OP_REQUIRES(
+        ctx, values->shape().dim_size(0) == indices->shape().dim_size(0),
+        errors::InvalidArgument(
+            "Number of values must match first dimension of indices. ", "Got ",
+            values->shape().dim_size(0),
+            " values, indices shape: ", indices->shape().DebugString()));
+    OP_REQUIRES(
+        ctx, dense_shape->shape().dim_size(0) == indices->shape().dim_size(1),
+        errors::InvalidArgument(
+            "Number of dimensions must match second dimension of indices. ",
+            "Got ", dense_shape->shape().dim_size(0),
+            " dimensions, indices shape: ", indices->shape().DebugString()));
+    OP_REQUIRES(ctx, dense_shape->NumElements() > 0,
                 errors::InvalidArgument(
-                    "Input shape should be a vector but received shape ",
-                    dense_shape->shape().DebugString()));
+                    "The shape argument requires at least one element."));
 
     // We currently ensure that `sparse_tensor` is ordered in the
     // batch dimension.
     // TODO(mrry): Investigate ways to avoid this unconditional check
     // if we can be sure that the sparse tensor was produced in an
diff --git a/tensorflow/python/data/kernel_tests/from_sparse_tensor_slices_test.py b/tensorflow/python/data/kernel_tests/from_sparse_tensor_slices_test.py
index 04f2c27bb9fe7..3d8ee43272fbf 100644
--- a/tensorflow/python/data/kernel_tests/from_sparse_tensor_slices_test.py
+++ b/tensorflow/python/data/kernel_tests/from_sparse_tensor_slices_test.py
@@ -132,10 +132,29 @@ class FromSparseTensorSlicesTest(test_base.DatasetTestBase,
                                                     empty_dense_shape)
       # Here, we expect the test to fail when running the feed.
       with self.assertRaises(errors.InvalidArgumentError):
         sess.run(init_op, feed_dict={st: sparse_feed})
 
+  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
+  def testEmptySparseTensorSlicesInvalid2(self):
+    """Test a dataset based on invalid `tf.sparse.SparseTensor`."""
+    st = array_ops.sparse_placeholder(dtypes.float64)
+    iterator = dataset_ops.make_initializable_iterator(
+        dataset_ops.Dataset.from_sparse_tensor_slices(st))
+    init_op = iterator.initializer
+
+    with self.cached_session() as sess:
+      # Test with an empty sparse tensor but with non empty values.
+      empty_indices = [[]]
+      empty_values = []
+      dense_shape = [1, 1]
+      sparse_feed = sparse_tensor.SparseTensorValue(empty_indices, empty_values,
+                                                    dense_shape)
+      # Here, we expect the test to fail when running the feed.
+      with self.assertRaises(errors.InvalidArgumentError):
+        sess.run(init_op, feed_dict={st: sparse_feed})
+
   @combinations.generate(combinations.combine(tf_api_version=2, mode=["eager"]))
   def testFromSparseTensorSlicesError(self):
     with self.assertRaises(AttributeError):
       dataset_ops.Dataset.from_sparse_tensor_slices(None)
 

Vulnerable code:                 errors::InvalidArgument(
                    "Input indices should be a matrix but received shape ",
                    indices->shape().DebugString()));

    const auto num_indices = indices->NumElements();
    const auto num_values = values->NumElements();
    if (num_indices == 0 || num_values == 0) {
      OP_REQUIRES(ctx, num_indices == num_values,
                  errors::InvalidArgument(
                      "If indices or values are empty, the other one must also "
                      "be. Got indices of shape ",
                      indices->shape().DebugString(), " and values of shape ",
                      values->shape().DebugString()));
    }
                errors::InvalidArgument(
                    "Input values should be a vector but received shape ",
                    indices->shape().DebugString()));
                    "Input shape should be a vector but received shape ",
                    dense_shape->shape().DebugString()));
Fixed code:                 errors::InvalidArgument("Input indices must be a matrix. Got: ",
                                        indices->shape().DebugString()));
                errors::InvalidArgument("Input values must be a vector. Got: ",
                                        values->shape().DebugString()));
                errors::InvalidArgument("Input shape must be a vector. Got: ",
                                        dense_shape->shape().DebugString()));
    OP_REQUIRES(
        ctx, values->shape().dim_size(0) == indices->shape().dim_size(0),
        errors::InvalidArgument(
            "Number of values must match first dimension of indices. ", "Got ",
            values->shape().dim_size(0),
            " values, indices shape: ", indices->shape().DebugString()));
    OP_REQUIRES(
        ctx, dense_shape->shape().dim_size(0) == indices->shape().dim_size(1),
        errors::InvalidArgument(
            "Number of dimensions must match second dimension of indices. ",
            "Got ", dense_shape->shape().dim_size(0),
            " dimensions, indices shape: ", indices->shape().DebugString()));
    OP_REQUIRES(ctx, dense_shape->NumElements() > 0,
                    "The shape argument requires at least one element."));
  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
  def testEmptySparseTensorSlicesInvalid2(self):
    """Test a dataset based on invalid `tf.sparse.SparseTensor`."""
    st = array_ops.sparse_placeholder(dtypes.float64)
    iterator = dataset_ops.make_initializable_iterator(
        dataset_ops.Dataset.from_sparse_tensor_slices(st))
    init_op = iterator.initializer

    with self.cached_session() as sess:
      # Test with an empty sparse tensor but with non empty values.
      empty_indices = [[]]
      empty_values = []
      dense_shape = [1, 1]
      sparse_feed = sparse_tensor.SparseTensorValue(empty_indices, empty_values,
                                                    dense_shape)
      # Here, we expect the test to fail when running the feed.
      with self.assertRaises(errors.InvalidArgumentError):
        sess.run(init_op, feed_dict={st: sparse_feed})