CVE-2022-21731: Access of Resource Using Incompatible Type ('Type Confusion') (CWE-843)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of shape inference for `ConcatV2` can be used to trigger a denial of service attack via a segfault caused by a type confusion. The `axis` argument is translated into `concat_dim` in the `ConcatShapeHelper` helper function. Then, a value for `min_rank` is computed based on `concat_dim`. This is then used to validate that the `values` tensor has at least the required rank. However, `WithRankAtLeast` receives the lower bound as a 64-bits value and then compares it against the maximum 32-bits integer value that could be represented. Due to the fact that `min_rank` is a 32-bits value and the value of `axis`, the `rank` argument is a negative value, so the error check is bypassed. 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 allocates or initializes a resource such as a pointer, object, or variable using one type, but it later accesses that resource using a type that is incompatible with the original type.
CVSS v3 base score: 6.5
Published: 2022-02-03T12:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/framework/common_shape_fns.cc, tensorflow/python/kernel_tests/array_ops/concat_op_test.py
Commit date: 2021-11-24 13:14:19 -0800
Commit message: Fix Segfault in Concat V2 shape function.

PiperOrigin-RevId: 412120654
Change-Id: I3ff915faea694f9ad8b00024e9af2de9909011be

Diff: diff --git a/tensorflow/core/framework/common_shape_fns.cc b/tensorflow/core/framework/common_shape_fns.cc
index 95f2670e0571d..4a5ce9ae661ff 100644
--- a/tensorflow/core/framework/common_shape_fns.cc
+++ b/tensorflow/core/framework/common_shape_fns.cc
@@ -2003,11 +2003,11 @@ Status ConcatShapeHelper(InferenceContext* c, int start_value_index,
   } else {
     concat_dim = concat_dim_t->flat<int64_t>()(0);
   }
 
   // Minimum required number of dimensions.
-  const int min_rank = concat_dim < 0 ? -concat_dim : concat_dim + 1;
+  const int64 min_rank = concat_dim < 0 ? -concat_dim : concat_dim + 1;
 
   ShapeHandle output_before;
   ShapeHandle output_after;
 
   ShapeHandle input = c->input(end_value_index - 1);
diff --git a/tensorflow/python/kernel_tests/array_ops/concat_op_test.py b/tensorflow/python/kernel_tests/array_ops/concat_op_test.py
index a957665903176..a5460861f9dd8 100644
--- a/tensorflow/python/kernel_tests/array_ops/concat_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/concat_op_test.py
@@ -14,10 +14,11 @@
 # ==============================================================================
 """Functional tests for Concat Op."""
 
 import numpy as np
 
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -568,10 +569,21 @@ class ConcatOpTest(test.TestCase):
       with test_util.use_gpu():
         t1 = [1]
         t2 = [2]
         gen_array_ops.concat_v2([t1, t2], 1).eval()
 
+  def testConcatInvalidAxisInTfFunction(self):
+
+    @def_function.function
+    def concat_wrapper():
+      y = gen_array_ops.concat_v2(
+          values=[[1, 2, 3], [4, 5, 6]], axis=0xb500005b)
+      return y
+
+    with self.assertRaises(ValueError):
+      concat_wrapper()
+
   def testConcatNegativeAxis(self):
     with test_util.use_gpu():
       t1 = [[1, 2, 3], [4, 5, 6]]
       t2 = [[7, 8, 9], [10, 11, 12]]
 

Vulnerable code:   const int min_rank = concat_dim < 0 ? -concat_dim : concat_dim + 1;
Fixed code:   const int64 min_rank = concat_dim < 0 ? -concat_dim : concat_dim + 1;
from tensorflow.python.eager import def_function
  def testConcatInvalidAxisInTfFunction(self):

    @def_function.function
    def concat_wrapper():
      y = gen_array_ops.concat_v2(
          values=[[1, 2, 3], [4, 5, 6]], axis=0xb500005b)
      return y

    with self.assertRaises(ValueError):
      concat_wrapper()