From b95613787a95be255565495bd7cffdf07f0b365f Mon Sep 17 00:00:00 2001
From: Aiken Harris <harraiken91@gmail.com>
Date: Tue, 24 Mar 2026 08:39:47 +0100
Subject: [PATCH] Strip MM_POOL_PROTECTED flag to maintain NT compatibility and
 ensure correct pool tracking hash lookups

---
 sdk/xtdk/mmtypes.h   |  1 +
 xtoskrnl/mm/alloc.cc | 20 +++++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/sdk/xtdk/mmtypes.h b/sdk/xtdk/mmtypes.h
index eea5ccf..f1e4199 100644
--- a/sdk/xtdk/mmtypes.h
+++ b/sdk/xtdk/mmtypes.h
@@ -35,6 +35,7 @@
 
 /* Pool flags */
 #define MM_POOL_BIG_ALLOCATIONS_ENTRY_FREE         0x1
+#define MM_POOL_PROTECTED                          0x80000000
 #define MM_POOL_RAISE_EXCEPTION                    0x10
 
 /* Number of reserved zeroed PTEs */
diff --git a/xtoskrnl/mm/alloc.cc b/xtoskrnl/mm/alloc.cc
index 6c77316..37bc957 100644
--- a/xtoskrnl/mm/alloc.cc
+++ b/xtoskrnl/mm/alloc.cc
@@ -1053,7 +1053,12 @@ MM::Allocator::FreePool(IN PVOID VirtualAddress,
 
         /* Retrieve original metadata while removing the allocation from the tracking table */
         Tag = UnregisterBigAllocationTag(VirtualAddress, &PageCount, PoolType);
-        if(!Tag)
+        if(Tag & MM_POOL_PROTECTED)
+        {
+            /* Strip the protected pool bit */
+            Tag &= ~MM_POOL_PROTECTED;
+        }
+        else if(!Tag)
         {
             /* Fallback to a default tag */
             Tag = SIGNATURE32('B', 'i', 'g', 'A');
@@ -1100,6 +1105,13 @@ MM::Allocator::FreePool(IN PVOID VirtualAddress,
     Tag = PoolEntry->PoolTag;
     Combined = FALSE;
 
+    /* Check if the allocation tag carries the protected pool modifier */
+    if(Tag & MM_POOL_PROTECTED)
+    {
+        /* Strip the protected pool bit */
+        Tag &= ~MM_POOL_PROTECTED;
+    }
+
     /* Remove the allocation from the tracking table */
     UnregisterAllocationTag(Tag, BlockSize * MM_POOL_BLOCK_SIZE, (MMPOOL_TYPE)(PoolEntry->PoolType - 1));
 
@@ -1467,6 +1479,9 @@ MM::Allocator::RegisterAllocationTag(IN ULONG Tag,
     Processor = KE::Processor::GetCurrentProcessorNumber();
     CpuTable = TagTables[Processor];
 
+    /* Strip the protected pool bit */
+    Tag &= ~MM_POOL_PROTECTED;
+
     /* Compute the initial hash index */
     Hash = ComputeHash(Tag, AllocationsTrackingTableMask);
     Index = Hash;
@@ -1861,6 +1876,9 @@ MM::Allocator::UnregisterAllocationTag(IN ULONG Tag,
     Processor = KE::Processor::GetCurrentProcessorNumber();
     CpuTable = TagTables[Processor];
 
+    /* Strip the protected pool bit */
+    Tag &= ~MM_POOL_PROTECTED;
+
     /* Compute the initial hash index */
     Hash = ComputeHash(Tag, AllocationsTrackingTableMask);
     Index = Hash;