CVE-2022-29206: NULL Pointer Dereference (CWE-476)
value: TensorFlow is an open source platform for machine learning. Prior to versions 2.9.0, 2.8.1, 2.7.2, and 2.6.4, the implementation of `tf.raw_ops.SparseTensorDenseAdd` does not fully validate the input arguments. In this case, a reference gets bound to a `nullptr` during kernel execution. This is undefined behavior. Versions 2.9.0, 2.8.1, 2.7.2, and 2.6.4 contain a patch for this issue.
CWE description: The product dereferences a pointer that it expects to be valid but is NULL.
CVSS v3 base score: 5.5
Published: 2022-05-20T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/sparse_tensor_dense_add_op.cc, tensorflow/python/kernel_tests/sparse_ops/sparse_add_op_test.py, tensorflow/python/kernel_tests/sparse_ops/sparse_ops_test.py
Commit date: 2022-05-03 07:54:57 -0700
Commit message: Fix UB in SparseTensorDenseAdd

Added more input validation to avoid nullptr dereferencing and array index
out of bounds issues.

PiperOrigin-RevId: 446192704

Diff: diff --git a/tensorflow/core/kernels/sparse_tensor_dense_add_op.cc b/tensorflow/core/kernels/sparse_tensor_dense_add_op.cc
index 48803e4b93980..6d6b05bf70f30 100644
--- a/tensorflow/core/kernels/sparse_tensor_dense_add_op.cc
+++ b/tensorflow/core/kernels/sparse_tensor_dense_add_op.cc
@@ -16,10 +16,11 @@ limitations under the License.
 #define EIGEN_USE_THREADS
 
 #include "tensorflow/core/kernels/sparse_tensor_dense_add_op.h"
 
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_util.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/util/sparse/sparse_tensor.h"
@@ -45,10 +46,21 @@ Status ValidateInputs(const Tensor *a_indices, const Tensor *a_values,
         "Inputs a_values and a_shape should be vectors "
         "but received shapes: ",
         a_values->shape().DebugString(), " and ",
         a_shape->shape().DebugString());
   }
+  int64_t nnz = a_indices->dim_size(0);
+  int64_t ndims = a_indices->dim_size(1);
+  if (a_values->dim_size(0) != nnz) {
+    return errors::InvalidArgument("Dimensions ", nnz, " and ",
+                                   a_values->dim_size(0),
+                                   " are not compatible");
+  }
+  if (a_shape->dim_size(0) != ndims) {
+    return errors::InvalidArgument("Dimensions ", ndims, " and ",
+                                   a_shape->dim_size(0), " are not compatible");
+  }
   if (a_shape->NumElements() != b->dims()) {
     return errors::InvalidArgument(
         "Two operands have different ranks; received: ", a_shape->NumElements(),
         " and ", b->dims());
   }
@@ -59,10 +71,28 @@ Status ValidateInputs(const Tensor *a_indices, const Tensor *a_values,
           "Dimension ", i,
           " does not equal (no broadcasting is supported): sparse side ",
           a_shape_flat(i), " vs dense side ", b->dim_size(i));
     }
   }
+
+  // Check for invalid indices.
+  const auto a_indices_mat = a_indices->flat_inner_dims<Index>();
+
+  for (int64_t zidx = 0; zidx < nnz; ++zidx) {
+    for (int64_t didx = 0; didx < ndims; ++didx) {
+      const Index idx = a_indices_mat(zidx, didx);
+      if (idx < 0 || idx >= a_shape_flat(didx)) {
+        return errors::InvalidArgument(
+            "Sparse tensor has an invalid index on dimension ", didx,
+            ": "
+            "a_indices(",
+            zidx, ",", didx, ") = ", idx,
+            ", dense tensor shape: ", a_shape_flat);
+      }
+    }
+  }
+
   return Status::OK();
 }
 
 }  // namespace
 
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_add_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_add_op_test.py
index 61ad45fb5e273..821184af5b569 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_add_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_add_op_test.py
@@ -187,11 +187,10 @@ class SparseAddTest(test.TestCase):
 
       err = gradient_checker.compute_gradient_error([sparse.values, dense],
                                                     [(nnz,), (n, m)], s, (n, m))
       self.assertLess(err, 1e-3)
 
-  @test_util.run_deprecated_v1
   def testInvalidSparseTensor(self):
     with test_util.force_cpu():
       shape = [2, 2]
       val = [0]
       dense = constant_op.constant(np.zeros(shape, dtype=np.int32))
@@ -199,16 +198,53 @@ class SparseAddTest(test.TestCase):
       for bad_idx in [
           [[-1, 0]],  # -1 is invalid.
           [[1, 3]],  # ...so is 3.
       ]:
         sparse = sparse_tensor.SparseTensorValue(bad_idx, val, shape)
-        s = sparse_ops.sparse_add(sparse, dense)
-
-        with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
-                                    "invalid index"):
+        with self.assertRaisesRegex(
+            (ValueError, errors_impl.InvalidArgumentError), "invalid index"):
+          s = sparse_ops.sparse_add(sparse, dense)
           self.evaluate(s)
 
+  def _testSparseDenseInvalidInputs(self,
+                                    a_indices,
+                                    a_values,
+                                    a_shape,
+                                    b,
+                                    expected_error=""):
+    # Public API call to sparse-dense add.
+    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
+                                expected_error):
+      a = sparse_tensor.SparseTensor(a_indices, a_values, a_shape)
+      self.evaluate(sparse_ops.sparse_add(a, b))
+    # Directly call generated kernel, by-passing SparseTensor validation.
+    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
+                                expected_error):
+      self.evaluate(
+          sparse_ops.gen_sparse_ops.sparse_tensor_dense_add(
+              a_indices, a_values, a_shape, b))
+
+  def testSparseDenseInvalidInputs(self):
+    self._testSparseDenseInvalidInputs(
+        a_indices=constant_op.constant(0, shape=[17, 2], dtype=dtypes.int64),
+        a_values=constant_op.constant(0, shape=[5], dtype=dtypes.float32),
+        a_shape=constant_op.constant([3, 4], dtype=dtypes.int64),
+        b=constant_op.constant(1, shape=[3, 4], dtype=dtypes.float32),
+        expected_error="Dimensions 17 and 5 are not compatible")
+    self._testSparseDenseInvalidInputs(
+        a_indices=constant_op.constant(0, shape=[17, 4], dtype=dtypes.int64),
+        a_values=constant_op.constant(0, shape=[17], dtype=dtypes.float32),
+        a_shape=constant_op.constant([3, 4], dtype=dtypes.int64),
+        b=constant_op.constant(1, shape=[3, 4], dtype=dtypes.float32),
+        expected_error="Dimensions 4 and 2 are not compatible")
+    self._testSparseDenseInvalidInputs(
+        a_indices=constant_op.constant(7, shape=[17, 2], dtype=dtypes.int64),
+        a_values=constant_op.constant(0, shape=[17], dtype=dtypes.float32),
+        a_shape=constant_op.constant([3, 4], dtype=dtypes.int64),
+        b=constant_op.constant(1, shape=[3, 4], dtype=dtypes.float32),
+        expected_error="invalid index")
+
 ######################## Benchmarking code
 
 
 def _s2d_add_vs_sparse_add(sparsity, n, m, num_iters=50):
   np.random.seed(1618)
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_ops_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_ops_test.py
index 4972d1d25f08f..684d1f98432b5 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_ops_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_ops_test.py
@@ -663,11 +663,11 @@ class SparseFillEmptyRowsTest(test_util.TensorFlowTestCase):
 
 
 class SparseAddTest(test_util.TensorFlowTestCase):
 
   def testValuesInVariable(self):
-    indices = constant_op.constant([[1]], dtype=dtypes.int64)
+    indices = constant_op.constant([[0]], dtype=dtypes.int64)
     values = variables.Variable([1], trainable=False, dtype=dtypes.float32)
     shape = constant_op.constant([1], dtype=dtypes.int64)
 
     sp_input = sparse_tensor.SparseTensor(indices, values, shape)
     sp_output = sparse_ops.sparse_add(sp_input, sp_input)

Vulnerable code:   @test_util.run_deprecated_v1
        s = sparse_ops.sparse_add(sparse, dense)

        with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
                                    "invalid index"):
    indices = constant_op.constant([[1]], dtype=dtypes.int64)
Fixed code: #include "tensorflow/core/framework/op_requires.h"
  int64_t nnz = a_indices->dim_size(0);
  int64_t ndims = a_indices->dim_size(1);
  if (a_values->dim_size(0) != nnz) {
    return errors::InvalidArgument("Dimensions ", nnz, " and ",
                                   a_values->dim_size(0),
                                   " are not compatible");
  }
  if (a_shape->dim_size(0) != ndims) {
    return errors::InvalidArgument("Dimensions ", ndims, " and ",
                                   a_shape->dim_size(0), " are not compatible");
  }

  // Check for invalid indices.
  const auto a_indices_mat = a_indices->flat_inner_dims<Index>();

  for (int64_t zidx = 0; zidx < nnz; ++zidx) {
    for (int64_t didx = 0; didx < ndims; ++didx) {
      const Index idx = a_indices_mat(zidx, didx);
      if (idx < 0 || idx >= a_shape_flat(didx)) {
        return errors::InvalidArgument(
            "Sparse tensor has an invalid index on dimension ", didx,
            ": "
            "a_indices(",
            zidx, ",", didx, ") = ", idx,
            ", dense tensor shape: ", a_shape_flat);
      }
    }
  }

        with self.assertRaisesRegex(
            (ValueError, errors_impl.InvalidArgumentError), "invalid index"):
          s = sparse_ops.sparse_add(sparse, dense)
  def _testSparseDenseInvalidInputs(self,
                                    a_indices,
                                    a_values,
                                    a_shape,
                                    b,
                                    expected_error=""):
    # Public API call to sparse-dense add.
    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
                                expected_error):
      a = sparse_tensor.SparseTensor(a_indices, a_values, a_shape)
      self.evaluate(sparse_ops.sparse_add(a, b))
    # Directly call generated kernel, by-passing SparseTensor validation.
    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
                                expected_error):
      self.evaluate(
          sparse_ops.gen_sparse_ops.sparse_tensor_dense_add(
              a_indices, a_values, a_shape, b))

  def testSparseDenseInvalidInputs(self):
    self._testSparseDenseInvalidInputs(
        a_indices=constant_op.constant(0, shape=[17, 2], dtype=dtypes.int64),
        a_values=constant_op.constant(0, shape=[5], dtype=dtypes.float32),
        a_shape=constant_op.constant([3, 4], dtype=dtypes.int64),
        b=constant_op.constant(1, shape=[3, 4], dtype=dtypes.float32),
        expected_error="Dimensions 17 and 5 are not compatible")
    self._testSparseDenseInvalidInputs(
        a_indices=constant_op.constant(0, shape=[17, 4], dtype=dtypes.int64),
        a_values=constant_op.constant(0, shape=[17], dtype=dtypes.float32),
        a_shape=constant_op.constant([3, 4], dtype=dtypes.int64),
        b=constant_op.constant(1, shape=[3, 4], dtype=dtypes.float32),
        expected_error="Dimensions 4 and 2 are not compatible")
    self._testSparseDenseInvalidInputs(
        a_indices=constant_op.constant(7, shape=[17, 2], dtype=dtypes.int64),
        a_values=constant_op.constant(0, shape=[17], dtype=dtypes.float32),
        a_shape=constant_op.constant([3, 4], dtype=dtypes.int64),
        b=constant_op.constant(1, shape=[3, 4], dtype=dtypes.float32),
        expected_error="invalid index")

    indices = constant_op.constant([[0]], dtype=dtypes.int64)