Commit message: Fix FPE issue with `tf.raw_ops.Reverse`.

PiperOrigin-RevId: 371176973
Change-Id: Ic6d483bfc95313ec2299c2d1c956cfe96c96626c
CVE: CVE-2021-29556
CWE: CWE-369
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. An attacker can cause a denial of service via a FPE runtime error in `tf.raw_ops.Reverse`. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/36229ea9e9451dac14a8b1f4711c435a1d84a594/tensorflow/core/kernels/reverse_op.cc#L75-L76) performs a division based on the first dimension of the tensor argument. 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: reverse_op.cc

void HandleReverseCase(OpKernelContext* context,
                       typename TTypes<bool, 1>::ConstTensor dims,
                       Tensor* result) {
  const Tensor& input = context->input(0);

  // Use optimized reverse if possible.
  if (NDIMS == 3 && std::is_same<Device, CPUDevice>::value &&
      data_type_can_memcpy<T>::value && (!dims(0) && dims(1) && !dims(2))) {
    if (input.dim_size(2) == 3) {
      DoHandleReverseCase<T, 3>(context, input, result);
    } else {
      DoHandleReverseCase<T, -1>(context, input, result);
    }
    return;
  }
  typename Eigen::array<bool, NDIMS> axes_di;
  for (int i = 0; i < NDIMS; i++) {
    axes_di[i] = dims(i);
  }
  functor::Reverse<Device, T, NDIMS>()(context->eigen_device<Device>(),
                                       input.tensor<T, NDIMS>(), axes_di,
                                       result->tensor<T, NDIMS>());
}

Project: tensorflow
File: reverse_op.cc

  void Compute(OpKernelContext* context) override {
    const Tensor& input = context->input(0);
    const Tensor& sparse_dims = context->input(1);

    if (TensorShapeUtils::IsScalar(input.shape()) || input.NumElements() == 0) {
      context->set_output(0, input);
    } else {
      const int input_dims = input.dims();
      const TensorShape& sparse_dims_shape = sparse_dims.shape();
      const auto& axes_sparse_flat = sparse_dims.flat<Tidx>();

      OP_REQUIRES(context, TensorShapeUtils::IsVector(sparse_dims_shape),
                  errors::InvalidArgument("'dims' must be 1-dimension, not ",
                                          sparse_dims.dims()));
      gtl::InlinedVector<bool, 8> axes_dense(input_dims, false);
      for (int dummy = 0; dummy < axes_sparse_flat.size(); dummy++) {
        Tidx axis = internal::SubtleMustCopy<Tidx>(axes_sparse_flat(dummy));
        Tidx canonical_axis = axis < 0 ? input_dims + axis : axis;
        OP_REQUIRES(context, canonical_axis >= 0 && canonical_axis < input_dims,
                    errors::InvalidArgument("'axis'[", dummy, "] = ", axis,
                                            " is out of valid range [", 0, ", ",
                                            input_dims - 1));
        OP_REQUIRES(context, !axes_dense[canonical_axis],
                    errors::InvalidArgument("axis ", canonical_axis,
                                            " specified more than once."));
        axes_dense[canonical_axis] = true;
      }

      OP_REQUIRES(context, input_dims <= 8,
                  errors::Unimplemented(
                      "reverse is not implemented for tensors of rank > 8."));

      Tensor* output = nullptr;
      OP_REQUIRES_OK(context,
                     context->allocate_output(0, input.shape(), &output));

      // TODO(cwhipkey): we can do dimension folding to reduce, e.g., a reverse
      // of a single dimension to the dims=3 or dims=2 case, regardless of the
      // number of dimensions in the tensor. This would let some ops use faster
      // lower-dimension code (and use optimized versions).

#define HANDLE_REVERSE(NDIMS)                                           \
  case NDIMS:                                                           \
    HandleReverseV2Case<Device, T, NDIMS>(context, axes_dense, output); \
    return;

      switch (input_dims) {
        HANDLE_REVERSE(0);
        HANDLE_REVERSE(1);
        HANDLE_REVERSE(2);
        HANDLE_REVERSE(3);
        HANDLE_REVERSE(4);
        HANDLE_REVERSE(5);
        HANDLE_REVERSE(6);
        HANDLE_REVERSE(7);
        HANDLE_REVERSE(8);
      }
#undef HANDLE_REVERSE
    }
  }

Project: tensorflow
File: reverse_op.cc

  explicit ReverseOp(OpKernelConstruction* context) : OpKernel(context) {}

Project: tensorflow
File: reverse_op.cc

DoHandleReverseCase(OpKernelContext* context, const Tensor& input,
                    Tensor* result) {}

Project: tensorflow
File: reverse_op.cc

  void Compute(OpKernelContext* context) override {
    const Tensor& input = context->input(0);
    // If input is provided, check to make sure the first dimension is valid.
    if (input.dims() > 0) {
      OP_REQUIRES(
          context, input.dim_size(0) != 0,
          errors::InvalidArgument("Invalid input first dimension. Found 0."));
    }
    const Tensor& dims = context->input(1);

    if (TensorShapeUtils::IsScalar(input.shape())) {
      context->set_output(0, input);
    } else {
      const int input_dims = input.dims();
      OP_REQUIRES(context, TensorShapeUtils::IsVector(dims.shape()),
                  errors::InvalidArgument("'dims' must be 1-dimension, not ",
                                          dims.dims()));

      OP_REQUIRES(
          context, input_dims == dims.dim_size(0),
          errors::InvalidArgument(
              "'dims' must have the same number of values as 'input' has "
              "dimensions. 'input' has ",
              input_dims, "'dims' has ", dims.dim_size(0), " values"));
      OP_REQUIRES(context, input_dims <= 8,
                  errors::Unimplemented(
                      "reverse is not implemented for tensors of rank > 8."));

      Tensor* output = nullptr;
      OP_REQUIRES_OK(context,
                     context->allocate_output(0, input.shape(), &output));

#define HANDLE_REVERSE(NDIMS)                                               \
  case NDIMS:                                                               \
    HandleReverseCase<Device, T, NDIMS>(context, dims.vec<bool>(), output); \
    return;

      switch (input_dims) {
        HANDLE_REVERSE(0);
        HANDLE_REVERSE(1);
        HANDLE_REVERSE(2);
        HANDLE_REVERSE(3);
        HANDLE_REVERSE(4);
        HANDLE_REVERSE(5);
        HANDLE_REVERSE(6);
        HANDLE_REVERSE(7);
        HANDLE_REVERSE(8);
      }
#undef HANDLE_REVERSE
    }
  }

Project: tensorflow
File: reverse_op.cc

void ReverseRows(OpKernelContext* context, const Tensor& input,
                 Tensor* result) {
  auto work = [&input, result](int64 start, int64 end) {
    const int64 inner_size =
        NUM_CHANNELS > 0 ? NUM_CHANNELS : input.dim_size(2);
    const int64 middle_size = input.dim_size(1);
    const int64 row_size = inner_size * middle_size;
    DCHECK_EQ(input.dim_size(2), inner_size);

    const T* in_ptr = input.bit_casted_tensor<T, 3>().data();
    T* out_ptr = result->bit_casted_tensor<T, 3>().data();

    in_ptr += start * row_size;
    out_ptr += start * row_size;

    for (int outer_dim = start; outer_dim < end; ++outer_dim) {
      out_ptr += row_size;
      int remaining = middle_size;
      while (remaining > 0) {
        out_ptr -= inner_size;
        memcpy(out_ptr, in_ptr, inner_size * sizeof(T));
        in_ptr += inner_size;
        --remaining;
      }

      out_ptr += row_size;
    }
  };

  // Shard across outer dimension.
  const int64 N = input.dim_size(0);
  const int64 cost_per_unit = input.NumElements() / N;
  auto worker_threads = context->device()->tensorflow_cpu_worker_threads();
  Shard(worker_threads->num_threads, worker_threads->workers, N, cost_per_unit,
        std::move(work));
}

Project: tensorflow
File: reverse_op.cc

  explicit ReverseV2Op(OpKernelConstruction* context) : OpKernel(context) {}

Project: tensorflow
File: reverse_op.cc

DoHandleReverseCase(OpKernelContext* context, const Tensor& input,
                    Tensor* result) {
  if (sizeof(T) == 1) {
    static_assert(sizeof(uint8) == 1, "uint8 must be 1 byte.");
    ReverseRows<uint8, NUM_CHANNELS>(context, input, result);
  } else if (sizeof(T) == 2) {
    static_assert(sizeof(uint16) == 2, "uint16 must be 2 bytes");
    ReverseRows<uint16, NUM_CHANNELS>(context, input, result);
  } else if (sizeof(T) == 4) {
    static_assert(sizeof(uint32) == 4, "uint32 must be 4 bytes");
    ReverseRows<uint32, NUM_CHANNELS>(context, input, result);
  } else if (sizeof(T) == 8) {
    static_assert(sizeof(uint64) == 8, "uint64 must be 8 bytes");
    ReverseRows<uint64, NUM_CHANNELS>(context, input, result);
  } else if (sizeof(T) == 16) {
    static_assert(sizeof(complex128) == 16, "complex128 must be 16 bytes");
    ReverseRows<complex128, NUM_CHANNELS>(context, input, result);
  } else {
    context->CtxFailure(errors::InvalidArgument(DataTypeString(input.dtype()),
                                                " has unexpected size of ",
                                                sizeof(T), " bytes"));
  }
}

Project: tensorflow
File: reverse_op.cc

void HandleReverseV2Case(OpKernelContext* context,
                         const gtl::ArraySlice<bool> axes, Tensor* result) {
  const Tensor& input = context->input(0);

  // Use optimized reverse if possible.
  if (NDIMS == 3 && std::is_same<Device, CPUDevice>::value &&
      data_type_can_memcpy<T>::value && (!axes[0] && axes[1] && !axes[2])) {
    if (input.dim_size(2) == 3) {
      DoHandleReverseCase<T, 3>(context, input, result);
    } else {
      DoHandleReverseCase<T, -1>(context, input, result);
    }
    return;
  }

  typename Eigen::array<bool, NDIMS> axes_di;
  for (int i = 0; i < NDIMS; i++) {
    axes_di[i] = axes[i];
  }
  functor::Reverse<Device, T, NDIMS>()(context->eigen_device<Device>(),
                                       input.tensor<T, NDIMS>(), axes_di,
                                       result->tensor<T, NDIMS>());
}