fix(cgif_rgb): add NULL checks after malloc/realloc and fix integer overflow

Author: uwezkhan

src/cgif_rgb.c

@@ -129,14 +129,25 @@ static int32_t col_hash_collision_count(const uint8_t* rgb, const uint8_t* hashT
 /* initialize hash table storing colors and their frequency */
 static colHashTable* init_col_hash_table(uint32_t tableSize){
   colHashTable* colhash = malloc(sizeof(colHashTable));
+  if(colhash == NULL) return NULL;
   colhash->tableSize = getNextPrimePower2(tableSize); // increase table size to next prime number
-  colhash->frequ = malloc(sizeof(uint32_t) * colhash->tableSize);
+  colhash->frequ     = malloc(sizeof(uint32_t) * colhash->tableSize);
   colhash->hashTable = malloc(3 * colhash->tableSize);
   colhash->indexUsed = malloc(colhash->tableSize);
-  colhash->pPalette = malloc(3 * colhash->tableSize);
-  colhash->colIdx = malloc(sizeof(uint32_t)*colhash->tableSize);
+  colhash->pPalette  = malloc(3 * colhash->tableSize);
+  colhash->colIdx    = malloc(sizeof(uint32_t) * colhash->tableSize);
+  if(colhash->frequ == NULL || colhash->hashTable == NULL || colhash->indexUsed == NULL ||
+     colhash->pPalette == NULL || colhash->colIdx == NULL) {
+    free(colhash->frequ);
+    free(colhash->hashTable);
+    free(colhash->indexUsed);
+    free(colhash->pPalette);
+    free(colhash->colIdx);
+    free(colhash);
+    return NULL;
+  }
   colhash->cnt = 0;                                      // no colors initially
-  memset(colhash->pPalette, 0, 3 * colhash->tableSize);  // unused part of color table is uninitialized otheriwse
+  memset(colhash->pPalette, 0, 3 * colhash->tableSize);  // unused part of color table is uninitialized otherwise
   memset(colhash->indexUsed, 0, colhash->tableSize);     // initially no entry in hash-table is used
   return colhash;
 }
@@ -154,12 +165,25 @@ static void free_col_hash_table(colHashTable* colhash){
 /* increase the size of the color hash table */
 static void resize_col_hash_table(colHashTable* colhash){
   uint32_t tableSizeNew;
-  tableSizeNew  = getNextPrimePower2(colhash->tableSize); // increase table size to the next prime number above the next power of two
-  colhash->pPalette      = realloc(colhash->pPalette, 3 * tableSizeNew);
-  colhash->colIdx        = realloc(colhash->colIdx, sizeof(uint32_t) * tableSizeNew);
-  uint8_t* hashTable_new = malloc(3 * tableSizeNew);
-  uint8_t* indexUsed_new = malloc(tableSizeNew);
-  uint32_t* frequ_new = malloc(sizeof(uint32_t) * tableSizeNew);
+  tableSizeNew = getNextPrimePower2(colhash->tableSize); // increase table size to the next prime number above the next power of two
+  uint8_t*  pPalette_new  = realloc(colhash->pPalette, 3 * tableSizeNew);
+  uint32_t* colIdx_new    = realloc(colhash->colIdx, sizeof(uint32_t) * tableSizeNew);
+  uint8_t*  hashTable_new = malloc(3 * tableSizeNew);
+  uint8_t*  indexUsed_new = malloc(tableSizeNew);
+  uint32_t* frequ_new     = malloc(sizeof(uint32_t) * tableSizeNew);
+  if(pPalette_new == NULL || colIdx_new == NULL || hashTable_new == NULL ||
+     indexUsed_new == NULL || frequ_new == NULL) {
+    // realloc failure: original pointers may be freed or unchanged; clean up and leave colhash intact
+    free(hashTable_new);
+    free(indexUsed_new);
+    free(frequ_new);
+    // if realloc returned NULL, original pointer is still valid
+    if(pPalette_new != NULL) colhash->pPalette = pPalette_new;
+    if(colIdx_new   != NULL) colhash->colIdx   = colIdx_new;
+    return;
+  }
+  colhash->pPalette = pPalette_new;
+  colhash->colIdx   = colIdx_new;
   memset(indexUsed_new, 0, tableSizeNew);
   colhash->cnt = 0;
   for(uint32_t j = 0; j < colhash->tableSize; ++j) { // TBD (no improvement when tested): easier to loop over pPalette and also leave pPalette in place?, if indexUsed is also unnecessary then
@@ -176,15 +200,16 @@ static void resize_col_hash_table(colHashTable* colhash){
   colhash->tableSize = tableSizeNew;
   free(colhash->hashTable); // free part of old hash table that is not used anymore
   free(colhash->indexUsed); // free part of old hash table that is not used anymore
-  free(colhash->frequ); // free part of old hash table that is not used anymore
+  free(colhash->frequ);     // free part of old hash table that is not used anymore
   colhash->hashTable = hashTable_new; // pass pointer to new hash table
   colhash->indexUsed = indexUsed_new; // pass pointer to new hash table
-  colhash->frequ = frequ_new; // pass pointer to new hash table
+  colhash->frequ     = frequ_new;     // pass pointer to new hash table
 }
 
 /* take frequ indexed by hash(rgb) and return corresponding dense array */
 static uint32_t* hash_to_dense(colHashTable* colhash, cgif_chan_fmt fmtChan) {
   uint32_t* frequDense = malloc(sizeof(uint32_t) * colhash->cnt);
+  if(frequDense == NULL) return NULL;
   uint32_t h;
   (void)fmtChan;
   for(uint32_t i = 0; i < colhash->cnt; ++i) {
@@ -243,6 +268,7 @@ static treeNode* new_tree_node(uint8_t* pPalette, uint32_t* frequ, uint16_t* num
   float var[3];
 
   treeNode* node = malloc(sizeof(treeNode));
+  if(node == NULL) return NULL;
   node->idxMin   = idxMin; // minimum color in pPalette belonging to the node
   node->idxMax   = idxMax; // maximum color in pPalette belonging to the node
   get_variance(pPalette, frequ, idxMin, idxMax, var, node->mean);
@@ -488,10 +514,18 @@ static uint32_t quantize_and_dither(colHashTable* colhash, const uint8_t* pImage
   const uint16_t colMax = (1uL << depthMax) - 1; // maximum number of colors (-1 to leave space for transparency), disadvantage (TBD): quantization for static image with 256 colors and no alpha channel unnecessary
   if(colhash->cnt > colMax) { // color-quantization is needed
     uint32_t* pFrequDense = hash_to_dense(colhash, fmtChan);
+    if(pFrequDense == NULL) return 0;
     treeNode* root        = create_decision_tree(colhash->pPalette, pFrequDense, pPalette256, colhash->cnt, colMax, depthMax); // create decision tree (dynamic, splits along rgb-dimension with highest variance)
     free(pFrequDense);
-    float* pImageDataRGBfloat = malloc(fmtChan * numPixel * sizeof(float)); // TBD fmtChan + only when hasAlpha
-    for(uint32_t i = 0; i < fmtChan * numPixel; ++i){
+    if(root == NULL) return 0;
+    // Guard against integer overflow: fmtChan (uint8_t) * numPixel (uint32_t) can overflow uint32_t
+    size_t floatBufSize = (size_t)fmtChan * (size_t)numPixel * sizeof(float);
+    float* pImageDataRGBfloat = malloc(floatBufSize);
+    if(pImageDataRGBfloat == NULL) {
+      free_decision_tree(root);
+      return 0;
+    }
+    for(uint32_t i = 0; i < (size_t)fmtChan * numPixel; ++i){
       pImageDataRGBfloat[i] = pImageDataRGB[i];
     }
     uint8_t transIndex = colMax;
@@ -519,8 +553,9 @@ static uint32_t quantize_and_dither(colHashTable* colhash, const uint8_t* pImage
 CGIFrgb* cgif_rgb_newgif(const CGIFrgb_Config* pConfig) {
   CGIF_Config idxConfig = {0};
   CGIFrgb* pGIFrgb;
-  
+
   pGIFrgb = malloc(sizeof(CGIFrgb));
+  if(pGIFrgb == NULL) return NULL;
   memset(pGIFrgb, 0, sizeof(CGIFrgb));
   idxConfig.pWriteFn  = pConfig->pWriteFn;
   idxConfig.pContext  = pConfig->pContext;
@@ -536,7 +571,7 @@ CGIFrgb* cgif_rgb_newgif(const CGIFrgb_Config* pConfig) {
   }
   pGIFrgb->config    = *pConfig;
   pGIFrgb->curResult = CGIF_PENDING;
-  return pGIFrgb;  
+  return pGIFrgb;
 }
 
 cgif_result cgif_rgb_addframe(CGIFrgb* pGIF, const CGIFrgb_FrameConfig* pConfig) {
@@ -558,9 +593,18 @@ cgif_result cgif_rgb_addframe(CGIFrgb* pGIF, const CGIFrgb_FrameConfig* pConfig)
     return CGIF_ERROR;
   }
   pNewBef = malloc(pConfig->fmtChan * MULU16(imageWidth, imageHeight));
+  if(pNewBef == NULL) {
+    pGIF->curResult = CGIF_ERROR;
+    return CGIF_ERROR;
+  }
   memcpy(pNewBef, pConfig->pImageData, pConfig->fmtChan * MULU16(imageWidth, imageHeight));
   fConfig.pLocalPalette = aPalette;
   fConfig.pImageData    = malloc(pGIF->config.width * (uint32_t)pGIF->config.height);
+  if(fConfig.pImageData == NULL) {
+    free(pNewBef);
+    pGIF->curResult = CGIF_ERROR;
+    return CGIF_ERROR;
+  }
   fConfig.delay         = pConfig->delay;
   fConfig.attrFlags     = CGIF_FRAME_ATTR_USE_LOCAL_TABLE;
   if(pConfig->attrFlags & CGIF_RGB_FRAME_ATTR_INTERLACED) {