CVE-2022-36002: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. When `Unbatch` receives a nonscalar input `id`, it gives a `CHECK` fail that can trigger a denial of service attack. We have patched the issue in GitHub commit 4419d10d576adefa36b0e0a9425d2569f7c0189f. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: The product contains an assert() or similar statement that can be triggered by an attacker, which leads to an application exit or other behavior that is more severe than necessary.
CVSS v3 base score: 7.5
Published: 2022-09-16T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/batch_kernels.cc, tensorflow/python/ops/batch_ops_test.py
Commit date: 2022-07-18 10:52:10 -0700
Commit message: Fix check failure in Unbatch Op kernel by checking whether input argument is a scalar before trying to extract value.

PiperOrigin-RevId: 461660186

Diff: diff --git a/tensorflow/core/kernels/batch_kernels.cc b/tensorflow/core/kernels/batch_kernels.cc
index 7aba9db55ddc7..df84bb9d5e6a3 100644
--- a/tensorflow/core/kernels/batch_kernels.cc
+++ b/tensorflow/core/kernels/batch_kernels.cc
@@ -21,10 +21,11 @@ limitations under the License.
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_util.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h"
 #include "tensorflow/core/kernels/batching_util/batch_resource_base.h"
 #include "tensorflow/core/kernels/batching_util/bounded_executor.h"
@@ -652,10 +653,16 @@ class UnbatchResource : public ResourceBase {
           "Wrong shape for index tensor. Expected 1st dimension size to be 3 ; "
           "Got: ",
           batch_index_t.shape().dim_size(1), ".");
     }
 
+    if (!TensorShapeUtils::IsScalar(context->input(2).shape())) {
+      return errors::InvalidArgument(
+          "Input id should be scalar; "
+          "Got: ",
+          context->input(2).DebugString(), ".");
+    }
     const int64_t batch_key = context->input(2).scalar<int64_t>()();
     const bool nonempty_input = batch_index_t.dim_size(0) > 0;
 
     // If we have a non-empty tensor, slice it up.
     // (It is important to do this outside of the critical section below.)
diff --git a/tensorflow/python/ops/batch_ops_test.py b/tensorflow/python/ops/batch_ops_test.py
index 8fae09a791199..15a1a71a93bb6 100644
--- a/tensorflow/python/ops/batch_ops_test.py
+++ b/tensorflow/python/ops/batch_ops_test.py
@@ -234,10 +234,30 @@ class BatchOpsTest(test.TestCase):
       main_results = sess.run([result], feed_dict={inp: [2]})
       worker_thread.join()
       self.assertEqual(thread_results[0], [2])
       self.assertEqual(main_results[0], [3])
 
+  def testUnbatchInvalidIdArg(self):
+    """Tests that unbatch work together."""
+    if context.executing_eagerly():
+      batched_tensor = constant_op.constant(
+          value=np.random.random(size=(3, 3, 1)), dtype=dtypes.float64)
+      batched_index = constant_op.constant(
+          value=np.random.randint(0, 100, size=(3, 3, 1)), dtype=dtypes.int64)
+      arg_id = constant_op.constant(
+          value=np.random.randint(0, 100, size=(3, 3, 1)), dtype=dtypes.int64)
+
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  "Input id should be scalar;"):
+        batch_ops.unbatch(
+            batched_tensor=batched_tensor,
+            batch_index=batched_index,
+            id=arg_id,
+            timeout_micros=50,
+            container="",
+            shared_name="")
+
   def testBatchDecoratedWithCapturedInput(self):
     """Tests that the batch_function decorator works."""
     if context.executing_eagerly():
       return
     with self.cached_session() as sess:

Fixed code: #include "tensorflow/core/framework/tensor_shape.h"
    if (!TensorShapeUtils::IsScalar(context->input(2).shape())) {
      return errors::InvalidArgument(
          "Input id should be scalar; "
          "Got: ",
          context->input(2).DebugString(), ".");
    }
  def testUnbatchInvalidIdArg(self):
    """Tests that unbatch work together."""
    if context.executing_eagerly():
      batched_tensor = constant_op.constant(
          value=np.random.random(size=(3, 3, 1)), dtype=dtypes.float64)
      batched_index = constant_op.constant(
          value=np.random.randint(0, 100, size=(3, 3, 1)), dtype=dtypes.int64)
      arg_id = constant_op.constant(
          value=np.random.randint(0, 100, size=(3, 3, 1)), dtype=dtypes.int64)

      with self.assertRaisesRegex(errors.InvalidArgumentError,
                                  "Input id should be scalar;"):
        batch_ops.unbatch(
            batched_tensor=batched_tensor,
            batch_index=batched_index,
            id=arg_id,
            timeout_micros=50,
            container="",
            shared_name="")