aom_mem: fix overflow in alignment calculation

the calculation of the aligned allocation size could rollover if the
requested size is near SIZE_MAX

+ add a missing overflow check for aom_calloc()

Bug: aomedia:3195
Bug: aomedia:3223
Change-Id: I385de3344dadf9af56f0673b1e1a69b254e73c92
(cherry picked from commit 0929aa0af301888445539626bbc94c09ed0e28b7)

aom_mem/aom_mem.c

diff --git a/aom_mem/aom_mem.c b/aom_mem/aom_mem.c
index e977b01..f13ee2f 100644
--- a/aom_mem/aom_mem.c
+++ b/aom_mem/aom_mem.c
@@ -10,25 +10,31 @@
  */
 
 #include "aom_mem.h"
-#include <stdio.h>
+#include <assert.h>
 #include <stdlib.h>
 #include <string.h>
 #include "include/aom_mem_intrnl.h"
 #include "aom/aom_integer.h"
 
-#if defined(AOM_MAX_ALLOCABLE_MEMORY)
-// Returns 0 in case of overflow of nmemb * size.
-static int check_size_argument_overflow(uint64_t nmemb, uint64_t size) {
-  const uint64_t total_size = nmemb * size;
-  if (nmemb == 0) return 1;
-  if (size > AOM_MAX_ALLOCABLE_MEMORY / nmemb) return 0;
-  if (total_size != (size_t)total_size) return 0;
-  return 1;
+static size_t GetAllocationPaddingSize(size_t align) {
+  assert(align > 0);
+  assert(align < SIZE_MAX - ADDRESS_STORAGE_SIZE);
+  return align - 1 + ADDRESS_STORAGE_SIZE;
 }
-#endif
 
-static size_t GetAlignedMallocSize(size_t size, size_t align) {
-  return size + align - 1 + ADDRESS_STORAGE_SIZE;
+// Returns 0 in case of overflow of nmemb * size.
+static int check_size_argument_overflow(size_t nmemb, size_t size,
+                                        size_t align) {
+  if (nmemb == 0) return 1;
+  const size_t alloc_padding = GetAllocationPaddingSize(align);
+#if defined(AOM_MAX_ALLOCABLE_MEMORY)
+  assert(AOM_MAX_ALLOCABLE_MEMORY >= alloc_padding);
+  assert(AOM_MAX_ALLOCABLE_MEMORY <= SIZE_MAX);
+  if (size > (AOM_MAX_ALLOCABLE_MEMORY - alloc_padding) / nmemb) return 0;
+#else
+  if (size > (SIZE_MAX - alloc_padding) / nmemb) return 0;
+#endif
+  return 1;
 }
 
 static size_t *GetMallocAddressLocation(void *const mem) {
@@ -48,10 +54,8 @@
 
 void *aom_memalign(size_t align, size_t size) {
   void *x = NULL;
-  const size_t aligned_size = GetAlignedMallocSize(size, align);
-#if defined(AOM_MAX_ALLOCABLE_MEMORY)
-  if (!check_size_argument_overflow(1, aligned_size)) return NULL;
-#endif
+  if (!check_size_argument_overflow(1, size, align)) return NULL;
+  const size_t aligned_size = size + GetAllocationPaddingSize(align);
   void *const addr = malloc(aligned_size);
   if (addr) {
     x = aom_align_addr((unsigned char *)addr + ADDRESS_STORAGE_SIZE, align);
@@ -63,6 +67,7 @@
 void *aom_malloc(size_t size) { return aom_memalign(DEFAULT_ALIGNMENT, size); }
 
 void *aom_calloc(size_t num, size_t size) {
+  if (!check_size_argument_overflow(num, size, DEFAULT_ALIGNMENT)) return NULL;
   const size_t total_size = num * size;
   void *const x = aom_malloc(total_size);
   if (x) memset(x, 0, total_size);

test/aom_mem_test.cc

diff --git a/test/aom_mem_test.cc b/test/aom_mem_test.cc
new file mode 100644
index 0000000..849ba64
--- /dev/null
+++ b/test/aom_mem_test.cc
@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2021, Alliance for Open Media. All rights reserved
+ *
+ * This source code is subject to the terms of the BSD 2 Clause License and
+ * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
+ * was not distributed with this source code in the LICENSE file, you can
+ * obtain it at www.aomedia.org/license/software. If the Alliance for Open
+ * Media Patent License 1.0 was not distributed with this source code in the
+ * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
+ */
+
+#include "aom_mem/aom_mem.h"
+
+#include <cstdio>
+#include <cstddef>
+
+#include "third_party/googletest/src/googletest/include/gtest/gtest.h"
+
+TEST(AomMemTest, Overflow) {
+  // Allocations are aligned > 1 so SIZE_MAX should always fail.
+  ASSERT_EQ(aom_malloc(SIZE_MAX), nullptr);
+  ASSERT_EQ(aom_calloc(1, SIZE_MAX), nullptr);
+  ASSERT_EQ(aom_calloc(32, SIZE_MAX / 32), nullptr);
+  ASSERT_EQ(aom_calloc(SIZE_MAX, SIZE_MAX), nullptr);
+  ASSERT_EQ(aom_memalign(1, SIZE_MAX), nullptr);
+  ASSERT_EQ(aom_memalign(64, SIZE_MAX), nullptr);
+  ASSERT_EQ(aom_memalign(64, SIZE_MAX - 64), nullptr);
+  ASSERT_EQ(aom_memalign(64, SIZE_MAX - 64 - sizeof(size_t) + 2), nullptr);
+}
+
+TEST(AomMemTest, NullParams) {
+  ASSERT_EQ(aom_memset16(nullptr, 0, 0), nullptr);
+  aom_free(nullptr);
+}

test/test.cmake

diff --git a/test/test.cmake b/test/test.cmake
index 31ae14b..183aff2 100644
--- a/test/test.cmake
+++ b/test/test.cmake
@@ -128,6 +128,7 @@
 
 if(NOT BUILD_SHARED_LIBS)
   list(APPEND AOM_UNIT_TEST_COMMON_SOURCES
+              "${AOM_ROOT}/test/aom_mem_test.cc"
               "${AOM_ROOT}/test/av1_common_int_test.cc"
               "${AOM_ROOT}/test/cdef_test.cc"
               "${AOM_ROOT}/test/cfl_test.cc"