CVE-2022-21734: Access of Resource Using Incompatible Type ('Type Confusion') (CWE-843)
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.
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-03T13:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/map_stage_op.cc, tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
Commit date: 2021-12-02 19:02:16 -0800
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

Diff: diff --git a/tensorflow/core/kernels/map_stage_op.cc b/tensorflow/core/kernels/map_stage_op.cc
index aee53ae78d3af..463e2a8ced5c2 100644
--- a/tensorflow/core/kernels/map_stage_op.cc
+++ b/tensorflow/core/kernels/map_stage_op.cc
@@ -534,10 +534,15 @@ class MapStageOp : public OpKernel {
     OP_REQUIRES_OK(ctx, ctx->input("indices", &indices_tensor));
     OP_REQUIRES_OK(ctx, ctx->input_list("values", &values_tensor));
     OP_REQUIRES(ctx, key_tensor->NumElements() > 0,
                 errors::InvalidArgument("key must not be empty"));
 
+    OP_REQUIRES(ctx, key_tensor->NumElements() == 1,
+                errors::InvalidArgument(
+                    "key must be an int64 scalar, got tensor with shape: ",
+                    key_tensor->shape()));
+
     // Create copy for insertion into Staging Area
     Tensor key(*key_tensor);
 
     // Create the tuple to store
     for (std::size_t i = 0; i < values_tensor.size(); ++i) {
diff --git a/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py b/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
index 313b5244ee15c..8600ad1f8d726 100644
--- a/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
+++ b/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
@@ -10,12 +10,15 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # 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 errors
+import numpy as np
+
+from tensorflow.python.framework import constant_op
 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
@@ -26,11 +29,11 @@ TIMEOUT = 1
 
 class MapStageTest(test.TestCase):
 
   @test_util.run_deprecated_v1
   def testSimple(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
         gi = array_ops.placeholder(dtypes.int64)
         v = 2. * (array_ops.zeros([128, 128]) + x)
@@ -38,21 +41,21 @@ class MapStageTest(test.TestCase):
         stager = data_flow_ops.MapStagingArea([dtypes.float32])
         stage = stager.put(pi, [v], [0])
         k, y = stager.get(gi)
         y = math_ops.reduce_max(math_ops.matmul(y, y))
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       sess.run(stage, feed_dict={x: -1, pi: 0})
       for i in range(10):
         _, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
         self.assertAllClose(4 * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
 
   @test_util.run_deprecated_v1
   def testMultiple(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
         gi = array_ops.placeholder(dtypes.int64)
         v = 2. * (array_ops.zeros([128, 128]) + x)
@@ -60,51 +63,50 @@ class MapStageTest(test.TestCase):
         stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32])
         stage = stager.put(pi, [x, v], [0, 1])
         k, (z, y) = stager.get(gi)
         y = math_ops.reduce_max(z * math_ops.matmul(y, y))
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       sess.run(stage, feed_dict={x: -1, pi: 0})
       for i in range(10):
         _, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
         self.assertAllClose(
             4 * (i - 1) * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
 
   @test_util.run_deprecated_v1
   def testDictionary(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
         gi = array_ops.placeholder(dtypes.int64)
         v = 2. * (array_ops.zeros([128, 128]) + x)
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [dtypes.float32, dtypes.float32],
-            shapes=[[], [128, 128]],
-            names=['x', 'v'])
+        stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32],
+                                              shapes=[[], [128, 128]],
+                                              names=['x', 'v'])
         stage = stager.put(pi, {'x': x, 'v': v})
         key, ret = stager.get(gi)
         z = ret['x']
         y = ret['v']
         y = math_ops.reduce_max(z * math_ops.matmul(y, y))
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       sess.run(stage, feed_dict={x: -1, pi: 0})
       for i in range(10):
         _, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
         self.assertAllClose(
             4 * (i - 1) * (i - 1) * (i - 1) * 128, yval, rtol=1e-4)
 
   def testColocation(self):
     gpu_dev = test.gpu_device_name()
 
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         v = 2. * (array_ops.zeros([128, 128]) + x)
       with ops.device(gpu_dev):
         stager = data_flow_ops.MapStagingArea([dtypes.float32])
@@ -117,62 +119,60 @@ class MapStageTest(test.TestCase):
         _, z = stager.get()
         self.assertEqual(x[0].device, '/device:CPU:0')
         self.assertEqual(y.device, '/device:CPU:0')
         self.assertEqual(z[0].device, '/device:CPU:0')
 
-    G.finalize()
+    g.finalize()
 
   @test_util.run_deprecated_v1
   def testPeek(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.int32, name='x')
         pi = array_ops.placeholder(dtypes.int64)
         gi = array_ops.placeholder(dtypes.int64)
         p = array_ops.placeholder(dtypes.int32, name='p')
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [
-                dtypes.int32,
-            ], shapes=[[]])
+        stager = data_flow_ops.MapStagingArea([
+            dtypes.int32,
+        ], shapes=[[]])
         stage = stager.put(pi, [x], [0])
         peek = stager.peek(gi)
         size = stager.size()
 
-    G.finalize()
+    g.finalize()
 
     n = 10
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       for i in range(n):
         sess.run(stage, feed_dict={x: i, pi: i})
 
       for i in range(n):
-        self.assertTrue(sess.run(peek, feed_dict={gi: i})[0] == i)
+        self.assertEqual(sess.run(peek, feed_dict={gi: i})[0], i)
 
-      self.assertTrue(sess.run(size) == 10)
+      self.assertEqual(sess.run(size), 10)
 
   @test_util.run_deprecated_v1
   def testSizeAndClear(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32, name='x')
         pi = array_ops.placeholder(dtypes.int64)
         gi = array_ops.placeholder(dtypes.int64)
         v = 2. * (array_ops.zeros([128, 128]) + x)
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [dtypes.float32, dtypes.float32],
-            shapes=[[], [128, 128]],
-            names=['x', 'v'])
+        stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32],
+                                              shapes=[[], [128, 128]],
+                                              names=['x', 'v'])
         stage = stager.put(pi, {'x': x, 'v': v})
         size = stager.size()
         clear = stager.clear()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       sess.run(stage, feed_dict={x: -1, pi: 3})
       self.assertEqual(sess.run(size), 1)
       sess.run(stage, feed_dict={x: -1, pi: 1})
       self.assertEqual(sess.run(size), 2)
       sess.run(clear)
@@ -180,34 +180,35 @@ class MapStageTest(test.TestCase):
 
   @test_util.run_deprecated_v1
   def testCapacity(self):
     capacity = 3
 
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.int32, name='x')
         pi = array_ops.placeholder(dtypes.int64, name='pi')
         gi = array_ops.placeholder(dtypes.int64, name='gi')
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [
-                dtypes.int32,
-            ], capacity=capacity, shapes=[[]])
+        stager = data_flow_ops.MapStagingArea([
+            dtypes.int32,
+        ],
+                                              capacity=capacity,
+                                              shapes=[[]])
 
       stage = stager.put(pi, [x], [0])
       get = stager.get()
       size = stager.size()
 
-    G.finalize()
+    g.finalize()
 
     from six.moves import queue as Queue
     import threading
 
     queue = Queue.Queue()
     n = 8
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # Stage data in a separate thread which will block
       # when it hits the staging area's capacity and thus
       # not fill the queue with n tokens
       def thread_run():
         for i in range(n):
@@ -232,46 +233,46 @@ class MapStageTest(test.TestCase):
                   "Staging Area size is '{}' and configured "
                   "capacity is '{}'.".format(capacity, i, sess.run(size),
                                              capacity))
 
       # Should have capacity elements in the staging area
-      self.assertTrue(sess.run(size) == capacity)
+      self.assertEqual(sess.run(size), capacity)
 
       # Clear the staging area completely
       for i in range(n):
         sess.run(get)
 
-      self.assertTrue(sess.run(size) == 0)
+      self.assertEqual(sess.run(size), 0)
 
   @test_util.run_deprecated_v1
   def testMemoryLimit(self):
     memory_limit = 512 * 1024  # 512K
     chunk = 200 * 1024  # 256K
     capacity = memory_limit // chunk
 
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.uint8, name='x')
         pi = array_ops.placeholder(dtypes.int64, name='pi')
         gi = array_ops.placeholder(dtypes.int64, name='gi')
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [dtypes.uint8], memory_limit=memory_limit, shapes=[[]])
+        stager = data_flow_ops.MapStagingArea([dtypes.uint8],
+                                              memory_limit=memory_limit,
+                                              shapes=[[]])
         stage = stager.put(pi, [x], [0])
         get = stager.get()
         size = stager.size()
 
-    G.finalize()
+    g.finalize()
 
     from six.moves import queue as Queue
     import threading
-    import numpy as np
 
     queue = Queue.Queue()
     n = 8
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # Stage data in a separate thread which will block
       # when it hits the staging area's capacity and thus
       # not fill the queue with n tokens
       def thread_run():
         for i in range(n):
@@ -297,60 +298,61 @@ class MapStageTest(test.TestCase):
                   "Staging Area size is '{}' and configured "
                   "capacity is '{}'.".format(capacity, i, sess.run(size),
                                              capacity))
 
       # Should have capacity elements in the staging area
-      self.assertTrue(sess.run(size) == capacity)
+      self.assertEqual(sess.run(size), capacity)
 
       # Clear the staging area completely
       for i in range(n):
         sess.run(get)
 
-      self.assertTrue(sess.run(size) == 0)
+      self.assertEqual(sess.run(size), 0)
 
   @test_util.run_deprecated_v1
   def testOrdering(self):
     import six
     import random
 
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.int32, name='x')
         pi = array_ops.placeholder(dtypes.int64, name='pi')
         gi = array_ops.placeholder(dtypes.int64, name='gi')
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [
-                dtypes.int32,
-            ], shapes=[[]], ordered=True)
+        stager = data_flow_ops.MapStagingArea([
+            dtypes.int32,
+        ],
+                                              shapes=[[]],
+                                              ordered=True)
         stage = stager.put(pi, [x], [0])
         get = stager.get()
         size = stager.size()
 
-    G.finalize()
+    g.finalize()
 
     n = 10
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # Keys n-1..0
       keys = list(reversed(six.moves.range(n)))
 
       for i in keys:
         sess.run(stage, feed_dict={pi: i, x: i})
 
-      self.assertTrue(sess.run(size) == n)
+      self.assertEqual(sess.run(size), n)
 
       # Check that key, values come out in ascending order
       for i, k in enumerate(reversed(keys)):
         get_key, values = sess.run(get)
         self.assertTrue(i == k == get_key == values)
 
-      self.assertTrue(sess.run(size) == 0)
+      self.assertEqual(sess.run(size), 0)
 
   @test_util.run_deprecated_v1
   def testPartialDictInsert(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         f = array_ops.placeholder(dtypes.float32)
         v = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
@@ -364,53 +366,51 @@ class MapStageTest(test.TestCase):
         stage_v = stager.put(pi, {'v': v})
         key, ret = stager.get(gi)
         size = stager.size()
         isize = stager.incomplete_size()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # 0 complete and incomplete entries
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
       # Stage key 0, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Stage key 1, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 2])
+      self.assertEqual(sess.run([size, isize]), [0, 2])
 
       # Now complete key 0 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 0, v: 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
       # We can now obtain tuple associated with key 0
-      self.assertTrue(
-          sess.run([key, ret], feed_dict={
-              gi: 0
-          }) == [0, {
+      self.assertEqual(
+          sess.run([key, ret], feed_dict={gi: 0}),
+          [0, {
               'x': 1,
               'f': 2,
               'v': 1
           }])
 
       # 0 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Now complete key 1 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 1, v: 3})
       # We can now obtain tuple associated with key 1
-      self.assertTrue(
-          sess.run([key, ret], feed_dict={
-              gi: 1
-          }) == [1, {
+      self.assertEqual(
+          sess.run([key, ret], feed_dict={gi: 1}),
+          [1, {
               'x': 1,
               'f': 2,
               'v': 3
           }])
 
   @test_util.run_deprecated_v1
   def testPartialIndexInsert(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         f = array_ops.placeholder(dtypes.float32)
         v = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
@@ -422,39 +422,39 @@ class MapStageTest(test.TestCase):
         stage_v = stager.put(pi, [v], [1])
         key, ret = stager.get(gi)
         size = stager.size()
         isize = stager.incomplete_size()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # 0 complete and incomplete entries
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
       # Stage key 0, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Stage key 1, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 2])
+      self.assertEqual(sess.run([size, isize]), [0, 2])
 
       # Now complete key 0 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 0, v: 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
       # We can now obtain tuple associated with key 0
-      self.assertTrue(sess.run([key, ret], feed_dict={gi: 0}) == [0, [1, 1, 2]])
+      self.assertEqual(sess.run([key, ret], feed_dict={gi: 0}), [0, [1, 1, 2]])
 
       # 0 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Now complete key 1 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 1, v: 3})
       # We can now obtain tuple associated with key 1
-      self.assertTrue(sess.run([key, ret], feed_dict={gi: 1}) == [1, [1, 3, 2]])
+      self.assertEqual(sess.run([key, ret], feed_dict={gi: 1}), [1, [1, 3, 2]])
 
   @test_util.run_deprecated_v1
   def testPartialDictGetsAndPeeks(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         f = array_ops.placeholder(dtypes.float32)
         v = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
@@ -474,85 +474,79 @@ class MapStageTest(test.TestCase):
         pop_key_xf, pop_xf = stager.get(indices=['x', 'f'])
         pop_key_v, pop_v = stager.get(pi, ['v'])
         size = stager.size()
         isize = stager.incomplete_size()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # 0 complete and incomplete entries
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
       # Stage key 0, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Stage key 1, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 2])
+      self.assertEqual(sess.run([size, isize]), [0, 2])
 
       # Now complete key 0 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 0, v: 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
 
       # We can now peek at 'x' and 'f' values associated with key 0
-      self.assertTrue(sess.run(peek_xf, feed_dict={pei: 0}) == {'x': 1, 'f': 2})
+      self.assertEqual(sess.run(peek_xf, feed_dict={pei: 0}), {'x': 1, 'f': 2})
       # Peek at 'v' value associated with key 0
-      self.assertTrue(sess.run(peek_v, feed_dict={pei: 0}) == {'v': 1})
+      self.assertEqual(sess.run(peek_v, feed_dict={pei: 0}), {'v': 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
 
       # We can now obtain 'x' and 'f' values associated with key 0
-      self.assertTrue(
-          sess.run([key_xf, get_xf], feed_dict={
-              gi: 0
-          }) == [0, {
+      self.assertEqual(
+          sess.run([key_xf, get_xf], feed_dict={gi: 0}), [0, {
               'x': 1,
               'f': 2
           }])
       # Still have 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
 
       # We can no longer get 'x' and 'f' from key 0
       with self.assertRaises(errors.InvalidArgumentError) as cm:
         sess.run([key_xf, get_xf], feed_dict={gi: 0})
 
       exc_str = ("Tensor at index '0' for key '0' " 'has already been removed.')
 
-      self.assertTrue(exc_str in cm.exception.message)
+      self.assertIn(exc_str, cm.exception.message)
 
       # Obtain 'v' value associated with key 0
-      self.assertTrue(
-          sess.run([key_v, get_v], feed_dict={
-              gi: 0
-          }) == [0, {
+      self.assertEqual(
+          sess.run([key_v, get_v], feed_dict={gi: 0}), [0, {
               'v': 1
           }])
       # 0 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
 
       # Now complete key 1 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 1, v: 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 0])
+      self.assertEqual(sess.run([size, isize]), [1, 0])
 
       # Pop without key to obtain 'x' and 'f' values associated with key 1
-      self.assertTrue(sess.run([pop_key_xf, pop_xf]) == [1, {'x': 1, 'f': 2}])
+      self.assertEqual(sess.run([pop_key_xf, pop_xf]), [1, {'x': 1, 'f': 2}])
       # still 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 0])
+      self.assertEqual(sess.run([size, isize]), [1, 0])
       # We can now obtain 'x' and 'f' values associated with key 1
-      self.assertTrue(
-          sess.run([pop_key_v, pop_v], feed_dict={
-              pi: 1
-          }) == [1, {
+      self.assertEqual(
+          sess.run([pop_key_v, pop_v], feed_dict={pi: 1}), [1, {
               'v': 1
           }])
       # Nothing is left
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
 
   @test_util.run_deprecated_v1
   def testPartialIndexGets(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         f = array_ops.placeholder(dtypes.float32)
         v = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
@@ -566,31 +560,75 @@ class MapStageTest(test.TestCase):
         key_xf, get_xf = stager.get(gi, [0, 2])
         key_v, get_v = stager.get(gi, [1])
         size = stager.size()
         isize = stager.incomplete_size()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # Stage complete tuple
       sess.run(stage_xvf, feed_dict={pi: 0, x: 1, f: 2, v: 3})
 
-      self.assertTrue(sess.run([size, isize]) == [1, 0])
+      self.assertEqual(sess.run([size, isize]), [1, 0])
 
       # Partial get using indices
-      self.assertTrue(
-          sess.run([key_xf, get_xf], feed_dict={
-              gi: 0
-          }) == [0, [1, 2]])
+      self.assertEqual(
+          sess.run([key_xf, get_xf], feed_dict={gi: 0}), [0, [1, 2]])
 
       # Still some of key 0 left
-      self.assertTrue(sess.run([size, isize]) == [1, 0])
+      self.assertEqual(sess.run([size, isize]), [1, 0])
 
       # Partial get of remaining index
-      self.assertTrue(sess.run([key_v, get_v], feed_dict={gi: 0}) == [0, [3]])
+      self.assertEqual(sess.run([key_v, get_v], feed_dict={gi: 0}), [0, [3]])
 
       # All gone
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      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})
 
 
 if __name__ == '__main__':
   test.main()

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})