CVE-2023-25662: Integer Overflow or Wraparound (CWE-190)
value: TensorFlow is an open source platform for machine learning. Versions prior to 2.12.0 and 2.11.1 are vulnerable to integer overflow in EditDistance. A fix is included in TensorFlow version 2.12.0 and version 2.11.1.

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: 7.5
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/edit_distance_op_test.py
Commit date: 2023-01-24 13:57:30 -0800
Commit message: Fix security vulnerability in EditDistance op shape function.

PiperOrigin-RevId: 504367470

Diff: diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 13e24aa503d03..ef23f456533fa 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -23,10 +23,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/status.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"
@@ -1070,17 +1071,28 @@ REGISTER_OP("EditDistance")
       if (hypothesis_shape_t == nullptr || truth_shape_t == nullptr) {
         // We need to know the runtime shape of the two tensors,
         // or else the output shape is unknown.
         return shape_inference::UnknownShape(c);
       }
-
       if (hypothesis_shape_t->NumElements() != truth_shape_t->NumElements()) {
         return errors::InvalidArgument(
             "Num elements of hypothesis_shape does not match truth_shape: ",
             hypothesis_shape_t->NumElements(), " vs. ",
             truth_shape_t->NumElements());
       }
+      if (hypothesis_shape_t->NumElements() < 2) {
+        return errors::InvalidArgument(
+            "Input Hypothesis SparseTensors must have rank at least 2, but "
+            "hypothesis_shape rank is: ",
+            hypothesis_shape_t->NumElements());
+      }
+      if (truth_shape_t->NumElements() < 2) {
+        return errors::InvalidArgument(
+            "Input Truth SparseTensors must have rank at least 2, but "
+            "truth_shape rank is: ",
+            truth_shape_t->NumElements());
+      }
 
       auto h_values = hypothesis_shape_t->flat<int64_t>();
       auto t_values = truth_shape_t->flat<int64_t>();
       std::vector<DimensionHandle> dims(hypothesis_shape_t->NumElements() - 1);
       for (int i = 0; i < dims.size(); ++i) {
diff --git a/tensorflow/python/kernel_tests/array_ops/edit_distance_op_test.py b/tensorflow/python/kernel_tests/array_ops/edit_distance_op_test.py
index 8aa22cb812a0e..3dc2f4565b7e0 100644
--- a/tensorflow/python/kernel_tests/array_ops/edit_distance_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/edit_distance_op_test.py
@@ -13,12 +13,13 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for tensorflow.kernels.edit_distance_op."""
 
 import numpy as np
-
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import test
 
@@ -223,8 +224,68 @@ class EditDistanceTest(test.TestCase):
         expected_output=expected_output,
         expected_err_re=(r"inner product -\d+ which would require writing "
                          "to outside of the buffer for the output tensor|"
                          r"Dimension -\d+ must be >= 0"))
 
+  def testEmptyShapeWithEditDistanceRaisesError(self):
+    para = {
+        "hypothesis_indices": [[]],
+        "hypothesis_values": ["tmp/"],
+        "hypothesis_shape": [],
+        "truth_indices": [[]],
+        "truth_values": [""],
+        "truth_shape": [],
+        "normalize": False,
+    }
+
+    # Check edit distance raw op with empty shape in eager mode.
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError),
+        (
+            r"Input Hypothesis SparseTensors must have rank at least 2, but"
+            " hypothesis_shape rank is: 0|Input SparseTensors must have rank "
+            "at least 2, but truth_shape rank is: 0"
+        ),
+    ):
+      array_ops.gen_array_ops.EditDistance(**para)
+
+    # Check raw op with tf.function
+    @def_function.function
+    def TestFunction():
+      """Wrapper function for edit distance call."""
+      array_ops.gen_array_ops.EditDistance(**para)
+
+    with self.assertRaisesRegex(
+        ValueError,
+        (
+            "Input Hypothesis SparseTensors must have rank at least 2, but"
+            " hypothesis_shape rank is: 0"
+        ),
+    ):
+      TestFunction()
+
+    # Check with python wrapper API
+    hypothesis_indices = [[]]
+    hypothesis_values = [0]
+    hypothesis_shape = []
+    truth_indices = [[]]
+    truth_values = [1]
+    truth_shape = []
+    expected_output = []  # dummy ignored
+
+    with self.assertRaisesRegex(
+        ValueError,
+        (
+            "Input Hypothesis SparseTensors must have rank at least 2, but"
+            " hypothesis_shape rank is: 0"
+        ),
+    ):
+      self._testEditDistance(
+          hypothesis=(hypothesis_indices, hypothesis_values, hypothesis_shape),
+          truth=(truth_indices, truth_values, truth_shape),
+          normalize=False,
+          expected_output=expected_output,
+      )
+
 
 if __name__ == "__main__":
   test.main()

Fixed code: #include "tensorflow/core/platform/status.h"
      if (hypothesis_shape_t->NumElements() < 2) {
        return errors::InvalidArgument(
            "Input Hypothesis SparseTensors must have rank at least 2, but "
            "hypothesis_shape rank is: ",
            hypothesis_shape_t->NumElements());
      }
      if (truth_shape_t->NumElements() < 2) {
        return errors::InvalidArgument(
            "Input Truth SparseTensors must have rank at least 2, but "
            "truth_shape rank is: ",
            truth_shape_t->NumElements());
      }
from tensorflow.python.eager import def_function
from tensorflow.python.framework import errors
  def testEmptyShapeWithEditDistanceRaisesError(self):
    para = {
        "hypothesis_indices": [[]],
        "hypothesis_values": ["tmp/"],
        "hypothesis_shape": [],
        "truth_indices": [[]],
        "truth_values": [""],
        "truth_shape": [],
        "normalize": False,
    }

    # Check edit distance raw op with empty shape in eager mode.
    with self.assertRaisesRegex(
        (errors.InvalidArgumentError, ValueError),
        (
            r"Input Hypothesis SparseTensors must have rank at least 2, but"
            " hypothesis_shape rank is: 0|Input SparseTensors must have rank "
            "at least 2, but truth_shape rank is: 0"
        ),
    ):
      array_ops.gen_array_ops.EditDistance(**para)

    # Check raw op with tf.function
    @def_function.function
    def TestFunction():
      """Wrapper function for edit distance call."""
      array_ops.gen_array_ops.EditDistance(**para)

    with self.assertRaisesRegex(
        ValueError,
        (
            "Input Hypothesis SparseTensors must have rank at least 2, but"
            " hypothesis_shape rank is: 0"
        ),
    ):
      TestFunction()

    # Check with python wrapper API
    hypothesis_indices = [[]]
    hypothesis_values = [0]
    hypothesis_shape = []
    truth_indices = [[]]
    truth_values = [1]
    truth_shape = []
    expected_output = []  # dummy ignored

    with self.assertRaisesRegex(
        ValueError,
        (
            "Input Hypothesis SparseTensors must have rank at least 2, but"
            " hypothesis_shape rank is: 0"
        ),
    ):
      self._testEditDistance(
          hypothesis=(hypothesis_indices, hypothesis_values, hypothesis_shape),
          truth=(truth_indices, truth_values, truth_shape),
          normalize=False,
          expected_output=expected_output,
      )