Commit message: Prevent heap OOB in sparse reduction ops.

PiperOrigin-RevId: 387934524
Change-Id: I894aa30f1e454f09b471d565b4a325da49322c1a
CVE: CVE-2021-37635
CWE: CWE-125
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. In affected versions the implementation of sparse reduction operations in TensorFlow can trigger accesses outside of bounds of heap allocated data. The [implementation](https://github.com/tensorflow/tensorflow/blob/a1bc56203f21a5a4995311825ffaba7a670d7747/tensorflow/core/kernels/sparse_reduce_op.cc#L217-L228) fails to validate that each reduction group does not overflow and that each corresponding index does not point to outside the bounds of the input tensor. We have patched the issue in GitHub commit 87158f43f05f2720a374f3e6d22a7aaa3a33f750. 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_reduce_op.cc

  void Compute(OpKernelContext *ctx) override {
    const Tensor *indices_t, *values_t, *shape_t, *reduction_axes_t;
    OP_REQUIRES_OK(ctx, ctx->input("input_indices", &indices_t));
    OP_REQUIRES_OK(ctx, ctx->input("input_values", &values_t));
    OP_REQUIRES_OK(ctx, ctx->input("input_shape", &shape_t));
    OP_REQUIRES_OK(ctx, ctx->input("reduction_axes", &reduction_axes_t));

    OP_REQUIRES_OK(ctx, ValidateInputs(shape_t, reduction_axes_t));

    SparseTensor sp;
    OP_REQUIRES_OK(ctx, SparseTensor::Create(tensor::DeepCopy(*indices_t),
                                         tensor::DeepCopy(*values_t),
                    TensorShape(shape_t->vec<int64>()), &sp));
    ReduceDetails reduction = SparseTensorReduceHelper(
        sp, reduction_axes_t->flat<int32>(), keep_dims_);

    sp.Reorder<T>(reduction.reorder_dims);
    // Count nnzs in the output SparseTensor.
    int64_t nnz = 0;
    auto iter = sp.group(reduction.group_by_dims);
    for (auto it = iter.begin(); it != iter.end(); ++it) {
      nnz++;
    }

    Tensor *out_indices_t;
    OP_REQUIRES_OK(ctx,
                   ctx->allocate_output(
                       0, TensorShape({nnz, reduction.reduced_shape.dims()}),
                       &out_indices_t));
    typename TTypes<int64>::Matrix out_indices_mat =
        out_indices_t->matrix<int64>();
    // For keep_dims. We don't explicitly set dim fields for reduced dims below.
    out_indices_mat.setZero();

    Tensor *out_values_t;
    OP_REQUIRES_OK(ctx,
                   ctx->allocate_output(1, TensorShape({nnz}), &out_values_t));
    auto out_flat = out_values_t->flat<T>();

    Tensor tmp_reduced_val;
    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<T>::value,
                                           TensorShape({}), &tmp_reduced_val));
    auto reduced_val = tmp_reduced_val.scalar<T>();
    int64_t i = 0;
    for (const auto &g : sp.group(reduction.group_by_dims)) {
      Op::template Run<T>(ctx, reduced_val, g.template values<T>());
      std::vector<int64> group = g.group();
      for (int64_t j = 0; j < group.size(); j++) {
        if (keep_dims_) {
          out_indices_mat(i, reduction.group_by_dims[j]) = group[j];
        } else {
          out_indices_mat(i, j) = group[j];
        }
      }
      out_flat(i) = reduced_val();
      i++;
      VLOG(2) << "coords: " << absl::StrJoin(g.group(), ",")
              << "; group " << Op::Name() << ": "
              << reduced_val();
    }

    Tensor *out_shape_t;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(
                            2, TensorShape({reduction.reduced_shape.dims()}),
                            &out_shape_t));
    auto out_shape_flat = out_shape_t->flat<int64>();
    auto out_dim_sizes = reduction.reduced_shape.dim_sizes();
    if (!out_dim_sizes.empty()) {
      std::copy(out_dim_sizes.begin(), out_dim_sizes.end(), &out_shape_flat(0));
    }
  }

Project: tensorflow
File: sparse_reduce_op.cc

Status ValidateInputs(const Tensor *shape_t, const Tensor *reduction_axes_t) {
  // indices and values are validated in SparseTensor ctor.
  if (!TensorShapeUtils::IsVector(shape_t->shape())) {
    return errors::InvalidArgument(
        "Expected input_shape to be a vector; got shape: ",
        shape_t->shape().DebugString());
  }
  if (!TensorShapeUtils::IsScalar(reduction_axes_t->shape()) &&
      !TensorShapeUtils::IsVector(reduction_axes_t->shape())) {
    return errors::InvalidArgument(
        "Expected reduction_axes to be a scalar or a vector; got shape: ",
        reduction_axes_t->shape().DebugString());
  }

  const auto reduction_axes_flat = reduction_axes_t->flat<int32>();
  for (int64_t i = 0; i < reduction_axes_flat.size(); i++) {
    int32_t axis = reduction_axes_flat(i);
    if (axis < -shape_t->NumElements() || axis >= shape_t->NumElements()) {
      return errors::InvalidArgument("Invalid reduction dimension ", axis,
                                     ", for input with ",
                                     shape_t->NumElements(), " dimensions.");
    }
  }

  return Status::OK();
}

Project: tensorflow
File: sparse_reduce_op.cc

  static void Run(OpKernelContext *ctx, typename TTypes<T>::Scalar &s, const typename TTypes<T>::UnalignedVec &v) {
      s.device(ctx->eigen_cpu_device()) = v.maximum();
  }

Project: tensorflow
File: sparse_reduce_op.cc

  void Compute(OpKernelContext *ctx) override {
    const Tensor *indices_t, *values_t, *shape_t, *reduction_axes_t;
    OP_REQUIRES_OK(ctx, ctx->input("input_indices", &indices_t));
    OP_REQUIRES_OK(ctx, ctx->input("input_values", &values_t));
    OP_REQUIRES_OK(ctx, ctx->input("input_shape", &shape_t));
    OP_REQUIRES_OK(ctx, ctx->input("reduction_axes", &reduction_axes_t));

    OP_REQUIRES_OK(ctx, ValidateInputs(shape_t, reduction_axes_t));

    // TODO(zongheng): we will call Reorder() below, which will modify
    // in-place the underlying indices and values buffers.  To avoid
    // surprises of this kernel being stateful, we work around the above by
    // making deep copies here.  Remove this if/when we change Reorder()'s
    // semantics.
    const auto shape_vec = shape_t->vec<int64>();
    SparseTensor sp;
    OP_REQUIRES_OK(ctx, SparseTensor::Create(
        tensor::DeepCopy(*indices_t), tensor::DeepCopy(*values_t),
                    TensorShape(shape_vec), &sp));
    ReduceDetails reduction = SparseTensorReduceHelper(
        sp, reduction_axes_t->flat<int32>(), keep_dims_);

    Tensor *out_values;
    OP_REQUIRES_OK(
        ctx, ctx->allocate_output(0, reduction.reduced_shape, &out_values));
    auto out_flat = out_values->flat<T>();
    out_flat.setZero();

    Tensor tmp_reduced_val;
    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DataTypeToEnum<T>::value,
                                           TensorShape({}), &tmp_reduced_val));
    auto reduced_val = tmp_reduced_val.scalar<T>();

    // Compute strides, and use it to convert coords to flat index.  The
    // coordinates returned by .group() have the same ndims as group_by_dims.
    gtl::InlinedVector<int64, 8> output_strides(reduction.group_by_dims.size());
    if (!output_strides.empty()) {  // Do this iff we don't reduce all.
      output_strides.back() = 1;
      for (int d = output_strides.size() - 2; d >= 0; --d) {
        output_strides[d] =
            output_strides[d + 1] * shape_vec(reduction.group_by_dims[d + 1]);
      }
    }

    auto CoordinatesToFlatIndex = [](ArraySlice<int64> coords,
                                     ArraySlice<int64> strides) -> int64 {
      if (strides.empty()) {  // Reduce all.
        return 0;
      }
      CHECK_EQ(coords.size(), strides.size());
      int64_t idx = 0;
      for (int i = 0; i < coords.size(); ++i) {
        idx += coords[i] * strides[i];
      }
      return idx;
    };

    // Each group maps one-on-one onto a value in the reduced tensor.
    // g.group() provides the coordinates of a particular reduced value.
    sp.Reorder<T>(reduction.reorder_dims);
    for (const auto &g : sp.group(reduction.group_by_dims)) {
      Op::template Run<T>(ctx, reduced_val, g.template values<T>());
      OP_REQUIRES(ctx,
                  output_strides.empty() ||
                  (g.group().size() == output_strides.size()),
                  errors::Internal(
                      "Expected group size and output_strides size to match",
                      ", but got ", g.group().size(), " and ",
                      output_strides.size()));
      const int64_t idx = CoordinatesToFlatIndex(g.group(), output_strides);
      OP_REQUIRES(ctx,
                  idx >= 0 && idx < out_flat.size(),
                  errors::Internal(
                      "Obtained a write index of ", idx,
                      " which is outside of bounds of [0, ",
                      out_flat.size(), ")"));
      out_flat(idx) = reduced_val();
      VLOG(2) << "coords: " << absl::StrJoin(g.group(), ",")
              << "; idx: " << idx << "; group " << Op::Name() << ": "
              << reduced_val();
    }
  }

Project: tensorflow
File: sparse_reduce_op.cc

  static StringPiece Name() {
      return "sum";
  }

Project: tensorflow
File: sparse_reduce_op.cc

ReduceDetails SparseTensorReduceHelper(const SparseTensor &sp,
                                       gtl::ArraySlice<int32> axes_slice,
                                       bool keep_dims) {
  ReduceDetails reduction;

  std::vector<int32> reduction_axes(axes_slice.begin(), axes_slice.end());
  int ndims = sp.dims();
  for (int64_t i = 0; i < reduction_axes.size(); ++i) {
    reduction_axes[i] = (reduction_axes[i] + ndims) % ndims;
  }
  std::sort(reduction_axes.begin(), reduction_axes.end());

  // (0) Calculate the grouping dimensions:
  // group_by_dims == {0, .., NDIMS-1} \ reduction_axes.
  std::vector<int64> perm(ndims);
  std::iota(perm.begin(), perm.end(), 0);

  // Requires perm and reduction_axes_ be sorted; group_by_dims will be
  // sorted as well.
  std::set_difference(
      perm.begin(), perm.end(), reduction_axes.begin(), reduction_axes.end(),
      std::inserter(reduction.group_by_dims, reduction.group_by_dims.begin()));

  // Now append the rest of the axes (the complement of group_by_dims_);
  // result is used by Reorder().
  reduction.reorder_dims = reduction.group_by_dims;
  std::set_difference(perm.begin(), perm.end(), reduction.group_by_dims.begin(),
                      reduction.group_by_dims.end(),
                      std::back_inserter(reduction.reorder_dims));

  // (1) Calculate the shape after reduction.
  auto sp_shape = sp.shape();
  std::vector<int64> out_dim_sizes;
  if (keep_dims) {
    out_dim_sizes.reserve(ndims);
    auto beg = reduction.group_by_dims.begin();
    auto end = reduction.group_by_dims.end();
    for (int d = 0; d < ndims; ++d) {
      if (std::find(beg, end, d) == end) {
        out_dim_sizes.push_back(1);  // A reduced axis.
      } else {
        out_dim_sizes.push_back(sp_shape[d]);
      }
    }
  } else {
    out_dim_sizes = sp.PickDims(reduction.group_by_dims);
  }

  reduction.reduced_shape = TensorShape(out_dim_sizes);
  return reduction;
}

Project: tensorflow
File: sparse_reduce_op.cc

  static StringPiece Name() {
      return "max";
  }

Project: tensorflow
File: sparse_reduce_op.cc

  explicit SparseReduceOp(OpKernelConstruction *ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("keep_dims", &keep_dims_));
  }

Project: tensorflow
File: sparse_reduce_op.cc

  explicit SparseReduceSparseOp(OpKernelConstruction *ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("keep_dims", &keep_dims_));
  }

Project: tensorflow
File: sparse_reduce_op.cc

  static void Run(OpKernelContext *ctx, typename TTypes<T>::Scalar &s, const typename TTypes<T>::UnalignedVec &v) {
      s.device(ctx->eigen_cpu_device()) = v.sum();
  }