umalloc fixes

ulib/assert.h

@@ -1,13 +1,13 @@
 #ifndef ULIB_ASSERT_H_
 #define ULIB_ASSERT_H_
 
-#include <uprintf.h>
+#include <write/write.h>
 #include <util/util.h>
 
 #define ASSERT(X, fmt, ...) \
   do { \
     if (!(X)) { \
-      uprintf("ASSERT %s:%d in %s "fmt, __FILE__, __LINE__, __func__, ##__VA_ARGS__); \
+      writefmt("ASSERT {s}:{d} in {s} "fmt, __FILE__, __LINE__, __func__, ##__VA_ARGS__); \
       quit(); \
     } \
   } while(0)

ulib/string/string.c

@@ -1,6 +1,7 @@
 #include <stddef.h>
 #include <string/string.h>
 #include <umalloc/umalloc.h>
+#include <write/write.h>
 
 size_t string_len(const char *s) {
   size_t l = 0;
@@ -140,26 +141,29 @@ char *string_tokenizealloc(char *s, char *delim) {
     return NULL;
   }
 
+  while (s[curridx] && string_strchr(delim, s[curridx])) {
+    curridx++;
+  }
+  if (!s[curridx]) {
+    return NULL;
+  }
+
   char *w = (char *)umalloc(sizeof(char) * STRING_TOKENIZEALLOC_TOK_SIZE);
   string_memset(w, 0, sizeof(char) * STRING_TOKENIZEALLOC_TOK_SIZE);
-  int i = curridx, k = 0, j = 0;
 
-  while (s[i] != '\0') {
+  int k = 0;
-    j = 0;
+  while (s[curridx] && !string_strchr(delim, s[curridx])) {
-    while (delim[j] != '\0') {
+    if (k >= STRING_TOKENIZEALLOC_TOK_SIZE - 1) {
-      if (s[i] != delim[j]) {
+      break;
-        w[k] = s[i];
-      } else {
-        goto it;
-      }
-      j++;
     }
-    i++;
+    w[k++] = s[curridx++];
-    k++;
   }
-  it: {
+
-    w[i] = 0;
+  w[k] = '\0';
-    curridx = i + 1;
+
-    return w;
+  if (s[curridx]) {
+    curridx++;
   }
+
+  return w;
 }

ulib/ubsan/ubsan.c

@@ -3,7 +3,7 @@
 
 #include <stdint.h>
 #include <stddef.h>
-#include <uprintf.h>
+#include <write/write.h>
 
 struct source_location {
 	const char *file;
@@ -96,7 +96,7 @@ struct pointer_overflow_data {
 
 static void ubsan_print(struct source_location *loc, const char *message)
 {
-  uprintf("UBSAN error %s:%d:%d %s\n", loc->file, loc->line, loc->column, message);
+  writefmt("UBSAN error {s}:{d}:{d} {s}\n", loc->file, loc->line, loc->column, message);
 }
 
 void __ubsan_handle_add_overflow(struct overflow_data *data)

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;
+  int64_t balance;
-} UmArena;
+} __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;
+  uintptr_t addr = usable->mem + usable->cursor;
-  usable->cursor += size;
+  UmBlock *block = (UmBlock *)addr;
-  usable->balance += 1;
+  block->size = size;
-  return (void *)current;
+  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_;
+  if (ptr_ == NULL) {
-
-  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) {
     return;
   }
 
-  freeable->balance -= 1;
+  uintptr_t ptr = (uintptr_t)ptr_;
 
-  if (freeable->balance < 0) {
+  UmArena *freeable = NULL;
-    writefmt("umalloc: ufree() called more times than umalloc() ({d})\n", freeable->balance);
+  size_t size = 0;
-    quit();
-  }
 
-  if (freeable->balance == 0) {
+  UmArena *arena, *arenatmp;
-    LL_REMOVE(ARENAS, freeable);
+  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;
+    }
   }
 }
 

ulib/write/write.c

@@ -6,6 +6,9 @@
 #include <system/system.h>
 #include <sysdefs/devctl.h>
 #include <write/write.h>
+#include <umalloc/umalloc.h>
+#include <string/string.h>
+#include <util/util.h>
 
 char *utoa(uintmax_t v, int base, bool uppercase) {
   if (base < 2 || base > 36) {
@@ -76,7 +79,7 @@ char *itoa(intmax_t v, int base, bool uppercase) {
   return &buf[i];
 }
 
-#define FMTBUF_MAX 2048
+#define FMTBUF_MAX (1024 * 4)
 
 extern Dev_t termdev;
 
@@ -88,16 +91,15 @@ enum {
 const char *convstr1 = "0123456789abcdef";
 const char *convstr2 = "0123456789ABCDEF";
 
-void writefmt(char *fmt, ...) {
+size_t writefmt(char *fmt, ...) {
   va_list list;
   va_start(list, fmt);
-  writevfmt(fmt, list);
+  size_t count = writevfmt(fmt, list);
   va_end(list);
+  return count;
 }
 
-#define ITOA_BUF_SIZE 32
+size_t writevsfmt(char *buf, char *fmt, va_list list) {
-
-void writevsfmt(char *buf, char *fmt, va_list list) {
   size_t c = 0;
   int WRITE_STATE = WRITE_NORMALMODE;
 
@@ -118,44 +120,48 @@ void writevsfmt(char *buf, char *fmt, va_list list) {
           break;
         case 's': {
           char *string = va_arg(list, char *);
-          while (*string) { buf[c++] = *string; string++; }
+          while (string && *string) { buf[c++] = *string; string++; }
         } break;
         case 'd': {
           int int1 = va_arg(list, int);
           char *string = itoa(int1, 10, false);
-          while (*string) { buf[c++] = *string; string++; }
+          while (string && *string) { buf[c++] = *string; string++; }
         } break;
         case 'x': {
           unsigned int int1 = va_arg(list, unsigned int);
           char *string = utoa(int1, 16, false);
-          while (*string) { buf[c++] = *string; string++; }
+          while (string && *string) { buf[c++] = *string; string++; }
         } break;
         case 'X': {
           unsigned int int1 = va_arg(list, unsigned int);
           char *string = utoa(int1, 16, true);
-          while (*string) { buf[c++] = *string; string++; }
+          while (string && *string) { buf[c++] = *string; string++; }
         } break;
         case 'p': {
           uintptr_t ptr = (uintptr_t)va_arg(list, void *);
           char *string = utoa(ptr, 16, false);
-          while (*string) { buf[c++] = *string; string++; }
+          while (string && *string) { buf[c++] = *string; string++; }
         } break;
         case 'P': {
           uintptr_t ptr = (uintptr_t)va_arg(list, void *);
           char *string = utoa(ptr, 16, true);
-          while (*string) { buf[c++] = *string; string++; }
+          while (string && *string) { buf[c++] = *string; string++; }
         } break;
         case 'c': {
-          char c1 = (char)va_arg(list, char);
+          int c1 = va_arg(list, int);
-          buf[c++] = c1;
+          buf[c++] = (char)c1;
         } break;
       }
     }
   }
+
+  return c;
 }
 
-void writevfmt(char *fmt, va_list list) {
+size_t writevfmt(char *fmt, va_list list) {
   char buf[FMTBUF_MAX];
-  writevsfmt(buf, fmt, list); 
+  string_memset(buf, 0, FMTBUF_MAX);
-  devctl(&termdev, 0x00, buf, sizeof(buf), 0);
+  size_t count = writevsfmt(buf, fmt, list); 
+  devctl(&termdev, 0x00, buf, count, 0);
+  return count;
 }

ulib/write/write.h

@@ -3,8 +3,8 @@
 
 #include <stdarg.h>
 
-void writefmt(char *fmt, ...);
+size_t writefmt(char *fmt, ...);
-void writevfmt(char *fmt, va_list list);
+size_t writevfmt(char *fmt, va_list list);
-void writevsfmt(char *buf, char *fmt, va_list list);
+size_t writevsfmt(char *buf, char *fmt, va_list list);
 
 #endif // ULIB_WRITE_WRITE_H_