CVE-2022-21738: Integer Overflow or Wraparound (CWE-190)
value: Tensorflow is an Open Source Machine Learning Framework. The implementation of `SparseCountSparseOutput` can be made to crash a TensorFlow process by an integer overflow whose result is then used in a memory allocation. 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 performs a calculation that can
         produce an integer overflow or wraparound when the logic
         assumes that the resulting value will always be larger than
         the original value. This occurs when an integer value is
         incremented to a value that is too large to store in the
         associated representation. When this occurs, the value may
         become a very small or negative number.
CVSS v3 base score: 6.5
Published: 2022-02-03T14:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/core/kernels/count_ops.cc
Commit date: 2021-12-07 20:09:14 -0800
Commit message: Prevent crash due to integer overflow followed by allocating negative sized array.

PiperOrigin-RevId: 414891322
Change-Id: I5df390e0dc1d9f115209293708950cdf9306931c

Diff: diff --git a/tensorflow/core/kernels/count_ops.cc b/tensorflow/core/kernels/count_ops.cc
index cc101b66f8140..3f9df8f0d69d6 100644
--- a/tensorflow/core/kernels/count_ops.cc
+++ b/tensorflow/core/kernels/count_ops.cc
@@ -11,20 +11,25 @@ distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <limits>
+
 #include "absl/container/flat_hash_map.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
 
+// Don't allocate too large `BatchedMap<T>` objects
+static int kMaxBatches = std::numeric_limits<int>::max();
+
 template <class T>
 using BatchedMap = std::vector<absl::flat_hash_map<int64_t, T>>;
 
 namespace {
 // TODO(momernick): Extend this function to work with outputs of rank > 2.
@@ -233,10 +238,14 @@ class SparseCount : public OpKernel {
               "; values shape: ", values.shape().DebugString()));
     }
 
     bool is_1d = shape.NumElements() == 1;
     int num_batches = is_1d ? 1 : shape_vector(0);
+    OP_REQUIRES(
+        context, 0 < num_batches && num_batches < kMaxBatches,
+        errors::InvalidArgument("Cannot allocate ", num_batches,
+                                " batches, is the dense shape too wide?"));
 
     const auto values_values = values.flat<T>();
     const auto weight_values = weights.flat<W>();
 
     auto per_batch_counts = BatchedMap<W>(num_batches);

Fixed code: #include <limits>

// Don't allocate too large `BatchedMap<T>` objects
static int kMaxBatches = std::numeric_limits<int>::max();

    OP_REQUIRES(
        context, 0 < num_batches && num_batches < kMaxBatches,
        errors::InvalidArgument("Cannot allocate ", num_batches,
                                " batches, is the dense shape too wide?"));