From 6622dc584e92972dac6a87a572bea95e1829139f Mon Sep 17 00:00:00 2001
From: Matthew Parkinson <mjp41@users.noreply.github.com>
Date: Thu, 27 Feb 2025 08:52:50 +0000
Subject: [PATCH] Bug fixes (#752)

* Fix overflow by alignment

* Bug fix: Ensure bytes_free is the total

The bytes freed was not added to the total, but
overrode it.  This meant it never fired. This
commit fixes that.
---
 src/snmalloc/mem/corealloc.h      |  2 +-
 src/snmalloc/mem/sizeclasstable.h | 50 ++++++++++++++++++++++++-------
 src/test/func/malloc/malloc.cc    |  9 +++---
 3 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/src/snmalloc/mem/corealloc.h b/src/snmalloc/mem/corealloc.h
index 55fd5ef..000f592 100644
--- a/src/snmalloc/mem/corealloc.h
+++ b/src/snmalloc/mem/corealloc.h
@@ -800,7 +800,7 @@ namespace snmalloc
         meta->as_key_tweak(),
         domesticate);
 
-      bytes_freed = objsize * length;
+      bytes_freed += objsize * length;
 
       // Update the head and the next pointer in the free list.
       meta->free_queue.append_segment(
diff --git a/src/snmalloc/mem/sizeclasstable.h b/src/snmalloc/mem/sizeclasstable.h
index 4dd2eec..c9f70be 100644
--- a/src/snmalloc/mem/sizeclasstable.h
+++ b/src/snmalloc/mem/sizeclasstable.h
@@ -129,17 +129,6 @@ namespace snmalloc
 
   using sizeclass_compress_t = uint8_t;
 
-  constexpr SNMALLOC_FAST_PATH static size_t
-  aligned_size(size_t alignment, size_t size)
-  {
-    // Client responsible for checking alignment is not zero
-    SNMALLOC_ASSERT(alignment != 0);
-    // Client responsible for checking alignment is a power of two
-    SNMALLOC_ASSERT(bits::is_pow2(alignment));
-
-    return ((alignment - 1) | (size - 1)) + 1;
-  }
-
   /**
    * This structure contains the fields required for fast paths for sizeclasses.
    */
@@ -535,4 +524,43 @@ namespace snmalloc
       return 1;
     return bits::one_at_bit(bits::ctz(rsize));
   }
+
+  constexpr SNMALLOC_FAST_PATH static size_t
+  aligned_size(size_t alignment, size_t size)
+  {
+    // Client responsible for checking alignment is not zero
+    SNMALLOC_ASSERT(alignment != 0);
+    // Client responsible for checking alignment is a power of two
+    SNMALLOC_ASSERT(bits::is_pow2(alignment));
+
+    // There are a class of corner cases to consider
+    //    alignment = 0x8
+    //    size = 0xfff...fff7
+    // for this result will be 0.  This should fail an allocation, so we need to
+    // check for this overflow.
+    // However,
+    //    alignment = 0x8
+    //    size      = 0x0
+    // will also result in 0, but this should be allowed to allocate.
+    // So we need to check for overflow, and return SIZE_MAX in this first case,
+    // and 0 in the second.
+    size_t result = ((alignment - 1) | (size - 1)) + 1;
+    // The following code is designed to fuse well with a subsequent
+    // sizeclass calculation.  We use the same fast path constant to
+    // move the case where result==0 to the slow path, and then check for which
+    // case we are in.
+    constexpr size_t SmallSizeClassUpperBound =
+      sizeclass_to_size(NUM_SMALL_SIZECLASSES - 1) - 1;
+    if (SNMALLOC_LIKELY((result - 1) < SmallSizeClassUpperBound))
+      return result;
+
+    // We are in the slow path, so we need to check for overflow.
+    if (SNMALLOC_UNLIKELY(result == 0))
+    {
+      // Check for overflow and return the maximum size.
+      if (SNMALLOC_UNLIKELY(result < size))
+        return SIZE_MAX;
+    }
+    return result;
+  }
 } // namespace snmalloc
diff --git a/src/test/func/malloc/malloc.cc b/src/test/func/malloc/malloc.cc
index a2b5bff..630d723 100644
--- a/src/test/func/malloc/malloc.cc
+++ b/src/test/func/malloc/malloc.cc
@@ -15,11 +15,7 @@ constexpr int SUCCESS = 0;
 void check_result(size_t size, size_t align, void* p, int err, bool null)
 {
   bool failed = false;
-  EXPECT(
-    (errno == err) || (err == SUCCESS),
-    "Expected error: {} but got {}",
-    err,
-    errno);
+  EXPECT((errno == err), "Expected error: {} but got {}", err, errno);
   if (null)
   {
     EXPECT(p == nullptr, "Expected null but got {}", p);
@@ -307,6 +303,9 @@ int main(int argc, char** argv)
   for (size_t align = sizeof(uintptr_t); align < MAX_SMALL_SIZECLASS_SIZE * 8;
        align <<= 1)
   {
+    // Check overflow with alignment taking it round to 0.
+    test_memalign(1 - align, align, ENOMEM, true);
+
     for (smallsizeclass_t sc = 0; sc < NUM_SMALL_SIZECLASSES - 6; sc++)
     {
       const size_t size = sizeclass_to_size(sc);