CVE-2022-21737 (C++, published 2022-02-03T14:15Z)
Description: 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.
Severity: MEDIUM (CVSS v3 base score: 6.5)
Weakness: CWE-754 \u2014 Improper Check for Unusual or Exceptional Conditions
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.
Repository: https://github.com/tensorflow/tensorflow
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 message: Fix check-fail when bincount ops are passed invalid values.

PiperOrigin-RevId: 415063028
Change-Id: I20f8dc09933ddca1111c4efbf9a3a1e863215d02


Vulnerable code:

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