CVE-2021-41220 (Python, published 2021-11-05T23:15Z)
Description: 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.
Severity: MEDIUM (CVSS v3 base score: 7.8)
Weakness: CWE-416 \u2014 Use After Free
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.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/collective_ops.cc
tensorflow/python/kernel_tests/collective_ops_test.py
Commit message: Fix undefined behavior in CollectiveReduceV2 and others

We should not call done after it's moved.

PiperOrigin-RevId: 400838185
Change-Id: Ifc979740054b8f8c6f4d50acc89472fe60c4fdb1


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)