CVE-2022-29200: Improper Validation of Specified Quantity in Input (CWE-1284)
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.LSTMBlockCell` 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 does not validate the ranks of any of the arguments to this API call. This results in `CHECK`-failures when the elements of the tensor are accessed. 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 that is expected to specify a quantity (such as size or length), but it does not validate or incorrectly validates that the quantity has the required properties.
CVSS v3 base score: 5.5
Published: 2022-05-20T22:16Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/rnn/lstm_ops.cc, tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py
Commit date: 2022-05-02 14:28:15 -0700
Commit message: Fix security vulnerability with LSTMBlockCellOp

PiperOrigin-RevId: 446028341

Diff: diff --git a/tensorflow/core/kernels/rnn/lstm_ops.cc b/tensorflow/core/kernels/rnn/lstm_ops.cc
index 711fc8f08275d..ab4b9c695a569 100644
--- a/tensorflow/core/kernels/rnn/lstm_ops.cc
+++ b/tensorflow/core/kernels/rnn/lstm_ops.cc
@@ -414,10 +414,69 @@ class LSTMBlockCellOp : public OpKernel {
                                       TensorShape({batch_size, cell_size * 4}),
                                       &gates_tensor));
 
     const Device& device = ctx->eigen_device<Device>();
 
+    // Sanity check that each of the tensors have the required NDIMS.
+    OP_REQUIRES(ctx, x_tensor->dims() == 2,
+                errors::InvalidArgument("x_tensor must be rank 2 but is rank ",
+                                        x_tensor->dims(), "."));
+    OP_REQUIRES(
+        ctx, cs_prev_tensor->dims() == 2,
+        errors::InvalidArgument("cs_prev_tensor must be rank 2 but is rank ",
+                                cs_prev_tensor->dims(), "."));
+    OP_REQUIRES(
+        ctx, h_prev_tensor->dims() == 2,
+        errors::InvalidArgument("h_prev_tensor must be rank 2 but is rank ",
+                                h_prev_tensor->dims(), "."));
+    OP_REQUIRES(ctx, w_tensor->dims() == 2,
+                errors::InvalidArgument("w_tensor must be rank 2 but is rank ",
+                                        w_tensor->dims(), "."));
+    OP_REQUIRES(
+        ctx, wci_tensor->dims() == 1,
+        errors::InvalidArgument("wci_tensor must be rank 1 but is rank ",
+                                wci_tensor->dims(), "."));
+    OP_REQUIRES(
+        ctx, wcf_tensor->dims() == 1,
+        errors::InvalidArgument("wcf_tensor must be rank 1 but is rank ",
+                                wci_tensor->dims(), "."));
+    OP_REQUIRES(
+        ctx, wco_tensor->dims() == 1,
+        errors::InvalidArgument("wco_tensor must be rank 1 but is rank ",
+                                wco_tensor->dims(), "."));
+    OP_REQUIRES(ctx, b_tensor->dims() == 1,
+                errors::InvalidArgument("b_tensor must be rank 1 but is rank ",
+                                        b_tensor->dims(), "."));
+    OP_REQUIRES(ctx, xh_tensor.dims() == 2,
+                errors::InvalidArgument("xh_tensor must be rank 2 but is rank ",
+                                        xh_tensor.dims(), "."));
+    OP_REQUIRES(ctx, i_tensor->dims() == 2,
+                errors::InvalidArgument("i_tensor must be rank 2 but is rank ",
+                                        i_tensor->dims(), "."));
+    OP_REQUIRES(ctx, cs_tensor->dims() == 2,
+                errors::InvalidArgument("cs_tensor must be rank 2 but is rank ",
+                                        cs_tensor->dims(), "."));
+    OP_REQUIRES(ctx, f_tensor->dims() == 2,
+                errors::InvalidArgument("f_tensor must be rank 2 but is rank ",
+                                        f_tensor->dims(), "."));
+    OP_REQUIRES(ctx, o_tensor->dims() == 2,
+                errors::InvalidArgument("o_tensor must be rank 2 but is rank ",
+                                        o_tensor->dims(), "."));
+    OP_REQUIRES(ctx, ci_tensor->dims() == 2,
+                errors::InvalidArgument("ci_tensor must be rank 2 but is rank ",
+                                        ci_tensor->dims(), "."));
+    OP_REQUIRES(ctx, co_tensor->dims() == 2,
+                errors::InvalidArgument("co_tensor must be rank 2 but is rank ",
+                                        co_tensor->dims(), "."));
+    OP_REQUIRES(
+        ctx, gates_tensor.dims() == 2,
+        errors::InvalidArgument("gates_tensor must be rank 2 but is rank ",
+                                gates_tensor.dims(), "."));
+    OP_REQUIRES(ctx, h_tensor->dims() == 2,
+                errors::InvalidArgument("h_tensor must be rank 2 but is rank ",
+                                        h_tensor->dims(), "."));
+
     functor::LSTMBlockCellFprop<Device, T, USE_CUBLAS, gate_layout>(
         batch_size, input_size, cell_size)(
         ctx, device, forget_bias_, cell_clip_, use_peephole_,
         x_tensor->matrix<T>(), cs_prev_tensor->matrix<T>(),
         h_prev_tensor->matrix<T>(), w_tensor->matrix<T>(), wci_tensor->vec<T>(),
diff --git a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py
index 438211f0cb71f..bed3cbfd8aa2a 100644
--- a/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py
@@ -31,10 +31,11 @@ from tensorflow.python.framework import random_seed
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import   array_ops
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_rnn_ops
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import rnn
 from tensorflow.python.ops import rnn_cell
@@ -1321,10 +1322,40 @@ class LSTMTest(test.TestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def testDynamicEquivalentToStaticRNNWithSequenceLength(self):
     self._testDynamicEquivalentToStaticRNN(use_sequence_length=True)
 
+  @test_util.run_in_graph_and_eager_modes
+  def testLSTMBlockCellErrorHandling(self):
+    forget_bias = 1
+    cell_clip = 0
+    use_peephole = False
+    x = constant_op.constant(0.837607, shape=[28, 29], dtype=dtypes.float32)
+    cs_prev = constant_op.constant(0, shape=[28, 17], dtype=dtypes.float32)
+    h_prev = constant_op.constant(
+        0.592631638, shape=[28, 17], dtype=dtypes.float32)
+    w = constant_op.constant(0.887386262, shape=[46, 68], dtype=dtypes.float32)
+    wci = constant_op.constant(0, shape=[], dtype=dtypes.float32)
+    wcf = constant_op.constant(0, shape=[17], dtype=dtypes.float32)
+    wco = constant_op.constant(
+        0.592631638, shape=[28, 17], dtype=dtypes.float32)
+    b = constant_op.constant(0.75259006, shape=[68], dtype=dtypes.float32)
+    with self.assertRaises(errors_impl.InvalidArgumentError):
+      self.evaluate(
+          gen_rnn_ops.lstm_block_cell(
+              x=x,
+              cs_prev=cs_prev,
+              h_prev=h_prev,
+              w=w,
+              wci=wci,
+              wcf=wcf,
+              wco=wco,
+              b=b,
+              forget_bias=forget_bias,
+              cell_clip=cell_clip,
+              use_peephole=use_peephole))
+
 
 class BidirectionalRNNTest(test.TestCase):
 
   def setUp(self):
     self._seed = 23489

Fixed code:     // Sanity check that each of the tensors have the required NDIMS.
    OP_REQUIRES(ctx, x_tensor->dims() == 2,
                errors::InvalidArgument("x_tensor must be rank 2 but is rank ",
                                        x_tensor->dims(), "."));
    OP_REQUIRES(
        ctx, cs_prev_tensor->dims() == 2,
        errors::InvalidArgument("cs_prev_tensor must be rank 2 but is rank ",
                                cs_prev_tensor->dims(), "."));
    OP_REQUIRES(
        ctx, h_prev_tensor->dims() == 2,
        errors::InvalidArgument("h_prev_tensor must be rank 2 but is rank ",
                                h_prev_tensor->dims(), "."));
    OP_REQUIRES(ctx, w_tensor->dims() == 2,
                errors::InvalidArgument("w_tensor must be rank 2 but is rank ",
                                        w_tensor->dims(), "."));
    OP_REQUIRES(
        ctx, wci_tensor->dims() == 1,
        errors::InvalidArgument("wci_tensor must be rank 1 but is rank ",
                                wci_tensor->dims(), "."));
    OP_REQUIRES(
        ctx, wcf_tensor->dims() == 1,
        errors::InvalidArgument("wcf_tensor must be rank 1 but is rank ",
                                wci_tensor->dims(), "."));
    OP_REQUIRES(
        ctx, wco_tensor->dims() == 1,
        errors::InvalidArgument("wco_tensor must be rank 1 but is rank ",
                                wco_tensor->dims(), "."));
    OP_REQUIRES(ctx, b_tensor->dims() == 1,
                errors::InvalidArgument("b_tensor must be rank 1 but is rank ",
                                        b_tensor->dims(), "."));
    OP_REQUIRES(ctx, xh_tensor.dims() == 2,
                errors::InvalidArgument("xh_tensor must be rank 2 but is rank ",
                                        xh_tensor.dims(), "."));
    OP_REQUIRES(ctx, i_tensor->dims() == 2,
                errors::InvalidArgument("i_tensor must be rank 2 but is rank ",
                                        i_tensor->dims(), "."));
    OP_REQUIRES(ctx, cs_tensor->dims() == 2,
                errors::InvalidArgument("cs_tensor must be rank 2 but is rank ",
                                        cs_tensor->dims(), "."));
    OP_REQUIRES(ctx, f_tensor->dims() == 2,
                errors::InvalidArgument("f_tensor must be rank 2 but is rank ",
                                        f_tensor->dims(), "."));
    OP_REQUIRES(ctx, o_tensor->dims() == 2,
                errors::InvalidArgument("o_tensor must be rank 2 but is rank ",
                                        o_tensor->dims(), "."));
    OP_REQUIRES(ctx, ci_tensor->dims() == 2,
                errors::InvalidArgument("ci_tensor must be rank 2 but is rank ",
                                        ci_tensor->dims(), "."));
    OP_REQUIRES(ctx, co_tensor->dims() == 2,
                errors::InvalidArgument("co_tensor must be rank 2 but is rank ",
                                        co_tensor->dims(), "."));
    OP_REQUIRES(
        ctx, gates_tensor.dims() == 2,
        errors::InvalidArgument("gates_tensor must be rank 2 but is rank ",
                                gates_tensor.dims(), "."));
    OP_REQUIRES(ctx, h_tensor->dims() == 2,
                errors::InvalidArgument("h_tensor must be rank 2 but is rank ",
                                        h_tensor->dims(), "."));

from tensorflow.python.ops import gen_rnn_ops
  @test_util.run_in_graph_and_eager_modes
  def testLSTMBlockCellErrorHandling(self):
    forget_bias = 1
    cell_clip = 0
    use_peephole = False
    x = constant_op.constant(0.837607, shape=[28, 29], dtype=dtypes.float32)
    cs_prev = constant_op.constant(0, shape=[28, 17], dtype=dtypes.float32)
    h_prev = constant_op.constant(
        0.592631638, shape=[28, 17], dtype=dtypes.float32)
    w = constant_op.constant(0.887386262, shape=[46, 68], dtype=dtypes.float32)
    wci = constant_op.constant(0, shape=[], dtype=dtypes.float32)
    wcf = constant_op.constant(0, shape=[17], dtype=dtypes.float32)
    wco = constant_op.constant(
        0.592631638, shape=[28, 17], dtype=dtypes.float32)
    b = constant_op.constant(0.75259006, shape=[68], dtype=dtypes.float32)
    with self.assertRaises(errors_impl.InvalidArgumentError):
      self.evaluate(
          gen_rnn_ops.lstm_block_cell(
              x=x,
              cs_prev=cs_prev,
              h_prev=h_prev,
              w=w,
              wci=wci,
              wcf=wcf,
              wco=wco,
              b=b,
              forget_bias=forget_bias,
              cell_clip=cell_clip,
              use_peephole=use_peephole))