Check overflow in allocation (pytorch#12683)

Summary:

Overflow is unsafe as a user will get back much less memory then expected

Reviewed By: lucylq

Differential Revision: D78678540

Author: JacobSzwejbka

runtime/core/memory_allocator.h

@@ -12,12 +12,15 @@
 #include <cinttypes>
 #include <cstdint>
 
+#include <c10/util/safe_numerics.h>
+
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/platform/assert.h>
 #include <executorch/runtime/platform/compiler.h>
 #include <executorch/runtime/platform/log.h>
 #include <executorch/runtime/platform/profiler.h>
 
+
 namespace executorch {
 namespace runtime {
 
@@ -137,7 +140,17 @@ class MemoryAllocator {
     // Some users of this method allocate lists of pointers, causing the next
     // line to expand to `sizeof(type *)`, which triggers a clang-tidy warning.
     // NOLINTNEXTLINE(bugprone-sizeof-expression)
-    return static_cast<T*>(this->allocate(size * sizeof(T), alignment));
+    size_t bytes_size = 0;
+    bool overflow =
+        c10::mul_overflows(size, sizeof(T), &bytes_size);
+    if (overflow) {
+      ET_LOG(
+          Error,
+          "Failed to allocate list of type %zu: size * sizeof(T) overflowed",
+          size);
+      return nullptr;
+    }
+    return static_cast<T*>(this->allocate(bytes_size, alignment));
   }
 
   // Returns the allocator memory's base address.