CVE-2021-29521: Incorrect Calculation of Buffer Size (CWE-131)
value: TensorFlow is an end-to-end open source platform for machine learning. Specifying a negative dense shape in `tf.raw_ops.SparseCountSparseOutput` results in a segmentation fault being thrown out from the standard library as `std::vector` invariants are broken. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/8f7b60ee8c0206a2c99802e3a4d1bb55d2bc0624/tensorflow/core/kernels/count_ops.cc#L199-L213) assumes the first element of the dense shape is always positive and uses it to initialize a `BatchedMap<T>` (i.e., `std::vector<absl::flat_hash_map<int64,T>>`(https://github.com/tensorflow/tensorflow/blob/8f7b60ee8c0206a2c99802e3a4d1bb55d2bc0624/tensorflow/core/kernels/count_ops.cc#L27)) data structure. If the `shape` tensor has more than one element, `num_batches` is the first value in `shape`. Ensuring that the `dense_shape` argument is a valid tensor shape (that is, all elements are non-negative) solves this issue. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2 and TensorFlow 2.3.3.
CWE description: The product does not correctly calculate the size to be used when allocating a buffer, which could lead to a buffer overflow.
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-04-19 11:41:07 -0700
Commit message: Fix the segfault in `tf.raw_ops.SparseCountSparseOutput`.

PiperOrigin-RevId: 369264941
Change-Id: I23a96a15b8370c01ee21ba3841e1c7dcbf55e93d

Diff: diff --git a/tensorflow/core/kernels/count_ops.cc b/tensorflow/core/kernels/count_ops.cc
index d6ab68c2c70b..eeb0b853c7aa 100644
--- a/tensorflow/core/kernels/count_ops.cc
+++ b/tensorflow/core/kernels/count_ops.cc
@@ -195,13 +195,21 @@ class SparseCount : public OpKernel {
     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);
+    auto shape_vector = shape.flat<int64>();
+    int num_batches = is_1d ? 1 : shape_vector(0);
     int num_values = values.NumElements();
 
+    for (int b = 0; b < shape_vector.size(); b++) {
+      OP_REQUIRES(context, shape_vector(b) >= 0,
+                  errors::InvalidArgument(
+                      "Elements in dense_shape must be >= 0. Instead got:",
+                      shape.DebugString()));
+    }
+
     OP_REQUIRES(context, num_values == indices.shape().dim_size(0),
                 errors::InvalidArgument(
                     "Number of values must match first dimension of indices.",
                     "Got ", num_values,
                     " values, indices shape: ", indices.shape().DebugString()));

Vulnerable code:     int num_batches = is_1d ? 1 : shape.flat<int64>()(0);
Fixed code:     auto shape_vector = shape.flat<int64>();
    int num_batches = is_1d ? 1 : shape_vector(0);
    for (int b = 0; b < shape_vector.size(); b++) {
      OP_REQUIRES(context, shape_vector(b) >= 0,
                  errors::InvalidArgument(
                      "Elements in dense_shape must be >= 0. Instead got:",
                      shape.DebugString()));
    }