CVE-2021-29584: Integer Overflow or Wraparound (CWE-190)
value: TensorFlow is an end-to-end open source platform for machine learning. An attacker can trigger a denial of service via a `CHECK`-fail in caused by an integer overflow in constructing a new tensor shape. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/0908c2f2397c099338b901b067f6495a5b96760b/tensorflow/core/kernels/sparse_split_op.cc#L66-L70) builds a dense shape without checking that the dimensions would not result in overflow. The `TensorShape` constructor(https://github.com/tensorflow/tensorflow/blob/6f9896890c4c703ae0a0845394086e2e1e523299/tensorflow/core/framework/tensor_shape.cc#L183-L188) uses a `CHECK` operation which triggers when `InitDims`(https://github.com/tensorflow/tensorflow/blob/6f9896890c4c703ae0a0845394086e2e1e523299/tensorflow/core/framework/tensor_shape.cc#L212-L296) returns a non-OK status. This is a legacy implementation of the constructor and operations should use `BuildTensorShapeBase` or `AddDimWithStatus` to prevent `CHECK`-failures in the presence of overflows. 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.
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-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/sparse_split_op.cc
Commit date: 2021-05-06 16:01:09 -0700
Commit message: Prevent overflow in sparse op

PiperOrigin-RevId: 372442006
Change-Id: I60fe31cd7e56fb3501e97c63500caf902ddeee96

Diff: diff --git a/tensorflow/core/kernels/sparse_split_op.cc b/tensorflow/core/kernels/sparse_split_op.cc
index 5c0457aa956c8..3b88a8ca2bf6e 100644
--- a/tensorflow/core/kernels/sparse_split_op.cc
+++ b/tensorflow/core/kernels/sparse_split_op.cc
@@ -61,15 +61,22 @@ class SparseSplitOp : public OpKernel {
                 errors::InvalidArgument("Input num_split should be between 1 "
                                         "and the splitting dimension size (",
                                         input_shape.vec<int64>()(axis),
                                         "), got ", num_split_));
 
+    // Prevent overflow by constructing the dense shape separately
+    TensorShape dense_shape;
+    const auto input_shape_flat = input_shape.flat<int64>();
+    for (int i = 0; i < input_shape.NumElements(); i++) {
+      OP_REQUIRES_OK(context,
+                     dense_shape.AddDimWithStatus(input_shape_flat(i)));
+    }
+
     sparse::SparseTensor sparse_tensor;
     OP_REQUIRES_OK(context,
-                   sparse::SparseTensor::Create(
-                       input_indices, input_values,
-                       TensorShape(input_shape.vec<int64>()), &sparse_tensor));
+                   sparse::SparseTensor::Create(input_indices, input_values,
+                                                dense_shape, &sparse_tensor));
 
     std::vector<sparse::SparseTensor> outputs;
     OP_REQUIRES_OK(context, sparse::SparseTensor::Split<T>(
                                 sparse_tensor, axis, num_split_, &outputs));
 

Vulnerable code:                    sparse::SparseTensor::Create(
                       input_indices, input_values,
                       TensorShape(input_shape.vec<int64>()), &sparse_tensor));
Fixed code:     // Prevent overflow by constructing the dense shape separately
    TensorShape dense_shape;
    const auto input_shape_flat = input_shape.flat<int64>();
    for (int i = 0; i < input_shape.NumElements(); i++) {
      OP_REQUIRES_OK(context,
                     dense_shape.AddDimWithStatus(input_shape_flat(i)));
    }

                   sparse::SparseTensor::Create(input_indices, input_values,
                                                dense_shape, &sparse_tensor));