CVE-2022-21737: Improper Check for Unusual or Exceptional Conditions (CWE-754)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of `*Bincount` operations allows malicious users to cause denial of service by passing in arguments which would trigger a `CHECK`-fail. There are several conditions that the input arguments must satisfy. Some are not caught during shape inference and others are not caught during kernel implementation. This results in `CHECK` failures later when the output tensors get allocated. 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 does not check or incorrectly checks for unusual or exceptional conditions that are not expected to occur frequently during day to day operation of the product.
CVSS v3 base score: 6.5
Published: 2022-02-03T14:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/bincount_op.cc, tensorflow/core/ops/math_ops.cc, tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
Commit date: 2021-12-08 12:10:17 -0800
Commit message: Fix check-fail when bincount ops are passed invalid values.

PiperOrigin-RevId: 415063028
Change-Id: I20f8dc09933ddca1111c4efbf9a3a1e863215d02

Diff: diff --git a/tensorflow/core/kernels/bincount_op.cc b/tensorflow/core/kernels/bincount_op.cc
index 0f661dd9f201a..aad11a45a09b6 100644
--- a/tensorflow/core/kernels/bincount_op.cc
+++ b/tensorflow/core/kernels/bincount_op.cc
@@ -274,10 +274,13 @@ class DenseBincountOp : public OpKernel {
                     "Shape must be at most rank 2 but is rank ", data.dims()));
 
     const Tensor& size_t = ctx->input(1);
     const Tensor& weights = ctx->input(2);
 
+    OP_REQUIRES(ctx, size_t.dims() == 0,
+                errors::InvalidArgument("Shape must be rank 0 but is rank ",
+                                        size_t.dims()));
     Tidx size = size_t.scalar<Tidx>()();
     OP_REQUIRES(
         ctx, size >= 0,
         errors::InvalidArgument("size (", size, ") must be non-negative"));
 
@@ -370,10 +373,13 @@ class SparseBincountOp : public OpKernel {
     const Tensor& dense_shape = ctx->input(2);
     const Tensor& size_t = ctx->input(3);
     const auto weights = ctx->input(4).flat<T>();
     const int64_t weights_size = weights.size();
 
+    OP_REQUIRES(ctx, size_t.dims() == 0,
+                errors::InvalidArgument("Shape must be rank 0 but is rank ",
+                                        size_t.dims()));
     Tidx size = size_t.scalar<Tidx>()();
     OP_REQUIRES(
         ctx, size >= 0,
         errors::InvalidArgument("size (", size, ") must be non-negative"));
 
@@ -460,10 +466,13 @@ class RaggedBincountOp : public OpKernel {
     const auto values = ctx->input(1).flat<Tidx>();
     const Tensor& size_t = ctx->input(2);
     const auto weights = ctx->input(3).flat<T>();
     const int64_t weights_size = weights.size();
 
+    OP_REQUIRES(ctx, size_t.dims() == 0,
+                errors::InvalidArgument("Shape must be rank 0 but is rank ",
+                                        size_t.dims()));
     Tidx size = size_t.scalar<Tidx>()();
     OP_REQUIRES(
         ctx, size >= 0,
         errors::InvalidArgument("size (", size, ") must be non-negative"));
 
diff --git a/tensorflow/core/ops/math_ops.cc b/tensorflow/core/ops/math_ops.cc
index a23f159cbfc5d..3f653287fdb2f 100644
--- a/tensorflow/core/ops/math_ops.cc
+++ b/tensorflow/core/ops/math_ops.cc
@@ -1697,10 +1697,15 @@ REGISTER_OP("Bincount")
         // Return unknown shape if size is not known.
         c->set_output(0, c->UnknownShapeOfRank(1));
         return Status::OK();
       }
 
+      if (size_tensor->dims() != 0) {
+        return errors::InvalidArgument("Shape must be rank 0 but is rank ",
+                                       size_tensor->dims());
+      }
+
       // Return `[size]` shape if size is known.
       int32_t size_val = size_tensor->scalar<int32>()();
       if (size_val < 0) {
         return errors::InvalidArgument("size (", size_val,
                                        ") must be non-negative");
@@ -1728,10 +1733,14 @@ REGISTER_OP("DenseBincount")
       if (size_tensor == nullptr) {
         // Return unknown shape if size is not known.
         c->set_output(0, c->UnknownShape());
         return Status::OK();
       }
+      if (size_tensor->dims() != 0) {
+        return errors::InvalidArgument("Shape must be rank 0 but is rank ",
+                                       size_tensor->dims());
+      }
 
       int64_t size_val;
       DataType dtype;
       TF_RETURN_IF_ERROR(c->GetAttr("Tidx", &dtype));
       if (dtype == DT_INT32) {
@@ -1769,10 +1778,14 @@ REGISTER_OP("SparseBincount")
       if (size_tensor == nullptr) {
         // Return unknown shape if size is not known.
         c->set_output(0, c->UnknownShape());
         return Status::OK();
       }
+      if (size_tensor->dims() != 0) {
+        return errors::InvalidArgument("Shape must be rank 0 but is rank ",
+                                       size_tensor->dims());
+      }
 
       int64_t size_val;
       DataType dtype;
       TF_RETURN_IF_ERROR(c->GetAttr("Tidx", &dtype));
       if (dtype == DT_INT32) {
diff --git a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
index 3574fac5e1844..9161b3b082270 100644
--- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
@@ -342,10 +342,18 @@ class BincountOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       with test_util.use_gpu():
         self.evaluate(
             gen_math_ops.dense_bincount(
                 input=[[[1, 2, 3], [0, 3, 2]]], weights=[], size=10))
 
+  @test_util.run_in_graph_and_eager_modes
+  def test_size_is_not_scalar(self):  # b/206619828
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "Shape must be rank 0 but is rank 1"):
+      self.evaluate(
+          gen_math_ops.dense_bincount(
+              input=[0], size=[1, 1], weights=[3], binary_output=False))
+
 
 class SparseBincountOpTest(test_util.TensorFlowTestCase,
                            parameterized.TestCase):
 
   @parameterized.parameters([{
@@ -509,10 +517,23 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
                 dense_shape=inp_sparse.dense_shape,
                 size=size,
                 weights=[],
                 binary_output=True)))
 
+  @test_util.run_in_graph_and_eager_modes
+  def test_size_is_not_scalar(self):  # b/206619828
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "Shape must be rank 0 but is rank 1"):
+      self.evaluate(
+          gen_math_ops.sparse_bincount(
+              indices=[[0], [1]],
+              values=[0, 0],
+              dense_shape=[1, 1],
+              size=[1, 1],
+              weights=[0, 0],
+              binary_output=False))
+
 
 class RaggedBincountOpTest(test_util.TensorFlowTestCase,
                            parameterized.TestCase):
 
   @parameterized.parameters([{
@@ -648,8 +669,21 @@ class RaggedBincountOpTest(test_util.TensorFlowTestCase,
                 values=x.values,
                 weights=[],
                 size=size,
                 binary_output=True)))
 
+  @test_util.run_in_graph_and_eager_modes
+  def test_size_is_not_scalar(self):  # b/206619828
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "Shape must be rank 0 but is rank 1"):
+      self.evaluate(
+          gen_math_ops.ragged_bincount(
+              splits=[0, 0, 1],
+              values=[1],
+              size=[1, 1],
+              weights=[0, 0, 0],
+              binary_output=False,
+              name=None))
+
 
 if __name__ == "__main__":
   googletest.main()

Fixed code:     OP_REQUIRES(ctx, size_t.dims() == 0,
                errors::InvalidArgument("Shape must be rank 0 but is rank ",
                                        size_t.dims()));
    OP_REQUIRES(ctx, size_t.dims() == 0,
                errors::InvalidArgument("Shape must be rank 0 but is rank ",
                                        size_t.dims()));
    OP_REQUIRES(ctx, size_t.dims() == 0,
                errors::InvalidArgument("Shape must be rank 0 but is rank ",
                                        size_t.dims()));
      if (size_tensor->dims() != 0) {
        return errors::InvalidArgument("Shape must be rank 0 but is rank ",
                                       size_tensor->dims());
      }

      if (size_tensor->dims() != 0) {
        return errors::InvalidArgument("Shape must be rank 0 but is rank ",
                                       size_tensor->dims());
      }
      if (size_tensor->dims() != 0) {
        return errors::InvalidArgument("Shape must be rank 0 but is rank ",
                                       size_tensor->dims());
      }
  @test_util.run_in_graph_and_eager_modes
  def test_size_is_not_scalar(self):  # b/206619828
    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Shape must be rank 0 but is rank 1"):
      self.evaluate(
          gen_math_ops.dense_bincount(
              input=[0], size=[1, 1], weights=[3], binary_output=False))

  @test_util.run_in_graph_and_eager_modes
  def test_size_is_not_scalar(self):  # b/206619828
    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Shape must be rank 0 but is rank 1"):
      self.evaluate(
          gen_math_ops.sparse_bincount(
              indices=[[0], [1]],
              values=[0, 0],
              dense_shape=[1, 1],
              size=[1, 1],
              weights=[0, 0],
              binary_output=False))

  @test_util.run_in_graph_and_eager_modes
  def test_size_is_not_scalar(self):  # b/206619828
    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "Shape must be rank 0 but is rank 1"):
      self.evaluate(
          gen_math_ops.ragged_bincount(
              splits=[0, 0, 1],
              values=[1],
              size=[1, 1],
              weights=[0, 0, 0],
              binary_output=False,
              name=None))