diff --git a/src/mem.c b/src/mem.c index 5022d26c48..57fda23072 100644 --- a/src/mem.c +++ b/src/mem.c @@ -5,6 +5,13 @@ #include +// Enable this to catch use-after-free issues when debugging on Windows. +// #define DEBUG_USE_AFTER_FREE_ON_WINDOWS + +#if !defined(DEBUG_USE_AFTER_FREE_ON_WINDOWS) + +// The standard implementation, using malloc() and free(). + void * avifAlloc(size_t size) { void * out = malloc(size); @@ -18,3 +25,72 @@ void avifFree(void * p) { free(p); } + +#else + +// This implementation rounds up all memory allocations to the nearest 4k (page size), allocates +// using VirtualAlloc(), and then records the allocation in a linked list. When avifFree is called, +// instead of freeing the memory, it simply revokes all read and write access to that region +// forever. If any use-after-free issues are discovered (via fuzzing or the like), enabling this +// should immediately catch it in a debugger right when it happens. + +#include +#include + +typedef struct Allocation +{ + void * ptr; + size_t originalSize; + size_t size; + struct Allocation * next; +} Allocation; + +static Allocation * allocations = NULL; + +void * avifAlloc(size_t size) +{ + size_t originalSize = size; + size = (size + 4095) & ~(4095); + + void * out = VirtualAlloc(NULL, size, MEM_COMMIT, PAGE_READWRITE); + if (out == NULL) { + abort(); + } + + Allocation * a = (Allocation *)malloc(sizeof(Allocation)); + a->ptr = out; + a->originalSize = originalSize; + a->size = size; + a->next = allocations; + allocations = a; + + // printf("Alloc %p %zu (%zu)\n", a->ptr, a->originalSize, a->size); + return out; +} + +void avifFree(void * p) +{ + if (!p) { + return; + } + + Allocation * a = allocations; + for (; a != NULL; a = a->next) { + if (a->ptr == p) { + break; + } + } + + if (!a) { + abort(); + } + + DWORD old; + if (!VirtualProtect(a->ptr, a->size, PAGE_NOACCESS, &old)) { + abort(); + } + + // printf("Free %p %zu (%zu)\n", a->ptr, a->originalSize, a->size); +} + +#endif