From cda805574fa1c5c6d99f8a780e6cced2b5d48406 Mon Sep 17 00:00:00 2001 From: KK Date: Fri, 3 Jul 2026 20:50:13 +0200 Subject: [PATCH] Hardening the BlockAllocator --- src/basic_alloc.cpp | 148 +++++++++++++++++++++++++++++++++++++++++++- src/stb_sprintf.h | 24 ++++++- 2 files changed, 167 insertions(+), 5 deletions(-) diff --git a/src/basic_alloc.cpp b/src/basic_alloc.cpp index 9262302..0dfd6fd 100644 --- a/src/basic_alloc.cpp +++ b/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); + } } diff --git a/src/stb_sprintf.h b/src/stb_sprintf.h index cfecb9f..9c85df8 100644 --- a/src/stb_sprintf.h +++ b/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