CVE-2021-41195: Integer Overflow or Wraparound (CWE-190)
value: TensorFlow is an open source platform for machine learning. In affected versions the implementation of `tf.math.segment_*` operations results in a `CHECK`-fail related abort (and denial of service) if a segment id in `segment_ids` is large. This is similar to CVE-2021-29584 (and similar other reported vulnerabilities in TensorFlow, localized to specific APIs): the implementation (both on CPU and GPU) computes the output shape using `AddDim`. However, if the number of elements in the tensor overflows an `int64_t` value, `AddDim` results in a `CHECK` failure which provokes a `std::abort`. Instead, code should use `AddDimWithStatus`. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.
CWE description: The product performs a calculation that can
         produce an integer overflow or wraparound when the logic
         assumes that the resulting value will always be larger than
         the original value. This occurs when an integer value is
         incremented to a value that is too large to store in the
         associated representation. When this occurs, the value may
         become a very small or negative number.
CVSS v3 base score: 5.5
Published: 2021-11-05T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/segment_reduction_ops_impl.h, tensorflow/python/kernel_tests/segment_reduction_ops_test.py
Commit date: 2021-10-27 17:36:34 -0700
Commit message: Merge pull request #51733 from yongtang:46888-tf.math.segment_

PiperOrigin-RevId: 406020083
Change-Id: I179a9a8fe548ed324fc97363e81a46be28aa19b8

Diff: diff --cc tensorflow/core/kernels/segment_reduction_ops_impl.h
index e257321fdc65f,0fb5087c51743..a44ede6dbe531
--- a/tensorflow/core/kernels/segment_reduction_ops_impl.h
+++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h
@@@ -108,11 -108,11 +108,13 @@@ class SegmentReductionOp : public OpKer
  
      OP_REQUIRES(context, input.dims() >= 1,
                  errors::InvalidArgument("Shape must be at least rank 1"));
  
      TensorShape output_shape = input.shape();
-     output_shape.set_dim(0, output_rows);
++    // Since we're changing the first dimension of the shape, we need to make
++    // sure the new shape won't overflow.
+     OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(0, output_rows));
  
      // Note that we do not initialize the output buffer with a default value, so
      // we need to explicitly set missing indices to the default value.
      Tensor* output = nullptr;
      OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
@@@ -288,11 -288,11 +290,14 @@@ class SegmentReductionGPUOp : public As
        OP_REQUIRES_ASYNC(context, output_rows > 0,
                          errors::InvalidArgument("segment ids must be >= 0"),
                          done);
  
        TensorShape output_shape = input.shape();
--      output_shape.set_dim(0, output_rows);
++      // Since we're changing the first dimension of the shape, we need to make
++      // sure the new shape won't overflow.
++      OP_REQUIRES_OK_ASYNC(context,
++                           output_shape.SetDimWithStatus(0, output_rows), done);
  
        Tensor* output = nullptr;
        OP_REQUIRES_OK_ASYNC(
            context, context->allocate_output(0, output_shape, &output), done);
  
diff --cc tensorflow/python/kernel_tests/segment_reduction_ops_test.py
index 4c159a72817f1,98652b11515a4..8293dba130028
--- a/tensorflow/python/kernel_tests/segment_reduction_ops_test.py
+++ b/tensorflow/python/kernel_tests/segment_reduction_ops_test.py
@@@ -264,10 -265,24 +264,24 @@@ class SegmentReductionOpTest(SegmentRed
              "must be at least rank 1"):
            s = math_ops.segment_mean(
                data=np.uint16(10), segment_ids=np.array([]).astype("int64"))
            self.evaluate(s)
  
+   def testInvalidIds(self):
+     # Test case for GitHub issue 46888.
+     for op in [
+         math_ops.segment_max,
+         math_ops.segment_min,
+         math_ops.segment_mean,
+         math_ops.segment_sum,
+         math_ops.segment_prod,
+     ]:
 -      with self.cached_session():
++      with self.cached_session(use_gpu=False):
+         with self.assertRaises((ValueError, errors_impl.InternalError)):
+           s = op(data=np.ones((1, 10, 1)), segment_ids=[1676240524292489355])
+           self.evaluate(s)
+ 
  
  class UnsortedSegmentTest(SegmentReductionHelper):
  
    def __init__(self, methodName='runTest'):
      # Each item is np_op1, np_op2, tf_op, initial_value functor

Vulnerable code:      output_shape.set_dim(0, output_rows);
-      output_shape.set_dim(0, output_rows);
Fixed code: +    // Since we're changing the first dimension of the shape, we need to make
+    // sure the new shape won't overflow.
     OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(0, output_rows));
+      // Since we're changing the first dimension of the shape, we need to make
+      // sure the new shape won't overflow.
+      OP_REQUIRES_OK_ASYNC(context,
+                           output_shape.SetDimWithStatus(0, output_rows), done);
   def testInvalidIds(self):
     # Test case for GitHub issue 46888.
     for op in [
         math_ops.segment_max,
         math_ops.segment_min,
         math_ops.segment_mean,
         math_ops.segment_sum,
         math_ops.segment_prod,
     ]:
+      with self.cached_session(use_gpu=False):
         with self.assertRaises((ValueError, errors_impl.InternalError)):
           s = op(data=np.ones((1, 10, 1)), segment_ids=[1676240524292489355])
           self.evaluate(s)