CVE-2022-36004: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. When `tf.random.gamma` receives large input shape and rates, it gives a `CHECK` fail that can trigger a denial of service attack. We have patched the issue in GitHub commit 552bfced6ce4809db5f3ca305f60ff80dd40c5a3. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: The product contains an assert() or similar statement that can be triggered by an attacker, which leads to an application exit or other behavior that is more severe than necessary.
CVSS v3 base score: 7.5
Published: 2022-09-16T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/random_op.cc, tensorflow/core/kernels/random_poisson_op.cc, tensorflow/python/kernel_tests/random/random_gamma_test.py, tensorflow/python/kernel_tests/random/random_poisson_test.py
Commit date: 2022-07-23 22:06:53 -0700
Commit message: Fix size check for large input shape and rates.

To address a check failure for exceedingly large output shapes, we need to
`AppendShapeWithStatus`.

PiperOrigin-RevId: 462885894

Diff: diff --git a/tensorflow/core/kernels/random_op.cc b/tensorflow/core/kernels/random_op.cc
index 8ec9ed7d24b08..b85c120eba9aa 100644
--- a/tensorflow/core/kernels/random_op.cc
+++ b/tensorflow/core/kernels/random_op.cc
@@ -164,11 +164,11 @@ class RandomGammaOp : public OpKernel {
       OP_REQUIRES_OK(ctx, TensorShapeUtils::MakeShape(vec.data(), vec.size(),
                                                       &samples_shape));
     }
     const int64_t samples_per_alpha = samples_shape.num_elements();
 
-    samples_shape.AppendShape(alpha_t.shape());
+    OP_REQUIRES_OK(ctx, samples_shape.AppendShapeWithStatus(alpha_t.shape()));
     // Allocate output samples.
     Tensor* samples_t = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, samples_shape, &samples_t));
 
     if (samples_shape.num_elements() == 0) return;
diff --git a/tensorflow/core/kernels/random_poisson_op.cc b/tensorflow/core/kernels/random_poisson_op.cc
index b4c4d5d95c188..a14bee790753c 100644
--- a/tensorflow/core/kernels/random_poisson_op.cc
+++ b/tensorflow/core/kernels/random_poisson_op.cc
@@ -294,12 +294,12 @@ class RandomPoissonOp : public OpKernel {
     const Tensor& rate_t = ctx->input(1);
 
     TensorShape samples_shape;
     OP_REQUIRES_OK(ctx, tensor::MakeShape(shape_t, &samples_shape));
     const int64_t num_samples = samples_shape.num_elements();
+    OP_REQUIRES_OK(ctx, samples_shape.AppendShapeWithStatus(rate_t.shape()));
 
-    samples_shape.AppendShape(rate_t.shape());
     // Allocate output samples.
     Tensor* samples_t = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, samples_shape, &samples_t));
     if (num_samples == 0) return;
 
diff --git a/tensorflow/python/kernel_tests/random/random_gamma_test.py b/tensorflow/python/kernel_tests/random/random_gamma_test.py
index 71b06da183f33..4d3ea4fe5284e 100644
--- a/tensorflow/python/kernel_tests/random/random_gamma_test.py
+++ b/tensorflow/python/kernel_tests/random/random_gamma_test.py
@@ -14,11 +14,14 @@
 # ==============================================================================
 """Tests for tensorflow.ops.random_ops.random_gamma."""
 
 import numpy as np
 
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
 from tensorflow.python.framework import test_util
 from tensorflow.python.kernel_tests.random import util
 from tensorflow.python.ops import array_ops
@@ -214,8 +217,18 @@ class RandomGammaTest(test.TestCase):
       with self.cached_session():
         x = random_ops.random_gamma(shape=[n], alpha=0.001, dtype=dt, seed=0)
         self.assertEqual(0, math_ops.reduce_sum(math_ops.cast(
             math_ops.less_equal(x, 0.), dtype=dtypes.int64)).eval())
 
+  def testSizeTooLarge(self):
+    # Grappler asserts on size overflow, so this error is only caught when
+    # running eagerly.
+    if context.executing_eagerly():
+      with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                  "overflow"):
+        rate = constant_op.constant(1.0, shape=(4, 4, 4, 4, 4))
+        self.evaluate(
+            random_ops.random_gamma(
+                shape=[46902, 51188, 34063, 59195], alpha=rate))
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/kernel_tests/random/random_poisson_test.py b/tensorflow/python/kernel_tests/random/random_poisson_test.py
index 9f21f91ed5728..c0470e6029aa9 100644
--- a/tensorflow/python/kernel_tests/random/random_poisson_test.py
+++ b/tensorflow/python/kernel_tests/random/random_poisson_test.py
@@ -15,10 +15,11 @@
 """Tests for tensorflow.ops.random_ops.random_poisson."""
 import numpy as np
 
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.kernel_tests.random import util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import random_ops
@@ -169,8 +170,16 @@ class RandomPoissonTest(test.TestCase):
   @test_util.run_deprecated_v1
   def testInfRate(self):
     sample = random_ops.random_poisson(shape=[2], lam=np.inf)
     self.assertAllEqual([np.inf, np.inf], self.evaluate(sample))
 
+  def testSizeTooLarge(self):
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "overflow"):
+      rate = constant_op.constant(1.0, shape=(4, 4, 4, 4, 4))
+      self.evaluate(
+          random_ops.random_poisson(
+              shape=[46902, 51188, 34063, 59195], lam=rate))
+
 
 if __name__ == "__main__":
   test.main()

Vulnerable code:     samples_shape.AppendShape(alpha_t.shape());
    samples_shape.AppendShape(rate_t.shape());
Fixed code:     OP_REQUIRES_OK(ctx, samples_shape.AppendShapeWithStatus(alpha_t.shape()));
    OP_REQUIRES_OK(ctx, samples_shape.AppendShapeWithStatus(rate_t.shape()));
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors
  def testSizeTooLarge(self):
    # Grappler asserts on size overflow, so this error is only caught when
    # running eagerly.
    if context.executing_eagerly():
      with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                  "overflow"):
        rate = constant_op.constant(1.0, shape=(4, 4, 4, 4, 4))
        self.evaluate(
            random_ops.random_gamma(
                shape=[46902, 51188, 34063, 59195], alpha=rate))
from tensorflow.python.framework import errors
  def testSizeTooLarge(self):
    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "overflow"):
      rate = constant_op.constant(1.0, shape=(4, 4, 4, 4, 4))
      self.evaluate(
          random_ops.random_poisson(
              shape=[46902, 51188, 34063, 59195], lam=rate))