CVE-2022-35987: Reachable Assertion (CWE-617)
value: TensorFlow is an open source platform for machine learning. `DenseBincount` assumes its input tensor `weights` to either have the same shape as its input tensor `input` or to be length-0. A different `weights` shape will trigger a `CHECK` fail that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit bf4c14353c2328636a18bfad1e151052c81d5f43. 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-16T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/compiler/tf2xla/kernels/bincount_op.cc, tensorflow/core/kernels/bincount_op.cc, tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
Commit date: 2022-07-13 16:18:10 -0700
Commit message: Fix security vulnerability with DenseBincountOp

PiperOrigin-RevId: 460826735

Diff: diff --git a/tensorflow/compiler/tf2xla/kernels/bincount_op.cc b/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
index e7fafcd5bc303..2d4bba2623931 100644
--- a/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
@@ -78,10 +78,21 @@ class DenseBincountOp : public XlaOpKernel {
     xla::XlaOp weights = ctx->Input(2);
     StatusOr<xla::Shape> weights_shape_or = ctx->builder()->GetShape(weights);
     OP_REQUIRES_OK(ctx, weights_shape_or.status());
 
     auto weights_shape = weights_shape_or.ValueOrDie();
+    OP_REQUIRES(ctx,
+                xla::ShapeUtil::CompatibleIgnoringElementType(weights_shape,
+                                                              input_shape) ||
+                    (weights_shape.dimensions_size() > 0 &&
+                     weights_shape.dimensions(0) == 0),
+                errors::InvalidArgument(
+                    "`weights` must be the same shape as `arr` or a length-0 "
+                    "`Tensor`, in which case it acts as all weights equal to "
+                    "1. Received ",
+                    weights_shape.DebugString()));
+
     auto weights_size = weights_shape.dimensions(0);
     bool has_weights = false;
     if (weights_size) {
       has_weights = true;
     }
diff --git a/tensorflow/core/kernels/bincount_op.cc b/tensorflow/core/kernels/bincount_op.cc
index 4bef8c5be2252..a45692814c7e3 100644
--- a/tensorflow/core/kernels/bincount_op.cc
+++ b/tensorflow/core/kernels/bincount_op.cc
@@ -278,10 +278,18 @@ class DenseBincountOp : public OpKernel {
     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()));
+    OP_REQUIRES(ctx,
+                weights.shape() == data.shape() || weights.NumElements() == 0,
+                errors::InvalidArgument(
+                    "`weights` must be the same shape as `arr` or a length-0 "
+                    "`Tensor`, in which case it acts as all weights equal to "
+                    "1. Received ",
+                    weights.shape().DebugString()));
+
     Tidx size = size_t.scalar<Tidx>()();
     OP_REQUIRES(
         ctx, size >= 0,
         errors::InvalidArgument("size (", size, ") must be non-negative"));
 
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 44213ca405146..a000d777681b9 100644
--- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
@@ -22,10 +22,11 @@ from tensorflow.python.framework import errors_impl
 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 bincount_ops
 from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import sparse_ops
 from tensorflow.python.ops.ragged import ragged_factory_ops
 from tensorflow.python.ops.ragged import ragged_tensor
 from tensorflow.python.platform import googletest
 
@@ -150,10 +151,35 @@ class BincountTest(test_util.TensorFlowTestCase):
     with ops.Graph().as_default():
       s = array_ops.placeholder(dtype=dtypes.int32)
       v2 = gen_math_ops.bincount([1, 2, 3, 1, 6, 8], s, [])
       self.assertAllEqual(v2.get_shape().as_list(), [None])
 
+  @test_util.run_in_graph_and_eager_modes
+  def test_invalid_inputs(self):
+    binary_output = True
+    inp = random_ops.random_uniform(
+        shape=[10, 10],
+        minval=-10000,
+        maxval=10000,
+        dtype=dtypes.int32,
+        seed=-2460)
+    size = random_ops.random_uniform(
+        shape=[], minval=-10000, maxval=10000, dtype=dtypes.int32, seed=-10000)
+    weights = random_ops.random_uniform(
+        shape=[],
+        minval=-10000,
+        maxval=10000,
+        dtype=dtypes.float32,
+        seed=-10000)
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(
+          gen_math_ops.dense_bincount(
+              input=inp,
+              size=size,
+              weights=weights,
+              binary_output=binary_output))
+
 
 class BincountOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
 
   @parameterized.parameters([{
       "dtype": np.int32,

Fixed code:     OP_REQUIRES(ctx,
                xla::ShapeUtil::CompatibleIgnoringElementType(weights_shape,
                                                              input_shape) ||
                    (weights_shape.dimensions_size() > 0 &&
                     weights_shape.dimensions(0) == 0),
                errors::InvalidArgument(
                    "`weights` must be the same shape as `arr` or a length-0 "
                    "`Tensor`, in which case it acts as all weights equal to "
                    "1. Received ",
                    weights_shape.DebugString()));

    OP_REQUIRES(ctx,
                weights.shape() == data.shape() || weights.NumElements() == 0,
                errors::InvalidArgument(
                    "`weights` must be the same shape as `arr` or a length-0 "
                    "`Tensor`, in which case it acts as all weights equal to "
                    "1. Received ",
                    weights.shape().DebugString()));

from tensorflow.python.ops import random_ops
  @test_util.run_in_graph_and_eager_modes
  def test_invalid_inputs(self):
    binary_output = True
    inp = random_ops.random_uniform(
        shape=[10, 10],
        minval=-10000,
        maxval=10000,
        dtype=dtypes.int32,
        seed=-2460)
    size = random_ops.random_uniform(
        shape=[], minval=-10000, maxval=10000, dtype=dtypes.int32, seed=-10000)
    weights = random_ops.random_uniform(
        shape=[],
        minval=-10000,
        maxval=10000,
        dtype=dtypes.float32,
        seed=-10000)
    with self.assertRaises(errors.InvalidArgumentError):
      self.evaluate(
          gen_math_ops.dense_bincount(
              input=inp,
              size=size,
              weights=weights,
              binary_output=binary_output))