CVE-2022-21727: Integer Overflow or Wraparound (CWE-190)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of shape inference for `Dequantize` is vulnerable to an integer overflow weakness. The `axis` argument can be `-1` (the default value for the optional argument) or any other positive value at most the number of dimensions of the input. Unfortunately, the upper bound is not checked, and, since the code computes `axis + 1`, an attacker can trigger an integer overflow. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.
CWE description: The product performs a calculation that can
         produce an integer overflow or wraparound when the logic
         assumes that the resulting value will always be larger than
         the original value. This occurs when an integer value is
         incremented to a value that is too large to store in the
         associated representation. When this occurs, the value may
         become a very small or negative number.
CVSS v3 base score: 8.8
Published: 2022-02-03T11:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/ops/array_ops.cc, tensorflow/python/kernel_tests/array_ops/array_ops_test.py
Commit date: 2021-11-24 13:25:41 -0800
Commit message: Fix Integer overflow error in Dequantize op shape function, by adding a bound check on axis.

PiperOrigin-RevId: 412121389
Change-Id: I3088dbad9e90f9998d406b618c16694388a9dfb4

Diff: diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 1a88b16b865dd..8c76caf4f0fed 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -22,10 +22,11 @@ limitations under the License.
 #include "tensorflow/core/framework/shape_inference.h"
 #include "tensorflow/core/framework/tensor.pb.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/mirror_pad_mode.h"
 #include "tensorflow/core/util/padding.h"
 #include "tensorflow/core/util/strided_slice_op.h"
 #include "tensorflow/core/util/tensor_format.h"
 
@@ -3026,17 +3027,30 @@ REGISTER_OP("Dequantize")
       }
       if (axis < -1) {
         return errors::InvalidArgument("axis should be at least -1, got ",
                                        axis);
       }
+      auto input_dims = c->Rank(c->input(0));
+      if (axis > input_dims) {
+        return errors::InvalidArgument(
+            "Axis must be less than input dimension(", input_dims, "), got ",
+            axis);
+      }
       const int minmax_rank = (axis == -1) ? 0 : 1;
       TF_RETURN_IF_ERROR(shape_inference::UnchangedShape(c));
       ShapeHandle minmax;
       TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
       TF_RETURN_IF_ERROR(c->WithRank(c->input(2), minmax_rank, &minmax));
       if (axis != -1) {
         ShapeHandle input;
+        if (axis >= kint32max) {
+          // Check int32 max bound for a corner case to prevent integer flow
+          // when input actually has kint32max rank and above bound check is not
+          // triggered.
+          return errors::InvalidArgument(
+              "Axis cannot be >= kint32max value, got ", axis);
+        }
         TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
         DimensionHandle depth;
         TF_RETURN_IF_ERROR(
             c->Merge(c->Dim(minmax, 0), c->Dim(input, axis), &depth));
       }
diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
index f3f2cca468ab7..9634184039e1d 100644
--- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
@@ -1702,10 +1702,25 @@ class QuantizeAndDequantizeTest(test_util.TensorFlowTestCase):
             input_max=max_threshold,
             range_given=True)
       output_grad = gradient_checker_v2.compute_gradient(f, [input_tensor])
       self.assertAllClose(output_grad[0], np.zeros([1, 4, 4]))
 
+  def testOutOfBoundAxis(self):
+    input_tensor = constant_op.constant([1., 1.])
+    input_min = [0]
+    input_max = [1]
+    q_input, _, _ = array_ops.quantize(input_tensor, 0, 1, dtypes.qint32)
+    error = (errors.InvalidArgumentError, ValueError)
+    with self.assertRaisesRegex(error,
+                                r".*Axis must be less than input dimension.*"):
+      self.evaluate(
+          gen_array_ops.dequantize(
+              input=q_input,
+              min_range=input_min,
+              max_range=input_max,
+              axis=2**31 - 1))
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class SortedSearchTest(test_util.TensorFlowTestCase):
 
   def testUpperBoundFloatHandCoded(self):

Fixed code: #include "tensorflow/core/platform/types.h"
      auto input_dims = c->Rank(c->input(0));
      if (axis > input_dims) {
        return errors::InvalidArgument(
            "Axis must be less than input dimension(", input_dims, "), got ",
            axis);
      }
        if (axis >= kint32max) {
          // Check int32 max bound for a corner case to prevent integer flow
          // when input actually has kint32max rank and above bound check is not
          // triggered.
          return errors::InvalidArgument(
              "Axis cannot be >= kint32max value, got ", axis);
        }
  def testOutOfBoundAxis(self):
    input_tensor = constant_op.constant([1., 1.])
    input_min = [0]
    input_max = [1]
    q_input, _, _ = array_ops.quantize(input_tensor, 0, 1, dtypes.qint32)
    error = (errors.InvalidArgumentError, ValueError)
    with self.assertRaisesRegex(error,
                                r".*Axis must be less than input dimension.*"):
      self.evaluate(
          gen_array_ops.dequantize(
              input=q_input,
              min_range=input_min,
              max_range=input_max,
              axis=2**31 - 1))