CVE-2021-29619: Improper Handling of Exceptional Conditions (CWE-755)
value: TensorFlow is an end-to-end open source platform for machine learning. Passing invalid arguments (e.g., discovered via fuzzing) to `tf.raw_ops.SparseCountSparseOutput` results in segfault. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.
CWE description: The product does not handle or incorrectly handles an exceptional condition.
CVSS v3 base score: 5.5
Published: 2021-05-14T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/count_ops.cc
Commit date: 2021-03-02 17:06:27 -0800
Commit message: Fix segfaults in `tf.raw_ops.SparseCountSparseOutput`.

PiperOrigin-RevId: 360547563
Change-Id: I781c7af4b54a63d867c6e18d43a44d64a5c4e7c9

Diff: diff --git a/tensorflow/core/kernels/count_ops.cc b/tensorflow/core/kernels/count_ops.cc
index 087deef0812f0..d6ab68c2c70bd 100644
--- a/tensorflow/core/kernels/count_ops.cc
+++ b/tensorflow/core/kernels/count_ops.cc
@@ -190,10 +190,14 @@ class SparseCount : public OpKernel {
               "Weights and values must have the same shape. Weight shape: ",
               weights.shape().DebugString(),
               "; values shape: ", values.shape().DebugString()));
     }
 
+    OP_REQUIRES(context, shape.NumElements() != 0,
+                errors::InvalidArgument(
+                    "The shape argument requires at least one element."));
+
     bool is_1d = shape.NumElements() == 1;
     int num_batches = is_1d ? 1 : shape.flat<int64>()(0);
     int num_values = values.NumElements();
 
     OP_REQUIRES(context, num_values == indices.shape().dim_size(0),
@@ -210,10 +214,18 @@ class SparseCount : public OpKernel {
 
     T max_value = 0;
 
     for (int idx = 0; idx < num_values; ++idx) {
       int batch = is_1d ? 0 : indices_values(idx, 0);
+      if (batch >= num_batches) {
+        OP_REQUIRES(context, batch < num_batches,
+                    errors::InvalidArgument(
+                        "Indices value along the first dimension must be ",
+                        "lower than the first index of the shape.", "Got ",
+                        batch, " as batch and ", num_batches,
+                        " as the first dimension of the shape."));
+      }
       const auto& value = values_values(idx);
       if (value >= 0 && (maxlength_ <= 0 || value < maxlength_)) {
         if (binary_output_) {
           per_batch_counts[batch][value] = 1;
         } else if (use_weights) {

Fixed code:     OP_REQUIRES(context, shape.NumElements() != 0,
                errors::InvalidArgument(
                    "The shape argument requires at least one element."));

      if (batch >= num_batches) {
        OP_REQUIRES(context, batch < num_batches,
                    errors::InvalidArgument(
                        "Indices value along the first dimension must be ",
                        "lower than the first index of the shape.", "Got ",
                        batch, " as batch and ", num_batches,
                        " as the first dimension of the shape."));
      }