Strip MM_POOL_PROTECTED flag to maintain NT compatibility and ensure correct pool tracking hash lookups

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 */

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;