Randomizer for kmalloc

7 files changed, 56 insertions, 4 deletions

diff --git a/kern/Makefile b/kern/Makefile
index f987b2a..139efc4 100644
--- a/kern/Makefile
+++ b/kern/Makefile
@@ -16,6 +16,10 @@ include ../defs.mk
 
 CFLAGS += $(DEPCFLAGS)
 
+ifneq ($(RANDOMIZE),)
+CFLAGS += -DRANDOMIZE_KMALLOC
+endif
+
 INCLUDE = -I$(KERN) -I$(PIOS)
 CFLAGS += $(INCLUDE)
 ASFLAGS += $(INCLUDE)
diff --git a/kern/boot.c b/kern/boot.c
index 605f4a1..08b6c52 100644
--- a/kern/boot.c
+++ b/kern/boot.c
@@ -4,6 +4,7 @@
 #include "kern.h"
 #include "sys.h"
 #include "vm.h"
+#include "timer.h"
 
 static pte_1mb_t __attribute__((aligned(1 << 14))) kernel_pagetable[4096];
 
@@ -70,6 +71,7 @@ static void vm_kernel_enable() {
 }
 
 void cstart() {
+    timer_init();
     extern int _kbss_start, _kbss_end;
     int *bss = (int *) ka2pa((uintptr_t) &_kbss_start);
     int *bss_end = (int *) ka2pa((uintptr_t) &_kbss_end);
diff --git a/kern/kern.c b/kern/kern.c
index bfea720..ed98e97 100644
--- a/kern/kern.c
+++ b/kern/kern.c
@@ -11,6 +11,7 @@
 #include "timer.h"
 #include "uart.h"
 #include "vm.h"
+#include "libc/rand.h"
 
 void reboot() {
     printf("DONE!!!\n");
@@ -36,10 +37,11 @@ void kernel_start() {
 #endif
 
     sys_enable_cache();
-
     uart_init(115200);
     init_printf(NULL, uart_putc);
 
+    rand_seed(timer_cycles());
+
     kmalloc_init();
 
     irq_init();
diff --git a/kern/kmalloc.c b/kern/kmalloc.c
index 6ace9b4..76b61af 100644
--- a/kern/kmalloc.c
+++ b/kern/kmalloc.c
@@ -1,6 +1,7 @@
 #include <stdbool.h>
 #include <stdint.h>
 
+#include "libc/rand.h"
 #include "kern.h"
 #include "kmalloc.h"
 #include "ksan.h"
@@ -112,7 +113,7 @@ void kmalloc_init() {
 
 // Allocate and returns a pointer to at least `sz` contiguous bytes of memory.
 // Returns `NULL` if `sz == 0` or on failure.
-void *kmalloc(size_t sz) {
+static void *kmalloc_internal(size_t sz) {
     if (sz == 0) {
         return NULL;
     }
@@ -172,7 +173,7 @@ void *kmalloc(size_t sz) {
 
 // Free a pointer previously returned by `kalloc`. Does nothing if `ptr ==
 // NULL`.
-void kfree(void *ptr) {
+static void kfree_internal(void *ptr) {
     if (!ptr || (char *) ptr < &_kheap_start) {
         return;
     }
@@ -213,3 +214,26 @@ void kfree(void *ptr) {
 
     ll_free_insert(pn_to_free(pn), order);
 }
+
+#ifdef RANDOMIZE_KMALLOC
+void *kmalloc(size_t size) {
+    unsigned r = rand16() % 8;
+    void *ptrs[r];
+    for (unsigned i = 0; i < r; i++) {
+        ptrs[i] = kmalloc_internal(size);
+    }
+    void *p = kmalloc_internal(size);
+    for (unsigned i = 0; i < r; i++) {
+        kfree(ptrs[i]);
+    }
+    return p;
+}
+#else
+void *kmalloc(size_t size) {
+    return kmalloc_internal(size);
+}
+#endif
+
+void kfree(void *p) {
+    kfree_internal(p);
+}
diff --git a/kern/kmalloc.h b/kern/kmalloc.h
index ddf503b..2133a43 100644
--- a/kern/kmalloc.h
+++ b/kern/kmalloc.h
@@ -7,7 +7,6 @@
 
 void kmalloc_init();
 void *kmalloc(size_t size);
-void *kmalloc_aligned(size_t size, size_t align);
 void kfree(void *p);
 
 static inline uintptr_t pagenum(uintptr_t pa) {
diff --git a/libc/rand.c b/libc/rand.c
new file mode 100644
index 0000000..9e6fdd0
--- /dev/null
+++ b/libc/rand.c
@@ -0,0 +1,18 @@
+// TODO improve rand implementation, this one is from libpi
+
+static unsigned short lfsr = 0xACE1u;
+static unsigned bit;
+
+unsigned short rand16() {
+    bit  = ((lfsr >> 0) ^ (lfsr >> 2) ^ (lfsr >> 3) ^ (lfsr >> 5) ) & 1;
+    lfsr =  (lfsr >> 1) | (bit << 15);
+    return lfsr;
+}
+
+unsigned long rand32() {
+    return (rand16() << 16) | rand16();
+}
+
+void rand_seed(unsigned seed) {
+    lfsr = seed;
+}
diff --git a/libc/rand.h b/libc/rand.h
new file mode 100644
index 0000000..6ce3bed
--- /dev/null
+++ b/libc/rand.h
@@ -0,0 +1,3 @@
+void rand_seed(unsigned seed);
+unsigned long rand32(void);
+unsigned short rand16(void);