Fix UB in SparseTensorDenseAdd

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

PiperOrigin-RevId: 446192704

tensorflow/core/kernels/sparse_tensor_dense_add_op.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #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"
@@ -47,6 +48,17 @@ Status ValidateInputs(const Tensor *a_indices, const Tensor *a_values,
         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(),
@@ -61,6 +73,24 @@ Status ValidateInputs(const Tensor *a_indices, const Tensor *a_values,
           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();
 }
 

tensorflow/python/kernel_tests/sparse_ops/sparse_add_op_test.py

@@ -189,7 +189,6 @@ def testSparseTensorDenseAddGradients(self):
                                                     [(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]
@@ -201,12 +200,49 @@ def testInvalidSparseTensor(self):
           [[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
 
 

tensorflow/python/kernel_tests/sparse_ops/sparse_ops_test.py

@@ -665,7 +665,7 @@ def testInvalidIndices(self):
 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)