pool.c

Go to the documentation of this file.

/*++
 
Module Name:
 
    pool.c
 
Purpose:
 
    This translation unit contains the implementation of pool allocations in the kernel.
 
Author:
 
    slep (Matanel) 2025.
 
Revision History:
 
--*/
 
#include "../../includes/mm.h"
#include "../../includes/me.h"
#include "../../includes/mg.h"
#include "../../assert.h"
 
// Can hold any size. (NonPaged)
POOL_DESCRIPTOR GlobalPool;
 
#define POOL_TYPE_GLOBAL 9999
#define POOL_TYPE_PAGED  1234
 
uintptr_t MmNonPagedPoolStart = 0;
uintptr_t MmNonPagedPoolEnd = 0;
uintptr_t MmPagedPoolStart = 0;
uintptr_t MmPagedPoolEnd = 0;
 
MTSTATUS
MiInitializePoolSystem(
    void
)
 
/*++
 
    Routine description:
 
        Initializes the Pool Allocation System of the kernel.
 
    Arguments:
 
        None.
 
    Return Values:
 
        MTSTATUS Status Code.
 
--*/
 
{
    PPROCESSOR cpu = MeGetCurrentProcessor();
    if (!cpu) return MT_NOT_FOUND;
    size_t base = 32; // Start size
 
    // Initialize normal NonPagedPool.
    for (int i = 0; i < MAX_POOL_DESCRIPTORS; i++) {
        PPOOL_DESCRIPTOR desc = &cpu->LookasidePools[i];
 
        // Would grow in binary (32,64,128,256... + sizeof(POOL_HEADER)) (max would be 2048)
        size_t blockSize = (base << i) + sizeof(POOL_HEADER);
        desc->BlockSize = blockSize;
        desc->FreeCount = 0;
        desc->FreeListHead.Next = NULL;
        desc->TotalBlocks = 0;
        desc->PoolLock.locked = 0;
        desc->PoolType = NonPagedPool;
    }
 
    // Initialize NonPagedPoolNx (no-execute)
    base = 32;
    for (int i = 0; i < MAX_POOL_DESCRIPTORS; i++) {
        PPOOL_DESCRIPTOR desc = &cpu->LookasidePoolsNx[i];
 
        size_t blockSize = (base << i) + sizeof(POOL_HEADER);
        desc->BlockSize = blockSize;
        desc->FreeCount = 0;
        desc->FreeListHead.Next = NULL;
        desc->TotalBlocks = 0;
        desc->PoolLock.locked = 0;
        desc->PoolType = NonPagedPoolNx;
    }
 
    // NPG and NPGNx pools reside in the same VA space.
    MmNonPagedPoolStart = MI_NONPAGED_POOL_BASE;
    MmNonPagedPoolEnd = MI_NONPAGED_POOL_END;
    MmPagedPoolStart = MI_PAGED_POOL_BASE;
    MmPagedPoolEnd = MI_PAGED_POOL_END;
    return MT_SUCCESS;
}
 
static
bool
MiRefillPool(
    PPOOL_DESCRIPTOR Desc,
    size_t PoolIndex
)
 
/*++
 
    Routine description:
 
        Refills the specified pool with a block of its size.
 
    Arguments:
 
        [IN]    PPOOL_DESCRIPTOR Desc - Pointer to descriptor.
        [IN]    size_T PoolIndex - Index of the slab in the CPU Lookaside buffer.
 
    Return Values:
 
        True or False based if allocation and/or refill succeeded.
 
--*/
 
{
    // Before allocating a va or another PFN, lets check the global pool first, see if we have a free 4KiB block.
    //IRQL oldIrql;
    uintptr_t PageVa = 0;
    size_t HeaderBlockSize = 0;
    //size_t Iterations = 0;
    /* If I ever return back freeing to global pool, I should check that Desc->BlockSize == VirtualPageSize, else it wont use it.
    * Since we have memory corruptions for larger block sizes..
    // Acquire the spinlock for atomicity.
    MsAcquireSpinlock(&GlobalPool.PoolLock, &oldIrql);
 
    // Initialize the local list so that we push back to not memory leak blocks from the global pol.
    SINGLE_LINKED_LIST localList;
    localList.Next = NULL;
 
    while (GlobalPool.FreeCount) {
        // As long as we have a free block in the global pool, we check it.
        PSINGLE_LINKED_LIST list = GlobalPool.FreeListHead.Next;
        if (list == NULL) break; // FreeCount was wrong, but that's ok
        GlobalPool.FreeListHead.Next = list->Next;
        PPOOL_HEADER header = CONTAINING_RECORD(list, POOL_HEADER, Metadata.FreeListEntry);
        GlobalPool.FreeCount--;
        
        if (header->PoolCanary != 'BEKA') {
            MeBugCheckEx(
                MEMORY_CORRUPT_HEADER,
                (void*)header,
                (void*)__read_rip(),
                NULL,
                NULL
            );
        }
 
        if (Desc->BlockSize > header->Metadata.BlockSize) {
            // If the block gotten from the global pool is smaller than the required refill size, we continue and add this block to the list (to push back later)
            header->Metadata.FreeListEntry.Next = localList.Next;
            localList.Next = &header->Metadata.FreeListEntry;
            Iterations++;
            continue;
        }
 
        // The block is good! The loop that refills the desc wil overwrite this header data. (this is why we dont add sizeof)
        PageVa = (uintptr_t)header;
        HeaderBlockSize = header->Metadata.BlockSize;
        break;
    }
 
    // Refill back the pool.
    while (Iterations--) {
        PSINGLE_LINKED_LIST entryToPushBack = localList.Next;
        if (entryToPushBack == NULL) {
            // Shouldn't happen if iterations is correct, but I always admire checking.
            break;
        }
 
        localList.Next = entryToPushBack->Next;
        entryToPushBack->Next = GlobalPool.FreeListHead.Next;
        GlobalPool.FreeListHead.Next = entryToPushBack;
 
        GlobalPool.FreeCount++;
    }
 
    // Release global pool lock.
    MsReleaseSpinlock(&GlobalPool.PoolLock, oldIrql);
    */
    if (!PageVa) {
        // The global pool is empty... lets allocate.
        // Allocate a 4KiB virtual address.
        PageVa = MiAllocatePoolVa(NonPagedPool, VirtualPageSize);
        if (!PageVa) return false; // Out of VA Space.
 
        // Allocate a 4KiB Physical page.
        PAGE_INDEX pfn = MiRequestPhysicalPage(PfnStateZeroed);
        if (pfn == PFN_ERROR) {
            MiFreePoolVaContiguous(PageVa, VirtualPageSize, NonPagedPool);
            return false;
        }
 
        // Map the page permanently.
        PMMPTE pte = MiGetPtePointer((uintptr_t)PageVa);
        if (!pte) {
            MiFreePoolVaContiguous(PageVa, VirtualPageSize, NonPagedPool);
            MiReleasePhysicalPage(pfn);
            return false;
        }
 
        uint64_t PteFlags = PAGE_PRESENT | PAGE_RW;
 
        // If the descriptor is a NonPagedPoolNx type, we add the NX bit.
        if (Desc->PoolType == NonPagedPoolNx) {
            PteFlags |= PAGE_NX;
        }
 
        // Get the PFN Physical address.
        uint64_t phys = PPFN_TO_PHYSICAL_ADDRESS(INDEX_TO_PPFN(pfn));
 
        MI_WRITE_PTE(pte, PageVa, phys, PteFlags);
 
        // Set block size.
        HeaderBlockSize = VirtualPageSize;
    }
 
    // Reaching here means we got a valid page (either from the global pool, or allocation), now we must carve it up to the appropriate slab's size.
    // Acquire the spinlock for the descriptor given before modifying its data.
    IRQL descIrql;
    MsAcquireSpinlock(&Desc->PoolLock, &descIrql);
 
    // Loop from the start to the end of the page, stepping by the small block size.
    for (size_t offset = 0; (offset + Desc->BlockSize) <= HeaderBlockSize; offset += Desc->BlockSize) {
        // newBlock points to the start of this Desc->BlockSize chunk.
        PPOOL_HEADER newBlock = (PPOOL_HEADER)((uint8_t*)PageVa + offset);
 
        // Set its header metadata.
        newBlock->Metadata.BlockSize = Desc->BlockSize;
        newBlock->Metadata.PoolIndex = PoolIndex;
        newBlock->PoolCanary = 'BEKA'; // Pool Canary
        newBlock->PoolTag = 'ADIR'; // Default Tag
        
        // Add this block to the list of the descriptor.
        newBlock->Metadata.FreeListEntry.Next = Desc->FreeListHead.Next;
        Desc->FreeListHead.Next = &newBlock->Metadata.FreeListEntry;
        Desc->TotalBlocks++;
        Desc->FreeCount++;
    }
 
    MsReleaseSpinlock(&Desc->PoolLock, descIrql);
    return true;
}
 
static
void*
MiAllocateLargePool(
    enum _POOL_TYPE PoolType,
    size_t NumberOfBytes,
    uint32_t Tag
)
 
/*++
 
    Routine description:
 
        Allocates a PoolType pool, and returns a pointer to start of region.
 
    Arguments:
 
        [IN]    enum _POOL_TYPE PoolType - The type of pool to allocate for. (must be NonPagedXXX variant)
        [IN]    size_t NumberOfBytes - Number of bytes needed to allocate.
        [IN]    uint32_t Tag - A 4 byte integer that signifies the current allocation, in little endian (e.g 'TSET' -> 'TEST')
 
    Return Values:
 
        Pointer to start of allocated region
 
--*/
 
{
    // Calculate required size.
    size_t RequiredSize = NumberOfBytes + sizeof(POOL_HEADER);
 
    // Convert to pages.
    size_t neededPages = BYTES_TO_PAGES(RequiredSize);
 
    // Allocate contiguous VAs.
    uintptr_t pageVa = MiAllocatePoolVa(NonPagedPool, RequiredSize);
    
    if (!pageVa) {
        // We don't have enough VA space to allocate required pool.
        return NULL;
    }
 
    // Now lets loop to request PFNs, we also create a safeguard to unroll the loop if failure happens.
    bool failure = false;
    size_t Iterations = 0;
 
    for (size_t i = 0; i < neededPages; i++) {
        // Increment by 4KiB each time.
        uint8_t* currVa = (uint8_t*)pageVa + (i * VirtualPageSize);
 
        PAGE_INDEX pfn = MiRequestPhysicalPage(PfnStateFree);
 
        if (pfn == PFN_ERROR) {
            // Allocation for a physical page failed, free the VA allocated, and unroll the loop (see code below loop)
            MiFreePoolVaContiguous(pageVa, RequiredSize, NonPagedPool);
            failure = true;
            break;
        }
 
        // Map the page.
        PMMPTE pte = MiGetPtePointer((uintptr_t)currVa);
        uint64_t phys = PPFN_TO_PHYSICAL_ADDRESS(INDEX_TO_PPFN(pfn));
 
        uint64_t PteFlags = PAGE_PRESENT | PAGE_RW;
 
        // If its NX pool, we add the NX bit.
        if (PoolType == NonPagedPoolNx) {
            PteFlags |= PAGE_NX;
        }
 
        MI_WRITE_PTE(pte, currVa, phys, PAGE_PRESENT | PAGE_RW);
        
        // Update PFN metadata.
        PPFN_ENTRY ppfn = INDEX_TO_PPFN(pfn);
        ppfn->State = PfnStateActive;
        ppfn->Flags = PFN_FLAG_NONPAGED;
        ppfn->Descriptor.Mapping.PteAddress = pte;
 
        // Update iterations for loop unroll (if failure)
        Iterations++;
    }
 
    // If failure, unroll PFNs using Iterations.
    if (failure) {
        for (size_t j = 0; j < Iterations; j++) {
            uint8_t* vaToFree = (uint8_t*)pageVa + (j * VirtualPageSize);
            PMMPTE pte = MiGetPtePointer((uintptr_t)vaToFree);
            PAGE_INDEX pfn = MiTranslatePteToPfn(pte);
            MiUnmapPte(pte);
            MiReleasePhysicalPage(pfn);
        }
        return NULL;
    }
 
    // Success! Initialize the block and return the pointer to caller.
    PPOOL_HEADER newHeader = (PPOOL_HEADER)pageVa;
    newHeader->PoolCanary = 'BEKA';
    newHeader->PoolTag = Tag;
    newHeader->Metadata.BlockSize = neededPages * VirtualPageSize; // Store allocated size.
    newHeader->Metadata.PoolIndex = POOL_TYPE_GLOBAL;
 
    void* UserAddress = (void*)((uint8_t*)newHeader + sizeof(POOL_HEADER));
    // Set to zero (to avoid kernel issues)
    kmemset(UserAddress, 0, NumberOfBytes);
    // Return the pointer (exclude metadata start).
    return UserAddress;
}
 
static
void*
MiAllocatePagedPool(
    IN  size_t NumberOfBytes,
    IN  uint32_t Tag
)
 
/*++
 
    Routine description:
 
        Allocates a paged pool.
 
    Arguments:
 
        [IN]    size_t NumberOfBytes - Number of bytes needed to allocate.
        [IN]    uint32_t Tag - A 4 byte integer that signifies the current allocation, in big endian (e.g 'TSET' - "TEST")
 
    Return Values:
 
        Pointer to allocated region, or NULL on failure.
 
    Notes:
 
        This function AND access to it's pool contents MUST be with IRQL < DISPATCH_LEVEL.
 
--*/
 
{
    size_t ActualSize = NumberOfBytes + sizeof(POOL_HEADER);
    uintptr_t PagedVa = MiAllocatePoolVa(PagedPool, ActualSize);
    if (!PagedVa) return NULL;
 
    PPOOL_HEADER header = (PPOOL_HEADER)PagedVa;
    // The first page will not be resident in memory anymore, it should be paged in as well, that means that pool free routines will need to be executed at < DISPATCH_LEVEL too.
    // Set each page to be a demand lazy allocation.
    size_t NumberOfPages = BYTES_TO_PAGES(ActualSize);
    size_t currVa = PagedVa;
 
    size_t i = 0;
    bool failure = false;
    for (; i < NumberOfPages; i++) {
        PMMPTE tmpPte = MiGetPtePointer(currVa);
        if (!tmpPte) { failure = true; break; }
        MMPTE TempPte = *tmpPte;
        // Set the PTE as demand zero. (paged pool is also NoExecute)
        MM_SET_DEMAND_ZERO_PTE(TempPte, PROT_KERNEL_READ | PROT_KERNEL_WRITE | PROT_KERNEL_NOEXECUTE, false);
        // Atomically exchange new value.
        MiAtomicExchangePte(tmpPte, TempPte.Value);
        // Invalidate the VA.
        MiInvalidateTlbForVa((void*)currVa);
        currVa += VirtualPageSize;
    }
 
    // Check if we failed getting a PTE, if so, unroll loop.
    if (failure) {
        while (i-- > 0) {
            // Go back to the page we processed before.
            currVa -= VirtualPageSize;
 
            // Acquire its PTE.
            PMMPTE Pte = MiGetPtePointer(currVa);
            if (!Pte) {
                assert(false, "Acquiring PTE pointer after it was valid before resulted in NULL, severe bug.");
                break;
            }
 
            // Clear PTE.
            MiUnmapPte(Pte);
        }
 
        // Return NULL, allocation failure.
        return NULL;
    }
 
    // Set metadata. (header should get paged in now).
    header->PoolCanary = 'BEKA';
    header->PoolTag = Tag;
    header->Metadata.BlockSize = ActualSize;
    header->Metadata.PoolIndex = POOL_TYPE_PAGED;
 
    // Return VA.
    return (void*)((uint8_t*)PagedVa + sizeof(POOL_HEADER));
}
 
void*
MmAllocatePoolWithTag(
    IN enum _POOL_TYPE PoolType,
    IN  size_t NumberOfBytes,
    IN  uint32_t Tag
)
 
/*++
 
    Routine description:
 
        Allocates a pool block of the specified type and returns a pointer to allocated block.
        On any allocation, the returned block(s) is/are zeroed, no exceptions.
 
    Arguments:
 
        [IN]    enum _POOL_TYPE - POOL_TYPE Enumerator, specifying the type of pool that will be allocated.
        [IN]    size_t NumberOfBytes - Number of bytes needed to allocate.
        [IN]    uint32_t Tag - A 4 byte integer that signifies the current allocation, in big endian (e.g 'TSET' - "TEST")
 
    Return Values:
 
        Pointer to allocated region, or NULL on failure.
 
    Notes:
 
        PagedPool allocations CANNOT happen at IRQL => DISPATCH_LEVEL.
        NonPagedPool allocations CANNOT happen at IRQL > DISPATCH_LEVEL.
 
        NonPagedPoolCachedXxX pool allocations are not supported.
 
    Todos:
 
        Implement POOL_TAGGING (global doubly linked list of all pool allocs, so we can have a poolmon of our own!)
 
--*/
 
{
    // Declarations
    IRQL oldIrql;
    size_t ActualSize;
    PPROCESSOR cpu;
    PPOOL_DESCRIPTOR Desc;
    PPOOL_HEADER header;
    PSINGLE_LINKED_LIST list;
    size_t Index;
 
    // Runtime assertions
    //assert((NumberOfBytes) != 0); Better to use the if statement, supplies retaddr.
    assert((Tag) != 0);
 
    if (NumberOfBytes == 0) {
        // Bad pool caller.
        MeBugCheckEx(
            BAD_POOL_CALLER,
            RETADDR(0),
            NULL,
            NULL,
            NULL
        );
    }
 
    IRQL currIrql = MeGetCurrentIrql();
 
    // IRQL Must be less or equal to DISPATCH_LEVEL if allocating with NonPagedPool.
    // IRQL Must be LESS than DISPATCH_LEVEL if allocating with PagedPool.
    if (currIrql <= DISPATCH_LEVEL) {
        if (PoolType == PagedPool && currIrql == DISPATCH_LEVEL) {
            MeBugCheckEx(
                BAD_POOL_CALLER,
                (void*)&MmAllocatePoolWithTag,
                (void*)MeGetCurrentIrql(),
                (void*)8,
                (void*)__builtin_return_address(0)
            );
        }
    }
    // IRQL Must NOT be greater than DISPATCH_LEVEL at any allocation.
    else {
        MeBugCheckEx(
            BAD_POOL_CALLER,
            (void*)&MmAllocatePoolWithTag,
            (void*)MeGetCurrentIrql(),
            (void*)8,
            (void*)__builtin_return_address(0)
        );
    }
 
    if (PoolType == PagedPool) {
        // Use the internal paged pool allocator.
        return MiAllocatePagedPool(NumberOfBytes, Tag);
    }
 
    ActualSize = NumberOfBytes + sizeof(POOL_HEADER);
    cpu = MeGetCurrentProcessor();
 
 
    // It's NonPagedPool OR NonPagedPooLNx. Find the correct slab.
    PPOOL_DESCRIPTOR TypeDescriptor = NULL;
 
    if (PoolType == NonPagedPool) {
        // Normal
        TypeDescriptor = cpu->LookasidePools;
    }
    else if (PoolType == NonPagedPoolNx) {
        // Nx
        TypeDescriptor = cpu->LookasidePoolsNx;
    }
    else {
        // Pool type is not supported.
        return NULL;
    }
 
    Desc = NULL;
    for (int i = 0; i < MAX_POOL_DESCRIPTORS; i++) {
        PPOOL_DESCRIPTOR currentSlab = &TypeDescriptor[i];
        if (ActualSize <= currentSlab->BlockSize) {
            Desc = currentSlab;
            Index = i;
            break; // Found the best-fit slab
        }
    }
 
    if (Desc == NULL) {
        // Allocation is larger than 2048 bytes, use the large pool allocator.
        return MiAllocateLargePool(PoolType, NumberOfBytes, Tag);
    }
    
    MsAcquireSpinlock(&Desc->PoolLock, &oldIrql);
    assert((Desc->FreeCount) != UINT64_T_MAX);
 
    if (Desc->FreeCount == 0) {
        // Looks like the pool is empty, refill all empty pools.
        // First, release the spinlock.
        MsReleaseSpinlock(&Desc->PoolLock, oldIrql);
        if (!MiRefillPool(Desc, Index)) {
            // If we failed allocation, act on failure.
            return NULL;
        }
 
        // Retry allocation.
        return MmAllocatePoolWithTag(PoolType, NumberOfBytes, Tag);
    }
 
    // Looks like we have a block to return! Return its PTR.
    // First, acquire it. (we are under spinlock, no need for interlocked pop)
    list = Desc->FreeListHead.Next;
    assert((list) != NULL, "Pool is nullptr even though freecount isn't zero.");
    Desc->FreeListHead.Next = list->Next; // Finish the pop
    header = CONTAINING_RECORD(list, POOL_HEADER, Metadata.FreeListEntry);
 
    // We must restore the metadata because the linked list pointer 
    // overwrote it while the block was sitting in the free list.
    header->Metadata.PoolIndex = (uint16_t)Index;
    header->Metadata.BlockSize = (uint16_t)Desc->BlockSize;
 
    // First check if the canary is wrong.
    if (header->PoolCanary != 'BEKA') {
        MeBugCheckEx(
            MEMORY_CORRUPT_HEADER,
            (void*)header,
            (void*)__read_rip(),
            NULL,
            NULL
        );
    }
 
    // Rewrite its tag.
    header->PoolTag = Tag;
    // Decrement descriptor free count.
    Desc->FreeCount--;
    assert((Desc->FreeCount) != SIZE_T_MAX); // Check for underflow.
    // Release spinlock.
    MsReleaseSpinlock(&Desc->PoolLock, oldIrql);
    void* UserAddress = (void*)((uint8_t*)header + sizeof(POOL_HEADER));
 
    // Set to zero (to avoid kernel issues)
    // If this is ever removed, massive kernel bugs will appear with uninitialized memory.
    // So we kinda depend on it now.
    // brilliant engineering, brilliant (its sarcasm)
    // (NOTE: If pool is allocated freshly by MiRefillPool (needed loop to find out) (or double counters),
    // it usually (USUALLY, maybe it changed) came from acquiring a physical page with a zeroed pfn state value
    // so the page comes zeroed, which means the memset below can be skipped) (PERFORMANCE TODO)
    kmemset(UserAddress, 0, NumberOfBytes);
 
    // Return the pointer (exclude metadata start).
    return UserAddress;
}
 
void
MmFreePool(
    IN  void* buf
)
 
/*++
 
    Routine description:
 
        Deallocates the buffer allocated by the MmAllocatePoolWithTag routine (or other pool allocation routines if present).
 
    Arguments:
 
        [IN]    void* buf - The pointer given by the routine. (start of allocated region)
 
    Return Values:
 
        None.
 
    Notes:
 
        Memory is freed here, do not use the pointer after freeing the pool.
        The Pool tag will not be modified when freeing, useful for debugging.
 
        Memory allocated with PagedPoolXxX type must be deallocated at IRQL < DISPATCH_LEVEL.
 
--*/
 
{
    if (!buf) return;
    assert(MeGetCurrentIrql() <= DISPATCH_LEVEL);
 
    // Convert the buffer to the header.
    PPOOL_HEADER header = (PPOOL_HEADER)((uint8_t*)buf - sizeof(POOL_HEADER));
 
    gop_printf(COLOR_YELLOW, "MmFreePool called with IRQL: %d | Header: %p\n", MeGetCurrentIrql(), header);
 
    if (header->PoolCanary != 'BEKA') {
        MeBugCheckEx(
            MEMORY_CORRUPT_HEADER,
            (void*)header,
            (void*)RETADDR(0),
            NULL,
            NULL
        );
    }
 
    // Obtain the pool index to free the region back into.
    uint16_t PoolIndex = header->Metadata.PoolIndex;
 
    if (PoolIndex == POOL_TYPE_GLOBAL) {
        // We destroy global pool allocations and free them back to main memory.
        size_t BlockSize = header->Metadata.BlockSize;
        size_t NumberOfPages = BYTES_TO_PAGES(BlockSize);
        uintptr_t CurrentVA = (uintptr_t)header;
 
        // Release physical pages and unmap PTEs.
        for (size_t i = 0; i < NumberOfPages; i++) {
            PMMPTE pte = MiGetPtePointer(CurrentVA);
 
            // Check if PTE is present, it should be though.
            if (pte && pte->Hard.Present) {
                PAGE_INDEX Pfn = MiTranslatePteToPfn(pte);
 
                // Invalidate PTE.
                MiUnmapPte(pte);
                // Free the PFN back to db.
                MiReleasePhysicalPage(Pfn);
            }
            else {
                MeBugCheckEx(MEMORY_CORRUPT_HEADER, (void*)CurrentVA, 0, 0, 0);
            }
 
            CurrentVA += VirtualPageSize;
        }
 
        // Free VA space given.
        MiFreePoolVaContiguous((uintptr_t)header, BlockSize, NonPagedPool);
 
        return;
    }
 
    if (PoolIndex == POOL_TYPE_PAGED) {
        // For a paged pool allocation, we just free every PTE, then returned the VA space consumed.
        // The BlockSize field in a PagedPool allocation is how many bytes were requested + sizeof(POOL_HEADER)
        size_t NumberOfPages = BYTES_TO_PAGES(header->Metadata.BlockSize);
 
        // Loop over the amount, if the PTE is present, unmap it and clear the demand page.
        uintptr_t CurrentVA = (uintptr_t)header;
 
        for (size_t i = 0; i < NumberOfPages; i++) {
            PMMPTE pte = MiGetPtePointer(CurrentVA);
            if (unlikely(!pte)) goto advance;
 
            // Check if the PTE is present, if it is, the demand zero page has been consumed, we deallocate, and unset the demand zero.
            if (pte->Hard.Present) {
                // It has a PFN.
                PAGE_INDEX Pfn = MiTranslatePteToPfn(pte);
                // Unmap PTE, free PFN.
                MiUnmapPte(pte);
                MiReleasePhysicalPage(Pfn);
            }
            else {
                // If the PTE isnt present, it still must contain a demand zero bit.
                assert(MM_IS_DEMAND_ZERO_PTE(*pte) == true);
                MMPTE TempPte = *pte;
                MM_UNSET_DEMAND_ZERO_PTE(TempPte);
                // Flip the demand zero bit.
                MiAtomicExchangePte(pte, TempPte.Value);
            }
 
            // Invalidate the VA. (only necessary for else though, as in MiUnmapPte it does invalidate TLB, but im so scared of bugs ill leave this here)
            MiInvalidateTlbForVa((void*)CurrentVA);
 
            advance:
            CurrentVA += VirtualPageSize;
        }
        return;
    }
 
    //
    // Nonpaged pool allocation
    //
 
    PPROCESSOR cpu = MeGetCurrentProcessor();
    PPOOL_DESCRIPTOR Desc = &cpu->LookasidePools[PoolIndex];
 
    // Acquire the same lock used by the allocator
    IRQL oldIrql;
    MsAcquireSpinlock(&Desc->PoolLock, &oldIrql);
 
    // Push the entry back onto the list (it's protected by the lock)
    header->Metadata.FreeListEntry.Next = Desc->FreeListHead.Next;
    Desc->FreeListHead.Next = &header->Metadata.FreeListEntry;
 
    // Increment the free count.
    Desc->FreeCount++;
 
    // Release the lock
    MsReleaseSpinlock(&Desc->PoolLock, oldIrql);
}