From 6b0a47b53604d8cdd8f44dcf7bc4bc69fb5fa6f5 Mon Sep 17 00:00:00 2001
From: Wan-Teh Chang <wtc@google.com>
Date: Wed, 31 Jul 2024 18:22:21 -0700
Subject: [PATCH] Fix overflows in avifRGBImageAllocatePixels()

Calculate rowBytes in uint32_t. Only the allocation size needs to be
size_t.

Also make sure it is safe to cast various rowBytes fields to ptrdiff_t.
We need to do this when subtracting rowBytes from a pointer to go back
one row.

Part of the fix to https://github.com/AOMediaCodec/libavif/issues/2271.
---
 src/avif.c | 41 ++++++++++++++++++++++++++++++-----------
 1 file changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/avif.c b/src/avif.c
index 00dc4c8fab..3d3511b0b8 100644
--- a/src/avif.c
+++ b/src/avif.c
@@ -352,15 +352,21 @@ avifResult avifImageAllocatePlanes(avifImage * image, avifPlanesFlags planes)
     if (image->width == 0 || image->height == 0) {
         return AVIF_RESULT_INVALID_ARGUMENT;
     }
-    const size_t channelSize = avifImageUsesU16(image) ? 2 : 1;
-    if (image->width > SIZE_MAX / channelSize) {
+    const uint32_t channelSize = avifImageUsesU16(image) ? 2 : 1;
+    if (image->width > UINT32_MAX / channelSize) {
         return AVIF_RESULT_INVALID_ARGUMENT;
     }
-    const size_t fullRowBytes = channelSize * image->width;
-    if ((fullRowBytes > UINT32_MAX) || (image->height > SIZE_MAX / fullRowBytes)) {
+    const uint32_t fullRowBytes = channelSize * image->width;
+#if UINT32_MAX > PTRDIFF_MAX
+    // Make sure it is safe to cast image->yuvRowBytes[i] or image->alphaRowBytes to ptrdiff_t.
+    if (fullRowBytes > PTRDIFF_MAX) {
         return AVIF_RESULT_INVALID_ARGUMENT;
     }
-    const size_t fullSize = fullRowBytes * image->height;
+#endif
+    if (image->height > SIZE_MAX / fullRowBytes) {
+        return AVIF_RESULT_INVALID_ARGUMENT;
+    }
+    const size_t fullSize = (size_t)fullRowBytes * image->height;
 
     if ((planes & AVIF_PLANES_YUV) && (image->yuvFormat != AVIF_PIXEL_FORMAT_NONE)) {
         avifPixelFormatInfo info;
@@ -368,11 +374,11 @@ avifResult avifImageAllocatePlanes(avifImage * image, avifPlanesFlags planes)
 
         image->imageOwnsYUVPlanes = AVIF_TRUE;
         if (!image->yuvPlanes[AVIF_CHAN_Y]) {
-            image->yuvRowBytes[AVIF_CHAN_Y] = (uint32_t)fullRowBytes;
             image->yuvPlanes[AVIF_CHAN_Y] = (uint8_t *)avifAlloc(fullSize);
             if (!image->yuvPlanes[AVIF_CHAN_Y]) {
                 return AVIF_RESULT_OUT_OF_MEMORY;
             }
+            image->yuvRowBytes[AVIF_CHAN_Y] = fullRowBytes;
         }
 
         if (!info.monochrome) {
@@ -381,16 +387,16 @@ avifResult avifImageAllocatePlanes(avifImage * image, avifPlanesFlags planes)
             const uint32_t shiftedH = (uint32_t)(((uint64_t)image->height + info.chromaShiftY) >> info.chromaShiftY);
 
             // These are less than or equal to fullRowBytes/fullSize. No need to check overflows.
-            const size_t uvRowBytes = channelSize * shiftedW;
-            const size_t uvSize = uvRowBytes * shiftedH;
+            const uint32_t uvRowBytes = channelSize * shiftedW;
+            const size_t uvSize = (size_t)uvRowBytes * shiftedH;
 
             for (int uvPlane = AVIF_CHAN_U; uvPlane <= AVIF_CHAN_V; ++uvPlane) {
                 if (!image->yuvPlanes[uvPlane]) {
-                    image->yuvRowBytes[uvPlane] = (uint32_t)uvRowBytes;
                     image->yuvPlanes[uvPlane] = (uint8_t *)avifAlloc(uvSize);
                     if (!image->yuvPlanes[uvPlane]) {
                         return AVIF_RESULT_OUT_OF_MEMORY;
                     }
+                    image->yuvRowBytes[uvPlane] = uvRowBytes;
                 }
             }
         }
@@ -398,11 +404,11 @@ avifResult avifImageAllocatePlanes(avifImage * image, avifPlanesFlags planes)
     if (planes & AVIF_PLANES_A) {
         image->imageOwnsAlphaPlane = AVIF_TRUE;
         if (!image->alphaPlane) {
-            image->alphaRowBytes = (uint32_t)fullRowBytes;
             image->alphaPlane = (uint8_t *)avifAlloc(fullSize);
             if (!image->alphaPlane) {
                 return AVIF_RESULT_OUT_OF_MEMORY;
             }
+            image->alphaRowBytes = fullRowBytes;
         }
     }
     return AVIF_RESULT_OK;
@@ -626,7 +632,20 @@ void avifRGBImageSetDefaults(avifRGBImage * rgb, const avifImage * image)
 avifResult avifRGBImageAllocatePixels(avifRGBImage * rgb)
 {
     avifRGBImageFreePixels(rgb);
-    const uint32_t rowBytes = rgb->width * avifRGBImagePixelSize(rgb);
+    const uint32_t pixelSize = avifRGBImagePixelSize(rgb);
+    if (rgb->width > UINT32_MAX / pixelSize) {
+        return AVIF_RESULT_INVALID_ARGUMENT;
+    }
+    const uint32_t rowBytes = rgb->width * pixelSize;
+#if UINT32_MAX > PTRDIFF_MAX
+    // Make sure it is safe to cast rgb->rowBytes to ptrdiff_t.
+    if (rowBytes > PTRDIFF_MAX) {
+        return AVIF_RESULT_INVALID_ARGUMENT;
+    }
+#endif
+    if (rgb->height > SIZE_MAX / rowBytes) {
+        return AVIF_RESULT_INVALID_ARGUMENT;
+    }
     rgb->pixels = (uint8_t *)avifAlloc((size_t)rowBytes * rgb->height);
     AVIF_CHECKERR(rgb->pixels, AVIF_RESULT_OUT_OF_MEMORY);
     rgb->rowBytes = rowBytes;