CVE-2022-29206: Improper Input Validation (CWE-20)
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 receives input or data, but it does
        not validate or incorrectly validates that the input has the
        properties that are required to process the data safely and
        correctly.
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)