From ba9e0f1bff82e1cbe8461a7383a0a56196850b4d Mon Sep 17 00:00:00 2001
From: Mathieu Maret <mathieu.maret@gmail.com>
Date: Fri, 21 Aug 2020 00:05:53 +0200
Subject: [PATCH] protect slub access from multithread

---
 core/alloc.c | 39 ++++++++++++++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/core/alloc.c b/core/alloc.c
index 09202b7..d2032da 100644
--- a/core/alloc.c
+++ b/core/alloc.c
@@ -2,6 +2,7 @@
 //#define DEBUG
 #include "alloc.h"
 #include "errno.h"
+#include "irq.h"
 #include "klibc.h"
 #include "list.h"
 #include "math.h"
@@ -46,9 +47,13 @@ int allocBookSlab(size_t size, int selfContained)
     struct slabDesc *slab = NULL;
     int slabIdx;
     int ret;
+    int flags;
+
+    disable_IRQs(flags);
     list_foreach(slub, slab, slabIdx)
     {
         if (slab->size == size) {
+            restore_IRQs(flags);
             return -EEXIST;
         }
         if (slab->size > size) {
@@ -56,13 +61,16 @@ int allocBookSlab(size_t size, int selfContained)
         }
     }
     struct slabDesc *newSlab;
-    if ((ret = allocSlab(&newSlab, size, selfContained)))
+    if ((ret = allocSlab(&newSlab, size, selfContained))) {
+        restore_IRQs(flags);
         return ret;
+    }
     if (list_foreach_early_break(slub, slab, slabIdx)) {
         list_insert_before(slub, slab, newSlab);
     } else {
         list_add_tail(slub, newSlab);
     }
+    restore_IRQs(flags);
     return 0;
 }
 
@@ -148,12 +156,17 @@ static void *allocFromSlab(struct slabEntry *slab)
 
 void *malloc(size_t size)
 {
+    int flags;
+    void *ret;
+
     if (size > (1U << SLUB_SIZE)) {
         printf("implement malloc for big size\n");
         return NULL;
     }
     struct slabDesc *slab;
     uint slubIdx;
+
+    disable_IRQs(flags);
     list_foreach(slub, slab, slubIdx)
     {
         if (size <= slab->size)
@@ -166,21 +179,27 @@ void *malloc(size_t size)
         if (!slabEntry->full) {
             // pr_devel("found place in slub %d at idx %d for size %d\n", slubIdx,
             //         slabIdx, size);
-            return allocFromSlab(slabEntry);
+            ret = allocFromSlab(slabEntry);
+            restore_IRQs(flags);
+            return ret;
         }
     }
 
     // No room found
     struct slabEntry *newSlabEntry;
     struct slabEntry *slabList = &slab->slab;
-    int ret;
-    if ((ret = allocSlabEntry(&newSlabEntry, slab->size, IS_SELF_CONTAINED(&slab->slab)))) {
-        pr_devel("Fail to allocSlabEntry %d\n", ret);
+    int retSlab;
+    if ((retSlab =
+             allocSlabEntry(&newSlabEntry, slab->size, IS_SELF_CONTAINED(&slab->slab)))) {
+        pr_devel("Fail to allocSlabEntry %d\n", retSlab);
+        restore_IRQs(flags);
         return NULL;
     }
     pr_devel("Allocate new slab for object of size %d\n", slab->size);
     list_add_tail(slabList, newSlabEntry);
-    return allocFromSlab(newSlabEntry);
+    ret = allocFromSlab(newSlabEntry);
+    restore_IRQs(flags);
+    return ret;
 }
 
 int freeFromSlab(void *ptr, struct slabEntry *slab)
@@ -211,15 +230,21 @@ void free(void *ptr)
 
     struct slabDesc *slab;
     int slabIdx;
+    int flags;
+
+    disable_IRQs(flags);
     list_foreach(slub, slab, slabIdx)
     {
         struct slabEntry *slabEntry;
         int entryIdx;
         list_foreach(&slab->slab, slabEntry, entryIdx)
         {
-            if (freeFromSlab(ptr, slabEntry))
+            if (freeFromSlab(ptr, slabEntry)) {
+                restore_IRQs(flags);
                 return;
+            }
         }
     }
+    restore_IRQs(flags);
     pr_devel("free: slab not found\n");
 }