CVE-2022-21734 (Python, published 2022-02-03T13:15Z)
Description: value: Tensorflow is an Open Source Machine Learning Framework. The implementation of `MapStage` is vulnerable a `CHECK`-fail if the key tensor is not a scalar. 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.
Severity: MEDIUM (CVSS v3 base score: 6.5)
Weakness: CWE-843 \u2014 Access of Resource Using Incompatible Type ('Type Confusion')
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.
Repository: https://github.com/tensorflow/tensorflow
Files changed:
tensorflow/core/kernels/map_stage_op.cc
tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
Commit message: Add a check for Key being scalar tensor for MapStage and OrderedMapStage ops.

According to documentation[1][2], key must be int64 value, but this wasn't enforced and the ops would fail with check failure for non-scalar key value.

[1]https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/ordered-map-stage
[2]https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/map-stage

PiperOrigin-RevId: 413822112
Change-Id: I9d118faf990e6361900aa32272eff486ad9f0e2e


Vulnerable code:
from tensorflow.python.framework import errors
    with ops.Graph().as_default() as G:
    G.finalize()
    with self.session(graph=G) as sess:
    with ops.Graph().as_default() as G:
    G.finalize()
    with self.session(graph=G) as sess:
    with ops.Graph().as_default() as G:
        stager = data_flow_ops.MapStagingArea(
            [dtypes.float32, dtypes.float32],
            shapes=[[], [128, 128]],
            names=['x', 'v'])
    G.finalize()
    with self.session(graph=G) as sess:
    with ops.Graph().as_default() as G:
    G.finalize()
    with ops.Graph().as_default() as G:
        stager = data_flow_ops.MapStagingArea(
            [
                dtypes.int32,
            ], shapes=[[]])
    G.finalize()
    with self.session(graph=G) as sess:
        self.assertTrue(sess.run(peek, feed_dict={gi: i})[0] == i)
      self.assertTrue(sess.run(size) == 10)
    with ops.Graph().as_default() as G:
        stager = data_flow_ops.MapStagingArea(
            [dtypes.float32, dtypes.float32],
            shapes=[[], [128, 128]],
            names=['x', 'v'])
    G.finalize()
    with self.session(graph=G) as sess:
    with ops.Graph().as_default() as G:
        stager = data_flow_ops.MapStagingArea(
            [
                dtypes.int32,
            ], capacity=capacity, shapes=[[]])
    G.finalize()
    with self.session(graph=G) as sess:
      self.assertTrue(sess.run(size) == capacity)
      self.assertTrue(sess.run(size) == 0)
    with ops.Graph().as_default() as G:
        stager = data_flow_ops.MapStagingArea(
            [dtypes.uint8], memory_limit=memory_limit, shapes=[[]])
    G.finalize()
    import numpy as np
    with self.session(graph=G) as sess:
      self.assertTrue(sess.run(size) == capacity)
      self.assertTrue(sess.run(size) == 0)
    with ops.Graph().as_default() as G:
        stager = data_flow_ops.MapStagingArea(
            [
                dtypes.int32,
            ], shapes=[[]], ordered=True)
    G.finalize()
    with self.session(graph=G) as sess:
      self.assertTrue(sess.run(size) == n)
      self.assertTrue(sess.run(size) == 0)
    with ops.Graph().as_default() as G:
    G.finalize()
    with self.session(graph=G) as sess:
      self.assertTrue(sess.run([size, isize]) == [0, 0])
      self.assertTrue(sess.run([size, isize]) == [0, 1])
      self.assertTrue(sess.run([size, isize]) == [0, 2])
      self.assertTrue(sess.run([size, isize]) == [1, 1])
      self.assertTrue(
          sess.run([key, ret], feed_dict={
              gi: 0
          }) == [0, {
      self.assertTrue(sess.run([size, isize]) == [0, 1])
      self.assertTrue(
          sess.run([key, ret], feed_dict={
              gi: 1
          }) == [1, {
    with ops.Graph().as_default() as G:
    G.finalize()
    with self.session(graph=G) as sess:
      self.assertTrue(sess.run([size, isize]) == [0, 0])
      self.assertTrue(sess.run([size, isize]) == [0, 1])
      self.assertTrue(sess.run([size, isize]) == [0, 2])
      self.assertTrue(sess.run([size, isize]) == [1, 1])
      self.assertTrue(sess.run([key, ret], feed_dict={gi: 0}) == [0, [1, 1, 2]])
      self.assertTrue(sess.run([size, isize]) == [0, 1])
      self.assertTrue(sess.run([key, ret], feed_dict={gi: 1}) == [1, [1, 3, 2]])
    with ops.Graph().as_default() as G:
    G.finalize()
    with self.session(graph=G) as sess:
      self.assertTrue(sess.run([size, isize]) == [0, 0])
      self.assertTrue(sess.run([size, isize]) == [0, 1])
      self.assertTrue(sess.run([size, isize]) == [0, 2])
      self.assertTrue(sess.run([size, isize]) == [1, 1])
      self.assertTrue(sess.run(peek_xf, feed_dict={pei: 0}) == {'x': 1, 'f': 2})
      self.assertTrue(sess.run(peek_v, feed_dict={pei: 0}) == {'v': 1})
      self.assertTrue(sess.run([size, isize]) == [1, 1])
      self.assertTrue(
          sess.run([key_xf, get_xf], feed_dict={
              gi: 0
          }) == [0, {
      self.assertTrue(sess.run([size, isize]) == [1, 1])
      self.assertTrue(exc_str in cm.exception.message)
      self.assertTrue(
          sess.run([key_v, get_v], feed_dict={
              gi: 0
          }) == [0, {
      self.assertTrue(sess.run([size, isize]) == [0, 1])
      self.assertTrue(sess.run([size, isize]) == [1, 0])
      self.assertTrue(sess.run([pop_key_xf, pop_xf]) == [1, {'x': 1, 'f': 2}])
      self.assertTrue(sess.run([size, isize]) == [1, 0])
      self.assertTrue(
          sess.run([pop_key_v, pop_v], feed_dict={
              pi: 1
          }) == [1, {
      self.assertTrue(sess.run([size, isize]) == [0, 0])
    with ops.Graph().as_default() as G:
    G.finalize()
    with self.session(graph=G) as sess:
      self.assertTrue(sess.run([size, isize]) == [1, 0])
      self.assertTrue(
          sess.run([key_xf, get_xf], feed_dict={
              gi: 0
          }) == [0, [1, 2]])
      self.assertTrue(sess.run([size, isize]) == [1, 0])
      self.assertTrue(sess.run([key_v, get_v], feed_dict={gi: 0}) == [0, [3]])
      self.assertTrue(sess.run([size, isize]) == [0, 0])

Fixed code:
    OP_REQUIRES(ctx, key_tensor->NumElements() == 1,
                errors::InvalidArgument(
                    "key must be an int64 scalar, got tensor with shape: ",
                    key_tensor->shape()));

import numpy as np

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import errors
    with ops.Graph().as_default() as g:
    g.finalize()
    with self.session(graph=g) as sess:
    with ops.Graph().as_default() as g:
    g.finalize()
    with self.session(graph=g) as sess:
    with ops.Graph().as_default() as g:
        stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32],
                                              shapes=[[], [128, 128]],
                                              names=['x', 'v'])
    g.finalize()
    with self.session(graph=g) as sess:
    with ops.Graph().as_default() as g:
    g.finalize()
    with ops.Graph().as_default() as g:
        stager = data_flow_ops.MapStagingArea([
            dtypes.int32,
        ], shapes=[[]])
    g.finalize()
    with self.session(graph=g) as sess:
        self.assertEqual(sess.run(peek, feed_dict={gi: i})[0], i)
      self.assertEqual(sess.run(size), 10)
    with ops.Graph().as_default() as g:
        stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32],
                                              shapes=[[], [128, 128]],
                                              names=['x', 'v'])
    g.finalize()
    with self.session(graph=g) as sess:
    with ops.Graph().as_default() as g:
        stager = data_flow_ops.MapStagingArea([
            dtypes.int32,
        ],
                                              capacity=capacity,
                                              shapes=[[]])
    g.finalize()
    with self.session(graph=g) as sess:
      self.assertEqual(sess.run(size), capacity)
      self.assertEqual(sess.run(size), 0)
    with ops.Graph().as_default() as g:
        stager = data_flow_ops.MapStagingArea([dtypes.uint8],
                                              memory_limit=memory_limit,
                                              shapes=[[]])
    g.finalize()
    with self.session(graph=g) as sess:
      self.assertEqual(sess.run(size), capacity)
      self.assertEqual(sess.run(size), 0)
    with ops.Graph().as_default() as g:
        stager = data_flow_ops.MapStagingArea([
            dtypes.int32,
        ],
                                              shapes=[[]],
                                              ordered=True)
    g.finalize()
    with self.session(graph=g) as sess:
      self.assertEqual(sess.run(size), n)
      self.assertEqual(sess.run(size), 0)
    with ops.Graph().as_default() as g:
    g.finalize()
    with self.session(graph=g) as sess:
      self.assertEqual(sess.run([size, isize]), [0, 0])
      self.assertEqual(sess.run([size, isize]), [0, 1])
      self.assertEqual(sess.run([size, isize]), [0, 2])
      self.assertEqual(sess.run([size, isize]), [1, 1])
      self.assertEqual(
          sess.run([key, ret], feed_dict={gi: 0}),
          [0, {
      self.assertEqual(sess.run([size, isize]), [0, 1])
      self.assertEqual(
          sess.run([key, ret], feed_dict={gi: 1}),
          [1, {
    with ops.Graph().as_default() as g:
    g.finalize()
    with self.session(graph=g) as sess:
      self.assertEqual(sess.run([size, isize]), [0, 0])
      self.assertEqual(sess.run([size, isize]), [0, 1])
      self.assertEqual(sess.run([size, isize]), [0, 2])
      self.assertEqual(sess.run([size, isize]), [1, 1])
      self.assertEqual(sess.run([key, ret], feed_dict={gi: 0}), [0, [1, 1, 2]])
      self.assertEqual(sess.run([size, isize]), [0, 1])
      self.assertEqual(sess.run([key, ret], feed_dict={gi: 1}), [1, [1, 3, 2]])
    with ops.Graph().as_default() as g:
    g.finalize()
    with self.session(graph=g) as sess:
      self.assertEqual(sess.run([size, isize]), [0, 0])
      self.assertEqual(sess.run([size, isize]), [0, 1])
      self.assertEqual(sess.run([size, isize]), [0, 2])
      self.assertEqual(sess.run([size, isize]), [1, 1])
      self.assertEqual(sess.run(peek_xf, feed_dict={pei: 0}), {'x': 1, 'f': 2})
      self.assertEqual(sess.run(peek_v, feed_dict={pei: 0}), {'v': 1})
      self.assertEqual(sess.run([size, isize]), [1, 1])
      self.assertEqual(
          sess.run([key_xf, get_xf], feed_dict={gi: 0}), [0, {
      self.assertEqual(sess.run([size, isize]), [1, 1])
      self.assertIn(exc_str, cm.exception.message)
      self.assertEqual(
          sess.run([key_v, get_v], feed_dict={gi: 0}), [0, {
      self.assertEqual(sess.run([size, isize]), [0, 1])
      self.assertEqual(sess.run([size, isize]), [1, 0])
      self.assertEqual(sess.run([pop_key_xf, pop_xf]), [1, {'x': 1, 'f': 2}])
      self.assertEqual(sess.run([size, isize]), [1, 0])
      self.assertEqual(
          sess.run([pop_key_v, pop_v], feed_dict={pi: 1}), [1, {
      self.assertEqual(sess.run([size, isize]), [0, 0])
    with ops.Graph().as_default() as g:
    g.finalize()
    with self.session(graph=g) as sess:
      self.assertEqual(sess.run([size, isize]), [1, 0])
      self.assertEqual(
          sess.run([key_xf, get_xf], feed_dict={gi: 0}), [0, [1, 2]])
      self.assertEqual(sess.run([size, isize]), [1, 0])
      self.assertEqual(sess.run([key_v, get_v], feed_dict={gi: 0}), [0, [3]])
      self.assertEqual(sess.run([size, isize]), [0, 0])

  @test_util.run_deprecated_v1
  def testNonScalarKeyOrderedMap(self):
    with ops.Graph().as_default() as g:
      x = array_ops.placeholder(dtypes.float32)
      v = 2. * (array_ops.zeros([128, 128]) + x)
      t = data_flow_ops.gen_data_flow_ops.ordered_map_stage(
          key=constant_op.constant(value=[1], shape=(1, 3), dtype=dtypes.int64),
          indices=np.array([[6]]),
          values=[x, v],
          dtypes=[dtypes.int64],
          capacity=0,
          memory_limit=0,
          container='container1',
          shared_name='',
          name=None)

    g.finalize()

    with self.session(graph=g) as sess:
      with self.assertRaisesRegex(errors.InvalidArgumentError,
                                  'key must be an int64 scalar'):
        sess.run(t, feed_dict={x: 1})

  @test_util.run_deprecated_v1
  def testNonScalarKeyUnorderedMap(self):
    with ops.Graph().as_default() as g:
      x = array_ops.placeholder(dtypes.float32)
      v = 2. * (array_ops.zeros([128, 128]) + x)
      t = data_flow_ops.gen_data_flow_ops.map_stage(
          key=constant_op.constant(value=[1], shape=(1, 3), dtype=dtypes.int64),
          indices=np.array([[6]]),
          values=[x, v],
          dtypes=[dtypes.int64],
          capacity=0,
          memory_limit=0,
          container='container1',
          shared_name='',
          name=None)

    g.finalize()

    with self.session(graph=g) as sess:
      with self.assertRaisesRegex(errors.InvalidArgumentError,
                                  'key must be an int64 scalar'):
        sess.run(t, feed_dict={x: 1})