Commit message: Validate that a and b are proper sparse tensors

PiperOrigin-RevId: 373274848
Change-Id: I3a665ac3a29dee9fb69bdf408a939330cb93ea75
CVE: CVE-2021-29607
CWE: CWE-369
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. Incomplete validation in `SparseAdd` results in allowing attackers to exploit undefined behavior (dereferencing null pointers) as well as write outside of bounds of heap allocated data. The implementation(https://github.com/tensorflow/tensorflow/blob/656e7673b14acd7835dc778867f84916c6d1cac2/tensorflow/core/kernels/sparse_sparse_binary_op_shared.cc) has a large set of validation for the two sparse tensor inputs (6 tensors in total), but does not validate that the tensors are not empty or that the second dimension of `*_indices` matches the size of corresponding `*_shape`. This allows attackers to send tensor triples that represent invalid sparse tensors to abuse code assumptions that are not protected by validation. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.

Project: tensorflow
File: sparse_sparse_binary_op_shared.cc

  void Compute(OpKernelContext *ctx) override {
    const Tensor *a_indices_t, *a_values_t, *a_shape_t, *b_indices_t,
        *b_values_t, *b_shape_t;
    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_indices", &b_indices_t));
    OP_REQUIRES_OK(ctx, ctx->input("b_values", &b_values_t));
    OP_REQUIRES_OK(ctx, ctx->input("b_shape", &b_shape_t));

    // Validations.
    OP_REQUIRES(
        ctx,
        TensorShapeUtils::IsMatrix(a_indices_t->shape()) &&
            TensorShapeUtils::IsMatrix(b_indices_t->shape()),
        errors::InvalidArgument("Inputs a_indices and b_indices should be "
                                "matrices but received shapes: ",
                                a_indices_t->shape().DebugString(), ", ",
                                b_indices_t->shape().DebugString()));
    OP_REQUIRES(ctx,
                TensorShapeUtils::IsVector(a_values_t->shape()) &&
                    TensorShapeUtils::IsVector(b_values_t->shape()),
                errors::InvalidArgument(
                    "Inputs a_values and b_values should be vectors "
                    "but received shapes: ",
                    a_values_t->shape().DebugString(), " and ",
                    b_values_t->shape().DebugString()));

    const int64 a_nnz = a_indices_t->dim_size(0);
    const int64 b_nnz = b_indices_t->dim_size(0);

    const auto a_values = a_values_t->vec<T>();
    const auto b_values = b_values_t->vec<T>();

    OP_REQUIRES(
        ctx, a_values.size() == a_nnz && b_values.size() == b_nnz,
        errors::InvalidArgument("Expected ", a_nnz, " and ", b_nnz,
                                " non-empty input values, got ",
                                a_values.size(), " and ", b_values.size()));

    OP_REQUIRES(ctx,
                TensorShapeUtils::IsVector(a_shape_t->shape()) &&
                    TensorShapeUtils::IsVector(b_shape_t->shape()),
                errors::InvalidArgument(
                    "Input shapes should be a vector but received shapes ",
                    a_shape_t->shape().DebugString(), " and ",
                    b_shape_t->shape().DebugString()));
    const int num_dims = a_indices_t->dim_size(1);
    OP_REQUIRES(
        ctx, a_shape_t->NumElements() == num_dims,
        errors::InvalidArgument("Second dimension of a_indices and length of "
                                "a_shape must match, got ",
                                num_dims, " and ", a_shape_t->NumElements()));
    OP_REQUIRES(ctx, num_dims > 0,
                errors::InvalidArgument("Tensors must not be empty"));
    OP_REQUIRES(ctx, a_shape_t->IsSameSize(*b_shape_t),
                errors::InvalidArgument(
                    "Operands do not have the same ranks; got shapes: ",
                    a_shape_t->SummarizeValue(10), " and ",
                    b_shape_t->SummarizeValue(10)));
    const auto a_shape = a_shape_t->flat<int64>();
    const auto b_shape = b_shape_t->flat<int64>();
    for (int i = 0; i < a_shape_t->NumElements(); ++i) {
      OP_REQUIRES(ctx, a_shape(i) == b_shape(i),
                  errors::InvalidArgument("Operands' shapes do not match: got ",
                                          a_shape(i), " and ", b_shape(i),
                                          " for dimension ", i));
    }

    const auto a_indices_mat = a_indices_t->matrix<int64>();
    const auto b_indices_mat = b_indices_t->matrix<int64>();
    std::vector<T> a_augmented_values, b_augmented_values;
    std::vector<std::pair<bool, int64>> entries_to_copy;  // from_a?, idx
    UnionSparseIndicesAndValues(a_indices_mat, a_values, a_nnz, b_indices_mat,
                                b_values, b_nnz, num_dims, &a_augmented_values,
                                &b_augmented_values, &entries_to_copy);

    // Allocates and fills output tensors.
    const int64 sum_nnz = a_augmented_values.size();
    Tensor *output_indices_t, *output_values_t;
    OP_REQUIRES_OK(ctx,
                   ctx->allocate_output(0, TensorShape({sum_nnz, num_dims}),
                                        &output_indices_t));
    OP_REQUIRES_OK(
        ctx, ctx->allocate_output(1, TensorShape({sum_nnz}), &output_values_t));
    auto output_indices_mat = output_indices_t->matrix<int64>();

    for (int64 i = 0; i < sum_nnz; ++i) {
      const bool from_a = entries_to_copy[i].first;
      const int64 idx = entries_to_copy[i].second;
      output_indices_mat.chip<0>(i) =
          from_a ? a_indices_mat.chip<0>(idx) : b_indices_mat.chip<0>(idx);
    }

    // Performs the functor operation using Eigen.
    //
    // Note that the two stack-allocated std::vector's may not be aligned. Using
    // allocate_temp() would've given us aligned storage, but we do not know
    // their sizes in advance, so we couldn't use allocate_temp() anyway.
    //
    // TODO(zongheng): measure if it's worthwhile to somehow force alignment.
    using UnalignedTensorMap =
        Eigen::TensorMap<Eigen::Tensor<const T, 1, Eigen::RowMajor>,
                         Eigen::Unaligned>;
    auto a_augmented_values_t =
        UnalignedTensorMap(a_augmented_values.data(), sum_nnz);
    auto b_augmented_values_t =
        UnalignedTensorMap(b_augmented_values.data(), sum_nnz);
    output_values_t->flat<T>().device(ctx->eigen_device<Device>()) =
        a_augmented_values_t.binaryExpr(b_augmented_values_t,
                                        typename Functor::func());
  }

Project: tensorflow
File: sparse_sparse_binary_op_shared.cc

void UnionSparseIndicesAndValues(
    typename TTypes<int64>::ConstMatrix a_indices_mat,
    typename TTypes<T>::ConstFlat a_values, int64 a_nnz,
    typename TTypes<int64>::ConstMatrix b_indices_mat,
    typename TTypes<T>::ConstFlat b_values, int64 b_nnz, int num_dims,
    std::vector<T> *a_augmented_values, std::vector<T> *b_augmented_values,
    std::vector<std::pair<bool, int64>> *entries_to_copy) {
  entries_to_copy->reserve(a_nnz + b_nnz);
  a_augmented_values->reserve(a_nnz);
  b_augmented_values->reserve(b_nnz);

  int64 i = 0, j = 0;
  const T kZero = T(0);
  while (i < a_nnz && j < b_nnz) {
    switch (sparse::DimComparator::cmp(a_indices_mat, b_indices_mat, i, j,
                                       num_dims)) {
      case -1:
        entries_to_copy->emplace_back(true, i);
        a_augmented_values->push_back(a_values(i));
        b_augmented_values->push_back(kZero);
        ++i;
        break;
      case 0:
        entries_to_copy->emplace_back(true, i);
        a_augmented_values->push_back(a_values(i));
        b_augmented_values->push_back(b_values(j));
        ++i;
        ++j;
        break;
      case 1:
        entries_to_copy->emplace_back(false, j);
        a_augmented_values->push_back(kZero);
        b_augmented_values->push_back(b_values(j));
        ++j;
        break;
    }
  }
  // Handles leftovers; at most one loop runs.
  while (i < a_nnz) {
    entries_to_copy->emplace_back(/* is_a */ true, i);
    a_augmented_values->push_back(a_values(i++));
    b_augmented_values->push_back(kZero);
  }
  while (j < b_nnz) {
    entries_to_copy->emplace_back(/* is_a */ false, j);
    a_augmented_values->push_back(kZero);
    b_augmented_values->push_back(b_values(j++));
  }
}

Project: tensorflow
File: sparse_sparse_binary_op_shared.cc

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