CVE-2022-29198: 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.SparseTensorToCSRSparseMatrix` does not fully validate the input arguments. This results in a `CHECK`-failure which can be used to trigger a denial of service attack. The code assumes `dense_shape` is a vector and `indices` is a matrix (as part of requirements for sparse tensors) but there is no validation for this. 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-20T22:16Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/sparse/sparse_tensor_to_csr_sparse_matrix_op.cc, tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_ops_test.py
Commit date: 2022-05-02 16:16:54 -0700
Commit message: Fix failed check in SparseTensorToCSRSparseMatrix

Security vulnerability fix. A `CHECK` fails if inputing either an empty `dense_shape`,
or a non-rank-2 `indices`. Added appropriate checks and tests.

PiperOrigin-RevId: 446053984

Diff: diff --git a/tensorflow/core/kernels/sparse/sparse_tensor_to_csr_sparse_matrix_op.cc b/tensorflow/core/kernels/sparse/sparse_tensor_to_csr_sparse_matrix_op.cc
index 36c9f85e197d3..acb201893ce65 100644
--- a/tensorflow/core/kernels/sparse/sparse_tensor_to_csr_sparse_matrix_op.cc
+++ b/tensorflow/core/kernels/sparse/sparse_tensor_to_csr_sparse_matrix_op.cc
@@ -65,10 +65,17 @@ class SparseTensorToCSRSparseMatrixCPUOp : public OpKernel {
   void Compute(OpKernelContext* ctx) final {
     const Tensor& indices = ctx->input(0);
     const Tensor& values = ctx->input(1);
     const Tensor& dense_shape = ctx->input(2);
     const int rank = dense_shape.NumElements();
+    OP_REQUIRES(
+        ctx, TensorShapeUtils::IsVector(dense_shape.shape()),
+        errors::InvalidArgument("dense_shape must be rank 1 but got rank",
+                                dense_shape.shape().dims()));
+    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(indices.shape()),
+                errors::InvalidArgument("indices must be rank 2 but got rank",
+                                        indices.shape().dims()));
     OP_REQUIRES(ctx, rank == 2 || rank == 3,
                 errors::InvalidArgument("SparseTensor must have rank 2 or 3; ",
                                         "but indices has rank: ", rank));
     auto dense_shape_vec = dense_shape.vec<int64_t>();
     const int64_t batch_size = (rank == 2) ? 1 : dense_shape_vec(0);
diff --git a/tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_ops_test.py b/tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_ops_test.py
index 0e72a00ef1c5e..035791ce0a5bf 100644
--- a/tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_ops_test.py
+++ b/tensorflow/python/kernel_tests/linalg/sparse/csr_sparse_matrix_ops_test.py
@@ -166,10 +166,29 @@ class CSRSparseMatrixOpsTest(test.TestCase):
 
     self.assertAllEqual(a_indices, a_st_rt_value.indices)
     self.assertAllClose(a_values, a_st_rt_value.values)
     self.assertAllEqual(a_dense_shape, a_st_rt_value.dense_shape)
 
+  def testSparseTensorConversionInvalidInputShapes(self):
+    values = constant_op.constant(
+        0.554979503, shape=[5], dtype=dtypes.float32)
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 1"):
+      indices = constant_op.constant(0, shape=[5, 2], dtype=dtypes.int64)
+      dense_shape = constant_op.constant(53, shape=[], dtype=dtypes.int64)
+      csr = sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
+          indices=indices, values=values, dense_shape=dense_shape)
+      self.evaluate(csr)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 2"):
+      indices = constant_op.constant(0, shape=[5], dtype=dtypes.int64)
+      dense_shape = constant_op.constant(53, shape=[1], dtype=dtypes.int64)
+      csr = sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
+          indices=indices, values=values, dense_shape=dense_shape)
+      self.evaluate(csr)
+
   # TODO(b/139491352): Add handle_data propagation to array_ops.identity.
   @test_util.run_deprecated_v1
   def testCSRSparseMatrixResourceVariable(self):
     if not self._gpu_available:
       return

Fixed code:     OP_REQUIRES(
        ctx, TensorShapeUtils::IsVector(dense_shape.shape()),
        errors::InvalidArgument("dense_shape must be rank 1 but got rank",
                                dense_shape.shape().dims()));
    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(indices.shape()),
                errors::InvalidArgument("indices must be rank 2 but got rank",
                                        indices.shape().dims()));
  def testSparseTensorConversionInvalidInputShapes(self):
    values = constant_op.constant(
        0.554979503, shape=[5], dtype=dtypes.float32)
    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 1"):
      indices = constant_op.constant(0, shape=[5, 2], dtype=dtypes.int64)
      dense_shape = constant_op.constant(53, shape=[], dtype=dtypes.int64)
      csr = sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
          indices=indices, values=values, dense_shape=dense_shape)
      self.evaluate(csr)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 2"):
      indices = constant_op.constant(0, shape=[5], dtype=dtypes.int64)
      dense_shape = constant_op.constant(53, shape=[1], dtype=dtypes.int64)
      csr = sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
          indices=indices, values=values, dense_shape=dense_shape)
      self.evaluate(csr)