CVE-2019-16778 (C++, published 2019-12-16T21:15Z)
Description: value: In TensorFlow before 1.15, a heap buffer overflow in UnsortedSegmentSum can be produced when the Index template argument is int32. In this case data_size and num_segments fields are truncated from int64 to int32 and can produce negative numbers, resulting in accessing out of bounds heap memory. This is unlikely to be exploitable and was detected and fixed internally in TensorFlow 1.15 and 2.0.
Severity: HIGH (CVSS v3 base score: 9.8)
Weakness: CWE-681 \u2014 Incorrect Conversion between Numeric Types
When converting from one data type to another, such as long to integer, data can be omitted or translated in a way that produces unexpected values. If the resulting values are used in a sensitive context, then dangerous behaviors may occur.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/segment_reduction_ops.cc
tensorflow/core/kernels/segment_reduction_ops.h
tensorflow/core/kernels/segment_reduction_ops_gpu.cu.cc
Commit message: Fix heap buffer overflow in UnsortedSegmentSum.

When Index=int32, data_size and num_segments were truncated from int64 to int32. This truncation can produce negative numbers, which causes UnsortedSegmentFunctor to access out of bounds memory.

Also:
- Switches some indexing calculations to int64 to avoid signed integer overflow when either the input or output tensors have more than 2**31 - 1 elements.
- Fixes a range check error in the GPU kernel. The segment ID was checked against an upper bound measured in elements, not segments.
PiperOrigin-RevId: 256451663


Vulnerable code:
  void operator()(OpKernelContext* ctx, const Index num_segments,
                  const TensorShape& segment_ids_shape,
                  const Index data_size, const T* data,
    if (data_size == 0) {
    auto data_flat = typename TTypes<T, 2>::ConstTensor(data, N, data_size / N);
      reduction(data_flat.template chip<0>(i), output.template chip<0>(j));
    const Index output_rows = internal::SubtleMustCopy(static_cast<Index>(
    auto data_ptr = data.template flat<T>().data();
    reduction_functor_(context, output_rows, segment_ids.shape(), segment_flat,
                       data.NumElements(), data_ptr, output_flat);
  void operator()(OpKernelContext* ctx, const Index num_segments,
                  const TensorShape& segment_ids_shape,
                  const Index data_size, const T* data,
__global__ void UnsortedSegmentCustomKernel(const Index input_outer_dim_size,
                                            const Index inner_dim_size,
                                            const Index output_outer_dim_size,
  const Index input_total_size = input_outer_dim_size * inner_dim_size;
  const Index output_total_size = output_outer_dim_size * inner_dim_size;
  for (int input_index : GpuGridRangeX(input_total_size)) {
    const Index input_segment_index = input_index / inner_dim_size;
    const Index segment_offset = input_index % inner_dim_size;
    if (output_segment_index < 0 || output_segment_index >= output_total_size) {
    const Index output_index =
  void operator()(OpKernelContext* ctx, const Index num_segments,
                  const TensorShape& segment_ids_shape,
                  const Index data_size, const T* data,
    const Index input_outer_dim_size = segment_ids.dimension(0);
    const Index input_inner_dim_size = data_size / input_outer_dim_size;
    TF_CHECK_OK(
        GpuLaunchKernel(UnsortedSegmentCustomKernel<T, Index, ReductionF>,
                        config.block_count, config.thread_per_block, 0,
                        d.stream(), input_outer_dim_size, input_inner_dim_size,
                        num_segments, segment_ids.data(), data, output.data()));

Fixed code:
  void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape,
                  typename TTypes<T, 2>::ConstTensor data,
    if (data.size() == 0) {
    const int64 num_segments = output.dimension(0);
      reduction(data.template chip<0>(i), output.template chip<0>(j));
    const int64 output_rows = internal::SubtleMustCopy(static_cast<int64>(
    auto data_flat = data.flat_inner_outer_dims<T, 2>(segment_ids.dims() - 1);
    reduction_functor_(context, segment_ids.shape(), segment_flat, data_flat,
                       output_flat);
  void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape,
                  typename TTypes<T, 2>::ConstTensor data,
__global__ void UnsortedSegmentCustomKernel(const int64 input_outer_dim_size,
                                            const int64 inner_dim_size,
                                            const int64 output_outer_dim_size,
  const int64 input_total_size = input_outer_dim_size * inner_dim_size;
  for (int64 input_index : GpuGridRangeX(input_total_size)) {
    const int64 input_segment_index = input_index / inner_dim_size;
    const int64 segment_offset = input_index % inner_dim_size;
    if (output_segment_index < 0 ||
        output_segment_index >= output_outer_dim_size) {
    const int64 output_index =
  void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape,
                  typename TTypes<T, 2>::ConstTensor data,
    const int64 data_size = data.size();
    const int64 input_outer_dim_size = segment_ids.dimension(0);
    const int64 input_inner_dim_size = data.dimension(1);
    const int64 output_outer_dim_size = output.dimension(0);
    TF_CHECK_OK(GpuLaunchKernel(
        UnsortedSegmentCustomKernel<T, Index, ReductionF>, config.block_count,
        config.thread_per_block, 0, d.stream(), input_outer_dim_size,
        input_inner_dim_size, output_outer_dim_size, segment_ids.data(),
        data.data(), output.data()));