CVE-2022-29195: 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.StagePeek` 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 `index` is a scalar 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/stage_op.cc, tensorflow/python/kernel_tests/data_structures/stage_op_test.py, tensorflow/python/ops/data_flow_ops.py
Commit date: 2022-04-29 15:59:39 -0700
Commit message: Fix tf.raw_ops.StagePeek vulnerability with invalid `index`.

Check that input is actually a scalar before treating it as such.

PiperOrigin-RevId: 445524908

Diff: diff --git a/tensorflow/core/kernels/stage_op.cc b/tensorflow/core/kernels/stage_op.cc
index 55c9db22ddf52..f7bb42f9c52b7 100644
--- a/tensorflow/core/kernels/stage_op.cc
+++ b/tensorflow/core/kernels/stage_op.cc
@@ -256,10 +256,12 @@ class StagePeekOp : public OpKernel {
     Buffer* buf = nullptr;
     OP_REQUIRES_OK(ctx, GetBuffer(ctx, def(), &buf));
     core::ScopedUnref scope(buf);
     Buffer::Tuple tuple;
 
+    OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->input(0).shape()),
+                errors::InvalidArgument("index must be scalar"));
     std::size_t index = ctx->input(0).scalar<int>()();
 
     OP_REQUIRES_OK(ctx, buf->Peek(index, &tuple));
 
     OP_REQUIRES(
diff --git a/tensorflow/python/kernel_tests/data_structures/stage_op_test.py b/tensorflow/python/kernel_tests/data_structures/stage_op_test.py
index c720155f3b6c9..d12624f106592 100644
--- a/tensorflow/python/kernel_tests/data_structures/stage_op_test.py
+++ b/tensorflow/python/kernel_tests/data_structures/stage_op_test.py
@@ -11,10 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import data_flow_ops
 from tensorflow.python.ops import math_ops
@@ -132,10 +133,20 @@ class StageTest(test.TestCase):
         sess.run(stage, feed_dict={x: i})
 
       for i in range(10):
         self.assertTrue(sess.run(peek, feed_dict={p: i}) == [i])
 
+  def testPeekBadIndex(self):
+    stager = data_flow_ops.StagingArea([
+        dtypes.int32,
+    ], shapes=[[10]])
+    stager.put([array_ops.zeros([10], dtype=dtypes.int32)])
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                'must be scalar'):
+      self.evaluate(stager.peek([]))
+
   @test_util.run_deprecated_v1
   def testSizeAndClear(self):
     with ops.Graph().as_default() as G:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32, name='x')
diff --git a/tensorflow/python/ops/data_flow_ops.py b/tensorflow/python/ops/data_flow_ops.py
index 0076e54833de6..42fd28c8cc1e6 100644
--- a/tensorflow/python/ops/data_flow_ops.py
+++ b/tensorflow/python/ops/data_flow_ops.py
@@ -1735,11 +1735,11 @@ class BaseStagingArea:
         vals = [vals]
         indices = [0]
 
     # Sanity check number of values
     if not len(vals) <= len(self._dtypes):
-      raise ValueError(f"Unexpected number of inputs {len(vals)} vs"
+      raise ValueError(f"Unexpected number of inputs {len(vals)} vs "
                        f"{len(self._dtypes)}")
 
     tensors = []
 
     for val, i in zip(vals, indices):

Vulnerable code:       raise ValueError(f"Unexpected number of inputs {len(vals)} vs"
Fixed code:     OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->input(0).shape()),
                errors::InvalidArgument("index must be scalar"));
from tensorflow.python.framework import errors
  def testPeekBadIndex(self):
    stager = data_flow_ops.StagingArea([
        dtypes.int32,
    ], shapes=[[10]])
    stager.put([array_ops.zeros([10], dtype=dtypes.int32)])

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                'must be scalar'):
      self.evaluate(stager.peek([]))

      raise ValueError(f"Unexpected number of inputs {len(vals)} vs "