From 5b2741abc3550964a0313f421b1c9181eea2afdb Mon Sep 17 00:00:00 2001
From: Yannis Guyon <yguyon@google.com>
Date: Thu, 27 Jul 2023 11:32:10 +0200
Subject: [PATCH] Keep avifRWDataRealloc() input unchanged on error

Free resources in aviftest.c.
Change TODO syntax to match style guide.
---
 apps/shared/iccmaker.c |  6 ++----
 include/avif/avif.h    |  1 +
 src/mem.c              |  3 +--
 src/rawdata.c          | 21 ++++++---------------
 src/stream.c           |  3 +--
 tests/aviftest.c       | 13 +++++++++----
 6 files changed, 20 insertions(+), 27 deletions(-)

diff --git a/apps/shared/iccmaker.c b/apps/shared/iccmaker.c
index 278199a663..bcfec28448 100644
--- a/apps/shared/iccmaker.c
+++ b/apps/shared/iccmaker.c
@@ -444,10 +444,9 @@ avifBool avifGenerateRGBICC(avifRWData * icc, float gamma, const float primaries
     }
 
     computeMD5(buffer, sizeof(iccColorTemplate));
-    if (avifRWDataRealloc(icc, iccColorLength) != AVIF_RESULT_OK) {
+    if (avifRWDataSet(icc, buffer, iccColorLength) != AVIF_RESULT_OK) {
         return AVIF_FALSE;
     }
-    memcpy(icc->data, buffer, iccColorLength);
 
     return AVIF_TRUE;
 }
@@ -471,10 +470,9 @@ avifBool avifGenerateGrayICC(avifRWData * icc, float gamma, const float white[2]
     }
 
     computeMD5(buffer, sizeof(iccGrayTemplate));
-    if (avifRWDataRealloc(icc, iccGrayLength) != AVIF_RESULT_OK) {
+    if (avifRWDataSet(icc, buffer, iccGrayLength) != AVIF_RESULT_OK) {
         return AVIF_FALSE;
     }
-    memcpy(icc->data, buffer, iccGrayLength);
 
     return AVIF_TRUE;
 }
diff --git a/include/avif/avif.h b/include/avif/avif.h
index 3827477abd..d2caef9353 100644
--- a/include/avif/avif.h
+++ b/include/avif/avif.h
@@ -199,6 +199,7 @@ typedef struct avifRWData
 // clang-format on
 
 // The avifRWData input must be zero-initialized before being manipulated with these functions.
+// If AVIF_RESULT_OUT_OF_MEMORY is returned, raw is left unchanged.
 AVIF_API avifResult avifRWDataRealloc(avifRWData * raw, size_t newSize);
 AVIF_API avifResult avifRWDataSet(avifRWData * raw, const uint8_t * data, size_t len);
 AVIF_API void avifRWDataFree(avifRWData * raw);
diff --git a/src/mem.c b/src/mem.c
index ec9a2f480e..f26891a13d 100644
--- a/src/mem.c
+++ b/src/mem.c
@@ -9,8 +9,7 @@ void * avifAlloc(size_t size)
 {
     void * out = malloc(size);
     if (out == NULL) {
-        // TODO: https://github.com/AOMediaCodec/libavif/issues/820
-        //  - Remove once all calling sites propagate the error as AVIF_RESULT_OUT_OF_MEMORY.
+        // TODO(issue #820): Remove once all calling sites propagate the error as AVIF_RESULT_OUT_OF_MEMORY.
         abort();
     }
     return out;
diff --git a/src/rawdata.c b/src/rawdata.c
index 519cd9659a..e4653160c8 100644
--- a/src/rawdata.c
+++ b/src/rawdata.c
@@ -1,7 +1,6 @@
 // Copyright 2019 Joe Drago. All rights reserved.
 // SPDX-License-Identifier: BSD-2-Clause
 
-#include "avif/avif.h"
 #include "avif/internal.h"
 
 #include <string.h>
@@ -9,22 +8,14 @@
 avifResult avifRWDataRealloc(avifRWData * raw, size_t newSize)
 {
     if (raw->size != newSize) {
-        uint8_t * old = raw->data;
-        size_t oldSize = raw->size;
-        raw->data = avifAlloc(newSize);
-        if (!raw->data) {
-            // The alternative would be to keep old in raw->data but this avoids
-            // the need of calling avifRWDataFree() on error.
-            avifFree(old);
-            raw->size = 0;
-            return AVIF_RESULT_OUT_OF_MEMORY;
+        uint8_t * newData = avifAlloc(newSize);
+        AVIF_CHECKERR(newData, AVIF_RESULT_OUT_OF_MEMORY);
+        if (raw->size && newSize) {
+            memcpy(newData, raw->data, AVIF_MIN(raw->size, newSize));
         }
+        avifFree(raw->data);
+        raw->data = newData;
         raw->size = newSize;
-        if (oldSize) {
-            size_t bytesToCopy = (oldSize < raw->size) ? oldSize : raw->size;
-            memcpy(raw->data, old, bytesToCopy);
-            avifFree(old);
-        }
     }
     return AVIF_RESULT_OK;
 }
diff --git a/src/stream.c b/src/stream.c
index 498572df32..a4e9d39bdd 100644
--- a/src/stream.c
+++ b/src/stream.c
@@ -241,8 +241,7 @@ static void makeRoom(avifRWStream * stream, size_t size)
         newSize += AVIF_STREAM_BUFFER_INCREMENT;
     }
     if (avifRWDataRealloc(stream->raw, newSize) != AVIF_RESULT_OK) {
-        // TODO: https://github.com/AOMediaCodec/libavif/issues/820
-        //  - Return AVIF_RESULT_OUT_OF_MEMORY instead.
+        // TODO(issue #820): Return AVIF_RESULT_OUT_OF_MEMORY instead.
         abort();
     }
 }
diff --git a/tests/aviftest.c b/tests/aviftest.c
index ea71bfcc8f..1bac2c7406 100644
--- a/tests/aviftest.c
+++ b/tests/aviftest.c
@@ -187,7 +187,8 @@ static int runIOTests(const char * dataDir)
         size_t filenameLen = strlen(filename);
         if ((ioDirLen + filenameLen) > FILENAME_MAX_LENGTH) {
             printf("Path too long: %s\n", filename);
-            return 1;
+            retCode = 1;
+            break;
         }
         strcpy(fullFilename, ioDir);
         strcat(fullFilename, filename);
@@ -195,19 +196,23 @@ static int runIOTests(const char * dataDir)
         FILE * f = fopen(fullFilename, "rb");
         if (!f) {
             printf("Can't open for read: %s\n", filename);
-            return 1;
+            retCode = 1;
+            break;
         }
         fseek(f, 0, SEEK_END);
         size_t fileSize = ftell(f);
         fseek(f, 0, SEEK_SET);
         if (avifRWDataRealloc(&fileBuffer, fileSize) != AVIF_RESULT_OK) {
             printf("Out of memory when allocating buffer to read file: %s\n", filename);
-            return 1;
+            fclose(f);
+            retCode = 1;
+            break;
         }
         if (fread(fileBuffer.data, 1, fileSize, f) != fileSize) {
             printf("Can't read entire file: %s\n", filename);
             fclose(f);
-            return 1;
+            retCode = 1;
+            break;
         }
         fclose(f);