CVE-2021-41222: Insufficient Information (NVD-CWE-noinfo)
value: TensorFlow is an open source platform for machine learning. In affected versions the implementation of `SplitV` can trigger a segfault is an attacker supplies negative arguments. This occurs whenever `size_splits` contains more than one value and at least one value is negative. 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: There is insufficient information about the issue to classify it; details are unkown or unspecified.
CVSS v3 base score: 5.5
Published: 2021-11-05T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/split_v_op.cc, tensorflow/core/ops/array_ops.cc, tensorflow/python/kernel_tests/split_op_test.py
Commit date: 2021-10-05 11:41:35 -0700
Commit message: A negative size in one of the split sizes allowed the computed size of another
to exceed the total dimension, leading to a segfault and security vulnerability.
Adding a check for negative sizes prevents this.

PiperOrigin-RevId: 401035665
Change-Id: I79bbe329787dac82aa4bf60397a9129b716aedab

Diff: diff --git a/tensorflow/core/kernels/split_v_op.cc b/tensorflow/core/kernels/split_v_op.cc
index 24efc6eff0d67..d7711dad3be10 100644
--- a/tensorflow/core/kernels/split_v_op.cc
+++ b/tensorflow/core/kernels/split_v_op.cc
@@ -136,10 +136,17 @@ class SplitVOpBase : public OpKernel {
 
     if (neg_one_dim >= 0) {
       (*split_sizes_vec)[neg_one_dim] = input_size_split_dim - determined_size;
     }
 
+    for (int i = 0; i < split_sizes_vec->size(); ++i) {
+      const Tlen& split_size = (*split_sizes_vec)[i];
+      OP_REQUIRES(context, split_size >= Tlen(0),
+                  errors::InvalidArgument("Split size at index ", i,
+                                          " must be >= 0. Got: ", split_size));
+    }
+
     // Special case 2: split along the 1st dimension. The requirements are that
     // either we are splitting the outer dimension of two or more such that
     // every outer subpart is aligned or that the split sizes mean that they are
     // always aligned. In these cases, we can share the underlying buffer.
     //
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 34d09931c0408..775a4cc469664 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -679,10 +679,16 @@ REGISTER_OP("SplitV")
         for (int i = 0; i < num_outputs; ++i) {
           auto size = data[i];
           if (data[i] == -1 && c->ValueKnown(split_dim_size)) {
             size = split_dim_size - total_size;
           }
+          // If we have a negative known size (either explicit, or computed
+          // via -1), then the split sizes are invalid.
+          if (size < -1 || (size == -1 && c->ValueKnown(split_dim_size))) {
+            return errors::InvalidArgument("Split size at index ", i,
+                                           " must be >= 0. Got: ", size);
+          }
           TF_RETURN_IF_ERROR(
               c->ReplaceDim(input, split_dim, c->MakeDim(size), &output_shape));
           c->set_output(i, output_shape);
         }
         if (c->ValueKnown(split_dim_size)) {
diff --git a/tensorflow/python/kernel_tests/split_op_test.py b/tensorflow/python/kernel_tests/split_op_test.py
index 78eef5aadf8a7..a89715c743be1 100644
--- a/tensorflow/python/kernel_tests/split_op_test.py
+++ b/tensorflow/python/kernel_tests/split_op_test.py
@@ -382,8 +382,26 @@ class SplitOpTest(test.TestCase):
     with self.session() as sess:
       with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
                                   "must have exactly one element"):
         sess.run(y, {x: np.array([], dtype=np.int32), splits: [4, 11, 15]})
 
+  @test_util.run_in_graph_and_eager_modes
+  def testNegativeSizes(self):
+    x = constant_op.constant([1, 2, 3], dtypes.float32)
+    # A size of -1 signifies to determine size based on sum of other splits.
+    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
+                                "Split size at index 1 must be >= 0. Got: -2"):
+      splits = [-1, -2]
+      self.evaluate(array_ops.split(x, splits, axis=0))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testBadSplitSizes(self):
+    x = constant_op.constant([1, 2], dtypes.float32)
+    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
+                                "Determined shape must either match input"
+                                "|can't split axis"):
+      splits = [1, 2]
+      self.evaluate(array_ops.split(x, splits, axis=0))
+
 
 if __name__ == "__main__":
   test.main()

Fixed code:     for (int i = 0; i < split_sizes_vec->size(); ++i) {
      const Tlen& split_size = (*split_sizes_vec)[i];
      OP_REQUIRES(context, split_size >= Tlen(0),
                  errors::InvalidArgument("Split size at index ", i,
                                          " must be >= 0. Got: ", split_size));
    }

          // If we have a negative known size (either explicit, or computed
          // via -1), then the split sizes are invalid.
          if (size < -1 || (size == -1 && c->ValueKnown(split_dim_size))) {
            return errors::InvalidArgument("Split size at index ", i,
                                           " must be >= 0. Got: ", size);
          }
  @test_util.run_in_graph_and_eager_modes
  def testNegativeSizes(self):
    x = constant_op.constant([1, 2, 3], dtypes.float32)
    # A size of -1 signifies to determine size based on sum of other splits.
    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
                                "Split size at index 1 must be >= 0. Got: -2"):
      splits = [-1, -2]
      self.evaluate(array_ops.split(x, splits, axis=0))

  @test_util.run_in_graph_and_eager_modes
  def testBadSplitSizes(self):
    x = constant_op.constant([1, 2], dtypes.float32)
    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
                                "Determined shape must either match input"
                                "|can't split axis"):
      splits = [1, 2]
      self.evaluate(array_ops.split(x, splits, axis=0))