Commit message: Fix out of bounds read in `ragged_cross_op.cc`.

PiperOrigin-RevId: 369757702
Change-Id: Ie6e5d2c21513a8d56bf41fcf35960caf76e890f9
CVE: CVE-2021-29532
CWE: CWE-369
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. An attacker can force accesses outside the bounds of heap allocated arrays by passing in invalid tensor values to `tf.raw_ops.RaggedCross`. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/efea03b38fb8d3b81762237dc85e579cc5fc6e87/tensorflow/core/kernels/ragged_cross_op.cc#L456-L487) lacks validation for the user supplied arguments. Each of the above branches call a helper function after accessing array elements via a `*_list[next_*]` pattern, followed by incrementing the `next_*` index. However, as there is no validation that the `next_*` values are in the valid range for the corresponding `*_list` arrays, this results in heap OOB reads. 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: ragged_cross_op.cc

  static Status BuildSparseFeatureReader(const Tensor& indices,
                                         const Tensor& values, int64 batch_size,
                                         FeatureReaders* features) {
    if (values.dtype() == DT_INT64) {
      features->emplace_back(
          new SparseFeatureReader<int64>(indices, values, batch_size));
    } else if (values.dtype() == DT_STRING) {
      features->emplace_back(
          new SparseFeatureReader<tstring>(indices, values, batch_size));
    } else {
      return errors::InvalidArgument("Unexpected dtype for input ",
                                     (features->size() + 1), ": ",
                                     values.dtype());
    }
    return Status::OK();
  }

Project: tensorflow
File: ragged_cross_op.cc

  void Compute(OpKernelContext* context) override {
    OpInputList ragged_values_list;
    OpInputList ragged_splits_list;
    OpInputList sparse_indices_list;
    OpInputList sparse_values_list;
    OpInputList sparse_shape_list;
    OpInputList dense_list;
    OP_REQUIRES_OK(context,
                   context->input_list("ragged_values", &ragged_values_list));
    OP_REQUIRES_OK(
        context, context->input_list("ragged_row_splits", &ragged_splits_list));
    OP_REQUIRES_OK(context,
                   context->input_list("sparse_indices", &sparse_indices_list));
    OP_REQUIRES_OK(context,
                   context->input_list("sparse_values", &sparse_values_list));
    OP_REQUIRES_OK(context,
                   context->input_list("sparse_shape", &sparse_shape_list));
    OP_REQUIRES_OK(context, context->input_list("dense_inputs", &dense_list));
    OP_REQUIRES_OK(context,
                   ValidateInput(ragged_values_list, ragged_splits_list,
                                 sparse_indices_list, sparse_values_list,
                                 sparse_shape_list, dense_list));

    int64 batch_size =
        CalculateBatchSize(ragged_splits_list, sparse_shape_list, dense_list);

    FeatureReaders features;
    OP_REQUIRES_OK(context,
                   BuildFeatureReaders(ragged_values_list, ragged_splits_list,
                                       sparse_indices_list, sparse_values_list,
                                       dense_list, batch_size, &features));

    Tensor* values_out;
    Tensor* row_splits_out;
    OP_REQUIRES_OK(context, BuildOutputTensors(features, batch_size, context,
                                               &values_out, &row_splits_out));

    std::unique_ptr<OutputWriter> output_writer = MakeOutputWriter(
        features, num_buckets_, hash_key_, row_splits_out, values_out);

    auto do_work = [&output_writer](int64 begin, int64 end) {
      output_writer->WriteOutputSlice(begin, end);
    };

    // TODO(edloper): optimize cost_per_batch
    const int cost_per_batch = 5000 * ragged_values_list.size();
    auto thread_pool =
        context->device()->tensorflow_cpu_worker_threads()->workers;
    thread_pool->ParallelFor(batch_size, cost_per_batch, do_work);
  }

Project: tensorflow
File: ragged_cross_op.cc

void CopyToString(const tstring& src, tstring* dst) {
  if (src.type() == tstring::SMALL) {
    *dst = src;  // string buffer fits in the tstring object (under ~24 bytes)
  } else {
    dst->assign_as_view(src);
  }
}

Project: tensorflow
File: ragged_cross_op.cc

void CopyToFingerprint(int64 feature, uint64* dst) { *dst = feature; }

Project: tensorflow
File: ragged_cross_op.cc

  void ReadValue(int64 batch, int64 n, tstring* out) const override {
    CopyToString(values_(row_splits_(batch) + n), out);
  }

Project: tensorflow
File: ragged_cross_op.cc

  static Status BuildRaggedFeatureReader(const Tensor& values,
                                         const Tensor& splits,
                                         FeatureReaders* features) {
    if (values.dtype() != DT_INT64 && values.dtype() != DT_STRING) {
      return errors::InvalidArgument("Unexpected dtype for input ",
                                     (features->size() + 1), ": ",
                                     values.dtype());
    }
    if (splits.dtype() != DT_INT64 && splits.dtype() != DT_INT32) {
      return errors::InvalidArgument("Unexpected row_splits.dtype for input ",
                                     (features->size() + 1), ": ",
                                     values.dtype());
    }
    if (values.dtype() == DT_INT64) {
      if (splits.dtype() == DT_INT64) {
        features->emplace_back(
            new RaggedFeatureReader<int64, int64>(values, splits));
      } else {
        features->emplace_back(
            new RaggedFeatureReader<int64, int32>(values, splits));
      }
    } else {
      if (splits.dtype() == DT_INT64) {
        features->emplace_back(
            new RaggedFeatureReader<tstring, int64>(values, splits));
      } else {
        features->emplace_back(
            new RaggedFeatureReader<tstring, int32>(values, splits));
      }
    }
    return Status::OK();
  }

Project: tensorflow
File: ragged_cross_op.cc

  static Status BuildDenseFeatureReader(const Tensor& values,
                                        FeatureReaders* features) {
    if (values.dtype() == DT_INT64) {
      features->emplace_back(new DenseFeatureReader<int64>(values));
    } else if (values.dtype() == DT_STRING) {
      features->emplace_back(new DenseFeatureReader<tstring>(values));
    } else {
      return errors::InvalidArgument("Unexpected dtype for input ",
                                     (features->size() + 1), ": ",
                                     values.dtype());
    }
    return Status::OK();
  }

Project: tensorflow
File: ragged_cross_op.cc

  Status ValidateInput(const OpInputList& ragged_values_list,
                       const OpInputList& ragged_splits_list,
                       const OpInputList& sparse_indices_list,
                       const OpInputList& sparse_values_list,
                       const OpInputList& sparse_shape_list,
                       const OpInputList& dense_list) {
    const auto num_ragged = ragged_values_list.size();
    const auto num_sparse = sparse_indices_list.size();

    // Validate tensor shapes.
    for (int i = 0; i < num_ragged; ++i) {
      if (!TensorShapeUtils::IsVector(ragged_values_list[i].shape())) {
        return errors::InvalidArgument(
            "tf.ragged.cross only supports inputs with rank=2.");
      }
      if (!TensorShapeUtils::IsVector(ragged_splits_list[i].shape()) ||
          (ragged_splits_list[i].NumElements() == 0)) {
        return errors::InvalidArgument("Invalid RaggedTensor");
      }
    }
    for (int i = 0; i < num_sparse; ++i) {
      if (!TensorShapeUtils::IsMatrix(sparse_indices_list[i].shape()) ||
          !TensorShapeUtils::IsVector(sparse_values_list[i].shape()) ||
          !TensorShapeUtils::IsVector(sparse_shape_list[i].shape())) {
        return errors::InvalidArgument("Invalid SparseTensor ", i);
      }
      if (sparse_shape_list[i].NumElements() != 2) {
        return errors::InvalidArgument(
            "tf.ragged.cross only supports inputs with rank=2.");
      }
    }
    for (int i = 0; i < dense_list.size(); ++i) {
      if (!TensorShapeUtils::IsMatrix(dense_list[i].shape())) {
        return errors::InvalidArgument(
            "tf.ragged.cross only supports inputs with rank=2.");
      }
    }

    // Check that batch sizes are consistent.
    int64 batch_size =
        CalculateBatchSize(ragged_splits_list, sparse_shape_list, dense_list);
    for (int i = 0; i < num_ragged; ++i) {
      if (ragged_splits_list[i].NumElements() - 1 != batch_size) {
        return errors::InvalidArgument(
            "inputs must all have the same batch dimension size.");
      }
    }
    for (int i = 0; i < num_sparse; ++i) {
      if (sparse_shape_list[i].flat<int64>()(0) != batch_size) {
        return errors::InvalidArgument(
            "inputs must all have the same batch dimension size.");
      }
    }
    for (int i = 0; i < dense_list.size(); ++i) {
      if (dense_list[i].dim_size(0) != batch_size) {
        return errors::InvalidArgument(
            "inputs must all have the same batch dimension size.");
      }
    }

    return Status::OK();
  }

Project: tensorflow
File: ragged_cross_op.cc

  void ReadValue(int64 batch, int64 n, uint64* out) const override {
    CopyToFingerprint(values_(row_splits_[batch] + n), out);
  }

Project: tensorflow
File: ragged_cross_op.cc

  void ReadValue(int64 batch, int64 n, uint64* out) const override {
    CopyToFingerprint(values_(row_splits_(batch) + n), out);
  }

Project: tensorflow
File: ragged_cross_op.cc

void CopyToFingerprint(const tstring& feature, uint64* dst) {
  *dst = Fingerprint64(feature);
}

Project: tensorflow
File: ragged_cross_op.cc

  int64 FeatureCount(int64 batch) const override {
    return row_splits_[batch + 1] - row_splits_[batch];
  }

Project: tensorflow
File: ragged_cross_op.cc

  void WriteCombination(int64 batch_index, const std::vector<int>& combination,
                        tstring* out) {
    static const auto k_feature_separator = "_X_";
    gtl::InlinedVector<tstring, 6> cross_vec(features_.size());
    for (int i = 0; i < combination.size(); ++i) {
      features_[i]->ReadValue(batch_index, combination[i], &cross_vec[i]);
    }
    *out = absl::StrJoin(cross_vec, k_feature_separator);
  }

Project: tensorflow
File: ragged_cross_op.cc

  Status BuildOutputTensors(const FeatureReaders& features, int64 batch_size,
                            OpKernelContext* context, Tensor** values_out,
                            Tensor** row_splits_out) {
    // Allocate and populate the row_splits output tensor.
    TF_RETURN_IF_ERROR(context->allocate_output(
        1, TensorShape({batch_size + 1}), row_splits_out));
    auto flat_row_splits = (*row_splits_out)->flat<SplitsType>();
    int64 cross_count_total = 0;
    flat_row_splits(0) = 0;
    for (int64 b = 0; b < batch_size; b++) {
      cross_count_total += CrossCountByBatchIndex(features, b);
      flat_row_splits(b + 1) = cross_count_total;
    }

    // Allocate the values output tensor.
    TF_RETURN_IF_ERROR(context->allocate_output(
        0, TensorShape({cross_count_total}), values_out));

    return Status::OK();
  }

Project: tensorflow
File: ragged_cross_op.cc

std::unique_ptr<OutputWriter> MakeOutputWriter(const FeatureReaders& features,
                                               int64 num_buckets,
                                               uint64 hash_key,
                                               const Tensor* splits_out,
                                               Tensor* values_out) {
  if (values_out->dtype() == DT_INT64) {
    if (splits_out->dtype() == DT_INT64) {
      return std::make_unique<OutputWriterImpl<int64, int64>>(
          features, num_buckets, hash_key, splits_out, values_out);
    } else {
      return std::make_unique<OutputWriterImpl<int64, int32>>(
          features, num_buckets, hash_key, splits_out, values_out);
    }
  } else {
    if (splits_out->dtype() == DT_INT64) {
      return std::make_unique<OutputWriterImpl<tstring, int64>>(
          features, num_buckets, hash_key, splits_out, values_out);
    } else {
      return std::make_unique<OutputWriterImpl<tstring, int32>>(
          features, num_buckets, hash_key, splits_out, values_out);
    }
  }
}

Project: tensorflow
File: ragged_cross_op.cc

void CopyToString(int64 src, tstring* dst) { *dst = std::to_string(src); }

Project: tensorflow
File: ragged_cross_op.cc

  explicit DenseFeatureReader(const Tensor& tensor)
      : values_(tensor.matrix<ValuesType>()),
        feature_count_(tensor.dim_size(1)) {}

Project: tensorflow
File: ragged_cross_op.cc

  OutputWriterImpl(const FeatureReaders& features, int64 num_buckets,
                   uint64 hash_key, const Tensor* splits_out,
                   Tensor* values_out)
      : features_(features),
        num_buckets_(num_buckets),
        hash_key_(hash_key),
        splits_out_(splits_out->flat<SplitsType>()),
        values_out_(values_out->flat<ValuesType>()) {}

Project: tensorflow
File: ragged_cross_op.cc

  virtual ~FeatureReader() {}

Project: tensorflow
File: ragged_cross_op.cc

  SparseFeatureReader(const Tensor& indices_t, const Tensor& values_t,
                      int64 batch_size)
      : values_(values_t.flat<ValuesType>()) {
    row_splits_.reserve(batch_size + 1);
    row_splits_.push_back(0);
    auto indices = indices_t.matrix<int64>();
    int64 num_values = values_.size();
    int64 i = 0;  // value index
    for (int row = 0; row < batch_size; row++) {
      while (i < num_values && indices(i, 0) <= row) ++i;
      row_splits_.push_back(i);
    }
  }

Project: tensorflow
File: ragged_cross_op.cc

  void WriteOutputSlice(int64 begin, int64 end) override {
    std::vector<int> combination(features_.size(), 0);
    for (int64 b = begin; b < end; ++b) {
      auto row_start = splits_out_(b);
      auto row_limit = splits_out_(b + 1);
      for (auto i = row_start; i < row_limit; ++i) {
        WriteCombination(b, combination, &values_out_(i));
        NextCombination(b, &combination);
      }
      combination.assign(features_.size(), 0);  // reset for next batch.
    }
  }

Project: tensorflow
File: ragged_cross_op.cc

  RaggedFeatureReader(const Tensor& values, const Tensor& row_splits)
      : values_(values.flat<ValuesType>()),
        row_splits_(row_splits.flat<SplitsType>()) {}

Project: tensorflow
File: ragged_cross_op.cc

  Status BuildFeatureReaders(const OpInputList& ragged_values_list,
                             const OpInputList& ragged_splits_list,
                             const OpInputList& sparse_indices_list,
                             const OpInputList& sparse_values_list,
                             const OpInputList& dense_list, int64 batch_size,
                             FeatureReaders* features) {
    features->reserve(input_order_.size());

    int next_ragged = 0;
    int next_sparse = 0;
    int next_dense = 0;
    for (char c : input_order_) {
      if (c == 'R') {
        if (next_ragged >= ragged_values_list.size())
          return errors::InvalidArgument(
              "input_order \"", input_order_,
              "\" specifies reading a ragged tensor value at index ",
              next_ragged, " from a list of ", ragged_values_list.size(),
              " values.");
        if (next_ragged >= ragged_splits_list.size())
          return errors::InvalidArgument(
              "input_order \"", input_order_,
              "\" specifies reading a ragged tensor split at index ",
              next_ragged, " from a list of ", ragged_splits_list.size(),
              " splits.");
        TF_RETURN_IF_ERROR(BuildRaggedFeatureReader(
            ragged_values_list[next_ragged], ragged_splits_list[next_ragged],
            features));
        next_ragged++;
      } else if (c == 'S') {
        if (next_sparse >= sparse_values_list.size())
          return errors::InvalidArgument(
              "input_order \"", input_order_,
              "\" specifies reading a sparse tensor value at index ",
              next_sparse, " from a list of ", sparse_values_list.size(),
              " values.");
        if (next_sparse >= sparse_indices_list.size())
          return errors::InvalidArgument(
              "input_order \"", input_order_,
              "\" specifies reading a sparse tensor index at index ",
              next_sparse, " from a list of ", sparse_indices_list.size(),
              " indices.");
        TF_RETURN_IF_ERROR(BuildSparseFeatureReader(
            sparse_indices_list[next_sparse], sparse_values_list[next_sparse],
            batch_size, features));
        next_sparse++;
      } else if (c == 'D') {
        if (next_dense >= dense_list.size())
          return errors::InvalidArgument(
              "input_order \"", input_order_,
              "\" specifies reading a dense tensor at index ", next_dense,
              " from a list of ", dense_list.size(), " tensors.");
        TF_RETURN_IF_ERROR(
            BuildDenseFeatureReader(dense_list[next_dense++], features));
      } else {
        return errors::InvalidArgument("Unexpected input_order value.");
      }
    }

    return Status::OK();
  }

Project: tensorflow
File: ragged_cross_op.cc

  int64 CalculateBatchSize(const OpInputList& ragged_splits_list,
                           const OpInputList& sparse_shape_list,
                           const OpInputList& dense_list) {
    if (ragged_splits_list.size() > 0) {
      return ragged_splits_list[0].NumElements() - 1;
    } else if (dense_list.size() > 0) {
      return dense_list[0].dim_size(0);
    } else if (sparse_shape_list.size() > 0) {
      return sparse_shape_list[0].flat<int64>()(0);
    } else {
      return 0;
    }
  }

Project: tensorflow
File: ragged_cross_op.cc

  void ReadValue(int64 batch, int64 n, uint64* out) const override {
    CopyToFingerprint(values_(batch, n), out);
  }

Project: tensorflow
File: ragged_cross_op.cc

  void NextCombination(int64 batch_index, std::vector<int>* combination) const {
    bool carry = true;
    for (int i = combination->size() - 1; i >= 0; i--) {
      if (carry) {
        (*combination)[i] = (*combination)[i] + 1;
      }
      if ((*combination)[i] == features_[i]->FeatureCount(batch_index)) {
        (*combination)[i] = 0;
      } else {
        carry = false;
        break;
      }
    }
  }

Project: tensorflow
File: ragged_cross_op.cc

  void WriteCombination(int64 batch_index, const std::vector<int>& combination,
                        int64* out) {
    // Do the fingerprint concatenation on uint64.
    uint64 hashed_output = hash_key_;
    for (size_t i = 0; i < combination.size(); ++i) {
      uint64 hash_i;
      features_[i]->ReadValue(batch_index, combination[i], &hash_i);
      hashed_output = FingerprintCat64(hashed_output, hash_i);
    }
    // The return value is int64 based on the number of buckets.
    if (num_buckets_ > 0) {
      *out = hashed_output % num_buckets_;
    } else {
      // To prevent negative output we take modulo to max int64.
      *out = hashed_output % std::numeric_limits<int64>::max();
    }
  }

Project: tensorflow
File: ragged_cross_op.cc

  int64 FeatureCount(int64 batch) const override {
    return row_splits_(batch + 1) - row_splits_(batch);
  }

Project: tensorflow
File: ragged_cross_op.cc

  explicit RaggedCrossOp(OpKernelConstruction* context) : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("num_buckets", &num_buckets_));
    // Read signed_hash_key_ as int64 since uint64 attributes are not
    // supported by REGISTER_OP.
    int64 signed_hash_key_;
    OP_REQUIRES_OK(context, context->GetAttr("hash_key", &signed_hash_key_));
    hash_key_ = static_cast<uint64>(signed_hash_key_);

    int num_sparse;
    OP_REQUIRES_OK(context, context->GetAttr("Nsparse", &num_sparse));

    OP_REQUIRES_OK(context, context->GetAttr("ragged_values_types",
                                             &ragged_values_types_));
    OP_REQUIRES_OK(context, context->GetAttr("ragged_splits_types",
                                             &ragged_splits_types_));
    OP_REQUIRES_OK(context, context->GetAttr("sparse_values_types",
                                             &sparse_values_types_));
    OP_REQUIRES_OK(context, context->GetAttr("dense_types", &dense_types_));
    OP_REQUIRES_OK(context, context->GetAttr("input_order", &input_order_));
    OP_REQUIRES(context,
                ragged_values_types_.size() == ragged_splits_types_.size(),
                errors::InvalidArgument(
                    "ragged values and splits must have the same length"));
    OP_REQUIRES(context, num_sparse == sparse_values_types_.size(),
                errors::InvalidArgument(
                    "sparse indices and values must have the same length"));
    OP_REQUIRES(context,
                ragged_values_types_.size() + sparse_values_types_.size() +
                        dense_types_.size() ==
                    input_order_.size(),
                errors::InvalidArgument("Invalid length for input_order"));
  }

Project: tensorflow
File: ragged_cross_op.cc

  int64 FeatureCount(int64 batch) const override { return feature_count_; }

Project: tensorflow
File: ragged_cross_op.cc

  void ReadValue(int64 batch, int64 n, tstring* out) const override {
    CopyToString(values_(row_splits_[batch] + n), out);
  }

Project: tensorflow
File: ragged_cross_op.cc

  virtual ~OutputWriter() {}

Project: tensorflow
File: ragged_cross_op.cc

  void ReadValue(int64 batch, int64 n, tstring* out) const override {
    CopyToString(values_(batch, n), out);
  }

Project: tensorflow
File: ragged_cross_op.cc

  int64 CrossCountByBatchIndex(const FeatureReaders& features,
                               int batch_index) {
    int64 cross_count = 1;
    for (int i = 0; i < features.size(); ++i) {
      const auto feature_count = features[i]->FeatureCount(batch_index);
      if (feature_count == 0) return 0;
      cross_count *= feature_count;
    }
    return cross_count;
  }