CVE-2022-21740: Out-of-bounds Write (CWE-787)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of `SparseCountSparseOutput` is vulnerable to a heap overflow. 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 writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 8.8
Published: 2022-02-03T15:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/count_ops.cc
Commit date: 2021-12-07 19:49:35 -0800
Commit message: Further validate sparse tensor for `SparseCount`: indices must be valid within dense shape.

PiperOrigin-RevId: 414888122
Change-Id: I4552bd74c135ecd4bcb5448acc0a3ce9402d8286

Diff: diff --git a/tensorflow/core/kernels/count_ops.cc b/tensorflow/core/kernels/count_ops.cc
index 1f99e0783e26f..cc101b66f8140 100644
--- a/tensorflow/core/kernels/count_ops.cc
+++ b/tensorflow/core/kernels/count_ops.cc
@@ -204,10 +204,27 @@ class SparseCount : public OpKernel {
             "Got ", shape.shape().dim_size(0),
             " dimensions, indices shape: ", indices.shape().DebugString()));
     OP_REQUIRES(context, shape.NumElements() > 0,
                 errors::InvalidArgument(
                     "The shape argument requires at least one element."));
+    // Validate indices: each index must be valid for the corresponding
+    // dimension. This could be possibly done better.
+    const auto indices_values = indices.matrix<int64_t>();
+    const auto shape_vector = shape.vec<int64_t>();
+    int num_values = values.NumElements();  // same as first dim of indices
+    int rank = indices.shape().dim_size(1);
+    for (int i = 0; i < num_values; ++i) {
+      for (int j = 0; j < rank; ++j) {
+        OP_REQUIRES(
+            context,
+            indices_values(i, j) >= 0 && indices_values(i, j) < shape_vector(j),
+            errors::InvalidArgument(
+                "Invalid index value at ", i, ": dimension ", j, " has value ",
+                indices_values(i, j), " which is not in [0, ", shape_vector(j),
+                ") (as given by dense shape ", shape.DebugString()));
+      }
+    }
 
     if (use_weights) {
       OP_REQUIRES(
           context, weights.shape() == values.shape(),
           errors::InvalidArgument(
@@ -215,15 +232,12 @@ class SparseCount : public OpKernel {
               weights.shape().DebugString(),
               "; values shape: ", values.shape().DebugString()));
     }
 
     bool is_1d = shape.NumElements() == 1;
-    auto shape_vector = shape.flat<int64_t>();
     int num_batches = is_1d ? 1 : shape_vector(0);
-    int num_values = values.NumElements();
 
-    const auto indices_values = indices.matrix<int64_t>();
     const auto values_values = values.flat<T>();
     const auto weight_values = weights.flat<W>();
 
     auto per_batch_counts = BatchedMap<W>(num_batches);
 

Vulnerable code:     auto shape_vector = shape.flat<int64_t>();
    int num_values = values.NumElements();
    const auto indices_values = indices.matrix<int64_t>();
Fixed code:     // Validate indices: each index must be valid for the corresponding
    // dimension. This could be possibly done better.
    const auto indices_values = indices.matrix<int64_t>();
    const auto shape_vector = shape.vec<int64_t>();
    int num_values = values.NumElements();  // same as first dim of indices
    int rank = indices.shape().dim_size(1);
    for (int i = 0; i < num_values; ++i) {
      for (int j = 0; j < rank; ++j) {
        OP_REQUIRES(
            context,
            indices_values(i, j) >= 0 && indices_values(i, j) < shape_vector(j),
            errors::InvalidArgument(
                "Invalid index value at ", i, ": dimension ", j, " has value ",
                indices_values(i, j), " which is not in [0, ", shape_vector(j),
                ") (as given by dense shape ", shape.DebugString()));
      }
    }