CVE-2021-37647: NULL Pointer Dereference (CWE-476)
value: TensorFlow is an end-to-end open source platform for machine learning. When a user does not supply arguments that determine a valid sparse tensor, `tf.raw_ops.SparseTensorSliceDataset` implementation can be made to dereference a null pointer. The [implementation](https://github.com/tensorflow/tensorflow/blob/8d72537c6abf5a44103b57b9c2e22c14f5f49698/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc#L240-L251) has some argument validation but fails to consider the case when either `indices` or `values` are provided for an empty sparse tensor when the other is not. If `indices` is empty, then [code that performs validation](https://github.com/tensorflow/tensorflow/blob/8d72537c6abf5a44103b57b9c2e22c14f5f49698/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc#L260-L261) (i.e., checking that the indices are monotonically increasing) results in a null pointer dereference. If `indices` as provided by the user is empty, then `indices` in the C++ code above is backed by an empty `std::vector`, hence calling `indices->dim_size(0)` results in null pointer dereferencing (same as calling `std::vector::at()` on an empty vector). We have patched the issue in GitHub commit 02cc160e29d20631de3859c6653184e3f876b9d7. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.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: 5.5
Published: 2021-08-12T19: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-08-03 15:57:40 -0700
Commit message: Prevent nullptr deref in SparseTensorSliceDataset

The arguments must determine a valid sparse tensor. This means that when indices are empty then the values must be empty too (and the reverse).

Also added test, by modifying existing test with empty sparse tensor to now run with an invalid sparse tensor input.

PiperOrigin-RevId: 388562757
Change-Id: Id8b54cd7c2316025b4f9a77292c8fb5344d17609

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 00b71d41a7ca1..58bb4b0b6d806 100644
--- a/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc
+++ b/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc
@@ -239,10 +239,21 @@ class SparseTensorSliceDatasetOp : public DatasetOpKernel {
 
     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()));
+    }
     OP_REQUIRES(ctx, TensorShapeUtils::IsVector(values->shape()),
                 errors::InvalidArgument(
                     "Input values should be a vector but received shape ",
                     indices->shape().DebugString()));
     OP_REQUIRES(ctx, TensorShapeUtils::IsVector(dense_shape->shape()),
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 25ecbf20680b3..8f93530010cae 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
@@ -116,10 +116,30 @@ class FromSparseTensorSlicesTest(test_base.DatasetTestBase,
                                                     empty_dense_shape)
       sess.run(init_op, feed_dict={st: sparse_feed})
       with self.assertRaises(errors.OutOfRangeError):
         sess.run(get_next)
 
+  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
+  def testEmptySparseTensorSlicesInvalid(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 = np.empty((0, 4), dtype=np.int64)
+      non_empty_values = [1, 2, 3, 4]
+      empty_dense_shape = [0, 4, 37, 9]
+      sparse_feed = sparse_tensor.SparseTensorValue(empty_indices,
+                                                    non_empty_values,
+                                                    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=2, mode=["eager"]))
   def testFromSparseTensorSlicesError(self):
     with self.assertRaises(AttributeError):
       dataset_ops.Dataset.from_sparse_tensor_slices(None)
 

Fixed code: 
    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()));
    }
  @combinations.generate(combinations.combine(tf_api_version=1, mode=["graph"]))
  def testEmptySparseTensorSlicesInvalid(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 = np.empty((0, 4), dtype=np.int64)
      non_empty_values = [1, 2, 3, 4]
      empty_dense_shape = [0, 4, 37, 9]
      sparse_feed = sparse_tensor.SparseTensorValue(empty_indices,
                                                    non_empty_values,
                                                    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})