Commit message: Fix UB in SparseTensorDenseAdd

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

PiperOrigin-RevId: 446192704
CVE: CVE-2022-29206
CWE: CWE-20
Vulnerability description: 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.

Project: tensorflow
File: sparse_tensor_dense_add_op.cc

  explicit SparseTensorDenseAddOp(OpKernelConstruction *ctx) : OpKernel(ctx) {}

Project: tensorflow
File: sparse_tensor_dense_add_op.cc

Status ValidateInputs(const Tensor *a_indices, const Tensor *a_values,
                      const Tensor *a_shape, const Tensor *b) {
  if (!TensorShapeUtils::IsMatrix(a_indices->shape())) {
    return errors::InvalidArgument(
        "Input a_indices should be a matrix but received shape: ",
        a_indices->shape().DebugString());
  }
  if (!TensorShapeUtils::IsVector(a_values->shape()) ||
      !TensorShapeUtils::IsVector(a_shape->shape())) {
    return errors::InvalidArgument(
        "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());
  }
  const auto a_shape_flat = a_shape->flat<Index>();
  for (int i = 0; i < b->dims(); ++i) {
    if (a_shape_flat(i) != b->dim_size(i)) {
      return errors::InvalidArgument(
          "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();
}

Project: tensorflow
File: sparse_tensor_dense_add_op.cc

  void Compute(OpKernelContext *ctx) override {
    const Tensor *a_indices_t, *a_values_t, *a_shape_t, *b;
    OP_REQUIRES_OK(ctx, ctx->input("a_indices", &a_indices_t));
    OP_REQUIRES_OK(ctx, ctx->input("a_values", &a_values_t));
    OP_REQUIRES_OK(ctx, ctx->input("a_shape", &a_shape_t));
    OP_REQUIRES_OK(ctx, ctx->input("b", &b));
    OP_REQUIRES_OK(
        ctx, ValidateInputs<Index>(a_indices_t, a_values_t, a_shape_t, b));

    Tensor *out_t;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, b->shape(), &out_t));

    const int ndims = static_cast<int>(a_indices_t->dim_size(1));
    const auto a_indices_mat = a_indices_t->flat_inner_dims<Index>();
    const auto a_values_flat = a_values_t->flat<T>();

    switch (ndims) {
#define NDIMS_CASE(N)                                                     \
  case N: {                                                               \
    auto out_tensor = out_t->tensor<T, N>();                              \
    out_tensor.device(ctx->eigen_device<Device>()) = b->tensor<T, N>();   \
    const Index result =                                                  \
        functor::ScatterNdFunctor<Device, T, Index, N,                    \
                                  scatter_op::UpdateOp::ADD>()(           \
            ctx->eigen_device<Device>(), a_indices_mat, a_values_flat,    \
            out_tensor);                                                  \
    OP_REQUIRES(                                                          \
        ctx, result == -1,                                                \
        errors::InvalidArgument(                                          \
            "Sparse tensor has some invalid index on dimension ", result, \
            "; dense tensor shape: ", b->shape().DebugString()));         \
  } break;

      NDIMS_CASE(1);
      NDIMS_CASE(2);
      NDIMS_CASE(3);
      NDIMS_CASE(4);
      NDIMS_CASE(5);
      default:
        OP_REQUIRES(
            ctx, false,
            errors::InvalidArgument("Only tensors with ranks between 1 and 5 "
                                    "are currently supported.  Tensor rank: ",
                                    ndims));
#undef NDIMS_CASE
    }
  }

Project: tensorflow
File: sparse_tensor_dense_add_op.cc

  Index operator()(const CPUDevice &d,
                   typename TTypes<Index>::ConstMatrix indices,
                   typename TTypes<T>::ConstFlat updates,
                   typename TTypes<T, NDIMS>::Tensor out) {
    Eigen::array<Eigen::DenseIndex, NDIMS> idx;
    const int num_nnz = static_cast<int>(indices.dimension(0));
    for (int i = 0; i < num_nnz; ++i) {
      for (int d = 0; d < NDIMS; ++d) {
        idx[d] = internal::SubtleMustCopy(indices(i, d));
        if (!FastBoundsCheck(idx[d], out.dimension(d))) {
          return d;  // on failure: d nonnegative
        }
      }
      out(idx) += updates(i);
    }
    return -1;  // on success
  }