CVE-2019-16778: Incorrect Conversion between Numeric Types (CWE-681)
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.
CWE description: 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.
CVSS v3 base score: 9.8
Published: 2019-12-16T21:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
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 date: 2019-07-03 18:10:21 -0700
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

Diff: diff --git a/tensorflow/core/kernels/segment_reduction_ops.cc b/tensorflow/core/kernels/segment_reduction_ops.cc
index fabb7f6de1..e26bc595f5 100644
--- a/tensorflow/core/kernels/segment_reduction_ops.cc
+++ b/tensorflow/core/kernels/segment_reduction_ops.cc
@@ -374,32 +374,31 @@ namespace functor {
 
 // The ReductionFunctor implementation for CPU.
 template <typename T, typename Index, typename InitialValueF,
           typename ReductionF>
 struct UnsortedSegmentFunctor<CPUDevice, T, Index, InitialValueF, ReductionF> {
-  void operator()(OpKernelContext* ctx, const Index num_segments,
-                  const TensorShape& segment_ids_shape,
+  void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape,
                   typename TTypes<Index>::ConstFlat segment_ids,
-                  const Index data_size, const T* data,
+                  typename TTypes<T, 2>::ConstTensor data,
                   typename TTypes<T, 2>::Tensor output) {
     output.setConstant(InitialValueF()());
-    if (data_size == 0) {
+    if (data.size() == 0) {
       return;
     }
     const int64 N = segment_ids.dimension(0);
+    const int64 num_segments = output.dimension(0);
     ReductionF reduction;
-    auto data_flat = typename TTypes<T, 2>::ConstTensor(data, N, data_size / N);
     for (int64 i = 0; i < N; ++i) {
       Index j = internal::SubtleMustCopy(segment_ids(i));
       if (j < 0) {
         continue;
       }
       OP_REQUIRES(ctx, FastBoundsCheck(j, num_segments),
                   errors::InvalidArgument(
                       "segment_ids", SliceDebugString(segment_ids_shape, i),
                       " = ", j, " is out of range [0, ", num_segments, ")"));
-      reduction(data_flat.template chip<0>(i), output.template chip<0>(j));
+      reduction(data.template chip<0>(i), output.template chip<0>(j));
     }
   }
 };
 
 template <typename T>
@@ -483,11 +482,11 @@ class UnsortedSegmentReductionOp : public OpKernel {
     if (!UnsortedSegmentReductionDoValidation(this, context, data, segment_ids,
                                               num_segments)) {
       return;
     }
     const auto segment_flat = segment_ids.flat<Index>();
-    const Index output_rows = internal::SubtleMustCopy(static_cast<Index>(
+    const int64 output_rows = internal::SubtleMustCopy(static_cast<int64>(
         num_segments.dtype() == DT_INT32 ? num_segments.scalar<int32>()()
                                          : num_segments.scalar<int64>()()));
     OP_REQUIRES(context, output_rows >= 0,
                 errors::InvalidArgument("Input num_segments == ", output_rows,
                                         " must not be negative."));
@@ -497,13 +496,13 @@ class UnsortedSegmentReductionOp : public OpKernel {
       output_shape.AddDim(data.dim_size(i));
     }
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
     auto output_flat = output->flat_outer_dims<T>();
-    auto data_ptr = data.template flat<T>().data();
-    reduction_functor_(context, output_rows, segment_ids.shape(), segment_flat,
-                       data.NumElements(), data_ptr, output_flat);
+    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);
   }
 
  protected:
   DeviceReductionFunctor reduction_functor_;
 };
diff --git a/tensorflow/core/kernels/segment_reduction_ops.h b/tensorflow/core/kernels/segment_reduction_ops.h
index ee8c94e8d8..dc311ca9d3 100644
--- a/tensorflow/core/kernels/segment_reduction_ops.h
+++ b/tensorflow/core/kernels/segment_reduction_ops.h
@@ -57,14 +57,13 @@ struct SegmentSumFunctor {
 #endif
 
 template <typename Device, typename T, typename Index, typename InitialValueF,
           typename ReductionF>
 struct UnsortedSegmentFunctor {
-  void operator()(OpKernelContext* ctx, const Index num_segments,
-                  const TensorShape& segment_ids_shape,
+  void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape,
                   typename TTypes<Index>::ConstFlat segment_ids,
-                  const Index data_size, const T* data,
+                  typename TTypes<T, 2>::ConstTensor data,
                   typename TTypes<T, 2>::Tensor output);
 };
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 // reduction functors for the gpu
diff --git a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.cc b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.cc
index 0623df166f..490c3d2de8 100644
--- a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.cc
+++ b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.cc
@@ -104,25 +104,25 @@ __global__ void SortedSegmentSumCustomKernel(const Index input_outer_dim_size,
 
 // UnsortedSegmentSumKernel processes 'input_total_size' elements.
 // Each element is mapped from input to output by a combination of its
 // 'segment_ids' mapping and 'inner_dim_size'.
 template <typename T, typename Index, typename KernelReductionFunctor>
-__global__ void UnsortedSegmentCustomKernel(const Index input_outer_dim_size,
-                                            const Index inner_dim_size,
-                                            const Index output_outer_dim_size,
+__global__ void UnsortedSegmentCustomKernel(const int64 input_outer_dim_size,
+                                            const int64 inner_dim_size,
+                                            const int64 output_outer_dim_size,
                                             const Index* segment_ids,
                                             const T* input, T* output) {
-  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;
+  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;
     const Index output_segment_index = segment_ids[input_segment_index];
-    if (output_segment_index < 0 || output_segment_index >= output_total_size) {
+    if (output_segment_index < 0 ||
+        output_segment_index >= output_outer_dim_size) {
       continue;
     }
-    const Index output_index =
+    const int64 output_index =
         output_segment_index * inner_dim_size + segment_offset;
     KernelReductionFunctor()(output + output_index, ldg(input + input_index));
   }
 }
 
@@ -172,41 +172,42 @@ void SegmentSumFunctor<T, Index>::operator()(
 }
 
 template <typename T, typename Index, typename InitialValueF,
           typename ReductionF>
 struct UnsortedSegmentFunctor<GPUDevice, T, Index, InitialValueF, ReductionF> {
-  void operator()(OpKernelContext* ctx, const Index num_segments,
-                  const TensorShape& segment_ids_shape,
+  void operator()(OpKernelContext* ctx, const TensorShape& segment_ids_shape,
                   typename TTypes<Index>::ConstFlat segment_ids,
-                  const Index data_size, const T* data,
+                  typename TTypes<T, 2>::ConstTensor data,
                   typename TTypes<T, 2>::Tensor output) {
     if (output.size() == 0) {
       return;
     }
     // Set 'output' to initial value.
     GPUDevice d = ctx->template eigen_device<GPUDevice>();
     GpuLaunchConfig config = GetGpuLaunchConfig(output.size(), d);
     TF_CHECK_OK(GpuLaunchKernel(
         SetToValue<T>, config.block_count, config.thread_per_block, 0,
         d.stream(), output.size(), output.data(), InitialValueF()()));
+    const int64 data_size = data.size();
     if (data_size == 0 || segment_ids_shape.num_elements() == 0) {
       return;
     }
     // Launch kernel to compute unsorted segment reduction.
     // Notes:
     // *) 'data_size' is the total number of elements to process.
     // *) 'segment_ids.shape' is a prefix of data's shape.
     // *) 'input_outer_dim_size' is the total number of segments to process.
-    const Index input_outer_dim_size = segment_ids.dimension(0);
-    const Index input_inner_dim_size = data_size / input_outer_dim_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);
     config = GetGpuLaunchConfig(data_size, d);
 
-    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()));
+    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()));
   }
 };
 
 #define DEFINE_SORTED_GPU_SPECS_INDEX(T, Index) \
   template struct SegmentSumFunctor<T, Index>

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()));