CVE-2023-25668: Heap-based Buffer Overflow (CWE-122)
value: TensorFlow is an open source platform for machine learning. Attackers using Tensorflow prior to 2.12.0 or 2.11.1 can access heap memory which is not in the control of user, leading to a crash or remote code execution. The fix will be included in TensorFlow version 2.12.0 and will also cherrypick this commit on TensorFlow version 2.11.1.
CWE description: A heap overflow condition is a buffer overflow, where the buffer that can be overwritten is allocated in the heap portion of memory, generally meaning that the buffer was allocated using a routine such as malloc().
CVSS v3 base score: 9.8
Published: 2023-03-25T00: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: 2022-11-09 15:47:01 -0800
Commit message: Fix asan issue with QuantizeAndDequantizeV2/V3/V4/V4Grad shape inference functions.

PiperOrigin-RevId: 487366553

Diff: diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index b57a44f44fdd7..8e52dde528eaf 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -2877,10 +2877,14 @@ REGISTER_OP("QuantizeAndDequantizeV2")
       if (axis < -1) {
         return errors::InvalidArgument("axis should be at least -1, got ",
                                        axis);
       } else if (axis != -1) {
         ShapeHandle input;
+        if (axis >= kint32max) {
+          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));
       }
@@ -2912,10 +2916,14 @@ REGISTER_OP("QuantizeAndDequantizeV4")
       if (axis < -1) {
         return errors::InvalidArgument("axis should be at least -1, got ",
                                        axis);
       } else if (axis != -1) {
         ShapeHandle input;
+        if (axis >= kint32max) {
+          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));
       }
@@ -2943,10 +2951,14 @@ REGISTER_OP("QuantizeAndDequantizeV4Grad")
       if (axis < -1) {
         return errors::InvalidArgument("axis should be at least -1, got ",
                                        axis);
       } else if (axis != -1) {
         ShapeHandle input;
+        if (axis >= kint32max) {
+          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));
       }
@@ -2979,10 +2991,14 @@ REGISTER_OP("QuantizeAndDequantizeV3")
       if (axis < -1) {
         return errors::InvalidArgument("axis should be at least -1, got ",
                                        axis);
       } else if (axis != -1) {
         ShapeHandle input;
+        if (axis >= kint32max) {
+          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 184e9949186a1..35c1361a2b1b7 100644
--- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
@@ -1854,10 +1854,76 @@ class QuantizeAndDequantizeTest(test_util.TensorFlowTestCase):
               input=q_input,
               min_range=input_min,
               max_range=input_max,
               axis=2**31 - 1))
 
+  @test_util.run_v2_only
+  def testInvalidAxis(self):
+
+    @def_function.function
+    def test_quantize_and_dequantize_v2():
+      gen_array_ops.quantize_and_dequantize_v2(
+          input=[2.5],
+          input_min=[1.0],
+          input_max=[10.0],
+          signed_input=True,
+          num_bits=1,
+          range_given=True,
+          round_mode="HALF_TO_EVEN",
+          narrow_range=True,
+          axis=0x7fffffff)
+
+    @def_function.function
+    def test_quantize_and_dequantize_v3():
+      gen_array_ops.quantize_and_dequantize_v3(
+          input=[2.5],
+          input_min=[1.0],
+          input_max=[10.0],
+          num_bits=1,
+          signed_input=True,
+          range_given=True,
+          narrow_range=True,
+          axis=0x7fffffff)
+
+    @def_function.function
+    def test_quantize_and_dequantize_v4():
+      gen_array_ops.quantize_and_dequantize_v4(
+          input=[2.5],
+          input_min=[1.0],
+          input_max=[10.0],
+          signed_input=True,
+          num_bits=1,
+          range_given=True,
+          round_mode="HALF_TO_EVEN",
+          narrow_range=True,
+          axis=0x7fffffff)
+
+    @def_function.function
+    def test_quantize_and_dequantize_v4_grad():
+      gen_array_ops.quantize_and_dequantize_v4_grad(
+          gradients=[2.5],
+          input=[2.5],
+          input_min=[1.0],
+          input_max=[10.0],
+          axis=0x7fffffff)
+
+    with self.assertRaisesRegex(
+        ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
+      test_quantize_and_dequantize_v2()
+
+    with self.assertRaisesRegex(
+        ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
+      test_quantize_and_dequantize_v3()
+
+    with self.assertRaisesRegex(
+        ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
+      test_quantize_and_dequantize_v4()
+
+    with self.assertRaisesRegex(
+        ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
+      test_quantize_and_dequantize_v4_grad()
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class SortedSearchTest(test_util.TensorFlowTestCase):
 
   def testUpperBoundFloatHandCoded(self):

Fixed code:         if (axis >= kint32max) {
          return errors::InvalidArgument(
              "Axis cannot be >= kint32max value, got ", axis);
        }
        if (axis >= kint32max) {
          return errors::InvalidArgument(
              "Axis cannot be >= kint32max value, got ", axis);
        }
        if (axis >= kint32max) {
          return errors::InvalidArgument(
              "Axis cannot be >= kint32max value, got ", axis);
        }
        if (axis >= kint32max) {
          return errors::InvalidArgument(
              "Axis cannot be >= kint32max value, got ", axis);
        }
  @test_util.run_v2_only
  def testInvalidAxis(self):

    @def_function.function
    def test_quantize_and_dequantize_v2():
      gen_array_ops.quantize_and_dequantize_v2(
          input=[2.5],
          input_min=[1.0],
          input_max=[10.0],
          signed_input=True,
          num_bits=1,
          range_given=True,
          round_mode="HALF_TO_EVEN",
          narrow_range=True,
          axis=0x7fffffff)

    @def_function.function
    def test_quantize_and_dequantize_v3():
      gen_array_ops.quantize_and_dequantize_v3(
          input=[2.5],
          input_min=[1.0],
          input_max=[10.0],
          num_bits=1,
          signed_input=True,
          range_given=True,
          narrow_range=True,
          axis=0x7fffffff)

    @def_function.function
    def test_quantize_and_dequantize_v4():
      gen_array_ops.quantize_and_dequantize_v4(
          input=[2.5],
          input_min=[1.0],
          input_max=[10.0],
          signed_input=True,
          num_bits=1,
          range_given=True,
          round_mode="HALF_TO_EVEN",
          narrow_range=True,
          axis=0x7fffffff)

    @def_function.function
    def test_quantize_and_dequantize_v4_grad():
      gen_array_ops.quantize_and_dequantize_v4_grad(
          gradients=[2.5],
          input=[2.5],
          input_min=[1.0],
          input_max=[10.0],
          axis=0x7fffffff)

    with self.assertRaisesRegex(
        ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
      test_quantize_and_dequantize_v2()

    with self.assertRaisesRegex(
        ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
      test_quantize_and_dequantize_v3()

    with self.assertRaisesRegex(
        ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
      test_quantize_and_dequantize_v4()

    with self.assertRaisesRegex(
        ValueError, "Axis cannot be >= kint32max value, got 2147483647"):
      test_quantize_and_dequantize_v4_grad()