From ed2d4744d7b73f3a7692d89988a7197718715811 Mon Sep 17 00:00:00 2001
From: "Chloe M." <chloe@mensia.org>
Date: Wed, 24 Jun 2026 17:43:45 +0000
Subject: [PATCH] stos: pool: Protect pool structures with lock

Signed-off-by: Chloe M. <chloe@mensia.org>
---
 paw/stos/ex/pool.c      | 9 ++++++++-
 paw/stos/head/ex/pool.h | 3 +++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/paw/stos/ex/pool.c b/paw/stos/ex/pool.c
index 7c73f7b..59b17c6 100644
--- a/paw/stos/ex/pool.c
+++ b/paw/stos/ex/pool.c
@@ -179,6 +179,7 @@ ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
 {
     KPCR *ThisCore;
     MEMORY_POOL *AllocPool;
+    VOID *Base;
 
     if (ByteCount == 0) {
         return NULL;
@@ -195,6 +196,7 @@ ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
 
     ThisCore = HalKpcrCurrent();
     AllocPool = &ThisCore->AllocPool;
+    KeSpinLockAcquire(&AllocPool->Lock, false);
 
     /*
      * Ensure that the byte count is aligned by the minimum granularity
@@ -207,10 +209,14 @@ ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
             "allocation in pool #%d exceeds page size\n",
             AllocPool->Id
         );
+
+        KeSpinLockRelease(&AllocPool->Lock);
         return NULL;
     }
 
-    return AllocatePoolWithTag(AllocPool, Type, ByteCount, Tag);
+    Base = AllocatePoolWithTag(AllocPool, Type, ByteCount, Tag);
+    KeSpinLockRelease(&AllocPool->Lock);
+    return Base;
 }
 
 ST_STATUS
@@ -225,6 +231,7 @@ ExPoolInit(MEMORY_POOL *Pool)
     }
 
     Pool->Id = PoolCount;
+    KeSpinLockInit(&Pool->Lock, "ex-pool");
     DTRACE("bringing up pool #%d\n", Pool->Id);
 
     /* Initialize the pools */
diff --git a/paw/stos/head/ex/pool.h b/paw/stos/head/ex/pool.h
index b7cf672..dd19f68 100644
--- a/paw/stos/head/ex/pool.h
+++ b/paw/stos/head/ex/pool.h
@@ -11,6 +11,7 @@
 
 #include <hal/page.h>
 #include <stapi/status.h>
+#include <ke/spinlock.h>
 #include <stdef.h>
 
 /* Minimum log2 allocation size */
@@ -71,10 +72,12 @@ typedef struct {
  *
  * @BlockLevels: Block levels [granularity : GRAN(level)]
  * @Id:          Pool ID
+ * @Lock:        Lock protecting this pool
  */
 typedef struct {
     MEMORY_BLOCK BlockLevels[LEVEL_COUNT];
     USIZE Id;
+    SPINLOCK Lock;
 } MEMORY_POOL;
 
 /*