Commit message: Properly validate sparse tensor in `SparseTensorSliceDataset`

Existing validation was incomplete.

PiperOrigin-RevId: 415375048
Change-Id: I14cd18f29ede73286f3ffac35171bd15828997e9
CVE: CVE-2022-21736
CWE: CWE-476
Vulnerability description: Tensorflow is an Open Source Machine Learning Framework. The implementation of `SparseTensorSliceDataset` has an undefined behavior: under certain condition it can be made to dereference a `nullptr` value. The 3 input arguments to `SparseTensorSliceDataset` represent a sparse tensor. However, there are some preconditions that these arguments must satisfy but these are not validated in the implementation. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.

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_t>()(i) =
            params.dataset->sparse_tensor_.shape()[i + 1];
      }
    }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

    std::shared_ptr<model::Node> CreateNode(
        IteratorContext* ctx, model::Node::Args args) const override {
      return model::MakeSourceNode(std::move(args));
    }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  Status InputDatasets(std::vector<const DatasetBase*>* inputs) const override {
    return Status::OK();
  }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  int64_t CardinalityInternal() const override {
    return sparse_tensor_.shape()[0];
  }

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_t> 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

  explicit SparseTensorSliceDatasetOp(OpKernelConstruction* ctx)
      : DatasetOpKernel(ctx) {}

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  const std::vector<PartialTensorShape>& output_shapes() const override {
    return shapes_;
  }

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 must be a matrix. Got: ",
                                        indices->shape().DebugString()));
    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(values->shape()),
                errors::InvalidArgument("Input values must be a vector. Got: ",
                                        values->shape().DebugString()));
    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(dense_shape->shape()),
                errors::InvalidArgument("Input shape must be a vector. Got: ",
                                        dense_shape->shape().DebugString()));
    OP_REQUIRES(
        ctx, values->shape().dim_size(0) == indices->shape().dim_size(0),
        errors::InvalidArgument(
            "Number of values must match first dimension of indices. ", "Got ",
            values->shape().dim_size(0),
            " values, indices shape: ", indices->shape().DebugString()));
    OP_REQUIRES(
        ctx, dense_shape->shape().dim_size(0) == indices->shape().dim_size(1),
        errors::InvalidArgument(
            "Number of dimensions must match second dimension of indices. ",
            "Got ", dense_shape->shape().dim_size(0),
            " dimensions, indices shape: ", indices->shape().DebugString()));
    OP_REQUIRES(ctx, dense_shape->NumElements() > 0,
                errors::InvalidArgument(
                    "The shape argument requires at least one element."));

    // 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_t>()(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_t, 8> std_order(dense_shape->NumElements(), 0);
    sparse::SparseTensor tensor;
    OP_REQUIRES_OK(
        ctx, sparse::SparseTensor::Create(
                 *indices, *values, TensorShape(dense_shape->vec<int64_t>()),
                 std_order, &tensor));
    *output = new Dataset<T>(ctx, std::move(tensor));
  }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  explicit Dataset(OpKernelContext* ctx,
                   const sparse::SparseTensor& sparse_tensor)
      : DatasetBase(DatasetContext(ctx)),
        sparse_tensor_(sparse_tensor),
        dtypes_({DT_INT64, sparse_tensor.dtype(), DT_INT64}),
        shapes_({{-1, sparse_tensor.dims() - 1},
                 {-1},
                 {sparse_tensor.dims() - 1}}) {}

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  string DebugString() const override {
    return "SparseTensorSliceDatasetOp::Dataset";
  }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  const DataTypeVector& output_dtypes() const override { return dtypes_; }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  Status CheckExternalState() const override { return Status::OK(); }

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_t>();
        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 SaveInternal(SerializationContext* ctx,
                        IteratorStateWriter* writer) override {
      mutex_lock l(mu_);
      TF_RETURN_IF_ERROR(writer->WriteScalar(Iterator::full_name("i"), i_));
      TF_RETURN_IF_ERROR(
          writer->WriteScalar(Iterator::full_name("iter_loc"), iter_.loc()));
      TF_RETURN_IF_ERROR(writer->WriteScalar(
          Iterator::full_name("next_non_empty_i_"), next_non_empty_i_));
      if (i_ <= next_non_empty_i_) {
        TF_RETURN_IF_ERROR(writer->WriteTensor(
            Iterator::full_name("next_indices_"), next_indices_));
        TF_RETURN_IF_ERROR(writer->WriteTensor(
            Iterator::full_name("next_values_"), next_values_));
      }
      return Status::OK();
    }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

  std::unique_ptr<IteratorBase> MakeIteratorInternal(
      const string& prefix) const override {
    return absl::make_unique<Iterator>(typename Iterator::Params{
        this, strings::StrCat(prefix, "::SparseTensorSlice")});
  }

Project: tensorflow
File: sparse_tensor_slice_dataset_op.cc

    Status RestoreInternal(IteratorContext* ctx,
                           IteratorStateReader* reader) override {
      mutex_lock l(mu_);
      TF_RETURN_IF_ERROR(reader->ReadScalar(Iterator::full_name("i"), &i_));
      int64_t iter_loc;
      TF_RETURN_IF_ERROR(
          reader->ReadScalar(Iterator::full_name("iter_loc"), &iter_loc));
      iter_ = group_iterable_.at(iter_loc);
      TF_RETURN_IF_ERROR(reader->ReadScalar(
          Iterator::full_name("next_non_empty_i_"), &next_non_empty_i_));
      if (i_ <= next_non_empty_i_) {
        TF_RETURN_IF_ERROR(reader->ReadTensor(
            Iterator::full_name("next_indices_"), &next_indices_));
        TF_RETURN_IF_ERROR(reader->ReadTensor(
            Iterator::full_name("next_values_"), &next_values_));
      }
      return Status::OK();
    }