CVE-2022-21737 (published 2022-02-03T14:15Z)

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 v2 base score: 4.0 (AV:N/AC:L/Au:S/C:N/I:N/A:P)
CVSS v3 base score: 6.5, MEDIUM (CVSS:3.1/AV:N/AC:L/PR:L/UI:N/S:U/C:N/I:N/A:H)

Weakness information:
cwe_id: CWE-754
cwe_name: Improper Check for Unusual or Exceptional Conditions
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.
extended_description: 
            The programmer may assume that certain events or conditions will never occur or do not need to be worried about, such as low memory conditions, lack of access to resources due to restrictive permissions, or misbehaving clients or components. However, attackers may intentionally trigger these unusual conditions, thus violating the programmer's assumptions, possibly introducing instability, incorrect behavior, or a vulnerability.
            Note that this entry is not exclusively about the use of exceptions and exception handling, which are mechanisms for both checking and handling unusual or unexpected conditions.
         
         

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Edward Loper
author_date: 2021-12-08 12:01:53-08:00
msg: Fix check-fail when bincount ops are passed invalid values.

PiperOrigin-RevId: 415063028
Change-Id: I20f8dc09933ddca1111c4efbf9a3a1e863215d02
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: bincount_op.cc
change_type: ModificationType.MODIFY
diff: @@ -276,6 +276,9 @@ class DenseBincountOp : public OpKernel {
     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,
@@ -372,6 +375,9 @@ class SparseBincountOp : public OpKernel {
     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,
@@ -462,6 +468,9 @@ class RaggedBincountOp : public OpKernel {
     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,

programming_language: C++
filename: math_ops.cc
change_type: ModificationType.MODIFY
diff: @@ -1699,6 +1699,11 @@ REGISTER_OP("Bincount")
         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) {
@@ -1730,6 +1735,10 @@ REGISTER_OP("DenseBincount")
         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;
@@ -1771,6 +1780,10 @@ REGISTER_OP("SparseBincount")
         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;

programming_language: Rust
filename: bincount_op_test.py
change_type: ModificationType.MODIFY
diff: @@ -344,6 +344,14 @@ class BincountOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
             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):
@@ -511,6 +519,19 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
                 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):
@@ -650,6 +671,19 @@ class RaggedBincountOpTest(test_util.TensorFlowTestCase,
                 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()

programming_language: Python