CVE-2022-35982: Insufficient Information (NVD-CWE-noinfo)
value: TensorFlow is an open source platform for machine learning. If `SparseBincount` is given inputs for `indices`, `values`, and `dense_shape` that do not make a valid sparse tensor, it results in a segfault that can be used to trigger a denial of service attack. We have patched the issue in GitHub commit 40adbe4dd15b582b0210dfbf40c243a62f5119fa. 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: There is insufficient information about the issue to classify it; details are unkown or unspecified.
CVSS v3 base score: 7.5
Published: 2022-09-16T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Other
Files changed: tensorflow/core/kernels/BUILD, tensorflow/core/kernels/bincount_op.cc, tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
Commit date: 2022-07-12 16:07:03 -0700
Commit message: Add sparse tensor validation to SparseBincountOp.

Addresses a security issue.

PiperOrigin-RevId: 460573866

Diff: diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 60a688886eeab..8a7712ac6f339 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -4419,10 +4419,11 @@ tf_kernel_library(
     name = "bincount_op",
     prefix = "bincount_op",
     deps = [
         ":fill_functor",
         ":gpu_prim_hdrs",
+        ":sparse_utils",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//third_party/eigen3",
     ],
@@ -5005,10 +5006,11 @@ cc_library(
 )
 
 SPARSE_DEPS = [
     "//tensorflow/core:framework",
     "//tensorflow/core:lib",
+    ":sparse_utils",
 ]
 
 tf_kernel_library(
     name = "sparse_add_grad_op",
     prefix = "sparse_add_grad_op",
@@ -6478,10 +6480,11 @@ filegroup(
         "sparse_xent_op.h",
         "sparse_fill_empty_rows_op.h",
         "sparse_reorder_op.h",
         "sparse_slice_op.h",
         "sparse_tensor_dense_matmul_op.h",
+        "sparse_utils.h",
         "string_util.h",
         "string_to_hash_bucket_op.h",
         "string_to_hash_bucket_fast_op.h",
         "tensor_array.h",
         "tensor_list.h",
@@ -6716,10 +6719,11 @@ filegroup(
         "random_op.cc",
         "random_op_cpu.h",
         "random_ops_util.h",
         "random_poisson_op.cc",
         "shuffle_common.h",
+        "sparse_utils.cc",
         "random_shuffle_op.cc",
         "reduce_join_op.cc",
         "reduction_ops_all.cc",
         "reduction_ops_any.cc",
         "reduction_ops_common.cc",
diff --git a/tensorflow/core/kernels/bincount_op.cc b/tensorflow/core/kernels/bincount_op.cc
index 24d9c5b19d5e3..4bef8c5be2252 100644
--- a/tensorflow/core/kernels/bincount_op.cc
+++ b/tensorflow/core/kernels/bincount_op.cc
@@ -21,10 +21,11 @@ limitations under the License.
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/bincount_op.h"
 #include "tensorflow/core/kernels/fill_functor.h"
+#include "tensorflow/core/kernels/sparse_utils.h"
 #include "tensorflow/core/lib/core/threadpool.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/determinism.h"
 
 namespace tensorflow {
@@ -367,11 +368,12 @@ class SparseBincountOp : public OpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("binary_output", &binary_output_));
   }
 
   void Compute(OpKernelContext* ctx) override {
     const Tensor& indices = ctx->input(0);
-    const auto values = ctx->input(1).flat<Tidx>();
+    const Tensor& values = ctx->input(1);
+    const auto values_flat = values.flat<Tidx>();
     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();
 
@@ -380,10 +382,13 @@ class SparseBincountOp : public OpKernel {
                                         size_t.dims()));
     Tidx size = size_t.scalar<Tidx>()();
     OP_REQUIRES(
         ctx, size >= 0,
         errors::InvalidArgument("size (", size, ") must be non-negative"));
+    OP_REQUIRES_OK(
+        ctx, sparse_utils::ValidateSparseTensor<int64_t>(
+                 indices, values, dense_shape, /*validate_indices=*/true));
 
     bool is_1d = dense_shape.NumElements() == 1;
 
     Tensor* out_t;
     functor::SetZeroFunctor<Device, T> fill;
@@ -392,15 +397,15 @@ class SparseBincountOp : public OpKernel {
       auto out = out_t->flat<T>();
       fill(ctx->eigen_device<Device>(), out);
       if (binary_output_) {
         OP_REQUIRES_OK(ctx,
                        functor::BincountFunctor<Device, Tidx, T, true>::Compute(
-                           ctx, values, weights, out, size));
+                           ctx, values_flat, weights, out, size));
       } else {
         OP_REQUIRES_OK(
             ctx, functor::BincountFunctor<Device, Tidx, T, false>::Compute(
-                     ctx, values, weights, out, size));
+                     ctx, values_flat, weights, out, size));
       }
     } else {
       const auto shape = dense_shape.flat<int64_t>();
       const int64_t num_rows = shape(0);
       OP_REQUIRES_OK(
@@ -408,11 +413,11 @@ class SparseBincountOp : public OpKernel {
       const auto out = out_t->matrix<T>();
       fill(ctx->eigen_device<Device>(), out_t->flat<T>());
       const auto indices_mat = indices.matrix<int64_t>();
       for (int64_t i = 0; i < indices_mat.dimension(0); ++i) {
         const int64_t batch = indices_mat(i, 0);
-        const Tidx bin = values(i);
+        const Tidx bin = values_flat(i);
         OP_REQUIRES(
             ctx, batch < out.dimension(0),
             errors::InvalidArgument("Index out of bound. `batch` (", batch,
                                     ") must be less than the dimension size (",
                                     out.dimension(0), ")."));
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 9161b3b082270..fedb697221236 100644
--- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
@@ -364,11 +364,11 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
   def test_sparse_bincount_all_count(self, dtype):
     np.random.seed(42)
     num_rows = 128
     size = 1000
     n_elems = 4096
-    inp_indices = np.random.randint(0, num_rows, (n_elems,))
+    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
     inp_vals = np.random.randint(0, size, (n_elems,), dtype=dtype)
 
     np_out = np.bincount(inp_vals, minlength=size)
     self.assertAllEqual(
         np_out,
@@ -388,11 +388,11 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
   def test_sparse_bincount_all_count_with_weights(self, dtype):
     np.random.seed(42)
     num_rows = 128
     size = 1000
     n_elems = 4096
-    inp_indices = np.random.randint(0, num_rows, (n_elems,))
+    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
     inp_vals = np.random.randint(0, size, (n_elems,), dtype=dtype)
     inp_weight = np.random.random((n_elems,))
 
     np_out = np.bincount(inp_vals, minlength=size, weights=inp_weight)
     self.assertAllEqual(
@@ -413,11 +413,11 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
   def test_sparse_bincount_all_binary(self, dtype):
     np.random.seed(42)
     num_rows = 128
     size = 10
     n_elems = 4096
-    inp_indices = np.random.randint(0, num_rows, (n_elems,))
+    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
     inp_vals = np.random.randint(0, size, (n_elems,), dtype=dtype)
 
     np_out = np.ones((size,))
     self.assertAllEqual(
         np_out,
@@ -438,11 +438,11 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
   def test_sparse_bincount_all_binary_weights(self, dtype):
     np.random.seed(42)
     num_rows = 128
     size = 10
     n_elems = 4096
-    inp_indices = np.random.randint(0, num_rows, (n_elems,))
+    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
     inp_vals = np.random.randint(0, size, (n_elems,), dtype=dtype)
     inp_weight = np.random.random((n_elems,))
 
     np_out = np.ones((size,))
     self.assertAllEqual(
@@ -530,10 +530,31 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
               dense_shape=[1, 1],
               size=[1, 1],
               weights=[0, 0],
               binary_output=False))
 
+  def test_sparse_bincount_input_validation(self):
+    np.random.seed(42)
+    num_rows = 128
+    size = 1000
+    n_elems = 4096
+    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
+    inp_vals = np.random.randint(0, size, (n_elems,))
+
+    # Insert negative index.
+    inp_indices[10, 0] = -2
+
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                "out of bounds"):
+      self.evaluate(
+          gen_math_ops.sparse_bincount(
+              indices=inp_indices,
+              values=inp_vals,
+              dense_shape=[num_rows],
+              size=size,
+              weights=[]))
+
 
 class RaggedBincountOpTest(test_util.TensorFlowTestCase,
                            parameterized.TestCase):
 
   @parameterized.parameters([{

Vulnerable code:     const auto values = ctx->input(1).flat<Tidx>();
                           ctx, values, weights, out, size));
                     ctx, values, weights, out, size));
        const Tidx bin = values(i);
    inp_indices = np.random.randint(0, num_rows, (n_elems,))
    inp_indices = np.random.randint(0, num_rows, (n_elems,))
    inp_indices = np.random.randint(0, num_rows, (n_elems,))
    inp_indices = np.random.randint(0, num_rows, (n_elems,))
Fixed code:         ":sparse_utils",
    ":sparse_utils",
        "sparse_utils.h",
        "sparse_utils.cc",
#include "tensorflow/core/kernels/sparse_utils.h"
    const Tensor& values = ctx->input(1);
    const auto values_flat = values.flat<Tidx>();
    OP_REQUIRES_OK(
        ctx, sparse_utils::ValidateSparseTensor<int64_t>(
                 indices, values, dense_shape, /*validate_indices=*/true));
                           ctx, values_flat, weights, out, size));
                     ctx, values_flat, weights, out, size));
        const Tidx bin = values_flat(i);
    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
  def test_sparse_bincount_input_validation(self):
    np.random.seed(42)
    num_rows = 128
    size = 1000
    n_elems = 4096
    inp_indices = np.random.randint(0, num_rows, (n_elems, 1))
    inp_vals = np.random.randint(0, size, (n_elems,))

    # Insert negative index.
    inp_indices[10, 0] = -2

    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                "out of bounds"):
      self.evaluate(
          gen_math_ops.sparse_bincount(
              indices=inp_indices,
              values=inp_vals,
              dense_shape=[num_rows],
              size=size,
              weights=[]))