Add SafeMul and generalize SafeAdd

lib/extras/dec/apng.cc

@@ -1024,13 +1024,13 @@ Status DecodeImageAPNG(const Span<const uint8_t> bytes,
           };
           bytes_per_pixel =
               num_channels * (format.data_type == JXL_TYPE_UINT16 ? 2 : 1);
-          // TODO(eustas): ensure multiplication is safe
-          uint64_t row_bytes =
-              static_cast<uint64_t>(image_rect.xsize()) * bytes_per_pixel;
-          uint64_t max_rows = std::numeric_limits<size_t>::max() / row_bytes;
-          if (image_rect.ysize() > max_rows) {
+          size_t row_bytes;
+          size_t total_bytes;
+          if (!SafeMul(image_rect.xsize(), bytes_per_pixel, row_bytes) ||
+              !SafeMul(row_bytes, image_rect.ysize(), total_bytes)) {
             return JXL_FAILURE("Image too big.");
           }
+          (void)total_bytes;  // Calculated only for check.
           // TODO(eustas): drop frameRaw
           JXL_RETURN_IF_ERROR(
               ctx.frameRaw.Resize(row_bytes, image_rect.ysize()));

lib/extras/gain_map.cc

@@ -105,13 +105,13 @@ JXL_BOOL JxlGainMapWriteBundle(const JxlGainMapBundle* map_bundle,
   uint64_t cursor = 0;
   uint64_t next_cursor = 0;
 
-#define SAFE_CURSOR_UPDATE(n)                    \
-  do {                                           \
-    cursor = next_cursor;                        \
-    if (!jxl::SafeAdd(cursor, n, next_cursor) || \
-        next_cursor > output_buffer_size) {      \
-      return JXL_FALSE;                          \
-    }                                            \
+#define SAFE_CURSOR_UPDATE(n)                              \
+  do {                                                     \
+    cursor = next_cursor;                                  \
+    if (!jxl::SafeAdd<uint64_t>(cursor, n, next_cursor) || \
+        next_cursor > output_buffer_size) {                \
+      return JXL_FALSE;                                    \
+    }                                                      \
   } while (false)
 
   SAFE_CURSOR_UPDATE(1);
@@ -164,13 +164,13 @@ JXL_BOOL JxlGainMapReadBundle(JxlGainMapBundle* map_bundle,
   uint64_t cursor = 0;
   uint64_t next_cursor = 0;
 
-#define SAFE_CURSOR_UPDATE(n)                    \
-  do {                                           \
-    cursor = next_cursor;                        \
-    if (!jxl::SafeAdd(cursor, n, next_cursor) || \
-        next_cursor > input_buffer_size) {       \
-      return JXL_FALSE;                          \
-    }                                            \
+#define SAFE_CURSOR_UPDATE(n)                              \
+  do {                                                     \
+    cursor = next_cursor;                                  \
+    if (!jxl::SafeAdd<uint64_t>(cursor, n, next_cursor) || \
+        next_cursor > input_buffer_size) {                 \
+      return JXL_FALSE;                                    \
+    }                                                      \
   } while (false)
 
   // Read the version byte

lib/jxl/base/common.h

@@ -12,6 +12,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstdio>
+#include <limits>
 #include <memory>
 #include <string>
 #include <type_traits>
@@ -32,12 +33,23 @@ constexpr inline size_t RoundUpToBlockDim(size_t dim) {
   return (dim + 7) & ~static_cast<size_t>(7);
 }
 
-static inline bool JXL_MAYBE_UNUSED SafeAdd(const uint64_t a, const uint64_t b,
-                                            uint64_t& sum) {
+template <typename U,
+          class = typename std::enable_if<std::is_unsigned<U>::value>::type>
+static inline bool SafeAdd(const U a, const U b, U& sum) {
   sum = a + b;
   return sum >= a;  // no need to check b - either sum >= both or < both.
 }
 
+template <typename U,
+          class = typename std::enable_if<std::is_unsigned<U>::value>::type>
+static inline bool SafeMul(const U a, const U b, U& product) {
+  product = 0;
+  if (a == 0 || b == 0) return true;
+  if (b > (std::numeric_limits<U>::max() / a)) return false;
+  product = a * b;
+  return true;
+}
+
 template <typename T1, typename T2>
 constexpr inline T1 DivCeil(T1 a, T2 b) {
   return (a + b - 1) / b;

lib/jxl/fields.cc

@@ -237,7 +237,7 @@ class ReadVisitor : public VisitorBase {
     // Skip new fields this (old?) decoder didn't know about, if any.
     const size_t bits_read = reader_->TotalBitsConsumed();
     uint64_t end;
-    if (!SafeAdd(pos_after_ext_size_, total_extension_bits_, end)) {
+    if (!SafeAdd<uint64_t>(pos_after_ext_size_, total_extension_bits_, end)) {
       return JXL_FAILURE("Invalid extension size, caused overflow");
     }
     if (bits_read > end) {

lib/jxl/image.cc

@@ -81,14 +81,14 @@ Status PlaneBase::Allocate(JxlMemoryManager* memory_manager,
     return true;
   }
 
-  size_t max_y_size = std::numeric_limits<size_t>::max() / bytes_per_row_;
-  if (ysize_ > max_y_size) {
+  size_t total_bytes;
+  if (!SafeMul<size_t>(ysize_, bytes_per_row_, total_bytes)) {
     return JXL_FAILURE("Image dimensions are too large");
   }
 
-  JXL_ASSIGN_OR_RETURN(
-      bytes_, AlignedMemory::Create(memory_manager, bytes_per_row_ * ysize_,
-                                    pre_padding * sizeof_t_));
+  JXL_ASSIGN_OR_RETURN(bytes_,
+                       AlignedMemory::Create(memory_manager, total_bytes,
+                                             pre_padding * sizeof_t_));
 
   InitializePadding(*this, sizeof_t_);
 

lib/jxl/memory_manager_internal.cc

@@ -92,8 +92,10 @@ size_t BytesPerRow(const size_t xsize, const size_t sizeof_t) {
 StatusOr<AlignedMemory> AlignedMemory::Create(JxlMemoryManager* memory_manager,
                                               size_t size, size_t pre_padding) {
   JXL_ENSURE(pre_padding <= memory_manager_internal::kAlias);
-  size_t allocation_size = size + pre_padding + memory_manager_internal::kAlias;
-  if (size > allocation_size) {
+  size_t allocation_size;
+  if (!SafeAdd<size_t>(size, pre_padding, allocation_size) ||
+      !SafeAdd(allocation_size, memory_manager_internal::kAlias,
+               allocation_size)) {
     return JXL_FAILURE("Requested allocation is too large");
   }
   JXL_ENSURE(memory_manager);