Fixed diasterous bug with hfree.

Before this fix, hfree would neglect to set the prev pointer in the next
used block and such was leaving the prev pointers invalid after
coalescing frees.

1 files changed, 94 insertions, 20 deletions

diff --git a/02-usart/src/kern/mem.c b/02-usart/src/kern/mem.c
index 79bcabf..5234fff 100644
--- a/02-usart/src/kern/mem.c
+++ b/02-usart/src/kern/mem.c
@@ -26,6 +26,8 @@ void* memset(void* dest, int c, size_t n);
 
 typedef uint16_t halloc_off_t;
 
+#define CANARY 0x5a
+
 // The sizes will count the number of WORDS allocated.
 // Since there's a max size of 16k, only 12 bits will be
 // needed for this.
@@ -34,19 +36,25 @@ typedef struct HALLOC_NODE {
     uint32_t header;
     struct {
       /* Is this memory block currently in use (hasn't been hfree'd) */
-      bool     used:1;
+      uint8_t     used:1;
       /* Number of words allocated. Does not include the header. */
-      uint16_t size:15;
+      uint16_t size:12;
       /* The location of the previous block (in WORDS from offest) */
-      halloc_off_t prev;
+      halloc_off_t prev:12;
+      uint8_t canary:7;
     } PACKED;
   };
 
   uint8_t mem[]; /* The memory to use. */
 } halloc_node_t;
 
+static_assert(offsetof(halloc_node_t, mem) == 4, "Offset check failed.");
+
 halloc_node_t* halloc_start;
 
+#define halloc_node_out_of_range(node) \
+   ((uint8_t*) (node) == ((uint8_t*)&DATA_SEGMENT_STOP) + MAX_HEAP_SIZE)
+
 #define halloc_node_next(cur) \
   ((halloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size + 1) * 4)))
 
@@ -58,6 +66,9 @@ halloc_node_t* halloc_start;
 #define halloc_node_get_off(node) \
   (((uint32_t)(((uint8_t*)(node)) - ((uint8_t*)(halloc_start)))) / 4)
 
+#define get_halloc_node(mem) \
+  ((halloc_node_t*)(((uint8_t*)mem) - 4))
+
 #define size_for(n) \
   (((n) / 4) + ((n) % 4 != 0))
 
@@ -67,6 +78,7 @@ void* halloc(size_t size)
     halloc_start = (halloc_node_t*) DATA_SEGMENT_STOP_ADDR;
     memset(halloc_start, 0, sizeof(halloc_node_t));
     halloc_start->size = (MAX_HEAP_SIZE / 4) - 1;
+    halloc_start->canary = CANARY;
   }
 
   size_t realsz = size_for(size); /* Clip the size to the nearest word. */
@@ -85,6 +97,7 @@ void* halloc(size_t size)
         next->used = 0;
         next->size = orig_size - realsz - sizeof(halloc_node_t) / 4;
         next->prev = offset;
+        next->canary = CANARY;
 
         halloc_node_t* nextnext = halloc_node_next(next);
         if (halloc_node_get_off(nextnext) < (MAX_HEAP_SIZE / 4)) {
@@ -101,36 +114,64 @@ void* halloc(size_t size)
   return NULL;
 }
 
+
 /* Joins this node with the previous and next nodes if they're free. */
 static void coalesce(halloc_node_t* cur)
 {
-  if (cur->used) return;
+  halloc_node_t* orig = cur;
+  halloc_node_t* last_freed;
+  halloc_node_t* next_used;
 
-  halloc_node_t* next = halloc_node_next(cur);
-  if (!next->used) {
-    cur->size += next->size + sizeof(halloc_node_t) / 4;
+  /* Find the earliest contiguous free'd block. */
+  while (!cur->used && cur != halloc_start) {
+    cur = halloc_node_prev(cur);
   }
 
-  if (cur != halloc_start) {
-    coalesce(halloc_node_prev(cur));
+  if (cur == halloc_start && !cur->used) {
+    last_freed = cur;
+  } else {
+    last_freed = halloc_node_next(cur);
   }
-}
 
-void hfree(void* mem)
-{
-  /* TODO this should be a critical section. */
-  if (!mem) {
-    /* Consistent with C, freeing a NULL pointer does nothing. */
-    return;
+  /* Find the next used block. */
+  cur = orig;
+  while (!cur->used && !halloc_node_out_of_range(cur)) {
+    cur = halloc_node_next(cur);
   }
 
-  halloc_node_t* header = ((halloc_node_t*) mem) - 1;
+  next_used = cur;
+
+  if (!halloc_node_out_of_range(next_used)) {
+    next_used->prev = halloc_node_get_off(last_freed);
+  }
+
+  last_freed->size = ((uint8_t*) next_used - (last_freed->mem)) / 4;
+}
 
-  if (!header->used) {
 #ifdef FOR_TESTING
-    assert(header->used);
+#include <assert.h>
+#include <stdio.h>
+void panic(const char* x)
+{
+  fprintf(stderr, "%s\n", x);
+  assert(0);
+}
+#else
+void panic(const char* x)
+{
+  for(;;);
+}
 #endif
-    return; // TODO make this fail on ARM.
+
+void hfree(void* mem)
+{
+  /* Like normal free(), do nothing on free'ing NULL */
+  if (!mem) return;
+
+  halloc_node_t* header = get_halloc_node(mem);
+  if (!header->used) {
+    panic("Heap double free or corruption!\n");
+    return;
   }
 
   header->used = 0;
@@ -157,6 +198,18 @@ void* debug_halloc_get_prev_ptr(void* ptr)
   return prev->mem;
 }
 
+void debug_print_blocks()
+{
+  printf("------ Print Blocks -------\n");
+  halloc_node_t* cur = halloc_node_at_off(0);
+
+  while (!halloc_node_out_of_range(cur)) {
+    printf("header (%04x) {used=%d, size=%5d, prev=%04x, canary=%02x}\n",
+        halloc_node_get_off(cur), cur->used, cur->size, cur->prev, cur->canary);
+    cur = halloc_node_next(cur);
+  }
+}
+
 /* Tests that we can walk up and down the allocated blocks and that they
  * are properly aligned. */
 int debug_halloc_assert_consistency(char* error, size_t len)
@@ -164,8 +217,16 @@ int debug_halloc_assert_consistency(char* error, size_t len)
   halloc_node_t* cur = halloc_node_at_off(0);
   size_t total_size = 0;
   size_t loop_check = 0;
+  size_t n_blocks = 1;
+  size_t n_blocks_back = 1;
 
   while(1) {
+    if (cur->canary != CANARY) {
+      snprintf(error, len, "Node has corrupted canary. %02x vs expected %02x\n",
+          cur->canary, CANARY);
+      return 1;
+    }
+
     total_size += cur->size + 1;
 
     halloc_node_t* next = halloc_node_next(cur);
@@ -178,7 +239,9 @@ int debug_halloc_assert_consistency(char* error, size_t len)
           (void*)(DATA_SEGMENT_STOP_ADDR + MAX_HEAP_SIZE));
       return 1;
     }
+
     cur = next;
+    ++ n_blocks;
   }
 
   if (total_size * 4 != MAX_HEAP_SIZE) {
@@ -188,10 +251,21 @@ int debug_halloc_assert_consistency(char* error, size_t len)
     return 1;
   }
 
+  if (cur == halloc_start) {
+    return 0;
+  }
+
   while (loop_check < 10000) {
     halloc_node_t* prev = halloc_node_prev(cur);
+    ++ n_blocks_back;
 
     if (prev == halloc_start)  {
+      if (n_blocks != n_blocks_back) {
+        snprintf(
+            error, len, "Different number of blocks found on the way back. Found %lu on the way back vs %lu up.\n",
+            n_blocks_back, n_blocks);
+        return 1;
+      }
       return 0;
     }