CVE-2022-36005: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. When `tf.quantization.fake_quant_with_min_max_vars_gradient` receives input `min` or `max` that is nonscalar, it gives a `CHECK` fail that can trigger a denial of service attack. We have patched the issue in GitHub commit f3cf67ac5705f4f04721d15e485e192bb319feed. 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/fake_quant_ops.cc, tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
Commit date: 2022-07-21 22:52:41 -0700
Commit message: Add IsScalar / IsVector (rank) checks to input min/max tensors for FakeQuantWithMinMaxVarsPerChannelGradientOp and FakeQuantWithMinMaxVarsGradientOp.

PiperOrigin-RevId: 462542629

Diff: diff --git a/tensorflow/core/kernels/fake_quant_ops.cc b/tensorflow/core/kernels/fake_quant_ops.cc
index b8c31701a61c0..682459866e988 100644
--- a/tensorflow/core/kernels/fake_quant_ops.cc
+++ b/tensorflow/core/kernels/fake_quant_ops.cc
@@ -259,10 +259,16 @@ class FakeQuantWithMinMaxVarsGradientOp : public OpKernel {
     const Tensor& input = context->input(1);
     OP_REQUIRES(context, input.IsSameSize(gradient),
                 InvalidArgument("gradient and input must be the same size"));
     const Tensor& min = context->input(2);
     const Tensor& max = context->input(3);
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(min.shape()),
+        InvalidArgument("`min` must be rank 0 but is rank ", min.dims()));
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsScalar(max.shape()),
+        InvalidArgument("`max` must be rank 0 but is rank ", max.dims()));
 
     Tensor* grad_wrt_input;
     OP_REQUIRES_OK(context,
                    context->allocate_output(0, input.shape(), &grad_wrt_input));
 
@@ -412,14 +418,20 @@ class FakeQuantWithMinMaxVarsPerChannelGradientOp : public OpKernel {
     const Tensor& input = context->input(1);
     OP_REQUIRES(context, input.IsSameSize(gradient),
                 InvalidArgument("gradient and input must be the same size"));
     const int depth = input.dim_size(input.dims() - 1);  // last dimension size.
     const Tensor& min = context->input(2);
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsVector(min.shape()),
+        InvalidArgument("`min` must be rank 1 but is rank ", min.dims()));
     OP_REQUIRES(context, min.dim_size(0) == depth,
                 InvalidArgument("min has incorrect size, expected ", depth,
                                 " was ", min.dim_size(0)));
     const Tensor& max = context->input(3);
+    OP_REQUIRES(
+        context, TensorShapeUtils::IsVector(max.shape()),
+        InvalidArgument("`max` must be rank 1 but is rank ", max.dims()));
     OP_REQUIRES(context, max.dim_size(0) == depth,
                 InvalidArgument("max has incorrect size, expected ", depth,
                                 " was ", max.dim_size(0)));
 
     Tensor* grad_wrt_input;
diff --git a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
index 4c6d9974c1fc9..21059a72e7cc9 100644
--- a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -75,10 +75,75 @@ class FakeQuantWithMinMaxVarsPerChannelOpTest(test_util.TensorFlowTestCase):
       self.evaluate(
           array_ops.fake_quant_with_min_max_vars_per_channel(
               inputs=inputs, min=[0.0], max=[1.0, 1.1]))
 
 
+class FakeQuantWithMinMaxVarsGradientOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    gradients = constant_op.constant(
+        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)
+    inputs = constant_op.constant(
+        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be equal rank|must be rank 0"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_gradient(
+              gradients=gradients,
+              inputs=inputs,
+              min=0.0,
+              max=[[1.0], [2.0], [4.0]]))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "must be rank 0"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_gradient(
+              gradients=gradients,
+              inputs=inputs,
+              min=[[1.0], [2.0], [4.0]],
+              max=[[1.0], [2.0], [4.0]]))
+
+
+class FakeQuantWithMinMaxVarsPerChannelGradientOpTest(
+    test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    gradients = constant_op.constant(
+        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)
+    inputs = constant_op.constant(
+        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "Shapes must be equal rank|must be rank 1"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
+              gradients=gradients, inputs=inputs, min=[[0.0]], max=[1.0]))
+
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError),
+        "Dimension 0 in both shapes must be equal|incorrect size"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
+              gradients=gradients, inputs=inputs, min=[0.0, 0.1], max=[1.0]))
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "Shapes must be equal rank|must be rank 1"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
+              gradients=gradients, inputs=inputs, min=[1.0], max=[[1.0]]))
+
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError),
+        "Dimension 0 in both shapes must be equal|incorrect size"):
+      self.evaluate(
+          array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
+              gradients=gradients, inputs=inputs, min=[0.0], max=[1.0, 1.1]))
+
+
 class QuantizedBiasedAddTest(test_util.TensorFlowTestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def test_invalid_inputs(self):
     inputs = constant_op.constant(
@@ -335,13 +400,12 @@ class QuantizeDownAndShrinkRangeOpTest(test_util.TensorFlowTestCase):
         np.int32(0), shape=[3, 3, 3, 3], dtype=dtypes.qint32)
 
     with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                 "must be rank 0"):
       self.evaluate(
-          math_ops.quantize_down_and_shrink_range(input=inputs,
-                                                  input_min=[],
-                                                  input_max=4.0,
-                                                  out_type=dtypes.quint8))
+          math_ops.quantize_down_and_shrink_range(
+              input=inputs, input_min=[], input_max=4.0,
+              out_type=dtypes.quint8))
 
 
 if __name__ == "__main__":
   googletest.main()

Vulnerable code:           math_ops.quantize_down_and_shrink_range(input=inputs,
                                                  input_min=[],
                                                  input_max=4.0,
                                                  out_type=dtypes.quint8))
Fixed code:     OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(min.shape()),
        InvalidArgument("`min` must be rank 0 but is rank ", min.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsScalar(max.shape()),
        InvalidArgument("`max` must be rank 0 but is rank ", max.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(min.shape()),
        InvalidArgument("`min` must be rank 1 but is rank ", min.dims()));
    OP_REQUIRES(
        context, TensorShapeUtils::IsVector(max.shape()),
        InvalidArgument("`max` must be rank 1 but is rank ", max.dims()));
class FakeQuantWithMinMaxVarsGradientOpTest(test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    gradients = constant_op.constant(
        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)
    inputs = constant_op.constant(
        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be equal rank|must be rank 0"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_gradient(
              gradients=gradients,
              inputs=inputs,
              min=0.0,
              max=[[1.0], [2.0], [4.0]]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "must be rank 0"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_gradient(
              gradients=gradients,
              inputs=inputs,
              min=[[1.0], [2.0], [4.0]],
              max=[[1.0], [2.0], [4.0]]))


class FakeQuantWithMinMaxVarsPerChannelGradientOpTest(
    test_util.TensorFlowTestCase):

  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    gradients = constant_op.constant(
        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)
    inputs = constant_op.constant(
        value=[[1.0], [2.0], [4.0]], dtype=dtypes.float32)

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Shapes must be equal rank|must be rank 1"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
              gradients=gradients, inputs=inputs, min=[[0.0]], max=[1.0]))

    with self.assertRaisesRegex(
        (ValueError, errors.InvalidArgumentError),
        "Dimension 0 in both shapes must be equal|incorrect size"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
              gradients=gradients, inputs=inputs, min=[0.0, 0.1], max=[1.0]))

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Shapes must be equal rank|must be rank 1"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
              gradients=gradients, inputs=inputs, min=[1.0], max=[[1.0]]))

    with self.assertRaisesRegex(
        (ValueError, errors.InvalidArgumentError),
        "Dimension 0 in both shapes must be equal|incorrect size"):
      self.evaluate(
          array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
              gradients=gradients, inputs=inputs, min=[0.0], max=[1.0, 1.1]))


          math_ops.quantize_down_and_shrink_range(
              input=inputs, input_min=[], input_max=4.0,
              out_type=dtypes.quint8))