From 11ced8467eccad9c7cb94867708be8fa5c66c730 Mon Sep 17 00:00:00 2001 From: Antonio Sanchez Date: Tue, 3 May 2022 07:51:51 -0700 Subject: [PATCH] Fix UB in SparseTensorDenseAdd Added more input validation to avoid nullptr dereferencing and array index out of bounds issues. PiperOrigin-RevId: 446192704 --- .../kernels/sparse_tensor_dense_add_op.cc | 30 ++++++++++++ .../sparse_ops/sparse_add_op_test.py | 46 +++++++++++++++++-- .../sparse_ops/sparse_ops_test.py | 2 +- 3 files changed, 72 insertions(+), 6 deletions(-) diff --git a/tensorflow/core/kernels/sparse_tensor_dense_add_op.cc b/tensorflow/core/kernels/sparse_tensor_dense_add_op.cc index 48803e4b939800..6d6b05bf70f30a 100644 --- a/tensorflow/core/kernels/sparse_tensor_dense_add_op.cc +++ b/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(); + + 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(); } 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 61ad45fb5e273e..821184af5b5699 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 @@ -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 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 4972d1d25f08f1..684d1f98432b53 100644 --- a/tensorflow/python/kernel_tests/sparse_ops/sparse_ops_test.py +++ b/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)