Commit message: Fix overflow CHECK issue with `tf.raw_ops.UnsortedSegmentJoin`.

PiperOrigin-RevId: 370766155
Change-Id: I33e7c6626224e1060a8a4ab51ad5d861c6d4c63e
CVE: CVE-2021-29552
CWE: CWE-703
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. An attacker can cause a denial of service by controlling the values of `num_segments` tensor argument for `UnsortedSegmentJoin`. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/a2a607db15c7cd01d754d37e5448d72a13491bdb/tensorflow/core/kernels/unsorted_segment_join_op.cc#L92-L93) assumes that the `num_segments` tensor is a valid scalar. Since the tensor is empty the `CHECK` involved in `.scalar<T>()()` that checks that the number of elements is exactly 1 will be invalidated and this would result in process termination. 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: unsorted_segment_join_op.cc

std::pair<INDICES_TYPE, INDICES_TYPE> GetStrides(
    const TensorShape& input_shape, const TensorShape& segment_id_shape) {
  int64 small_stride = 1;
  int64 big_stride = 1;
  for (auto i = 0; i < input_shape.dims(); i++) {
    if (i < segment_id_shape.dims()) {
      small_stride *= segment_id_shape.dim_size(i);
    } else {
      big_stride *= input_shape.dim_size(i);
    }
  }
  return std::make_pair(big_stride, small_stride);
}

Project: tensorflow
File: unsorted_segment_join_op.cc

  void Compute(OpKernelContext* context) override {
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const int32 input_dims = input_shape.dims();

    const Tensor& segment_id = context->input(1);
    const TensorShape& segment_id_shape = segment_id.shape();
    const int32 segment_dims = segment_id_shape.dims();

    const Tensor& num_segments_tensor = context->input(2);
    OP_REQUIRES(context, num_segments_tensor.NumElements() != 0,
                errors::InvalidArgument("Number of segments cannot be empty."));
    auto num_segments = num_segments_tensor.scalar<NUM_SEGMENTS_TYPE>()();

    OP_REQUIRES(context, segment_dims != 0,
                errors::InvalidArgument("Segment_id cannot have rank 0"));

    OP_REQUIRES(
        context, segment_dims <= input_dims,
        errors::OutOfRange("Invalid segment_id rank ", segment_dims,
                           " for input with ", input_dims, " dimension(s)"));
    for (auto i = 0; i < segment_dims; i++) {
      OP_REQUIRES(
          context, segment_id_shape.dim_size(i) == input_shape.dim_size(i),
          errors::InvalidArgument(
              "Segment dimension is ", segment_id_shape.dim_size(i),
              " while input dimension is ", input_dims, " in rank ", i));
    }

    // Making output tensor.
    Tensor* output_tensor = nullptr;
    TensorShape output_shape =
        GetOutputShape(input_shape, segment_id_shape, num_segments);
    OP_REQUIRES_OK(context, context->allocate_output("output", output_shape,
                                                     &output_tensor));

    // Preparating flat tensors.
    auto output_flat = output_tensor->flat<tstring>();
    auto flat_segment_id = segment_id.flat<INDICES_TYPE>();
    auto flat_input = input.flat<tstring>();

    for (int i = 0; i < flat_segment_id.size(); i++) {
      OP_REQUIRES(
          context,
          ((flat_segment_id(i) < num_segments) && (flat_segment_id(i) >= 0)),
          errors::InvalidArgument(
              "segment_ids are not allowed to exceed num_segments or"
              " to have negative values."));
    }

    int64 big_stride;
    int64 small_stride;
    std::tie(big_stride, small_stride) =
        GetStrides<INDICES_TYPE>(input_shape, segment_id_shape);
    auto relative_offset_set =
        GetFlattenedRelativeOffsets<INDICES_TYPE>(small_stride, big_stride);
    for (auto start_offset = 0; start_offset < big_stride; start_offset++) {
      for (auto i = 0; i < relative_offset_set.size(); i++) {
        auto output_index = start_offset + flat_segment_id(i) * big_stride;
        auto offset = start_offset + relative_offset_set[i];
        if (output_flat(output_index).length() != 0)
          output_flat(output_index).append(separator_.c_str());
        output_flat(output_index).append(flat_input(offset));
      }
    }
  }

Project: tensorflow
File: unsorted_segment_join_op.cc

TensorShape GetOutputShape(const TensorShape& input_shape,
                           const TensorShape& segment_id_shape,
                           const int64 num_segments) {
  TensorShape output_shape;
  output_shape.AddDim(num_segments);
  for (size_t index = segment_id_shape.dims(); index < input_shape.dims();
       ++index) {
    output_shape.AddDim(input_shape.dim_size(index));
  }
  return output_shape;
}