commit 32ff5cf2a4638d207b3d54abb1356a822cb2bddd
author Sai Deng <sdeng@google.com> Fri Feb 26 20:36:30 2021 -0800
committer Sai Deng <sdeng@google.com> Wed Mar 03 20:20:01 2021 +0000
tree 0c58b78afb0af3edc71c919bff6f72ecb8700267
parent 3386e9fa635029e6c1da7f325a6009c14e36726a

Add checks to the mallocs in butteraugli.c

Change-Id: Ib707859e4d51ae2a8d27533f0341514f9c835d68

aom_dsp/butteraugli.c

diff --git a/aom_dsp/butteraugli.c b/aom_dsp/butteraugli.c
index 1b47aba..b9fc0b4 100644
--- a/aom_dsp/butteraugli.c
+++ b/aom_dsp/butteraugli.c
@@ -16,9 +16,9 @@
 #include "aom_mem/aom_mem.h"
 #include "third_party/libyuv/include/libyuv/convert_argb.h"
 
-void aom_calc_butteraugli(const YV12_BUFFER_CONFIG *source,
-                          const YV12_BUFFER_CONFIG *distorted, int bit_depth,
-                          float *dist_map) {
+int aom_calc_butteraugli(const YV12_BUFFER_CONFIG *source,
+                         const YV12_BUFFER_CONFIG *distorted, int bit_depth,
+                         float *dist_map) {
   (void)bit_depth;
   assert(bit_depth == 8);
   assert(source->y_width == source->uv_width * 2);
@@ -28,7 +28,12 @@
   size_t buffer_size = width * height * 3;
   uint8_t *src_rgb = (uint8_t *)aom_malloc(buffer_size);
   uint8_t *distorted_rgb = (uint8_t *)aom_malloc(buffer_size);
-  // TODO(sdeng): Convert them to sRGB.
+  if (!src_rgb || !distorted_rgb) {
+    aom_free(src_rgb);
+    aom_free(distorted_rgb);
+    return 0;
+  }
+
   I420ToRGB24Matrix(source->y_buffer, source->y_stride, source->u_buffer,
                     source->uv_stride, source->v_buffer, source->uv_stride,
                     src_rgb, width * 3, &kYuvH709Constants, width, height);
@@ -45,9 +50,16 @@
       api, width, height, &pixel_format, src_rgb, buffer_size, &pixel_format,
       distorted_rgb, buffer_size);
 
-  const float *distmap;
+  const float *distmap = NULL;
   uint32_t row_stride;
   JxlButteraugliResultGetDistmap(result, &distmap, &row_stride);
+  if (distmap == NULL) {
+    JxlButteraugliApiDestroy(api);
+    JxlButteraugliResultDestroy(result);
+    aom_free(src_rgb);
+    aom_free(distorted_rgb);
+    return 0;
+  }
 
   for (int j = 0; j < height; ++j) {
     for (int i = 0; i < width; ++i) {
@@ -55,7 +67,9 @@
     }
   }
 
+  JxlButteraugliApiDestroy(api);
   JxlButteraugliResultDestroy(result);
   aom_free(src_rgb);
   aom_free(distorted_rgb);
+  return 1;
 }

aom_dsp/butteraugli.h

diff --git a/aom_dsp/butteraugli.h b/aom_dsp/butteraugli.h
index 95314ce..06402aa 100644
--- a/aom_dsp/butteraugli.h
+++ b/aom_dsp/butteraugli.h
@@ -14,8 +14,8 @@
 
 #include "aom_scale/yv12config.h"
 
-void aom_calc_butteraugli(const YV12_BUFFER_CONFIG *source,
-                          const YV12_BUFFER_CONFIG *distorted, int bit_depth,
-                          float *dist_map);
+int aom_calc_butteraugli(const YV12_BUFFER_CONFIG *source,
+                         const YV12_BUFFER_CONFIG *distorted, int bit_depth,
+                         float *dist_map);
 
 #endif  // AOM_AOM_DSP_BUTTERAUGLI_H_

av1/encoder/tune_butteraugli.c

diff --git a/av1/encoder/tune_butteraugli.c b/av1/encoder/tune_butteraugli.c
index 4d5866a..c43930b 100644
--- a/av1/encoder/tune_butteraugli.c
+++ b/av1/encoder/tune_butteraugli.c
@@ -48,7 +48,10 @@
   const YV12_BUFFER_CONFIG *recon = &cpi->butteraugli_info.recon;
   float *diffmap;
   CHECK_MEM_ERROR(cm, diffmap, aom_malloc(width * height * sizeof(*diffmap)));
-  aom_calc_butteraugli(source, recon, bit_depth, diffmap);
+  if (!aom_calc_butteraugli(source, recon, bit_depth, diffmap)) {
+    aom_internal_error(&cm->error, AOM_CODEC_ERROR,
+                       "Failed to calculate Butteraugli distances.");
+  }
 
   const int block_size = BLOCK_16X16;
   const int num_mi_w = mi_size_wide[block_size];