CVE-2022-29199: Improper Input Validation (CWE-20)
value: TensorFlow is an open source platform for machine learning. Prior to versions 2.9.0, 2.8.1, 2.7.2, and 2.6.4, the implementation of `tf.raw_ops.LoadAndRemapMatrix does not fully validate the input arguments. This results in a `CHECK`-failure which can be used to trigger a denial of service attack. The code assumes `initializing_values` is a vector but there is no validation for this before accessing its value. Versions 2.9.0, 2.8.1, 2.7.2, and 2.6.4 contain a patch for this issue.
CWE description: The product receives input or data, but it does
        not validate or incorrectly validates that the input has the
        properties that are required to process the data safely and
        correctly.
CVSS v3 base score: 5.5
Published: 2022-05-20T22:16Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/load_and_remap_matrix_op.cc, tensorflow/python/kernel_tests/io_ops/checkpoint_ops_test.py
Commit date: 2022-04-29 15:46:05 -0700
Commit message: Fix tf.raw_ops.LoadAndRemapMatrix vulnerability with invalid `row_remapping`.

Check that `row_remapping` has the correct dims().

PiperOrigin-RevId: 445522800

Diff: diff --git a/tensorflow/core/kernels/load_and_remap_matrix_op.cc b/tensorflow/core/kernels/load_and_remap_matrix_op.cc
index 3fa753251c0f4..4276e16059a9c 100644
--- a/tensorflow/core/kernels/load_and_remap_matrix_op.cc
+++ b/tensorflow/core/kernels/load_and_remap_matrix_op.cc
@@ -72,10 +72,15 @@ class LoadAndRemapMatrixOp : public OpKernel {
     // be maps with key = old ID, value = new ID.
     std::unordered_map<int64_t, int64_t> old_row_to_new_row_map;
     std::vector<bool> row_id_present;
     const Tensor* row_remapping_t;
     OP_REQUIRES_OK(context, context->input("row_remapping", &row_remapping_t));
+    OP_REQUIRES(
+        context, row_remapping_t->dims() == 1,
+        errors::InvalidArgument("The `row_remapping` tensor must be 1-D, got "
+                                "a tensor of shape ",
+                                row_remapping_t->shape().DebugString()));
     const auto row_remapping = row_remapping_t->vec<int64_t>();
     OP_REQUIRES(context, row_remapping.size() == num_rows_,
                 errors::InvalidArgument(strings::StrCat(
                     "Size of row_remapping is ", row_remapping.size(),
                     " instead of being equal to num_rows=", num_rows_)));
diff --git a/tensorflow/python/kernel_tests/io_ops/checkpoint_ops_test.py b/tensorflow/python/kernel_tests/io_ops/checkpoint_ops_test.py
index f357de2de7c84..91618f974b31a 100644
--- a/tensorflow/python/kernel_tests/io_ops/checkpoint_ops_test.py
+++ b/tensorflow/python/kernel_tests/io_ops/checkpoint_ops_test.py
@@ -225,10 +225,36 @@ class LoadAndRemapMatrixTest(test.TestCase):
     with self.cached_session():
       self.assertAllClose(
           np.reshape(initializing_values, (num_rows, num_cols)),
           self.evaluate(remapped_matrix))
 
+  def test_load_and_remap_invalid_dims(self):
+    ckpt_path = constant_op.constant(
+        '/tmp/warm_starting_util_test5kl2a3pc/tmpph76tep2/model-0',
+        shape=[],
+        dtype=dtypes.string)
+    old_tensor_name = constant_op.constant(
+        '/tmp/warm_starting_util_test5kl2a3pc/tmpph76tep2/model-0',
+        shape=[],
+        dtype=dtypes.string)
+    row_remapping = constant_op.constant(0, shape=[], dtype=dtypes.int64)
+    col_remapping = constant_op.constant(3, shape=[3], dtype=dtypes.int64)
+    initializing_values = constant_op.constant([],
+                                               shape=[0, 1],
+                                               dtype=dtypes.float32)
+    with self.cached_session(), self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError), 'tensor must be 1-D'):
+      self.evaluate(
+          gen_checkpoint_ops.load_and_remap_matrix(
+              ckpt_path=ckpt_path,
+              old_tensor_name=old_tensor_name,
+              row_remapping=row_remapping,
+              col_remapping=col_remapping,
+              initializing_values=initializing_values,
+              num_rows=1,
+              num_cols=1))
+
   @test_util.run_deprecated_v1
   def test_load_and_remap_invalid_remapping(self):
     """Tests that errors are raised when an ID maps to multiple new IDs.
 
     (This should usually not happen when using public APIs).

Fixed code:     OP_REQUIRES(
        context, row_remapping_t->dims() == 1,
        errors::InvalidArgument("The `row_remapping` tensor must be 1-D, got "
                                "a tensor of shape ",
                                row_remapping_t->shape().DebugString()));
  def test_load_and_remap_invalid_dims(self):
    ckpt_path = constant_op.constant(
        '/tmp/warm_starting_util_test5kl2a3pc/tmpph76tep2/model-0',
        shape=[],
        dtype=dtypes.string)
    old_tensor_name = constant_op.constant(
        '/tmp/warm_starting_util_test5kl2a3pc/tmpph76tep2/model-0',
        shape=[],
        dtype=dtypes.string)
    row_remapping = constant_op.constant(0, shape=[], dtype=dtypes.int64)
    col_remapping = constant_op.constant(3, shape=[3], dtype=dtypes.int64)
    initializing_values = constant_op.constant([],
                                               shape=[0, 1],
                                               dtype=dtypes.float32)
    with self.cached_session(), self.assertRaisesRegex(
        (ValueError, errors.InvalidArgumentError), 'tensor must be 1-D'):
      self.evaluate(
          gen_checkpoint_ops.load_and_remap_matrix(
              ckpt_path=ckpt_path,
              old_tensor_name=old_tensor_name,
              row_remapping=row_remapping,
              col_remapping=col_remapping,
              initializing_values=initializing_values,
              num_rows=1,
              num_cols=1))