Avoid integer overflows in align_image_dimension()

Impose maximum values on the input parameters so that we can perform
arithmetic operations without worrying about overflows.

Fix a bug (introduced in commit 7aa2edc) that the ~ operator is applied
to (stride_align - 1), which is unsigned int, and then the result is
converted to uint64_t.

Also change the AomImageTest.AomImgAllocHugeWidth test to write to the
first and last samples in the first row of the Y plane, so that the test
will crash if there is unsigned integer overflow in the calculation of
stride_in_bytes.

Bug: chromium:332382766
Change-Id: I634c38c35a296b5bbf3de7ddf10040e7ec5ee9a1
(cherry picked from commit 60653dff7f8ee3e769a0aeec5e210a4fc2687717)

3 files changed, 63 insertions, 13 deletions

diff --git a/aom/aom_image.h b/aom/aom_image.h
index 7f02df052..68fb31222 100644
--- a/aom/aom_image.h
+++ b/aom/aom_image.h
@@ -246,10 +246,13 @@ typedef struct aom_image {
  *                         is NULL, the storage for the descriptor will be
  *                         allocated on the heap.
  * \param[in]    fmt       Format for the image
- * \param[in]    d_w       Width of the image
- * \param[in]    d_h       Height of the image
+ * \param[in]    d_w       Width of the image. Must not exceed 0x08000000
+ *                         (2^27).
+ * \param[in]    d_h       Height of the image. Must not exceed 0x08000000
+ *                         (2^27).
  * \param[in]    align     Alignment, in bytes, of the image buffer and
- *                         each row in the image (stride).
+ *                         each row in the image (stride). Must not exceed
+ *                         65536.
  *
  * \return Returns a pointer to the initialized image descriptor. If the img
  *         parameter is non-null, the value of the img parameter will be
@@ -269,10 +272,12 @@ aom_image_t *aom_img_alloc(aom_image_t *img, aom_img_fmt_t fmt,
  *                         is NULL, the storage for the descriptor will be
  *                         allocated on the heap.
  * \param[in]    fmt       Format for the image
- * \param[in]    d_w       Width of the image
- * \param[in]    d_h       Height of the image
+ * \param[in]    d_w       Width of the image. Must not exceed 0x08000000
+ *                         (2^27).
+ * \param[in]    d_h       Height of the image. Must not exceed 0x08000000
+ *                         (2^27).
  * \param[in]    align     Alignment, in bytes, of each row in the image
- *                         (stride).
+ *                         (stride). Must not exceed 65536.
  * \param[in]    img_data  Storage to use for the image
  *
  * \return Returns a pointer to the initialized image descriptor. If the img
@@ -293,12 +298,17 @@ aom_image_t *aom_img_wrap(aom_image_t *img, aom_img_fmt_t fmt, unsigned int d_w,
  *                          is NULL, the storage for the descriptor will be
  *                          allocated on the heap.
  * \param[in]    fmt        Format for the image
- * \param[in]    d_w        Width of the image
- * \param[in]    d_h        Height of the image
+ * \param[in]    d_w        Width of the image. Must not exceed 0x08000000
+ *                          (2^27).
+ * \param[in]    d_h        Height of the image. Must not exceed 0x08000000
+ *                          (2^27).
  * \param[in]    align      Alignment, in bytes, of the image buffer and
- *                          each row in the image (stride).
+ *                          each row in the image (stride). Must not exceed
+ *                          65536.
  * \param[in]    size_align Alignment, in pixels, of the image width and height.
+ *                          Must not exceed 65536.
  * \param[in]    border     A border that is padded on four sides of the image.
+ *                          Must not exceed 65536.
  *
  * \return Returns a pointer to the initialized image descriptor. If the img
  *         parameter is non-null, the value of the img parameter will be
diff --git a/aom/src/aom_image.c b/aom/src/aom_image.c
index b68dc4c8f..e10b8a9ad 100644
--- a/aom/src/aom_image.c
+++ b/aom/src/aom_image.c
@@ -9,6 +9,7 @@
  * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
  */
 
+#include <assert.h>
 #include <limits.h>
 #include <stdlib.h>
 #include <string.h>
@@ -36,12 +37,20 @@ static aom_image_t *img_alloc_helper(
   /* NOTE: In this function, bit_depth is either 8 or 16 (if
    * AOM_IMG_FMT_HIGHBITDEPTH is set), never 10 or 12.
    */
-  unsigned int h, w, xcs, ycs, bps, bit_depth;
+  unsigned int xcs, ycs, bps, bit_depth;
 
   if (img != NULL) memset(img, 0, sizeof(aom_image_t));
 
   if (fmt == AOM_IMG_FMT_NONE) goto fail;
 
+  /* Impose maximum values on input parameters so that this function can
+   * perform arithmetic operations without worrying about overflows.
+   */
+  if (d_w > 0x08000000 || d_h > 0x08000000 || buf_align > 65536 ||
+      stride_align > 65536 || size_align > 65536 || border > 65536) {
+    goto fail;
+  }
+
   /* Treat align==0 like align==1 */
   if (!buf_align) buf_align = 1;
 
@@ -104,11 +113,13 @@ static aom_image_t *img_alloc_helper(
   }
 
   /* Calculate storage sizes given the chroma subsampling */
-  w = align_image_dimension(d_w, xcs, size_align);
-  h = align_image_dimension(d_h, ycs, size_align);
+  const unsigned int w = align_image_dimension(d_w, xcs, size_align);
+  assert(d_w <= w);
+  const unsigned int h = align_image_dimension(d_h, ycs, size_align);
+  assert(d_h <= h);
 
   uint64_t s = (fmt & AOM_IMG_FMT_PLANAR) ? w : (uint64_t)bps * w / bit_depth;
-  s = (s + 2 * border + stride_align - 1) & ~(stride_align - 1);
+  s = (s + 2 * border + stride_align - 1) & ~((uint64_t)stride_align - 1);
   s = s * bit_depth / 8;
   if (s > INT_MAX) goto fail;
   const int stride_in_bytes = (int)s;
diff --git a/test/aom_image_test.cc b/test/aom_image_test.cc
index 62f3c1274..0dfb91221 100644
--- a/test/aom_image_test.cc
+++ b/test/aom_image_test.cc
@@ -9,6 +9,8 @@
  * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
  */
 
+#include <climits>
+
 #include "aom/aom_image.h"
 #include "third_party/googletest/src/googletest/include/gtest/gtest.h"
 
@@ -81,6 +83,20 @@ TEST(AomImageTest, AomImgAllocHugeWidth) {
   image = aom_img_alloc(nullptr, AOM_IMG_FMT_I420, 0x80000000, 1, 1);
   ASSERT_EQ(image, nullptr);
 
+  // The aligned width (UINT_MAX + 1) would overflow unsigned int.
+  image = aom_img_alloc(nullptr, AOM_IMG_FMT_I420, UINT_MAX, 1, 1);
+  ASSERT_EQ(image, nullptr);
+
+  image = aom_img_alloc_with_border(nullptr, AOM_IMG_FMT_I422, 1, INT_MAX, 1,
+                                    0x40000000, 0);
+  if (image) {
+    uint16_t *y_plane =
+        reinterpret_cast<uint16_t *>(image->planes[AOM_PLANE_Y]);
+    y_plane[0] = 0;
+    y_plane[image->d_w - 1] = 0;
+    aom_img_free(image);
+  }
+
   image = aom_img_alloc(nullptr, AOM_IMG_FMT_I420, 0x7ffffffe, 1, 1);
   if (image) {
     aom_img_free(image);
@@ -101,8 +117,21 @@ TEST(AomImageTest, AomImgAllocHugeWidth) {
     aom_img_free(image);
   }
 
+  image = aom_img_alloc(nullptr, AOM_IMG_FMT_I42016, 65536, 2, 1);
+  if (image) {
+    uint16_t *y_plane =
+        reinterpret_cast<uint16_t *>(image->planes[AOM_PLANE_Y]);
+    y_plane[0] = 0;
+    y_plane[image->d_w - 1] = 0;
+    aom_img_free(image);
+  }
+
   image = aom_img_alloc(nullptr, AOM_IMG_FMT_I42016, 285245883, 2, 1);
   if (image) {
+    uint16_t *y_plane =
+        reinterpret_cast<uint16_t *>(image->planes[AOM_PLANE_Y]);
+    y_plane[0] = 0;
+    y_plane[image->d_w - 1] = 0;
     aom_img_free(image);
   }
 }