Commit message: Fix heap-buffer-overflow issue with `tf.raw_ops.SparseSplit`.

PiperOrigin-RevId: 371242872
Change-Id: I482bb3d12602c7c3cc9446f97fb9f584bb98e9a4
CVE: CVE-2021-29558
CWE: CWE-703
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. An attacker can cause a heap buffer overflow in `tf.raw_ops.SparseSplit`. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/699bff5d961f0abfde8fa3f876e6d241681fbef8/tensorflow/core/util/sparse/sparse_tensor.h#L528-L530) accesses an array element based on a user controlled offset. 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_tensor.h

inline bool SparseTensor::ValidateAndInitializeToDense(Tensor* out,
                                                       bool initialize) {
  DCHECK_EQ(DataTypeToEnum<T>::v(), dtype())
      << "ToDense requested with the wrong datatype";

  DCHECK_EQ(out->shape().dims(), dims_)
      << "Incompatible dimensions between SparseTensor and output";

  DCHECK_EQ(out->dtype(), DataTypeToEnum<T>::v())
      << "Output must be type: " << DataTypeToEnum<T>::v()
      << " but got: " << out->dtype();

  // Make sure the dense output is the same rank and has room
  // to hold the SparseTensor.
  const auto& out_shape = out->shape();
  if (shape_.size() != out_shape.dims()) return false;
  for (int d = 0; d < shape_.size(); ++d) {
    if (shape_[d] > out_shape.dim_size(d)) return false;
  }

  if (initialize) {
    auto out_t = out->flat<T>();
    out_t.setConstant(T());
  }

  return true;
}

Project: tensorflow
File: sparse_tensor.h

inline SparseTensor SparseTensor::Concat(
    const gtl::ArraySlice<SparseTensor>& tensors) {
  DCHECK_GE(tensors.size(), size_t{1}) << "Cannot concat 0 SparseTensors";
  const int dims = tensors[0].dims_;
  DCHECK_GE(dims, 1) << "Cannot concat 0-dimensional SparseTensors";
  auto order_0 = tensors[0].order();
  const int primary_dim = order_0[0];
  ShapeArray final_order(order_0.begin(), order_0.end());
  ShapeArray final_shape(tensors[0].shape().begin(), tensors[0].shape().end());
  final_shape[primary_dim] = 0;  // We'll build this up as we go along.
  int num_entries = 0;

  bool fully_ordered = true;
  for (const SparseTensor& st : tensors) {
    DCHECK_EQ(st.dims_, dims) << "All SparseTensors must have the same rank.";
    DCHECK_EQ(DataTypeToEnum<T>::v(), st.dtype())
        << "Concat requested with the wrong data type";
    DCHECK_GE(st.order()[0], 0) << "SparseTensor must be ordered";
    DCHECK_EQ(st.order()[0], primary_dim)
        << "All SparseTensors' order[0] must match.  This is the concat dim.";
    if (st.order() != final_order) fully_ordered = false;
    const VarDimArray& st_shape = st.shape();
    for (int d = 0; d < dims - 1; ++d) {
      const int cdim = (d < primary_dim) ? d : d + 1;
      DCHECK_EQ(final_shape[cdim], st_shape[cdim])
          << "All SparseTensors' shapes must match except on the concat dim.  "
          << "Concat dim: " << primary_dim
          << ", mismatched shape at dim: " << cdim
          << ".  Expecting shape like: [" << str_util::Join(final_shape, ",")
          << "] but saw shape: [" << str_util::Join(st_shape, ",") << "]";
    }

    // Update dimension of final shape
    final_shape[primary_dim] =
        (final_shape[primary_dim] + st_shape[primary_dim]);

    num_entries += st.num_entries();  // Update number of entries
  }

  // If nonconsistent ordering among inputs, set final order to -1s.
  if (!fully_ordered) {
    final_order = UndefinedOrder(final_shape);
  }

  Tensor output_ix(DT_INT64, TensorShape({num_entries, dims}));
  Tensor output_vals(DataTypeToEnum<T>::v(), TensorShape({num_entries}));

  TTypes<int64>::Matrix ix_t = output_ix.matrix<int64>();
  typename TTypes<T>::Vec vals_t = output_vals.vec<T>();

  Eigen::DenseIndex offset = 0;
  int64 shape_offset = 0;
  for (const SparseTensor& st : tensors) {
    const int st_num_entries = st.num_entries();

    // Fill in indices & values.
    if (st_num_entries > 0) {
      std::copy_n(&st.vals_.vec<T>()(0), st_num_entries, &vals_t(offset));

      const auto* st_ix = &st.ix_.matrix<int64>()(0, 0);
      auto* ix_out = &ix_t(offset, 0);
      for (std::size_t i = 0; i < st_num_entries * dims; ++i) {
        *ix_out++ = *st_ix++ + ((i % dims == primary_dim) ? shape_offset : 0);
      }
    }

    offset += st_num_entries;
    shape_offset += st.shape()[primary_dim];
  }

  return SparseTensor(output_ix, output_vals, final_shape, final_order);
}

Project: tensorflow
File: sparse_tensor.h

inline bool SparseTensor::ToDense(Tensor* out, bool initialize) {
  if (!ValidateAndInitializeToDense<T>(out, initialize)) return false;

  auto out_t = out->flat<T>();
  auto vals_t = vals_.vec<T>();
  auto ix_t = ix_.matrix<int64>();
  const int64* const ix_ptr = ix_t.data();

  if (dims_ == 1) {
    // Fast path for sparse vectors.
    const int64 out_length = out->shape().dim_size(0);
    for (int n = 0; n < vals_t.dimension(0); ++n) {
      const int64 index = internal::SubtleMustCopy(ix_ptr[n]);
      if (!FastBoundsCheck(index, out_length)) return false;
      out_t(index) = vals_t(n);
    }
    return true;
  } else if (dims_ == 2) {
    // Fast path for sparse matrices.
    const auto& out_shape = out->shape();
    const int64 out_rows = out_shape.dim_size(0);
    const int64 out_cols = out_shape.dim_size(1);
    for (int n = 0; n < vals_t.dimension(0); ++n) {
      const int64 row_index = internal::SubtleMustCopy(ix_ptr[n * 2]);
      const int64 col_index = internal::SubtleMustCopy(ix_ptr[n * 2 + 1]);
      if (!(FastBoundsCheck(row_index, out_rows) &&
            FastBoundsCheck(col_index, out_cols))) {
        return false;
      }
      out_t(row_index * out_cols + col_index) = vals_t(n);
    }
    return true;
  } else {
    // General path for N-dimensional sparse tensors.
    gtl::InlinedVector<int64, 4> strides(dims_);
    const auto& out_shape = out->shape().dim_sizes();
    if (dims_ > 0) {
      strides[dims_ - 1] = 1;
    }
    for (int d = dims_ - 2; d >= 0; --d) {
      strides[d] = strides[d + 1] * out_shape[d + 1];
    }

    for (int n = 0; n < vals_t.dimension(0); ++n) {
      bool invalid_dims = false;
      int64 ix = 0;
      for (int d = 0; d < dims_; ++d) {
        const int64 ix_n_d = internal::SubtleMustCopy(ix_ptr[n * dims_ + d]);
        if (!FastBoundsCheck(ix_n_d, out_shape[d])) {
          invalid_dims = true;
        }
        ix += strides[d] * ix_n_d;
      }
      if (invalid_dims) return false;
      out_t(ix) = vals_t(n);
    }
    return true;
  }
}

Project: tensorflow
File: sparse_tensor.h

  int dims() const { return shape_.size(); }

Project: tensorflow
File: sparse_tensor.h

  VarDimArray shape() const { return shape_; }

Project: tensorflow
File: sparse_tensor.h

  static inline int GetDimensionInSlice(const int dim, const int split_size,
                                        const int residual) {
    DCHECK_GT(split_size, 0);
    DCHECK_GE(dim, 0);
    if (residual == 0) return dim % split_size;
    const int offset = residual * (split_size + 1);
    if (dim < offset) {
      return dim % (split_size + 1);
    } else {
      return (dim - offset) % split_size;
    }
  }

Project: tensorflow
File: sparse_tensor.h

  static inline ShapeArray TensorShapeToVector(const TensorShape& shape) {
    ShapeArray vec(shape.dims());
    for (int i = 0; i < shape.dims(); ++i) vec[i] = shape.dim_size(i);
    return vec;
  }

Project: tensorflow
File: sparse_tensor.h

inline SparseTensor SparseTensor::Slice(const SparseTensor& input_tensor,
                                        const gtl::ArraySlice<int64>& start,
                                        const gtl::ArraySlice<int64>& size) {
  TensorShape output_shape(input_tensor.shape());

  const int dims = input_tensor.dims();
  for (int dim = 0; dim < dims; dim++) {
    // Determine the size of the result; if the selected slice goes beyond the
    // input boundary, the result will correspond to the size of the overlap
    // between the input and the selected slice.
    const int64 input_size = output_shape.dim_size(dim);
    const int64 start_index = start[dim];
    const int64 slice_size = size[dim];
    if (start_index + slice_size < input_size) {
      // The entire selection is within input boundaries.
      output_shape.set_dim(dim, slice_size);
    } else if (start_index < input_size) {
      // The selection starts within input boundaries, but goes beyond them.
      output_shape.set_dim(dim, input_size - start_index);
    } else {
      // The selection is entirely out of input boundaries.
      output_shape.set_dim(dim, 0);
    }
  }

  auto input_indices_t = input_tensor.indices().matrix<int64>();
  auto input_values_t = input_tensor.values().vec<T>();

  // Find the number of indices that fall inside start and size.
  int count = 0;
  for (int i = 0; i < input_tensor.indices().dim_size(0); i++) {
    // The following will check to see if an input is within the
    // range specified by start and size.
    // The for loop below iterates through all dimensions. In case
    // the index falls outside of the start and size at any dimension,
    // it will be considered as a "no hit" (hit = false). In this
    // case, it will not be counted as the index that fall inside
    // the range specified by start and size.
    bool hit = true;
    for (int dim = 0; dim < dims; dim++) {
      if (!(start[dim] <= input_indices_t(i, dim) &&
            input_indices_t(i, dim) < start[dim] + size[dim])) {
        hit = false;
        break;
      }
    }
    if (!hit) {
      continue;
    }
    count++;
  }

  Tensor output_values(DataTypeToEnum<T>::v(), TensorShape({count}));
  Tensor output_indices(DT_INT64, TensorShape({count, dims}));

  auto output_values_t = output_values.vec<T>();
  auto output_indices_t = output_indices.matrix<int64>();

  // Obtain the output indices that fall inside start and size.
  int index = 0;
  for (int i = 0; i < input_tensor.indices().dim_size(0) && index < count;
       i++) {
    // The logic here is similar as the above except that the above
    // only count the number of indices while here we actually generate
    // the output.
    bool hit = true;
    for (int dim = 0; dim < dims; dim++) {
      if (!(start[dim] <= input_indices_t(i, dim) &&
            input_indices_t(i, dim) < start[dim] + size[dim])) {
        hit = false;
        break;
      }
    }
    if (!hit) {
      continue;
    }
    output_values_t(index) = input_values_t(i);
    for (int dim = 0; dim < dims; dim++) {
      output_indices_t(index, dim) = input_indices_t(i, dim) - start[dim];
    }
    index++;
  }

  return SparseTensor(output_indices, output_values, output_shape);
}

Project: tensorflow
File: sparse_tensor.h

  DataType dtype() const { return vals_.dtype(); }

Project: tensorflow
File: sparse_tensor.h

  SparseTensor(const SparseTensor& other)
      : SparseTensor(other.ix_, other.vals_, other.shape_, other.order_) {}

Project: tensorflow
File: sparse_tensor.h

  SparseTensor& operator=(SparseTensor&& other) {
    ix_ = std::move(other.ix_);
    vals_ = std::move(other.vals_);
    shape_ = std::move(other.shape_);
    order_ = std::move(other.order_);
    dims_ = std::move(other.dims_);
    return *this;
  }

Project: tensorflow
File: sparse_tensor.h

  VarDimArray order() const { return order_; }

Project: tensorflow
File: sparse_tensor.h

  std::size_t num_entries() const { return ix_.dim_size(0); }

Project: tensorflow
File: sparse_tensor.h

  SparseTensor(SparseTensor&& other)
      : SparseTensor(std::move(other.ix_), std::move(other.vals_),
                     std::move(other.shape_), std::move(other.order_)) {}

Project: tensorflow
File: sparse_tensor.h

  SparseTensor& operator=(const SparseTensor& other) {
    ix_ = other.ix_;
    vals_ = other.vals_;
    shape_ = other.shape_;
    order_ = other.order_;
    dims_ = other.dims_;
    return *this;
  }

Project: tensorflow
File: sparse_tensor.h

  const Tensor& indices() const { return ix_; }

Project: tensorflow
File: sparse_tensor.h

  SparseTensor(Tensor ix, Tensor vals, const TensorShape& shape)
      : SparseTensor(std::move(ix), std::move(vals), TensorShapeToVector(shape),
                     UndefinedOrder(TensorShapeToVector(shape))) {}

Project: tensorflow
File: sparse_tensor.h

  std::vector<int64> PickDims(gtl::ArraySlice<int64> dim_indices) const {
    std::vector<int64> res(dim_indices.size());
    for (size_t i = 0; i < dim_indices.size(); ++i) {
      res[i] = shape_[dim_indices[i]];
    }
    return res;
  }

Project: tensorflow
File: sparse_tensor.h

  const Tensor& values() const { return vals_; }

Project: tensorflow
File: sparse_tensor.h

  GroupIterable group(const VarDimArray& group_ix) const {
    DCHECK_LE(group_ix.size(), dims_);
    for (std::size_t di = 0; di < group_ix.size(); ++di) {
      DCHECK_GE(group_ix[di], 0) << "Group dimension out of range";
      DCHECK_LT(group_ix[di], dims_) << "Group dimension out of range";
      DCHECK_EQ(group_ix[di], order_[di])
          << "Group dimension does not match sorted order";
    }
    return GroupIterable(ix_, vals_, dims_, group_ix);
  }

Project: tensorflow
File: sparse_tensor.h

  static inline ShapeArray UndefinedOrder(const VarDimArray shape) {
    return ShapeArray(shape.size(), -1);
  }

Project: tensorflow
File: sparse_tensor.h

inline void SparseTensor::Reorder(const VarDimArray& order) {
  DCHECK_EQ(DataTypeToEnum<T>::v(), dtype())
      << "Reorder requested with the wrong datatype";
  DCHECK_EQ(order.size(), dims_) << "Order length must be SparseTensor rank";
  auto ix_t = ix_.matrix<int64>();
  auto vals_t = vals_.vec<T>();

  std::vector<int64> reorder(num_entries());
  std::iota(reorder.begin(), reorder.end(), 0);

  // Sort to get order of indices
  switch (order.size()) {
#define CASE_SORT(ORDER_SIZE)                                    \
  case ORDER_SIZE: {                                             \
    FixedDimComparator<ORDER_SIZE> sorter(ix_t, order, shape()); \
    std::sort(reorder.begin(), reorder.end(), sorter);           \
    break;                                                       \
  }
    CASE_SORT(0);
    CASE_SORT(1);
    CASE_SORT(2);
    CASE_SORT(3);
    CASE_SORT(4);
    CASE_SORT(5);
#undef CASE_SORT
    default: {
      DimComparator sorter(ix_t, order, shape());
      std::sort(reorder.begin(), reorder.end(), sorter);
    }
  }

  // We have a forward reordering, but what we'll need is a
  // permutation (the inverse).  This can be calculated with O(1)
  // additional
  // and O(n) time (INVPERM) but we just do the simple thing here.
  std::vector<size_t> permutation(reorder.size());
  for (std::size_t n = 0; n < reorder.size(); ++n) {
    permutation[reorder[n]] = n;
  }

  // Update indices & values by converting the permutations to
  // a product of transpositions.  Iterate over the cycles in the
  // permutation, and convert each of those into a product of
  // transpositions (swaps):
  //   https://en.wikipedia.org/wiki/Cyclic_permutation
  // This is N swaps, 2*N comparisons.
  for (std::size_t n = 0; n + 1 < permutation.size(); ++n) {
    while (n != permutation[n]) {
      std::size_t r = permutation[n];
      std::swap_ranges(&(ix_t(n, 0)), &(ix_t(n + 1, 0)), &(ix_t(r, 0)));
      std::swap(vals_t(n), vals_t(r));
      std::swap(permutation[n], permutation[r]);
    }
  }

  order_ = ShapeArray(order.begin(), order.end());
}

Project: tensorflow
File: sparse_tensor.h

  SparseTensor(Tensor ix, Tensor vals, const TensorShape& shape,
               const VarDimArray order)
      : SparseTensor(std::move(ix), std::move(vals), TensorShapeToVector(shape),
                     order) {}

Project: tensorflow
File: sparse_tensor.h

  SparseTensor() : dims_(0) {}

Project: tensorflow
File: sparse_tensor.h

inline Status SparseTensor::Split(const SparseTensor& input_tensor,
                                  const int split_dim, const int num_split,
                                  std::vector<SparseTensor>* result) {
  std::vector<Tensor> output_indices;
  std::vector<Tensor> output_values;
  std::vector<TensorShape> output_shapes;
  output_indices.reserve(num_split);
  output_values.reserve(num_split);
  output_shapes.reserve(num_split);

  std::vector<typename TTypes<int64>::Matrix> output_indices_t;
  std::vector<typename TTypes<T>::Vec> output_values_t;
  output_indices_t.reserve(num_split);
  output_values_t.reserve(num_split);
  auto input_values_t = input_tensor.values().vec<T>();
  auto input_indices_t = input_tensor.indices().matrix<int64>();

  std::vector<int> num_values(num_split, 0);
  const int num_dim = input_tensor.shape().size();
  const int split_dim_size = input_tensor.shape()[split_dim];
  const int split_size = split_dim_size / num_split;

  if (!(num_split > 0 && num_split <= split_dim_size)) {
    return errors::InvalidArgument("num_split must be in the interval (0, ",
                                   split_dim_size, "]");
  }
  if (!(split_dim >= 0 && split_dim < num_dim)) {
    return errors::InvalidArgument("num_dim must be in the interval [0, ",
                                   num_dim, ")");
  }

  const int residual = split_dim_size % num_split;
  for (int i = 0; i < input_tensor.indices().dim_size(0); ++i) {
    const int dim = input_tensor.indices().matrix<int64>()(i, split_dim);
    int slice_index = GetSliceIndex(dim, split_size, residual);
    if (slice_index >= num_values.size()) {
      return errors::InvalidArgument("Slice index ", slice_index,
                                     " is larger than num_split.");
    }
    num_values[slice_index]++;
  }

  for (int i = 0; i < num_split; ++i) {
    // TODO(ataei): Pass an allocator to avoid allocating large memory buffer.
    output_indices.emplace_back(DT_INT64,
                                TensorShape({num_values[i], num_dim}));
    output_values.emplace_back(DataTypeToEnum<T>::v(),
                               TensorShape({num_values[i]}));
    output_shapes.emplace_back(input_tensor.shape());
    output_indices_t.emplace_back(output_indices[i].matrix<int64>());
    output_values_t.emplace_back(output_values[i].vec<T>());
    const int size = GetSliceShape(i, split_size, residual);
    output_shapes[i].set_dim(split_dim, size);
  }

  std::vector<int> values_inserted_in_slice(num_split, 0);
  for (int i = 0; i < input_tensor.indices().dim_size(0); ++i) {
    const int dim = input_indices_t(i, split_dim);
    const int slice_index = GetSliceIndex(dim, split_size, residual);
    const int slice_dim = values_inserted_in_slice[slice_index]++;
    output_values_t[slice_index](slice_dim) = input_values_t(i);
    for (int j = 0; j < num_dim; ++j) {
      const int64 original_dim = input_indices_t(i, j);
      output_indices_t[slice_index](slice_dim, j) =
          (j == split_dim)
              ? GetDimensionInSlice(original_dim, split_size, residual)
              : original_dim;
    }
  }

  result->clear();
  result->reserve(num_split);
  for (int i = 0; i < num_split; ++i) {
    SparseTensor tensor;
    Status create_status =
        Create(output_indices[i], output_values[i], output_shapes[i], &tensor);
    if (!create_status.ok()) {
      return create_status;
    }
    result->push_back(std::move(tensor));
  }
  return Status::OK();
}

Project: tensorflow
File: sparse_tensor.h

  static inline int GetSliceIndex(const int dim, const int split_size,
                                  const int residual) {
    DCHECK_GT(split_size, 0);
    DCHECK_GE(dim, 0);
    if (residual == 0) return dim / split_size;
    const int offset = residual * (split_size + 1);
    if (dim < offset) {
      return dim / (split_size + 1);
    } else {
      return residual + ((dim - offset) / split_size);
    }
  }

Project: tensorflow
File: sparse_tensor.h

  static inline int GetSliceShape(const int slice_index, const int split_size,
                                  const int residual) {
    DCHECK_GT(split_size, 0);
    DCHECK_GE(slice_index, 0);
    if (residual == 0) return split_size;
    if (slice_index < residual) {
      return split_size + 1;
    } else {
      return split_size;
    }
  }

Project: tensorflow
File: sparse_tensor.h

  SparseTensor(Tensor ix, Tensor vals, const VarDimArray shape)
      : SparseTensor(std::move(ix), std::move(vals), shape,
                     UndefinedOrder(shape)) {}