Commit message: Fix tf.raw_ops.UnsortedSegmentJoin vulnerability with invalid num_segments.

Check that input is actually a scalar before treating it as such.

PiperOrigin-RevId: 445206880
CVE: CVE-2022-29197
CWE: CWE-703
Vulnerability description: TensorFlow is an open source platform for machine learning. Prior to versions 2.9.0, 2.8.1, 2.7.2, and 2.6.4, the implementation of `tf.raw_ops.UnsortedSegmentJoin` does not fully validate the input arguments. This results in a `CHECK`-failure which can be used to trigger a denial of service attack. The code assumes `num_segments` is a scalar but there is no validation for this before accessing its value. Versions 2.9.0, 2.8.1, 2.7.2, and 2.6.4 contain a patch for this issue.

Project: tensorflow
File: unsorted_segment_join_op.cc

gtl::InlinedVector<INDICES_TYPE, 8> GetFlattenedRelativeOffsets(
    INDICES_TYPE small_stride, INDICES_TYPE big_stride) {
  gtl::InlinedVector<INDICES_TYPE, 8> flattened_offsets(small_stride);
  for (auto i = 0; i < small_stride; i++) {
    flattened_offsets[i] = i * big_stride;
  }
  return flattened_offsets;
}

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_t small_stride = 1;
  int64_t 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_t input_dims = input_shape.dims();

    const Tensor& segment_id = context->input(1);
    const TensorShape& segment_id_shape = segment_id.shape();
    const int32_t 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."));
    OP_REQUIRES(context,
                TensorShapeUtils::IsScalar(num_segments_tensor.shape()),
                errors::InvalidArgument("Number of segments must be a scalar"));
    auto num_segments = num_segments_tensor.scalar<NUM_SEGMENTS_TYPE>()();

    OP_REQUIRES(
        context, num_segments >= 0,
        errors::InvalidArgument(
            "Number of segments must be non-negative but got ", num_segments));
    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_t big_stride;
    int64_t 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

  explicit UnsortedSegmentJoinOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("separator", &separator_));
  }

Project: tensorflow
File: unsorted_segment_join_op.cc

TensorShape GetOutputShape(const TensorShape& input_shape,
                           const TensorShape& segment_id_shape,
                           const int64_t 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;
}