CVE-2021-41197: Integer Overflow or Wraparound (CWE-190)
value: TensorFlow is an open source platform for machine learning. In affected versions TensorFlow allows tensor to have a large number of dimensions and each dimension can be as large as desired. However, the total number of elements in a tensor must fit within an `int64_t`. If an overflow occurs, `MultiplyWithoutOverflow` would return a negative result. In the majority of TensorFlow codebase this then results in a `CHECK`-failure. Newer constructs exist which return a `Status` instead of crashing the binary. This is similar to CVE-2021-29584. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, 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: 5.5
Published: 2021-11-05T20:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/kernels/image/crop_and_resize_op.cc, tensorflow/python/ops/image_ops_test.py
Commit date: 2021-08-31 16:29:59 -0700
Commit message: PR #51732: Fix crash of tf.image.crop_and_resize when input is large number

Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/51732

This PR is part of the effort in #46890 where
tf.image.crop_and_resize will crash if shape consists of large number.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
Copybara import of the project:

--
c8d87055a56d8740d27ad8bdc74a7459ede6900e by Yong Tang <yong.tang.github@outlook.com>:

Fix crash of tf.image.crop_and_resize when input is large number

This PR is part of the effort in 46890 where
tf.image.crop_and_resize will crash if shape consists of large number.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/51732 from yongtang:46890-tf.image.crop_and_resize c8d87055a56d8740d27ad8bdc74a7459ede6900e
PiperOrigin-RevId: 394109830
Change-Id: If049dad0844df9353722029ee95bc76819eda1f4

Diff: diff --git a/tensorflow/core/kernels/image/crop_and_resize_op.cc b/tensorflow/core/kernels/image/crop_and_resize_op.cc
index b6c49335ebff7..0fd905acabda3 100644
--- a/tensorflow/core/kernels/image/crop_and_resize_op.cc
+++ b/tensorflow/core/kernels/image/crop_and_resize_op.cc
@@ -168,18 +168,19 @@ class CropAndResizeOp : public AsyncOpKernel {
     const int crop_width = internal::SubtleMustCopy(crop_size_vec(1));
     OP_REQUIRES_ASYNC(
         context, crop_height > 0 && crop_width > 0,
         errors::InvalidArgument("crop dimensions must be positive"), done);
 
+    TensorShape shape;
+    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(num_boxes), done);
+    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(crop_height), done);
+    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(crop_width), done);
+    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(depth), done);
     // Allocate output tensor.
     Tensor* output = nullptr;
-    OP_REQUIRES_OK_ASYNC(
-        context,
-        context->allocate_output(
-            0, TensorShape({num_boxes, crop_height, crop_width, depth}),
-            &output),
-        done);
+    OP_REQUIRES_OK_ASYNC(context, context->allocate_output(0, shape, &output),
+                         done);
 
     auto compute_callback = [this, context, output]() {
       const Tensor& image = context->input(0);
       const Tensor& boxes = context->input(1);
       const Tensor& box_index = context->input(2);
@@ -415,18 +416,19 @@ class CropAndResizeGradImageOp : public AsyncOpKernel {
               "Deterministic GPU implementation of CropAndResizeBackpropImage"
               " not available."),
           done);
     }
 
+    TensorShape shape;
+    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(batch_size), done);
+    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(image_height), done);
+    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(image_width), done);
+    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(depth), done);
     // Allocate output tensor.
     Tensor* output = nullptr;
-    OP_REQUIRES_OK_ASYNC(
-        context,
-        context->allocate_output(
-            0, TensorShape({batch_size, image_height, image_width, depth}),
-            &output),
-        done);
+    OP_REQUIRES_OK_ASYNC(context, context->allocate_output(0, shape, &output),
+                         done);
 
     auto compute_callback = [this, context, output]() {
       const Tensor& grads = context->input(0);
       const Tensor& boxes = context->input(1);
       const Tensor& box_index = context->input(2);
diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py
index 17f00822ede98..3e1e2da19f513 100644
--- a/tensorflow/python/ops/image_ops_test.py
+++ b/tensorflow/python/ops/image_ops_test.py
@@ -6073,10 +6073,20 @@ class DecodeImageTest(test_util.TensorFlowTestCase, parameterized.TestCase):
             boxes=[[1.0e+40, 0, 0, 0]],
             box_indices=[1],
             crop_size=[1, 1])
         self.evaluate(op)
 
+  def testImageCropAndResizeWithInvalidInput(self):
+    with self.session():
+      with self.assertRaises((errors.InternalError, ValueError)):
+        op = image_ops_impl.crop_and_resize_v2(
+            image=np.ones((1, 1, 1, 1)),
+            boxes=np.ones((11, 4)),
+            box_indices=np.ones((11)),
+            crop_size=[2065374891, 1145309325])
+        self.evaluate(op)
+
   @parameterized.named_parameters(
       ("_jpeg", "JPEG", "jpeg_merge_test1.jpg"),
       ("_png", "PNG", "lena_rgba.png"),
       ("_gif", "GIF", "scan.gif"),
   )

Vulnerable code:     OP_REQUIRES_OK_ASYNC(
        context,
        context->allocate_output(
            0, TensorShape({num_boxes, crop_height, crop_width, depth}),
            &output),
        done);
    OP_REQUIRES_OK_ASYNC(
        context,
        context->allocate_output(
            0, TensorShape({batch_size, image_height, image_width, depth}),
            &output),
        done);
Fixed code:     TensorShape shape;
    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(num_boxes), done);
    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(crop_height), done);
    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(crop_width), done);
    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(depth), done);
    OP_REQUIRES_OK_ASYNC(context, context->allocate_output(0, shape, &output),
                         done);
    TensorShape shape;
    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(batch_size), done);
    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(image_height), done);
    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(image_width), done);
    OP_REQUIRES_OK_ASYNC(context, shape.AddDimWithStatus(depth), done);
    OP_REQUIRES_OK_ASYNC(context, context->allocate_output(0, shape, &output),
                         done);
  def testImageCropAndResizeWithInvalidInput(self):
    with self.session():
      with self.assertRaises((errors.InternalError, ValueError)):
        op = image_ops_impl.crop_and_resize_v2(
            image=np.ones((1, 1, 1, 1)),
            boxes=np.ones((11, 4)),
            box_indices=np.ones((11)),
            crop_size=[2065374891, 1145309325])
        self.evaluate(op)