CVE-2021-41220: Use After Free (CWE-416)
value: TensorFlow is an open source platform for machine learning. In affected versions the async implementation of `CollectiveReduceV2` suffers from a memory leak and a use after free. This occurs due to the asynchronous computation and the fact that objects that have been `std::move()`d from are still accessed. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, as this version is the only one that is also affected.
CWE description: The product reuses or references memory after it has been freed. At some point afterward, the memory may be allocated again and saved in another pointer, while the original pointer references a location somewhere within the new allocation. Any operations using the original pointer are no longer valid because the memory "belongs" to the code that operates on the new pointer.
CVSS v3 base score: 7.8
Published: 2021-11-05T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/collective_ops.cc, tensorflow/python/kernel_tests/collective_ops_test.py
Commit date: 2021-10-04 16:15:01 -0700
Commit message: Fix undefined behavior in CollectiveReduceV2 and others

We should not call done after it's moved.

PiperOrigin-RevId: 400838185
Change-Id: Ifc979740054b8f8c6f4d50acc89472fe60c4fdb1

Diff: diff --git a/tensorflow/core/kernels/collective_ops.cc b/tensorflow/core/kernels/collective_ops.cc
index c59464cad7001..1824148a401a3 100644
--- a/tensorflow/core/kernels/collective_ops.cc
+++ b/tensorflow/core/kernels/collective_ops.cc
@@ -492,19 +492,21 @@ class CollectiveOpV2Kernel : public AsyncOpKernel {
   Status FillCollectiveParams(CollectiveParams* col_params,
                               CollectiveType collective_type,
                               const Tensor& group_size, const Tensor& group_key,
                               const Tensor& instance_key) {
     if (group_size.dims() > 0) {
-      return errors::Internal("Unexpected dimensions on input group_size, got ",
-                              group_size.shape().DebugString());
+      return errors::InvalidArgument(
+          "Unexpected dimensions on input group_size, got ",
+          group_size.shape().DebugString());
     }
     if (group_key.dims() > 0) {
-      return errors::Internal("Unexpected dimensions on input group_key, got ",
-                              group_key.shape().DebugString());
+      return errors::InvalidArgument(
+          "Unexpected dimensions on input group_key, got ",
+          group_key.shape().DebugString());
     }
     if (instance_key.dims() > 0) {
-      return errors::Internal(
+      return errors::InvalidArgument(
           "Unexpected dimensions on input instance_key, got ",
           instance_key.shape().DebugString());
     }
     col_params->name = name_;
     col_params->group.device_type = device_type_;
@@ -623,11 +625,11 @@ class CollectiveReduceV2OpKernel : public CollectiveOpV2Kernel {
     OP_REQUIRES_OK_ASYNC(c,
                          FillCollectiveParams(col_params, REDUCTION_COLLECTIVE,
                                               /*group_size*/ c->input(1),
                                               /*group_key*/ c->input(2),
                                               /*instance_key*/ c->input(3)),
-                         done);
+                         done_with_cleanup);
     col_params->instance.shape = c->input(0).shape();
     col_params->merge_op = merge_op_.get();
     col_params->final_op = final_op_.get();
     VLOG(1) << "CollectiveReduceV2 group_size " << col_params->group.group_size
             << " group_key " << col_params->group.group_key << " instance_key "
@@ -853,18 +855,19 @@ class CollectiveInitializeCommunicatorOpKernel : public AsyncOpKernel {
     device_type_ = c->device_type();
   }
 
   Status CheckInputs(Tensor group_size_t, Tensor group_key_t) {
     if (group_size_t.dims() > 0) {
-      return errors::Internal(
+      return errors::InvalidArgument(
           "Unexpected dimensions on input group_size. "
           "It shoulbe a scalar, got tensor with shape ",
           group_size_t.shape().DebugString());
     }
     if (group_key_t.dims() > 0) {
-      return errors::Internal("Unexpected dimensions on input group_key, got ",
-                              group_key_t.shape().DebugString());
+      return errors::InvalidArgument(
+          "Unexpected dimensions on input group_key, got ",
+          group_key_t.shape().DebugString());
     }
 
     auto group_size = group_size_t.unaligned_flat<int32>()(0);
     if (group_size <= 0) {
       return errors::InvalidArgument(
@@ -1082,11 +1085,11 @@ class CollectiveReduceV3OpKernel : public CollectiveOpV3Kernel {
       done();
       col_params->Unref();
     };
     core::RefCountPtr<CollectiveGroupResource> resource;
     OP_REQUIRES_OK_ASYNC(c, LookupResource(c, HandleFromInput(c, 1), &resource),
-                         done);
+                         done_with_cleanup);
 
     Tensor group_assignment = c->input(2);
 
     OP_REQUIRES_OK_ASYNC(
         c,
@@ -1132,11 +1135,11 @@ class CollectiveAllToAllV3OpKernel : public CollectiveOpV3Kernel {
       done();
       col_params->Unref();
     };
     core::RefCountPtr<CollectiveGroupResource> resource;
     OP_REQUIRES_OK_ASYNC(c, LookupResource(c, HandleFromInput(c, 1), &resource),
-                         done);
+                         done_with_cleanup);
 
     Tensor group_assignment = c->input(2);
 
     OP_REQUIRES_OK_ASYNC(
         c,
diff --git a/tensorflow/python/kernel_tests/collective_ops_test.py b/tensorflow/python/kernel_tests/collective_ops_test.py
index 73e8110dbdc95..98b95c5666c80 100644
--- a/tensorflow/python/kernel_tests/collective_ops_test.py
+++ b/tensorflow/python/kernel_tests/collective_ops_test.py
@@ -1180,10 +1180,73 @@ class InputPipelineTest(test.TestCase):
       return value0, value1
 
     self.assertAllEqual(self.evaluate(f()), [[3.], [3.]])
 
 
+@combinations.generate(
+    combinations.times(
+        combinations.combine(collective_op=[
+            combinations.NamedObject('all_reduce_v2',
+                                     CollectiveOpsV2.all_reduce),
+            combinations.NamedObject('all_gather_v2',
+                                     CollectiveOpsV2.all_gather)
+        ]), device_combination))
+class InvalidInputTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    _setup_context()
+    super().setUp()
+
+  def testInvalidGroupKey(self, collective_op, device, communication):
+    dev0 = '/device:%s:0' % device
+    group_size = 2
+    group_key = [100]
+    instance_key = 100
+    in_tensor = constant_op.constant([1.])
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      with ops.device(dev0):
+        collective_op(
+            in_tensor,
+            group_size,
+            group_key,
+            instance_key,
+            communication_hint=communication)
+
+  def testInvalidGroupSize(self, collective_op, device, communication):
+    dev0 = '/device:%s:0' % device
+    group_size = -2
+    group_key = 100
+    instance_key = 100
+    in_tensor = constant_op.constant([1.])
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      with ops.device(dev0):
+        collective_op(
+            in_tensor,
+            group_size,
+            group_key,
+            instance_key,
+            communication_hint=communication)
+
+  def testInvalidInstanceKey(self, collective_op, device, communication):
+    dev0 = '/device:%s:0' % device
+    group_size = 2
+    group_key = 100
+    instance_key = [100]
+    in_tensor = constant_op.constant([1.])
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      with ops.device(dev0):
+        collective_op(
+            in_tensor,
+            group_size,
+            group_key,
+            instance_key,
+            communication_hint=communication)
+
+
 class CollectiveOpsV3Test(test.TestCase, parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
     _setup_context()

Vulnerable code:       return errors::Internal("Unexpected dimensions on input group_size, got ",
                              group_size.shape().DebugString());
      return errors::Internal("Unexpected dimensions on input group_key, got ",
                              group_key.shape().DebugString());
      return errors::Internal(
                         done);
      return errors::Internal(
      return errors::Internal("Unexpected dimensions on input group_key, got ",
                              group_key_t.shape().DebugString());
                         done);
                         done);
Fixed code:       return errors::InvalidArgument(
          "Unexpected dimensions on input group_size, got ",
          group_size.shape().DebugString());
      return errors::InvalidArgument(
          "Unexpected dimensions on input group_key, got ",
          group_key.shape().DebugString());
      return errors::InvalidArgument(
                         done_with_cleanup);
      return errors::InvalidArgument(
      return errors::InvalidArgument(
          "Unexpected dimensions on input group_key, got ",
          group_key_t.shape().DebugString());
                         done_with_cleanup);
                         done_with_cleanup);
@combinations.generate(
    combinations.times(
        combinations.combine(collective_op=[
            combinations.NamedObject('all_reduce_v2',
                                     CollectiveOpsV2.all_reduce),
            combinations.NamedObject('all_gather_v2',
                                     CollectiveOpsV2.all_gather)
        ]), device_combination))
class InvalidInputTest(test.TestCase, parameterized.TestCase):

  def setUp(self):
    _setup_context()
    super().setUp()

  def testInvalidGroupKey(self, collective_op, device, communication):
    dev0 = '/device:%s:0' % device
    group_size = 2
    group_key = [100]
    instance_key = 100
    in_tensor = constant_op.constant([1.])

    with self.assertRaises(errors.InvalidArgumentError):
      with ops.device(dev0):
        collective_op(
            in_tensor,
            group_size,
            group_key,
            instance_key,
            communication_hint=communication)

  def testInvalidGroupSize(self, collective_op, device, communication):
    dev0 = '/device:%s:0' % device
    group_size = -2
    group_key = 100
    instance_key = 100
    in_tensor = constant_op.constant([1.])

    with self.assertRaises(errors.InvalidArgumentError):
      with ops.device(dev0):
        collective_op(
            in_tensor,
            group_size,
            group_key,
            instance_key,
            communication_hint=communication)

  def testInvalidInstanceKey(self, collective_op, device, communication):
    dev0 = '/device:%s:0' % device
    group_size = 2
    group_key = 100
    instance_key = [100]
    in_tensor = constant_op.constant([1.])

    with self.assertRaises(errors.InvalidArgumentError):
      with ops.device(dev0):
        collective_op(
            in_tensor,
            group_size,
            group_key,
            instance_key,
            communication_hint=communication)