From 184ce5735e9c40e28c4471d83181a3fcfe722829 Mon Sep 17 00:00:00 2001
From: Aiken Harris <harraiken91@gmail.com>
Date: Mon, 16 Mar 2026 15:33:36 +0100
Subject: [PATCH] Add runlevel verification to memory pool allocations

---
 sdk/xtdk/mmtypes.h            |  4 +++
 xtoskrnl/includes/mm/alloc.hh |  3 ++
 xtoskrnl/mm/alloc.cc          | 65 +++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+)

diff --git a/sdk/xtdk/mmtypes.h b/sdk/xtdk/mmtypes.h
index 2cf29d7..7e81b0f 100644
--- a/sdk/xtdk/mmtypes.h
+++ b/sdk/xtdk/mmtypes.h
@@ -29,6 +29,10 @@
 /* Memory manager pool type mask definition */
 #define MM_POOL_TYPE_MASK                          1
 
+/* Bad pool caller reasons */
+#define MM_POOL_INVALID_ALLOC_RUNLEVEL             8
+#define MM_POOL_INVALID_FREE_RUNLEVEL              9
+
 /* Number of reserved zeroed PTEs */
 #define MM_RESERVED_ZERO_PTES                      32
 
diff --git a/xtoskrnl/includes/mm/alloc.hh b/xtoskrnl/includes/mm/alloc.hh
index 80a4651..1c82200 100644
--- a/xtoskrnl/includes/mm/alloc.hh
+++ b/xtoskrnl/includes/mm/alloc.hh
@@ -48,6 +48,9 @@ namespace MM
             STATIC XTAPI XTSTATUS FreeNonPagedPoolPages(IN PVOID VirtualAddress);
             STATIC XTAPI XTSTATUS FreePagedPoolPages(IN PVOID VirtualAddress);
             STATIC XTAPI VOID MapNonPagedPool(VOID);
+            STATIC XTINLINE VOID VerifyRunLevel(IN MMPOOL_TYPE PoolType,
+                                                IN SIZE_T Bytes,
+                                                IN PVOID Entry);
     };
 }
 
diff --git a/xtoskrnl/mm/alloc.cc b/xtoskrnl/mm/alloc.cc
index 26e6190..9f81436 100644
--- a/xtoskrnl/mm/alloc.cc
+++ b/xtoskrnl/mm/alloc.cc
@@ -312,6 +312,16 @@ MM::Allocator::AllocatePool(IN MMPOOL_TYPE PoolType,
 {
     UNIMPLEMENTED;
 
+    /* Verify run level for the specified pool */
+    VerifyRunLevel(PoolType, Bytes, NULLPTR);
+
+    /* Check if there are any bytes to allocate */
+    if(!Bytes)
+    {
+        /* Allocate at least a single byte */
+        Bytes = 1;
+    }
+
     /* Allocate pages */
     return AllocatePages(PoolType, Bytes, Memory);
 }
@@ -668,3 +678,58 @@ MM::Allocator::InitializePagedPool(VOID)
 {
     UNIMPLEMENTED;
 }
+
+/**
+ * Validates the run level for the specified pool. If the run level is invalid, the kernel panics.
+ *
+ * @param PoolType
+ *        Supplies the pool type.
+ *
+ * @param Bytes
+ *        Supplies the size of the allocation.
+ *
+ * @param Entry
+ *        Supplies a pointer to the allocation entry.
+ *
+ * @return This routine does not return any value.
+ *
+ * @since XT 1.0
+ */
+XTINLINE
+VOID
+MM::Allocator::VerifyRunLevel(IN MMPOOL_TYPE PoolType,
+                              IN SIZE_T Bytes,
+                              IN PVOID Entry)
+{
+    KRUNLEVEL RunLevel;
+
+    /* Get current run level */
+    RunLevel = KE::RunLevel::GetCurrentRunLevel();
+
+    /* Validate run level */
+    if((PoolType & MM_POOL_TYPE_MASK) == PagedPool)
+    {
+        /* Paged pool runs up to APC level */
+        if(RunLevel <= APC_LEVEL)
+        {
+            /* Run level is valid */
+            return;
+        }
+    }
+    else
+    {
+        /* Non-paged pool runs up to DISPATCH_LEVEL */
+        if(RunLevel <= DISPATCH_LEVEL)
+        {
+            /* Run level is valid */
+            return;
+        }
+    }
+
+    /* Invalid run level for specified pool, kernel panic */
+    KE::Crash::Panic(0xC2,
+                     (Entry ? MM_POOL_INVALID_FREE_RUNLEVEL : MM_POOL_INVALID_ALLOC_RUNLEVEL),
+                     RunLevel,
+                     PoolType,
+                     (Entry ? (ULONG_PTR)Entry : Bytes));
+}