Commit message: Prevent nullptr deref in SparseTensorSliceDataset

The arguments must determine a valid sparse tensor. This means that when indices are empty then the values must be empty too (and the reverse).

Also added test, by modifying existing test with empty sparse tensor to now run with an invalid sparse tensor input.

PiperOrigin-RevId: 388562757
Change-Id: Id8b54cd7c2316025b4f9a77292c8fb5344d17609
CVE: CVE-2021-37647
CWE: CWE-476
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. When a user does not supply arguments that determine a valid sparse tensor, `tf.raw_ops.SparseTensorSliceDataset` implementation can be made to dereference a null pointer. The [implementation](https://github.com/tensorflow/tensorflow/blob/8d72537c6abf5a44103b57b9c2e22c14f5f49698/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc#L240-L251) has some argument validation but fails to consider the case when either `indices` or `values` are provided for an empty sparse tensor when the other is not. If `indices` is empty, then [code that performs validation](https://github.com/tensorflow/tensorflow/blob/8d72537c6abf5a44103b57b9c2e22c14f5f49698/tensorflow/core/kernels/data/sparse_tensor_slice_dataset_op.cc#L260-L261) (i.e., checking that the indices are monotonically increasing) results in a null pointer dereference. If `indices` as provided by the user is empty, then `indices` in the C++ code above is backed by an empty `std::vector`, hence calling `indices->dim_size(0)` results in null pointer dereferencing (same as calling `std::vector::at()` on an empty vector). We have patched the issue in GitHub commit 02cc160e29d20631de3859c6653184e3f876b9d7. The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

    Status GetNextInternal(IteratorContext* ctx,
                           std::vector<Tensor>* out_tensors,
                           bool* end_of_sequence) override {
      mutex_lock l(mu_);
      if (i_ == num_elements_) {
        *end_of_sequence = true;
        return Status::OK();
      }

      out_tensors->clear();
      out_tensors->reserve(3);
      const int rank = Iterator::dataset()->sparse_tensor_.dims();

      if (i_ > next_non_empty_i_ && iter_ != group_iterable_.end()) {
        // We still have elements to consume from `group_iterable_`
        // and we have emitted all elements up to and including the
        // current position.
        sparse::Group group = *iter_;
        const auto indices = group.indices();
        const auto values = group.values<T>();
        const int64_t num_entries = values.size();
        next_non_empty_i_ = indices(0, 0);

        next_indices_ = Tensor(DT_INT64, {num_entries, rank - 1});
        next_values_ = Tensor(DataTypeToEnum<T>::value, {num_entries});

        auto next_indices_t = next_indices_.matrix<int64>();
        auto next_values_t = next_values_.vec<T>();

        for (int64_t i = 0; i < num_entries; ++i) {
          for (int d = 1; d < rank; ++d) {
            next_indices_t(i, d - 1) = indices(i, d);
          }
          next_values_t(i) = values(i);
        }

        ++iter_;
      }
      if (i_ == next_non_empty_i_) {
        // The current position is non-empty in the input
        // `SparseTensor`, and we have already read the value from the
        // `GroupIterable`.
        out_tensors->push_back(std::move(next_indices_));
        out_tensors->push_back(std::move(next_values_));
        out_tensors->push_back(dense_shape_);
        next_non_empty_i_ = kNextNonEmptyUnknown;
      } else {
        DCHECK(i_ < next_non_empty_i_ || iter_ == group_iterable_.end());
        // The current position is empty in the input `SparseTensor`,
        // so emit empty indices and values.
        out_tensors->push_back(Tensor(DT_INT64, TensorShape({0, rank - 1})));
        out_tensors->push_back(Tensor(DataTypeToEnum<T>::value, {0}));
        out_tensors->push_back(dense_shape_);
      }

      ++i_;
      *end_of_sequence = false;
      return Status::OK();
    }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  Status AsGraphDefInternal(SerializationContext* ctx,
                            DatasetGraphDefBuilder* b,
                            Node** output) const override {
    Node* indices_node;
    TF_RETURN_IF_ERROR(b->AddTensor(sparse_tensor_.indices(), &indices_node));
    Node* value_node;
    TF_RETURN_IF_ERROR(b->AddTensor(sparse_tensor_.values(), &value_node));
    Node* dense_shape_node;
    std::vector<int64> dense_shape;
    dense_shape.reserve(sparse_tensor_.shape().size());
    for (int i = 0; i < sparse_tensor_.shape().size(); i++)
      dense_shape.emplace_back(sparse_tensor_.shape()[i]);
    TF_RETURN_IF_ERROR(b->AddVector(dense_shape, &dense_shape_node));
    AttrValue val_dtype;
    b->BuildAttrValue(sparse_tensor_.dtype(), &val_dtype);
    TF_RETURN_IF_ERROR(
        b->AddDataset(this, {indices_node, value_node, dense_shape_node},
                      {{"Tvalues", val_dtype}}, output));
    return Status::OK();
  }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  void MakeDataset(OpKernelContext* ctx, DatasetBase** output) override {
    // Create a new SparseTensorSliceDatasetOp::Dataset, insert it in
    // the step container, and return it as the output.
    const Tensor* indices;
    OP_REQUIRES_OK(ctx, ctx->input("indices", &indices));
    const Tensor* values;
    OP_REQUIRES_OK(ctx, ctx->input("values", &values));
    const Tensor* dense_shape;
    OP_REQUIRES_OK(ctx, ctx->input("dense_shape", &dense_shape));

    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(indices->shape()),
                errors::InvalidArgument(
                    "Input indices should be a matrix but received shape ",
                    indices->shape().DebugString()));

    const auto num_indices = indices->NumElements();
    const auto num_values = values->NumElements();
    if (num_indices == 0 || num_values == 0) {
      OP_REQUIRES(ctx, num_indices == num_values,
                  errors::InvalidArgument(
                      "If indices or values are empty, the other one must also "
                      "be. Got indices of shape ",
                      indices->shape().DebugString(), " and values of shape ",
                      values->shape().DebugString()));
    }
    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(values->shape()),
                errors::InvalidArgument(
                    "Input values should be a vector but received shape ",
                    indices->shape().DebugString()));
    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(dense_shape->shape()),
                errors::InvalidArgument(
                    "Input shape should be a vector but received shape ",
                    dense_shape->shape().DebugString()));

    // We currently ensure that `sparse_tensor` is ordered in the
    // batch dimension.
    // TODO(mrry): Investigate ways to avoid this unconditional check
    // if we can be sure that the sparse tensor was produced in an
    // appropriate order (e.g. by `tf.parse_example()` or a Dataset
    // that batches elements into rows of a SparseTensor).
    int64_t previous_batch_index = -1;
    for (int64_t i = 0; i < indices->dim_size(0); ++i) {
      int64_t next_batch_index = indices->matrix<int64>()(i, 0);
      OP_REQUIRES(
          ctx, next_batch_index >= previous_batch_index,
          errors::Unimplemented("The SparseTensor must be ordered in the batch "
                                "dimension; handling arbitrarily ordered input "
                                "is not currently supported."));
      previous_batch_index = next_batch_index;
    }
    gtl::InlinedVector<int64, 8> std_order(dense_shape->NumElements(), 0);
    sparse::SparseTensor tensor;
    OP_REQUIRES_OK(
        ctx, sparse::SparseTensor::Create(
                 *indices, *values, TensorShape(dense_shape->vec<int64>()),
                 std_order, &tensor));
    *output = new Dataset<T>(ctx, std::move(tensor));
  }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  int64 Cardinality() const override { return sparse_tensor_.shape()[0]; }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

    explicit Iterator(const typename Iterator::Params& params)
        : DatasetIterator<Dataset<T>>(params),
          num_elements_(params.dataset->sparse_tensor_.shape()[0]),
          dense_shape_(DT_INT64, {params.dataset->sparse_tensor_.dims() - 1}),
          group_iterable_(params.dataset->sparse_tensor_.group({0})),
          iter_(group_iterable_.begin()) {
      for (size_t i = 0; i < dense_shape_.NumElements(); ++i) {
        dense_shape_.vec<int64>()(i) =
            params.dataset->sparse_tensor_.shape()[i + 1];
      }
    }