Fix kalloc bug.

Bitfields are officially stupid. Bizzarre behavior was found in how the
bitfields integers were overflowing and causing other members to change
value, causing really screwy behavior. In addition, with the discovery
of 48k being available to the heap, a 12-bit value was no longer
sufficient to define the size.

I rewrote parts of the kalloc code to allow a generic size for the
kalloc header because now it'll require 2 words per block allocated,
and who knows what size the header will be on different platforms
with more memory.

Unfortunately, the second word of the header consists only of the "used"
bool. Because I wish to keep alignmennt with 32-bit words, 31 bits are
"wasted." However, these bits are used as a canary value to detect
heap corruption, so they're not completely wasted.

Also, testing was broken since adding the huge amount of platform
dependent code for doing system calls. These dependent parts were
put under a macro guard so they don't interfere with the x86 testing.

7 files changed, 155 insertions, 83 deletions

diff --git a/include/kern/mem.h b/include/kern/mem.h
index 2af49cb..20b09bb 100644
--- a/include/kern/mem.h
+++ b/include/kern/mem.h
@@ -4,6 +4,8 @@
 #include "arch.h"
 #include <stddef.h>
 
+void kalloc_init();
+
 /* allocates memory on the head, which is stored in sram2 */
 void* kalloc(size_t n);
 
diff --git a/src/kern/init.c b/src/kern/init.c
index c776b4a..9869749 100644
--- a/src/kern/init.c
+++ b/src/kern/init.c
@@ -5,6 +5,13 @@
 #include "arch/stm32l4xxx/peripherals/system.h"
 #include "kern/log.h"
 
+static _no_init init_level_t initlevel;
+
+init_level_t get_system_init_level()
+{
+  return initlevel;
+}
+
 /* Forward-declare the main function. This is implemented in main.c. */
 int main();
 
@@ -29,13 +36,6 @@ extern uint32_t INIT_5_END;
 extern uint32_t INIT_6_END;
 extern uint32_t INIT_7_END;
 
-static _no_init init_level_t initlevel;
-
-init_level_t get_system_init_level()
-{
-  return initlevel;
-}
-
 init0()
 {
   /* Enable a higher clock speed. This is the first thing we do
diff --git a/src/kern/mem.c b/src/kern/mem.c
index 533b394..31756e5 100644
--- a/src/kern/mem.c
+++ b/src/kern/mem.c
@@ -28,79 +28,90 @@ void* memset(void* dest, int c, size_t n);
 
 typedef uint16_t kalloc_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.
 typedef struct KALLOC_NODE {
-  union {
-    uint32_t header;
-    struct {
-      /* Is this memory block currently in use (hasn't been kfree'd) */
-      uint8_t used : 1;
-      /* Number of words allocated. Does not include the header. */
-      uint16_t size : 12;
-      /* The location of the previous block (in WORDS from offest) */
-      kalloc_off_t prev : 12;
-      uint8_t canary : 7;
-    } PACKED;
-  };
+  uint16_t size_words;
+  kalloc_off_t prev;
+  /*
+   * LSB is whether this block is used.
+   *
+   * Rest of the bits must equal (0xdeadbeee >> 1)
+   */
+  uint32_t used_and_canary;
 
   uint8_t mem[]; /* The memory to use. */
-} kalloc_node_t;
-
-static_assert(offsetof(kalloc_node_t, mem) == 4, "Offset check failed.");
+} PACKED kalloc_node_t;
 
 kalloc_node_t* kalloc_start;
 
+#define CANARY 0xdeadbeee
+#define kalloc_node_in_use(node) ((node)->used_and_canary & 1)
+#define kalloc_node_get_canary(node) ((node)->used_and_canary & (~1))
+#define WORD_SIZE (sizeof(uint32_t))
+#define SIZEOF_KALLOC_NODE_WORDS (sizeof(kalloc_node_t) / WORD_SIZE)
 #define HEAP_START_ADDR ((ptrdiff_t)&HEAP_START)
 #define REAL_HEAP_START \
   (*((unsigned char*)((HEAP_START_ADDR & (~3)) + (HEAP_START_ADDR % 4 != 0))))
 #define MAX_HEAP_SIZE ((&HEAP_STOP - &REAL_HEAP_START))
+#define MAX_HEAP_SIZE_WORDS (MAX_HEAP_SIZE / WORD_SIZE)
 #define kalloc_node_out_of_range(node) ((void*)(node) >= (void*)&HEAP_STOP)
 
 #define kalloc_node_next(cur) \
-  ((kalloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size + 1) * 4)))
+  ((kalloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size_words + SIZEOF_KALLOC_NODE_WORDS) * WORD_SIZE)))
 
 #define kalloc_node_prev(cur) kalloc_node_at_off(cur->prev)
 
 #define kalloc_node_at_off(offset) \
-  ((kalloc_node_t*)(((uint8_t*)kalloc_start) + (offset)*4))
+  ((kalloc_node_t*)(((uint8_t*)kalloc_start) + (offset)*WORD_SIZE))
 
 #define kalloc_node_get_off(node) \
-  (((uint32_t)(((uint8_t*)(node)) - ((uint8_t*)(kalloc_start)))) / 4)
+  ((uint32_t)(((((uint8_t*)(node)) - ((uint8_t*)(kalloc_start)))) / WORD_SIZE))
 
-#define get_kalloc_node(mem) ((kalloc_node_t*)(((uint8_t*)mem) - 4))
+#define get_kalloc_node(mem) \
+  ((kalloc_node_t*)(((uint8_t*)mem) - SIZEOF_KALLOC_NODE_WORDS * WORD_SIZE))
 
 #define size_for(n) (((n) / 4) + ((n) % 4 != 0))
 
+void kalloc_init()
+{
+  kalloc_start = (kalloc_node_t*)&REAL_HEAP_START;
+  memset(kalloc_start, 0, sizeof(kalloc_node_t));
+  kalloc_start->size_words = MAX_HEAP_SIZE_WORDS - SIZEOF_KALLOC_NODE_WORDS;
+  kalloc_start->used_and_canary = CANARY;
+}
+
 void* kalloc(size_t size)
 {
   if (!kalloc_start) {
-    kalloc_start = (kalloc_node_t*)&REAL_HEAP_START;
-    memset(kalloc_start, 0, sizeof(kalloc_node_t));
-    kalloc_start->size = (MAX_HEAP_SIZE / 4) - 1;
-    kalloc_start->canary = CANARY;
+    kalloc_init();
   }
 
-  size_t realsz = size_for(size); /* Clip the size to the nearest word. */
+  size_t realsz_words = size_for(size); /* Clip the size to the nearest word. */
   kalloc_off_t offset = 0;
-  while (offset < (MAX_HEAP_SIZE / 4)) {
+  while (offset < (MAX_HEAP_SIZE_WORDS)) {
     kalloc_node_t* cur = kalloc_node_at_off(offset);
 
-    if (!cur->used && (cur->size >= realsz)) {
-      cur->used = true;
-      size_t orig_size = cur->size;
-      cur->size = realsz;
+    if (!kalloc_node_in_use(cur) && (cur->size_words >= realsz_words)) {
+      cur->used_and_canary |= 1;
+      size_t orig_size_words = cur->size_words;
+
+      if (orig_size_words < realsz_words + SIZEOF_KALLOC_NODE_WORDS * 2) {
+        /* If the original size is only slightly larger than the size to
+         * allocate, there's no point in saving the leftover space. */
+        realsz_words = orig_size_words;
+      }
+
+      cur->size_words = realsz_words;
 
-      if (orig_size > realsz) {
+      if (orig_size_words > realsz_words) {
         /* This kalloc node needs to split into two blocks. */
         kalloc_node_t* next = kalloc_node_next(cur);
-        next->used = 0;
-        next->size = orig_size - realsz - sizeof(kalloc_node_t) / 4;
+        next->used_and_canary = CANARY;
+        next->size_words =
+            orig_size_words - realsz_words - SIZEOF_KALLOC_NODE_WORDS;
         next->prev = offset;
-        next->canary = CANARY;
 
         kalloc_node_t* nextnext = kalloc_node_next(next);
         if (kalloc_node_get_off(nextnext) < (MAX_HEAP_SIZE / 4)) {
@@ -111,7 +122,7 @@ void* kalloc(size_t size)
       return (void*)cur->mem;
     }
 
-    offset += (sizeof(kalloc_node_t) / 4) + cur->size;
+    offset += SIZEOF_KALLOC_NODE_WORDS + cur->size_words;
   }
 
   return NULL;
@@ -125,11 +136,11 @@ static void coalesce(kalloc_node_t* cur)
   kalloc_node_t* next_used;
 
   /* Find the earliest contiguous free'd block. */
-  while (!cur->used && cur != kalloc_start) {
+  while (!kalloc_node_in_use(cur) && cur != kalloc_start) {
     cur = kalloc_node_prev(cur);
   }
 
-  if (cur == kalloc_start && !cur->used) {
+  if (cur == kalloc_start && !kalloc_node_in_use(cur)) {
     last_freed = cur;
   } else {
     last_freed = kalloc_node_next(cur);
@@ -137,7 +148,7 @@ static void coalesce(kalloc_node_t* cur)
 
   /* Find the next used block. */
   cur = orig;
-  while (!cur->used && !kalloc_node_out_of_range(cur)) {
+  while (!kalloc_node_in_use(cur) && !kalloc_node_out_of_range(cur)) {
     cur = kalloc_node_next(cur);
   }
 
@@ -147,7 +158,7 @@ static void coalesce(kalloc_node_t* cur)
     next_used->prev = kalloc_node_get_off(last_freed);
   }
 
-  last_freed->size = ((uint8_t*)next_used - (last_freed->mem)) / 4;
+  last_freed->size_words = ((uint8_t*)next_used - (last_freed->mem)) / WORD_SIZE;
 }
 
 void kfree(void* mem)
@@ -156,12 +167,12 @@ void kfree(void* mem)
   if (!mem) return;
 
   kalloc_node_t* header = get_kalloc_node(mem);
-  if (!header->used) {
+  if (!kalloc_node_in_use(header)) {
     panic("Heap double free or corruption!\n");
     return;
   }
 
-  header->used = 0;
+  header->used_and_canary &= ~1;
   coalesce(header);
 }
 
@@ -188,16 +199,22 @@ void* debug_kalloc_get_prev_ptr(void* ptr)
 void debug_print_blocks()
 {
   printf("------ Print Blocks -------\n");
+  printf("MAX_HEAP_SIZE      : %ld\n", MAX_HEAP_SIZE);
+  printf("MAX_HEAP_SIZE_WORDS: %ld\n", MAX_HEAP_SIZE_WORDS);
+  printf("HEAP_START         : %p\n", &HEAP_START);
+  printf("HEAP_STOP          : %p\n", &HEAP_STOP);
+  printf("---------------------------\n");
   kalloc_node_t* cur = kalloc_node_at_off(0);
 
   while (!kalloc_node_out_of_range(cur)) {
     printf(
-        "header (%04x) {used=%d, size=%5d, prev=%04x, canary=%02x}\n",
+        "header (%04x)@%p {used=%d, size=%5d, prev=%04x, canary=%04x}\n",
         kalloc_node_get_off(cur),
-        cur->used,
-        cur->size,
+        cur->mem,
+        kalloc_node_in_use(cur),
+        cur->size_words,
         cur->prev,
-        cur->canary);
+        kalloc_node_get_canary(cur));
     cur = kalloc_node_next(cur);
   }
 }
@@ -213,17 +230,17 @@ int debug_kalloc_assert_consistency(char* error, size_t len)
   size_t n_blocks_back = 1;
 
   while (1) {
-    if (cur->canary != CANARY) {
+    if (kalloc_node_get_canary(cur) != CANARY) {
       snprintf(
           error,
           len,
           "Node has corrupted canary. %02x vs expected %02x\n",
-          cur->canary,
+          kalloc_node_get_canary(cur),
           CANARY);
       return 1;
     }
 
-    total_size += cur->size + 1;
+    total_size += cur->size_words + SIZEOF_KALLOC_NODE_WORDS;
 
     kalloc_node_t* next = kalloc_node_next(cur);
     if ((void*)next == (void*)&HEAP_STOP) {
@@ -247,7 +264,7 @@ int debug_kalloc_assert_consistency(char* error, size_t len)
         error,
         len,
         "Total recorded size is inconsistent. %lu vs %lu\n",
-        total_size * 4,
+        total_size * WORD_SIZE,
         MAX_HEAP_SIZE);
     return 1;
   }
diff --git a/src/kern/priv.c b/src/kern/priv.c
index 2c45eca..ba25c26 100644
--- a/src/kern/priv.c
+++ b/src/kern/priv.c
@@ -4,6 +4,7 @@
 #include "kern/log.h"
 #include "kern/mem.h"
 
+#ifdef ARCH_STM32L4
 void set_control_register(uint32_t reg)
 {
   asm volatile("msr control, %0" : "=r"(reg) :);
@@ -38,3 +39,4 @@ void jump_to_user_mode()
       "dsb\n\t"
       "b usermode_start\n\t" : : "r"(new_stack));
 }
+#endif
diff --git a/src/kern/svc.c b/src/kern/svc.c
index 5527364..ceca5fa 100644
--- a/src/kern/svc.c
+++ b/src/kern/svc.c
@@ -24,6 +24,7 @@ void handle_svc_call(
   }
 }
 
+#ifdef ARCH_STM32L4
 /* The actual handling for the svc call. Overrides the weak
  * symbol on_svc in irq.h */
 asm("	.align	2\n"
@@ -49,3 +50,4 @@ asm("	.align	2\n"
     "	mrsne r12,psp\n"
     "	stm   r12,{r0-r3}\n"
     "	bx    lr\n");
+#endif
diff --git a/src/user/syscall.c b/src/user/syscall.c
index 3aa0707..601c7a3 100644
--- a/src/user/syscall.c
+++ b/src/user/syscall.c
@@ -1,3 +1,4 @@
+#include "arch.h"
 #include "user/syscall.h"
 
 #include <stdint.h>
@@ -6,7 +7,9 @@ void __attribute__ ((noinline)) do_syscall(
     volatile uint32_t id,
     volatile uint32_t arg)
 {
+#ifdef ARCH_STM32L4
   asm volatile("svc #0x04");
+#endif
 }
 
 #define SYSCALL(id, fn, kernfn, argt) \
diff --git a/tests/test_memory.c b/tests/test_memory.c
index a977a70..c85e228 100644
--- a/tests/test_memory.c
+++ b/tests/test_memory.c
@@ -14,26 +14,25 @@ struct TEST_STRUCT {
 };
 
 struct TEST_STRUCT2 {
-  uint32_t array[10];
+  uint32_t array[16];
 };
 
 /* Copy of the node structure. */
+// 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.
 typedef struct KALLOC_NODE {
-  union {
-    uint32_t header;
-    struct {
-      /* Is this memory block currently in use (hasn't been kfree'd) */
-      bool used : 1;
-      /* Number of words allocated. Does not include the header. */
-      uint16_t size : 12;
-      /* The location of the previous block (in WORDS from offest) */
-      uint16_t prev : 12;
-      uint8_t canary : 7;
-    } PACKED;
-  };
+  uint16_t size_words;
+  uint16_t prev;
+  /*
+   * LSB is whether this block is used.
+   *
+   * Rest of the bits must equal (0xdeadbeee >> 1)
+   */
+  uint32_t used_and_canary;
 
   uint8_t mem[]; /* The memory to use. */
-} kalloc_node_t;
+} PACKED kalloc_node_t;
 
 extern kalloc_node_t* kalloc_start;
 
@@ -65,7 +64,8 @@ static struct TEST_STRUCT2* new_test_struct2()
   return ret;
 }
 
-#define ASSERT_CHAIN(t1, t2) ASSERT_EQ(V(t1) + sizeof(*t1) + 4, V(t2))
+#define ASSERT_CHAIN(t1, t2) \
+  ASSERT_EQ(V(t1) + sizeof(*t1) + sizeof(kalloc_node_t), V(t2))
 
 TEST(memory, kalloc)
 {
@@ -127,15 +127,17 @@ TEST(memory, uneven_kalloc)
   struct UNEVEN_STRUCT* test1 = new_uneven_struct();
   struct UNEVEN_STRUCT* test2 = new_uneven_struct();
 
-  ASSERT_EQ(V(test1) + 12, test2);
+  ASSERT_EQ(V(test1) + 8 + sizeof(kalloc_node_t), test2);
 
   wipeout_kalloc();
 
   return 0;
 }
 
-#define ASSERT_KALLOC_EMPTY() \
-  ASSERT_EQ((void*)(kalloc_start + kalloc_start->size + 1), (void*)&HEAP_STOP)
+#define ASSERT_KALLOC_EMPTY()                                                                                \
+  ASSERT_EQ(                                                                                                 \
+      (void*)((uint8_t*)kalloc_start + kalloc_start->size_words * sizeof(uint32_t) + sizeof(kalloc_node_t)), \
+      (void*)&HEAP_STOP)
 
 TEST(memory, kalloc_free)
 {
@@ -143,6 +145,10 @@ TEST(memory, kalloc_free)
     wipeout_kalloc();
   }
 
+  kalloc_init();
+
+  ASSERT_KALLOC_EMPTY();
+
   struct TEST_STRUCT* test1 = new_test_struct();
   struct TEST_STRUCT2* test2 = new_test_struct2();
   struct TEST_STRUCT* test3 = new_test_struct();
@@ -248,6 +254,46 @@ TEST(memory, relink_backref_after_free)
   return 0;
 }
 
+TEST(memory, alloc1)
+{
+  char buf[1024];
+
+  kalloc(5);
+  kalloc(1);
+
+  if (debug_kalloc_assert_consistency(buf, 1024)) {
+    fprintf(
+        stderr,
+        "Consistency check failed. (%s:%d)\n",
+        __FILE__,
+        __LINE__);
+    fprintf(stderr, "%s", buf);
+    ASSERT_TRUE(false);
+  }
+
+  return 0;
+}
+
+TEST(memory, alloc0)
+{
+  char buf[1024];
+
+  kalloc(5);
+  kalloc(0);
+
+  if (debug_kalloc_assert_consistency(buf, 1024)) {
+    fprintf(
+        stderr,
+        "Consistency check failed. (%s:%d)\n",
+        __FILE__,
+        __LINE__);
+    fprintf(stderr, "%s", buf);
+    ASSERT_TRUE(false);
+  }
+
+  return 0;
+}
+
 TEST(memory, consistency_stress)
 {
 #define NRUNS 500
@@ -258,6 +304,7 @@ TEST(memory, consistency_stress)
   int i;
   void* allocd[NRUNS] = {0};
   char buf[1024];
+  void* tofree;
 
   for (i = 0; i < NRUNS; ++i) {
     size_t nalloc = rand() % 20;
@@ -266,10 +313,11 @@ TEST(memory, consistency_stress)
     if (debug_kalloc_assert_consistency(buf, 1024)) {
       fprintf(
           stderr,
-          "Consistency check failed. (At index=%d, %s:%d)\n",
+          "Consistency check failed. (At index=%d, %s:%d, nalloc=%ld)\n",
           i,
           __FILE__,
-          __LINE__);
+          __LINE__,
+          nalloc);
       fprintf(stderr, "%s", buf);
       ASSERT_TRUE(false);
     }
@@ -290,7 +338,8 @@ TEST(memory, consistency_stress)
       ASSERT_TRUE(false);
     }
 
-    kfree(allocd[idx]);
+    tofree = allocd[idx];
+    kfree(tofree);
     allocd[idx] = NULL;
 
     if (debug_kalloc_assert_consistency(buf, 1024)) {
@@ -305,7 +354,8 @@ TEST(memory, consistency_stress)
     }
 
     idx = rand() % NRUNS;
-    kfree(allocd[idx]);
+    tofree = allocd[idx];
+    kfree(tofree);
     allocd[idx] = NULL;
 
     if (debug_kalloc_assert_consistency(buf, 1024)) {
@@ -363,11 +413,7 @@ TEST(memory, kalloc_free_alloc)
 
   kfree(test2);
   struct TEST_STRUCT* test7 = new_test_struct();
-  struct TEST_STRUCT* test8 = new_test_struct();
 
-  // Test 2 was large enough to accomodate 3 smaller structs.
   ASSERT_EQ(V(test7), V(test2));
-  ASSERT_EQ(V(test8), V(test2) + sizeof(*test7) + 4);
-
   return 0;
 }