CVE-2021-41213 (Python, published 2021-11-05T23:15Z)
Description: value: TensorFlow is an open source platform for machine learning. In affected versions the code behind `tf.function` API can be made to deadlock when two `tf.function` decorated Python functions are mutually recursive. This occurs due to using a non-reentrant `Lock` Python object. Loading any model which contains mutually recursive functions is vulnerable. An attacker can cause denial of service by causing users to load such models and calling a recursive `tf.function`, although this is not a frequent scenario. 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.
Severity: MEDIUM (CVSS v3 base score: 5.5)
Weakness: CWE-662 \u2014 Improper Synchronization
The product utilizes multiple threads or processes to allow temporary access to a shared resource that can only be exclusive to one process at a time, but it does not properly synchronize these actions, which might cause simultaneous accesses of this resource by multiple threads or processes.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/python/eager/def_function.py
tensorflow/python/eager/def_function_test.py
tensorflow/python/eager/function.py
Commit message: Fix the deadlock issue of recursive tf.function.

Replace threading.Lock with threading.RLock to allow recursive tf.function.

PiperOrigin-RevId: 401282729
Change-Id: I3d10416f2eb2c15e2055bb4f4afee3d62bd6c428


Vulnerable code:
    self._lock = threading.Lock()
    self._lock = threading.Lock()
    self._lock = threading.Lock()

Fixed code:
    self._lock = threading.RLock()
    self._lock = threading.RLock()
from tensorflow.python.eager import backprop
from tensorflow.python.ops import cond_v2
  def test_recursive_tf_function(self):

    @def_function.function
    def recursive_fn(n):
      if n > 0:
        return recursive_fn(n - 1)
      return 1

    self.assertEqual(recursive_fn(5).numpy(), 1)

  def test_recursive_tf_function_with_gradients(self):

    @def_function.function
    def recursive_fn(n, x):
      if n > 0:
        return n * recursive_fn(n - 1, x)
      else:
        return x

    x = variables.Variable(1.0)
    with backprop.GradientTape() as tape:
      g = recursive_fn(5, x)

    dg_dx = tape.gradient(g, x)
    self.assertEqual(dg_dx.numpy(), 120)

  def test_recursive_python_function(self):

    def recursive_py_fn(n):
      if n > 0:
        return recursive_py_fn(n - 1)
      return 1

    @def_function.function
    def recursive_fn(n):
      return recursive_py_fn(n)

    self.assertEqual(recursive_fn(5).numpy(), 1)

  def test_recursive_python_function_with_gradients(self):

    def recursive_py_fn(n, x):
      if n > 0:
        return n * recursive_py_fn(n - 1, x)
      return x

    @def_function.function
    def recursive_fn(n, x):
      return recursive_py_fn(n, x)

    x = variables.Variable(1.0)
    with backprop.GradientTape() as tape:
      g = recursive_fn(5, x)

    dg_dx = tape.gradient(g, x)
    self.assertEqual(dg_dx.numpy(), 120)

  def test_recursive_tf_function_call_each_other(self):

    @def_function.function
    def recursive_fn1(n):
      if n <= 1:
        return 1
      return recursive_fn2(n - 1)

    @def_function.function
    def recursive_fn2(n):
      if n <= 1:
        return 2
      return recursive_fn1(n - 1)

    self.assertEqual(recursive_fn1(5).numpy(), 1)
    self.assertEqual(recursive_fn1(6).numpy(), 2)
    self.assertEqual(recursive_fn2(5).numpy(), 2)
    self.assertEqual(recursive_fn2(6).numpy(), 1)

  def test_recursive_tf_function_call_each_other_with_gradients(self):

    @def_function.function
    def recursive_fn1(n, x):
      if n <= 1:
        return x
      return n * recursive_fn2(n - 1, x)

    @def_function.function
    def recursive_fn2(n, x):
      if n <= 1:
        return 2 * x
      return n * recursive_fn1(n - 1, x)

    x = variables.Variable(1.0)
    with backprop.GradientTape() as tape:
      g1 = recursive_fn1(5, x)

    dg1_dx = tape.gradient(g1, x)
    self.assertEqual(dg1_dx.numpy(), 120)

    with backprop.GradientTape() as tape:
      g2 = recursive_fn2(5, x)

    dg2_dx = tape.gradient(g2, x)
    self.assertEqual(dg2_dx.numpy(), 240)

  def test_recursive_tf_function_with_cond(self):
    @def_function.function(autograph=False)
    def recursive_fn(n):
      return cond_v2.cond_v2(n > 0, recursive_fn(n - 1), 1)

    with self.assertRaises(RecursionError):
      recursive_fn(constant_op.constant(5))

    self._lock = threading.RLock()