umalloc fixes

ulib/umalloc/umalloc.c

@@ -8,15 +8,24 @@
 #include <string/string.h>
 #include <write/write.h>
 #include <util/util.h>
+#include <assert.h>
 
 #define ARENA_SIZE (1<<20)
+#define BLOCK_MAGIC 0xDEADBEEF
 
 typedef struct UmArena {
   struct UmArena *next;
   uintptr_t mem;
+  uintptr_t base;
   size_t cursor;
-  int balance;
-} UmArena;
+  int64_t balance;
+} __attribute__((aligned(16))) UmArena;
+
+typedef struct {
+  uint32_t magic;
+  size_t size;
+  uint8_t data[];
+} __attribute__((aligned(16))) UmBlock;
 
 static UmArena *ARENAS = NULL;
 
@@ -28,11 +37,12 @@ UmArena *um_newarena(void) {
   if (mem == NULL || err != E_OK) {
     return NULL;
   }
-  string_memset(mem, 0, arenasize);
 
   UmArena *arena = (UmArena *)mem;
-  arena->mem = (uintptr_t)mem + sizeof(UmArena);
+  arena->base = (uintptr_t)mem;
+  arena->mem = ((uintptr_t)mem + sizeof(UmArena) + 15) & ~0xF;
   arena->cursor = 0;
+  arena->balance = 0;
   LL_APPEND(ARENAS, arena);
 
   return arena;
@@ -41,9 +51,11 @@ UmArena *um_newarena(void) {
 void *umalloc(size_t size) {
   UmArena *usable = NULL;
 
+  size = (size + 15) & ~(size_t)0xF;
+
   UmArena *arena, *arenatmp;
   LL_FOREACH_SAFE(ARENAS, arena, arenatmp) {
-    if (arena->cursor + size <= ARENA_SIZE - sizeof(UmArena)) {
+    if (arena->cursor + sizeof(UmBlock) + size <= ARENA_SIZE - ((uintptr_t)arena->mem - arena->base)) {
       usable = arena;
       break;
     }
@@ -57,40 +69,53 @@ void *umalloc(size_t size) {
     return NULL;
   }
 
-  uintptr_t current = usable->mem + usable->cursor;
-  usable->cursor += size;
-  usable->balance += 1;
-  return (void *)current;
+  uintptr_t addr = usable->mem + usable->cursor;
+  UmBlock *block = (UmBlock *)addr;
+  block->size = size;
+  block->magic = BLOCK_MAGIC;
+
+  usable->cursor += sizeof(UmBlock) + size;
+  usable->balance += size;
+
+  string_memset(block->data, 0, size);
+  return block->data;
 }
 
 void ufree(void *ptr_) {
-  uintptr_t ptr = (uintptr_t)ptr_;
-
-  UmArena *freeable = NULL;
-
-  UmArena *arena, *arenatmp;
-  LL_FOREACH_SAFE(ARENAS, arena, arenatmp) {
-    if (ptr >= arena->mem && ptr < arena->mem + ARENA_SIZE - sizeof(UmArena)) {
-      freeable = arena;
-      break;
-    }
-  }
-
-  if (freeable == NULL) {
+  if (ptr_ == NULL) {
     return;
   }
 
-  freeable->balance -= 1;
+  uintptr_t ptr = (uintptr_t)ptr_;
 
-  if (freeable->balance < 0) {
-    writefmt("umalloc: ufree() called more times than umalloc() ({d})\n", freeable->balance);
-    quit();
-  }
+  UmArena *freeable = NULL;
+  size_t size = 0;
 
-  if (freeable->balance == 0) {
-    LL_REMOVE(ARENAS, freeable);
+  UmArena *arena, *arenatmp;
+  LL_FOREACH_SAFE(ARENAS, arena, arenatmp) {
+    if (ptr >= arena->mem && ptr < arena->mem + arena->cursor) {
+      UmBlock *block = (UmBlock *)(ptr - sizeof(UmBlock));
 
-    mman_unmap((void *)(freeable->mem - sizeof(UmArena)));
+      if (((uintptr_t)block->data != ptr)
+       || (block->magic != BLOCK_MAGIC)
+       || (block->size == 0)
+       || (block->size > ARENA_SIZE)) {
+        return;
+      }
+
+      string_memset(block->data, 0xDD, block->size);
+      block->magic = 0;
+
+      arena->balance -= block->size;
+      ASSERT(arena->balance >= 0, "umalloc: imbalance after free\n");
+
+      if (arena->balance == 0) {
+        LL_REMOVE(ARENAS, arena);
+        mman_unmap((void *)arena->base);
+      }
+
+      return;
+    }
   }
 }