Hardening the BlockAllocator

src/basic_alloc.cpp

@@ -284,7 +284,7 @@ API void Init(BlockArena *arena) {
 }
 
 API void *PushSize(BlockArena *arena, size_t size) {
-    if (size > (size_t)(arena->end - arena->start)) {
+    if (!arena->blocks || size > (size_t)(arena->end - arena->start)) {
         AddBlock(arena, size);
     }
     U8 *result = arena->start;
@@ -323,7 +323,7 @@ API void *BlockArenaAllocatorProc(void *object, int kind, void *p, size_t size)
     BlockArena *arena = (BlockArena *)object;
     if (kind == AllocatorKind_Allocate) {
         return PushSize(arena, size);
-    } else if (AllocatorKind_Deallocate) {
+    } else if (kind == AllocatorKind_Deallocate) {
     } else {
         Assert(!"invalid codepath");
     }
@@ -439,4 +439,148 @@ void RunArenaTest() {
         Release(&arena);
         Assert(MemoryTrackingRecord.len == 0);
     }
+
+    // Unwind-to-null must leave the arena reusable. This catches crashes/UB in
+    // temp-memory style "clear and then allocate again" usage.
+    {
+        BlockArena arena = {};
+        arena.allocator = memory_tracking_allocator;
+        Unwind(&arena, 0);
+        U8 *a = (U8 *)PushSize(&arena, 64);
+        for (int i = 0; i < 64; i += 1) a[i] = (U8)i;
+        Unwind(&arena, 0);
+        Assert(MemoryTrackingRecord.len == 0);
+        U8 *b = (U8 *)PushSize(&arena, 128);
+        for (int i = 0; i < 128; i += 1) b[i] = (U8)(255 - i);
+        Release(&arena);
+        Assert(MemoryTrackingRecord.len == 0);
+    }
+
+    // Nested temp checkpoints, including a child allocation that forces a new
+    // block. Older allocations must survive, newer blocks must be released, and
+    // the checkpoint address must be reused by the next allocation.
+    {
+        BlockArena arena = {};
+        arena.allocator = memory_tracking_allocator;
+        U8 *permanent = (U8 *)PushSize(&arena, 256);
+        for (int i = 0; i < 256; i += 1) permanent[i] = (U8)i;
+
+        U8 *checkpoint_a = arena.start;
+        U8 *temp_a = (U8 *)PushSize(&arena, KiB(64));
+        for (size_t i = 0; i < KiB(64); i += 1) temp_a[i] = 0xA1;
+
+        U8 *checkpoint_b = arena.start;
+        U8 *temp_b = (U8 *)PushSize(&arena, MiB(2));
+        for (size_t i = 0; i < MiB(2); i += 4096) temp_b[i] = 0xB2;
+        Assert(MemoryTrackingRecord.len == 2);
+
+        Unwind(&arena, checkpoint_b);
+        Assert(arena.start == checkpoint_b);
+        Assert(MemoryTrackingRecord.len == 1);
+        for (int i = 0; i < 256; i += 1) Assert(permanent[i] == (U8)i);
+        for (size_t i = 0; i < KiB(64); i += 1) Assert(temp_a[i] == 0xA1);
+        Assert(PushSize(&arena, 32) == checkpoint_b);
+
+        Unwind(&arena, checkpoint_a);
+        Assert(arena.start == checkpoint_a);
+        for (int i = 0; i < 256; i += 1) Assert(permanent[i] == (U8)i);
+        Assert(PushSize(&arena, 32) == checkpoint_a);
+
+        Release(&arena);
+        Assert(MemoryTrackingRecord.len == 0);
+    }
+
+    // Repeatedly grow through multiple blocks, unwind back through all of them,
+    // allocate again, and do it many times. This is the pattern a temp arena can
+    // hit if one outer checkpoint is reused around variable-size work.
+    {
+        BlockArena arena = {};
+        arena.allocator = memory_tracking_allocator;
+        U8 *guard = (U8 *)PushSize(&arena, 1024);
+        for (int i = 0; i < 1024; i += 1) guard[i] = (U8)i;
+        U8 *base_checkpoint = arena.start;
+
+        for (int cycle = 0; cycle < 32; cycle += 1) {
+            U8 *first_block_temp = (U8 *)PushSize(&arena, KiB(900));
+            first_block_temp[0] = (U8)cycle;
+            first_block_temp[KiB(900) - 1] = (U8)(cycle + 1);
+            U8 *mid_checkpoint = arena.start;
+
+            U8 *large_temp = (U8 *)PushSize(&arena, MiB(2) + (size_t)cycle * 8);
+            large_temp[0] = (U8)(cycle + 2);
+            large_temp[MiB(2) - 1] = (U8)(cycle + 3);
+            U8 *last_temp = (U8 *)PushSize(&arena, KiB(700));
+            last_temp[0] = (U8)(cycle + 4);
+            last_temp[KiB(700) - 1] = (U8)(cycle + 5);
+            Assert(MemoryTrackingRecord.len == 3);
+
+            Unwind(&arena, mid_checkpoint);
+            Assert(arena.start == mid_checkpoint);
+            Assert(MemoryTrackingRecord.len == 1);
+            Assert(first_block_temp[0] == (U8)cycle);
+            Assert(first_block_temp[KiB(900) - 1] == (U8)(cycle + 1));
+            for (int i = 0; i < 1024; i += 1) Assert(guard[i] == (U8)i);
+            Assert(PushSize(&arena, 16) == mid_checkpoint);
+
+            Unwind(&arena, base_checkpoint);
+            Assert(arena.start == base_checkpoint);
+            Assert(MemoryTrackingRecord.len == 1);
+            for (int i = 0; i < 1024; i += 1) Assert(guard[i] == (U8)i);
+            Assert(PushSize(&arena, 16) == base_checkpoint);
+            Unwind(&arena, base_checkpoint);
+        }
+
+        Release(&arena);
+        Assert(MemoryTrackingRecord.len == 0);
+    }
+
+    // Checkpoint exactly at the end of a block is legal; unwinding to it should
+    // drop later blocks and leave the full block as the current block.
+    {
+        BlockArena arena = {};
+        arena.allocator = memory_tracking_allocator;
+        PushSize(&arena, 1);
+        size_t remaining = (size_t)(arena.end - arena.start);
+        PushSize(&arena, remaining);
+        U8 *checkpoint = arena.start;
+        Assert(checkpoint == arena.end);
+        PushSize(&arena, 1);
+        Assert(MemoryTrackingRecord.len == 2);
+        Unwind(&arena, checkpoint);
+        Assert(arena.blocks);
+        Assert(arena.blocks->next == NULL);
+        Assert(arena.start == checkpoint);
+        Assert(arena.end == checkpoint);
+        PushSize(&arena, 1);
+        Assert(MemoryTrackingRecord.len == 2);
+        Release(&arena);
+        Assert(MemoryTrackingRecord.len == 0);
+    }
+
+    // Deterministic stack-style stress test: allocate temps at many nested
+    // checkpoints, unwind in reverse order, and verify surviving bytes.
+    {
+        BlockArena arena = {};
+        arena.allocator = memory_tracking_allocator;
+        U8 *checkpoints[128];
+        U8 *ptrs[128];
+        size_t sizes[128];
+        for (int i = 0; i < 128; i += 1) {
+            checkpoints[i] = arena.start;
+            sizes[i] = (size_t)(((i * 7919) % (64 * 1024)) + 1);
+            ptrs[i] = (U8 *)PushSize(&arena, sizes[i]);
+            for (size_t j = 0; j < sizes[i]; j += 4096) ptrs[i][j] = (U8)i;
+            ptrs[i][sizes[i] - 1] = (U8)i;
+        }
+        for (int i = 127; i >= 0; i -= 1) {
+            Assert(ptrs[i][0] == (U8)i);
+            Assert(ptrs[i][sizes[i] - 1] == (U8)i);
+            Unwind(&arena, checkpoints[i]);
+            Assert(arena.start == checkpoints[i]);
+        }
+        Unwind(&arena, 0);
+        Assert(MemoryTrackingRecord.len == 0);
+        Release(&arena);
+        Assert(MemoryTrackingRecord.len == 0);
+    }
 }

src/stb_sprintf.h

@@ -164,16 +164,34 @@ PERFORMANCE vs MSVC 2008 32-/64-bit (GCC is even slower than MSVC):
 #define STBSP__ASAN
 #endif
 
+#if defined(__clang__)
+ #if defined(__has_attribute)
+  #if __has_attribute(__no_sanitize__)
+   // stb_sprintf intentionally uses unaligned word loads/stores for speed.
+   // UBSan reports these on x86/x64 even though the hardware permits them.
+   #define STBSP__UBSAN __attribute__((__no_sanitize__("undefined")))
+  #endif
+ #endif
+#elif defined(__GNUC__) && (__GNUC__ >= 5)
+ #define STBSP__UBSAN __attribute__((__no_sanitize_undefined__))
+#endif
+
+#ifndef STBSP__UBSAN
+#define STBSP__UBSAN
+#endif
+
+#define STBSP__SANITIZE_OFF STBSP__ASAN STBSP__UBSAN
+
 #ifdef STB_SPRINTF_STATIC
 #define STBSP__PUBLICDEC static
-#define STBSP__PUBLICDEF static STBSP__ASAN
+#define STBSP__PUBLICDEF static STBSP__SANITIZE_OFF
 #else
 #ifdef __cplusplus
 #define STBSP__PUBLICDEC extern "C"
-#define STBSP__PUBLICDEF extern "C" STBSP__ASAN
+#define STBSP__PUBLICDEF extern "C" STBSP__SANITIZE_OFF
 #else
 #define STBSP__PUBLICDEC extern
-#define STBSP__PUBLICDEF STBSP__ASAN
+#define STBSP__PUBLICDEF STBSP__SANITIZE_OFF
 #endif
 #endif