// Kernel pool allocator and corruption diagnostics.

#include <stdbool.h>
#include <stddef.h>
#include <string.h>
#include <intrin.h>
#include <kernel/mm_pool.h>
#include <kernel/mm_pte.h>
#include <kernel/mm_pfn.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_spinlock.h>

// The boundary-tag allocator stores a header before each allocation and a
// footer after it. Mirrored size and state fields permit constant-time
// coalescing while retaining enough metadata for corruption diagnostics.

#define BLOCK_MAGIC  0xB10CB10CUL

#define BLOCK_STATE_USED        0u
#define BLOCK_STATE_FREE        1u
#define BLOCK_STATE_QUARANTINE  2u

#define BLOCK_FLAG_CANARY       0x00000001u
#define BLOCK_FLAG_POISONED     0x00000002u

#define POOL_TAIL_CANARY_BASE   0xC0FFEE1234ABCDEFULL
#define POOL_POISON_FREE        0xDD
#define POOL_QUARANTINE_SLOTS   256u
#define POOL_ALIGNMENT          16u
#define POOL_TAG_DEFAULT        0x6C6F6F70u  // 'pool'.

typedef struct BLOCK_HDR {
    uint32_t Magic;
    uint32_t Free;      // BLOCK_STATE_*
    size_t   Size;      // usable allocator bytes (excludes header + footer)
    struct BLOCK_HDR* PrevFree;
    struct BLOCK_HDR* NextFree;

    // Debug metadata. Present in all builds so boot.ini bits can enable
    // diagnostics without changing allocator layout or ABI.
    size_t   RequestedSize;
    uint32_t Tag;
    uint32_t Flags;
    uint64_t Sequence;
    uintptr_t AllocCaller;
    uintptr_t FreeCaller;
    uintptr_t FreeListInsertCaller;
    uintptr_t FreeListRemoveCaller;
    uintptr_t CoalesceCaller;
    uint64_t FreeListSequence;

    // Keep user data 16-byte aligned for MSVC optimized code and image constants.
    uint64_t AlignmentPadding;
} BLOCK_HDR;

typedef struct BLOCK_FTR {
    size_t Size;
    uint32_t Free;
    uint32_t Magic;
} BLOCK_FTR;

_Static_assert((sizeof(BLOCK_HDR) % POOL_ALIGNMENT) == 0, "pool header must preserve 16-byte user alignment");
_Static_assert((sizeof(BLOCK_FTR) % POOL_ALIGNMENT) == 0, "pool footer must preserve 16-byte block alignment");

#define HDR_SIZE  (sizeof(BLOCK_HDR))
#define FTR_SIZE  (sizeof(BLOCK_FTR))
#define OVERHEAD  (HDR_SIZE + FTR_SIZE)
#define MIN_BLOCK (OVERHEAD + sizeof(void*))

// Implements the hdr To Ftr operation.
static inline BLOCK_FTR* KiHdrToFtr(BLOCK_HDR* h) {
    return (BLOCK_FTR*)((char*)h + HDR_SIZE + h->Size);
}
// Implements the ftr To Hdr operation.
static inline BLOCK_HDR* KiFtrToHdr(BLOCK_FTR* f) {
    return (BLOCK_HDR*)((char*)f - HDR_SIZE - f->Size);
}
// Implements the next Block operation.
static inline BLOCK_HDR* KiNextBlock(BLOCK_HDR* h) {
    return (BLOCK_HDR*)((char*)h + HDR_SIZE + h->Size + FTR_SIZE);
}
// Implements the prev Ftr operation.
static inline BLOCK_FTR* KiPrevFtr(BLOCK_HDR* h) {
    return (BLOCK_FTR*)((char*)h - FTR_SIZE);
}

typedef struct HEAP HEAP;

struct HEAP {
    BLOCK_HDR* FreeList;    // doubly-linked list of free blocks
    uintptr_t   Start;       // virtual base of the arena
    uintptr_t   Brk;         // current committed end (page-aligned)
    uintptr_t   Limit;       // hard upper bound
    uint64_t    MapFlags;    // PAGE_* flags for new committed pages
    SPINLOCK    Lock;
    int         DecommitTailOnFree;

    uint64_t    NextSequence;
    uint64_t    NextFreeListSequence;
    BLOCK_HDR*  Quarantine[POOL_QUARANTINE_SLOTS];
    size_t      QuarantineCursor;
    size_t      QuarantineCount;
};

// Returns the active allocator diagnostic options.
static inline uint32_t KiPoolDebugFlags(void) {
    return KeGetDebugFlags() & (KE_DEBUG_POOL_CANARY |
                                KE_DEBUG_POOL_POISON |
                                KE_DEBUG_POOL_QUARANTINE |
                                KE_DEBUG_POOL_CALLERS);
}

// Derives a per-allocation tail canary from stable block metadata.
static inline uint64_t KiPoolTailCanaryValue(const BLOCK_HDR* blk) {
    return POOL_TAIL_CANARY_BASE ^
           ((uint64_t)(uintptr_t)blk << 13) ^
           ((uint64_t)blk->RequestedSize << 1) ^
           blk->Sequence;
}

// Returns the first caller-visible byte of an allocator block.
static inline uint8_t* KiBlockUserBase(BLOCK_HDR* blk) {
    return (uint8_t*)blk + HDR_SIZE;
}

// Returns the canary slot immediately after the requested payload.
static inline uint64_t* KiBlockTailCanary(BLOCK_HDR* blk) {
    return (uint64_t*)(void*)(KiBlockUserBase(blk) + blk->RequestedSize);
}

// Returns a diagnostic name for an allocator block state.
static const char* KiPoolStateName(uint32_t state) {
    switch (state) {
        case BLOCK_STATE_USED: return "used";
        case BLOCK_STATE_FREE: return "free";
        case BLOCK_STATE_QUARANTINE: return "quarantine";
        default: return "bad-state";
    }
}

// Stops with allocator metadata that identifies the corrupt block and callers.
static _Noreturn void KiPoolPanic(const char* reason, HEAP* h, BLOCK_HDR* blk, const void* ptr) {
    KePanic(KE_PANIC_CODE_GENERIC,
            "MM: pool corruption: %s ptr=0x%016llX hdr=0x%016llX heap=[0x%016llX,0x%016llX) magic=%08x state=%u(%s) size=%llu req=%llu tag=%08x flags=%08x seq=%llu alloc=0x%016llX free=0x%016llX fl_ins=0x%016llX fl_rem=0x%016llX coal=0x%016llX fl_seq=%llu\n",
            reason ? reason : "unknown",
            (unsigned long long)(uintptr_t)ptr,
            (unsigned long long)(uintptr_t)blk,
            (unsigned long long)(h ? h->Start : 0),
            (unsigned long long)(h ? h->Brk : 0),
            blk ? blk->Magic : 0,
            blk ? blk->Free : 0,
            blk ? KiPoolStateName(blk->Free) : "<null>",
            (unsigned long long)(blk ? blk->Size : 0),
            (unsigned long long)(blk ? blk->RequestedSize : 0),
            blk ? blk->Tag : 0,
            blk ? blk->Flags : 0,
            (unsigned long long)(blk ? blk->Sequence : 0),
            (unsigned long long)(blk ? blk->AllocCaller : 0),
            (unsigned long long)(blk ? blk->FreeCaller : 0),
            (unsigned long long)(blk ? blk->FreeListInsertCaller : 0),
            (unsigned long long)(blk ? blk->FreeListRemoveCaller : 0),
            (unsigned long long)(blk ? blk->CoalesceCaller : 0),
            (unsigned long long)(blk ? blk->FreeListSequence : 0));
}

// Tests whether an allocator header pointer is inside the committed heap.
static bool KiPoolHeaderInCommittedHeap(const HEAP* h, const BLOCK_HDR* blk) {
    if (!h || !blk)
        return false;
    uintptr_t addr = (uintptr_t)blk;
    return addr >= h->Start && addr < h->Brk;
}

// Validates the mirrored footer and committed-heap bounds for one block.
static void KiPoolValidateFooterLocked(HEAP* h, BLOCK_HDR* blk, const char* where) {
    if (!h || !blk)
        KiPoolPanic(where, h, blk, blk ? KiBlockUserBase(blk) : (void*)0);
    if (blk->Magic != BLOCK_MAGIC)
        KiPoolPanic(where, h, blk, KiBlockUserBase(blk));
    if ((uintptr_t)blk < h->Start || (uintptr_t)blk >= h->Brk)
        KiPoolPanic("block outside heap", h, blk, KiBlockUserBase(blk));
    if (blk->Size > (size_t)(h->Brk - (uintptr_t)blk - OVERHEAD))
        KiPoolPanic("block size outside heap", h, blk, KiBlockUserBase(blk));
    BLOCK_FTR* ftr = KiHdrToFtr(blk);
    if ((uintptr_t)ftr + FTR_SIZE > h->Brk)
        KiPoolPanic("footer outside heap", h, blk, KiBlockUserBase(blk));
    if (ftr->Magic != BLOCK_MAGIC || ftr->Size != blk->Size || ftr->Free != blk->Free)
        KiPoolPanic(where, h, blk, KiBlockUserBase(blk));
}

// Writes the optional tail canary after the requested payload.
static void KiPoolWriteTailCanary(BLOCK_HDR* blk) {
    if (!blk || !(blk->Flags & BLOCK_FLAG_CANARY))
        return;
    *KiBlockTailCanary(blk) = KiPoolTailCanaryValue(blk);
}

// Validates an enabled tail canary while the heap lock is held.
static void KiPoolValidateTailCanaryLocked(HEAP* h, BLOCK_HDR* blk, const char* where) {
    if (!blk || !(blk->Flags & BLOCK_FLAG_CANARY))
        return;
    if (blk->RequestedSize + sizeof(uint64_t) > blk->Size)
        KiPoolPanic("invalid canary layout", h, blk, KiBlockUserBase(blk));
    uint64_t actual = *KiBlockTailCanary(blk);
    uint64_t expect = KiPoolTailCanaryValue(blk);
    if (actual != expect)
        KiPoolPanic(where, h, blk, KiBlockUserBase(blk));
}

// Poisons the requested payload before a freed block becomes reusable.
static void KiPoolPoisonFreedPayload(BLOCK_HDR* blk) {
    if (!blk)
        return;
    size_t bytes = blk->RequestedSize;
    if (bytes > blk->Size)
        bytes = blk->Size;
    memset(KiBlockUserBase(blk), POOL_POISON_FREE, bytes);
    blk->Flags |= BLOCK_FLAG_POISONED;
    KiPoolWriteTailCanary(blk);
}

// Detects writes into a poisoned free or quarantined payload.
static void KiPoolValidateFreedPoisonLocked(HEAP* h, BLOCK_HDR* blk, const char* where) {
    if (!blk || !(blk->Flags & BLOCK_FLAG_POISONED))
        return;
    size_t bytes = blk->RequestedSize;
    if (bytes > blk->Size)
        bytes = blk->Size;
    const uint8_t* p = KiBlockUserBase(blk);
    for (size_t i = 0; i < bytes; i++) {
        if (p[i] != POOL_POISON_FREE)
            KiPoolPanic(where, h, blk, KiBlockUserBase(blk));
    }
}

// Rolls back a failed heap grow attempt. KiHeapGrow maps pages before the
// new free block is linked into the allocator; if a later frame allocation or
// page-table allocation fails, those earlier mappings must be returned or the
// pool silently leaks physical frames and leaves stale mappings below Brk.
static void KiHeapRollbackGrow(uintptr_t start, uintptr_t end) {
    for (uintptr_t va = start; va < end; va += PAGE_SIZE) {
        PHYSICAL_ADDRESS frame = MmPteTranslate(va);
        if (!frame)
            continue;
        if (KSUCCESS(MmPteUnmapPage(va)))
            (void)MmPfnFree(frame);
    }
}

static void KiFreeListInsertHeadLocked(HEAP* h, BLOCK_HDR* blk, uintptr_t caller);
static void KiFreeListRemoveLocked(HEAP* h, BLOCK_HDR* blk, uintptr_t caller);

// Extend heap by at least 'need' bytes (rounded up to page boundary).
static bool KiHeapGrow(HEAP* h, size_t need, uintptr_t caller) {
    size_t grow = (need + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
    if (h->Brk + grow > h->Limit) return false;

    uintptr_t start = h->Brk;
    uintptr_t mappedEnd = start;
    for (uintptr_t va = start; va < start + grow; va += PAGE_SIZE) {
        PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseKernel);
        if (!frame) {
            KiHeapRollbackGrow(start, mappedEnd);
            return false;
        }
        if (!KSUCCESS(MmPteMapPage(va, frame, h->MapFlags, MmPteMapFailIfPresent))) {
            MmPfnFree(frame);
            KiHeapRollbackGrow(start, mappedEnd);
            return false;
        }
        mappedEnd = va + PAGE_SIZE;
        memset((void*)va, 0, PAGE_SIZE);
    }

    // The grow path owns the only transition from reserved heap VA to mapped heap
    // pages, so a successful loop above is the heap's committed-frontier proof.
    // Re-querying every heap page on later allocations only burns page-table locks.
    // Missing pool pages are corruption now, not something the allocator repairs.

    uintptr_t oldBrk = h->Brk;
    uintptr_t newBrk = oldBrk + grow;

    // If the current tail block is already free, extend it in-place instead of
    // leaving adjacent free blocks. This matters once tail decommit is enabled:
    // a later grow should rebuild the same tail block, not fragment the arena
    // into many page-sized free runs.
    if (oldBrk > h->Start) {
        BLOCK_FTR* prevFtr = (BLOCK_FTR*)(oldBrk - FTR_SIZE);
        if (prevFtr->Magic == BLOCK_MAGIC && prevFtr->Free == BLOCK_STATE_FREE) {
            BLOCK_HDR* prev = KiFtrToHdr(prevFtr);
            if ((uintptr_t)prev >= h->Start && prev->Magic == BLOCK_MAGIC && prev->Free == BLOCK_STATE_FREE) {
                prev->CoalesceCaller = caller;
                prev->Size += grow;
                h->Brk = newBrk;
                KiHdrToFtr(prev)->Magic = BLOCK_MAGIC;
                KiHdrToFtr(prev)->Free  = BLOCK_STATE_FREE;
                KiHdrToFtr(prev)->Size  = prev->Size;
                return true;
            }
        }
    }

    // Make the new pages a single free block.
    BLOCK_HDR* blk = (BLOCK_HDR*)oldBrk;
    blk->Magic    = BLOCK_MAGIC;
    blk->Free     = BLOCK_STATE_FREE;
    blk->Size     = grow - OVERHEAD;
    blk->PrevFree = (BLOCK_HDR*)0;
    blk->NextFree = (BLOCK_HDR*)0;
    blk->RequestedSize = 0;
    blk->Tag = 0;
    blk->Flags = 0;
    blk->Sequence = 0;
    blk->AllocCaller = 0;
    blk->FreeCaller = 0;
    blk->FreeListInsertCaller = 0;
    blk->FreeListRemoveCaller = 0;
    blk->CoalesceCaller = 0;
    blk->FreeListSequence = 0;

    h->Brk = newBrk;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
    KiHdrToFtr(blk)->Size  = blk->Size;
    KiFreeListInsertHeadLocked(h, blk, caller);

    return true;
}

// Validates a free-list node before the allocator mutates list links.
static void KiFreeListValidateNodeLocked(HEAP* h, BLOCK_HDR* blk, const char* where) {
    KiPoolValidateFooterLocked(h, blk, where);
    if (blk->Free != BLOCK_STATE_FREE)
        KiPoolPanic("free list node is not free", h, blk, KiBlockUserBase(blk));

    if (blk->PrevFree) {
        if (!KiPoolHeaderInCommittedHeap(h, blk->PrevFree))
            KiPoolPanic("free list previous pointer outside heap", h, blk, KiBlockUserBase(blk));
        if (blk->PrevFree->Magic != BLOCK_MAGIC)
            KiPoolPanic("free list previous pointer bad magic", h, blk->PrevFree, KiBlockUserBase(blk->PrevFree));
        if (blk->PrevFree->NextFree != blk)
            KiPoolPanic("free list previous backlink corrupt", h, blk, KiBlockUserBase(blk));
    }

    if (blk->NextFree) {
        if (!KiPoolHeaderInCommittedHeap(h, blk->NextFree))
            KiPoolPanic("free list next pointer outside heap", h, blk, KiBlockUserBase(blk));
        if (blk->NextFree->Magic != BLOCK_MAGIC)
            KiPoolPanic("free list next pointer bad magic", h, blk->NextFree, KiBlockUserBase(blk->NextFree));
        if (blk->NextFree->PrevFree != blk)
            KiPoolPanic("free list next backlink corrupt", h, blk, KiBlockUserBase(blk));
    }
}

// Inserts a free block at the head of the free list and records the mutation site.
static void KiFreeListInsertHeadLocked(HEAP* h, BLOCK_HDR* blk, uintptr_t caller) {
    KiPoolValidateFooterLocked(h, blk, "free list insert block corrupt");
    if (blk->Free != BLOCK_STATE_FREE)
        KiPoolPanic("free list insert of non-free block", h, blk, KiBlockUserBase(blk));
    if (blk->PrevFree || blk->NextFree)
        KiPoolPanic("free list insert of linked block", h, blk, KiBlockUserBase(blk));

    blk->FreeListInsertCaller = caller;
    blk->FreeListSequence = ++h->NextFreeListSequence;
    blk->PrevFree = (BLOCK_HDR*)0;
    blk->NextFree = h->FreeList;

    if (h->FreeList) {
        KiFreeListValidateNodeLocked(h, h->FreeList, "free list head corrupt before insert");
        h->FreeList->PrevFree = blk;
    }
    h->FreeList = blk;
}

// Unlinks a free block from the free list and records the mutation site.
static void KiFreeListRemoveLocked(HEAP* h, BLOCK_HDR* blk, uintptr_t caller) {
    KiFreeListValidateNodeLocked(h, blk, "free list remove block corrupt");

    blk->FreeListRemoveCaller = caller;
    blk->FreeListSequence = ++h->NextFreeListSequence;

    if (blk->PrevFree) blk->PrevFree->NextFree = blk->NextFree;
    else               h->FreeList = blk->NextFree;
    if (blk->NextFree) blk->NextFree->PrevFree = blk->PrevFree;
    blk->PrevFree = blk->NextFree = (BLOCK_HDR*)0;
}

// Implements the heap Try Decommit Tail operation.
static void KiHeapTryDecommitTail(HEAP* h, BLOCK_HDR* blk) {
    if (!h || !blk || blk->Free == BLOCK_STATE_USED) return;

    // Do not eagerly decommit the kernel pool on every free. Large transient
    // syscall bounce buffers, especially SysBlitScreen's frame-sized buffer,
    // are often allocated/freed once per frame. If the tail block is unmapped
    // immediately, the next frame has to map the same pages again, producing
    // massive page-table churn and global TLB shootdowns. Keep the committed
    // pool high-water resident during normal operation; add an explicit idle or
    // pressure pool trim later if needed.
    if (!h->DecommitTailOnFree) {
        KiHdrToFtr(blk)->Size  = blk->Size;
        KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
        KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
        return;
    }

    uintptr_t blockBase = (uintptr_t)blk;
    uintptr_t oldBrk = h->Brk;

    // Only the last free block can return committed heap pages. Interior free
    // blocks must remain mapped because the allocator's boundary tags and free
    // list links are stored inside those blocks.
    if ((uintptr_t)KiNextBlock(blk) != oldBrk) {
        KiHdrToFtr(blk)->Size  = blk->Size;
        KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
        KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
        return;
    }

    // Keep enough mapped space for the tail block's header, footer, and a small
    // usable payload, then decommit whole pages above that. The footer moves to
    // the new Brk so future coalescing can still walk the boundary tag chain.
    uintptr_t minEnd = blockBase + MIN_BLOCK;
    uintptr_t newBrk = (minEnd + PAGE_SIZE - 1) & ~(uintptr_t)(PAGE_SIZE - 1);

    if (newBrk >= oldBrk) {
        KiHdrToFtr(blk)->Size  = blk->Size;
        KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
        KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
        return;
    }

    for (uintptr_t va = newBrk; va < oldBrk; va += PAGE_SIZE) {
        PHYSICAL_ADDRESS frame = MmPteTranslate(va);
        if (!frame)
            continue;
        if (KSUCCESS(MmPteUnmapPage(va)))
            (void)MmPfnFree(frame);
    }

    h->Brk = newBrk;
    blk->Size = newBrk - blockBase - OVERHEAD;
    KiHdrToFtr(blk)->Size  = blk->Size;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
}

// Releases a previously quarantined block into the real free list. Caller holds h->Lock.
static void KiHeapReleaseFreeBlockLocked(HEAP* h, BLOCK_HDR* blk, uintptr_t caller) {
    if (!h || !blk) return;

    KiPoolValidateFooterLocked(h, blk, "release free footer corrupt");
    KiPoolValidateTailCanaryLocked(h, blk, "release free tail canary corrupt");
    KiPoolValidateFreedPoisonLocked(h, blk, "use-after-free payload write detected");

    blk->Free = BLOCK_STATE_FREE;
    blk->RequestedSize = 0;
    blk->Tag = 0;
    blk->Flags = 0;
    blk->AllocCaller = 0;
    // Keep FreeCaller/Sequence for diagnostics until the block is reused.
    KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Size  = blk->Size;

    // Coalesce with next block if it is really free. Quarantined neighbors are
    // intentionally not coalesced; their payload must remain intact for UAF checks.
    BLOCK_HDR* next = KiNextBlock(blk);
    if ((uintptr_t)next < h->Brk) {
        KiPoolValidateFooterLocked(h, next, "next footer corrupt while freeing");
        if (next->Free == BLOCK_STATE_FREE) {
            KiFreeListRemoveLocked(h, next, caller);
            blk->CoalesceCaller = caller;
            next->CoalesceCaller = caller;
            blk->Size += OVERHEAD + next->Size;
            KiHdrToFtr(blk)->Size  = blk->Size;
            KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
            KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
        }
    }

    // Coalesce with previous block if it is really free.
    if ((uintptr_t)blk > h->Start) {
        BLOCK_FTR* prevFtr = KiPrevFtr(blk);
        if (prevFtr->Magic == BLOCK_MAGIC) {
            BLOCK_HDR* prev = KiFtrToHdr(prevFtr);
            KiPoolValidateFooterLocked(h, prev, "previous footer corrupt while freeing");
            if (prev->Free == BLOCK_STATE_FREE) {
                KiFreeListRemoveLocked(h, prev, caller);
                prev->CoalesceCaller = caller;
                blk->CoalesceCaller = caller;
                prev->Size += OVERHEAD + blk->Size;
                KiHdrToFtr(prev)->Size  = prev->Size;
                KiHdrToFtr(prev)->Magic = BLOCK_MAGIC;
                KiHdrToFtr(prev)->Free  = BLOCK_STATE_FREE;
                blk = prev;
            }
        }
    }

    KiFreeListInsertHeadLocked(h, blk, caller);

    KiHdrToFtr(blk)->Free  = BLOCK_STATE_FREE;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Size  = blk->Size;

    KiHeapTryDecommitTail(h, blk);
}

// Releases one quarantine slot into the normal free-list path.
static void KiHeapReleaseQuarantineSlotLocked(HEAP* h, size_t index, uintptr_t caller) {
    if (!h || index >= POOL_QUARANTINE_SLOTS)
        return;
    BLOCK_HDR* old = h->Quarantine[index];
    if (!old)
        return;
    h->Quarantine[index] = (BLOCK_HDR*)0;
    if (h->QuarantineCount)
        h->QuarantineCount--;
    if (old->Free != BLOCK_STATE_QUARANTINE)
        KiPoolPanic("quarantine slot state corrupt", h, old, KiBlockUserBase(old));
    KiHeapReleaseFreeBlockLocked(h, old, caller);
}

// Drains every quarantine slot when allocation pressure requires reuse.
static void KiHeapDrainQuarantineLocked(HEAP* h, uintptr_t caller) {
    if (!h)
        return;
    for (size_t i = 0; i < POOL_QUARANTINE_SLOTS; i++)
        KiHeapReleaseQuarantineSlotLocked(h, i, caller);
    h->QuarantineCursor = 0;
}

// Places a freed block in the bounded quarantine before it can be reused.
static void KiHeapQuarantineBlockLocked(HEAP* h, BLOCK_HDR* blk, uintptr_t caller) {
    if (!h || !blk)
        return;
    size_t idx = h->QuarantineCursor % POOL_QUARANTINE_SLOTS;
    KiHeapReleaseQuarantineSlotLocked(h, idx, caller);

    blk->Free = BLOCK_STATE_QUARANTINE;
    KiHdrToFtr(blk)->Free  = BLOCK_STATE_QUARANTINE;
    KiHdrToFtr(blk)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(blk)->Size  = blk->Size;

    h->Quarantine[idx] = blk;
    h->QuarantineCursor = (idx + 1u) % POOL_QUARANTINE_SLOTS;
    h->QuarantineCount++;
}

// Allocates a block from the heap.
static void* KiHeapAlloc(HEAP* h, size_t requestedSize, uint32_t tag, uintptr_t caller) {
    if (!requestedSize) return (void*)0;

    uint32_t dbg = KiPoolDebugFlags();
    size_t size = requestedSize;
    if (dbg & KE_DEBUG_POOL_CANARY) {
        if (size > (size_t)-1 - sizeof(uint64_t))
            return (void*)0;
        size += sizeof(uint64_t);
    }
    if (size > (size_t)-1 - (POOL_ALIGNMENT - 1u))
        return (void*)0;
    size = (size + (POOL_ALIGNMENT - 1u)) & ~(size_t)(POOL_ALIGNMENT - 1u);

    uintptr_t flags = KeAcquireSpinlock(&h->Lock);

    // Best-fit search.
    BLOCK_HDR* best = (BLOCK_HDR*)0;
    for (BLOCK_HDR* cur = h->FreeList; cur; cur = cur->NextFree) {
        KiPoolValidateFooterLocked(h, cur, "free list block corrupt");
        if (cur->Free != BLOCK_STATE_FREE)
            KiPoolPanic("free list contains non-free block", h, cur, KiBlockUserBase(cur));
        if (cur->Size >= size) {
            if (!best || cur->Size < best->Size) best = cur;
            if (best->Size == size) break;
        }
    }

    if (!best && h->QuarantineCount) {
        KiHeapDrainQuarantineLocked(h, caller);
        for (BLOCK_HDR* cur = h->FreeList; cur; cur = cur->NextFree) {
            KiPoolValidateFooterLocked(h, cur, "free list block corrupt after quarantine drain");
            if (cur->Size >= size) {
                if (!best || cur->Size < best->Size) best = cur;
                if (best->Size == size) break;
            }
        }
    }

    if (!best) {
        if (!KiHeapGrow(h, size + OVERHEAD, caller)) {
            KeReleaseSpinlock(&h->Lock, flags);
            return (void*)0;
        }
        // Re-search; KiHeapGrow inserted at the front.
        best = h->FreeList;
        while (best && best->Size < size) best = best->NextFree;
        if (!best) {
            KeReleaseSpinlock(&h->Lock, flags);
            return (void*)0;
        }
        KiPoolValidateFooterLocked(h, best, "grown free block corrupt");
    }

    KiFreeListRemoveLocked(h, best, caller);

    // Split the block if the remainder is big enough.
    if (best->Size >= size + OVERHEAD + POOL_ALIGNMENT) {
        BLOCK_HDR* rem = (BLOCK_HDR*)((char*)best + HDR_SIZE + size + FTR_SIZE);
        rem->Magic     = BLOCK_MAGIC;
        rem->Free      = BLOCK_STATE_FREE;
        rem->Size      = best->Size - size - OVERHEAD;
        rem->PrevFree  = (BLOCK_HDR*)0;
        rem->NextFree  = (BLOCK_HDR*)0;
        rem->RequestedSize = 0;
        rem->Tag = 0;
        rem->Flags = 0;
        rem->Sequence = 0;
        rem->AllocCaller = 0;
        rem->FreeCaller = 0;
        rem->FreeListInsertCaller = 0;
        rem->FreeListRemoveCaller = 0;
        rem->CoalesceCaller = 0;
        rem->FreeListSequence = 0;
        KiHdrToFtr(rem)->Magic = BLOCK_MAGIC;
        KiHdrToFtr(rem)->Free  = BLOCK_STATE_FREE;
        KiHdrToFtr(rem)->Size  = rem->Size;
        KiFreeListInsertHeadLocked(h, rem, caller);

        best->Size = size;
    }

    best->Free = BLOCK_STATE_USED;
    best->RequestedSize = requestedSize;
    best->Tag = tag ? tag : POOL_TAG_DEFAULT;
    best->Flags = 0;
    best->Sequence = ++h->NextSequence;
    best->AllocCaller = (dbg & KE_DEBUG_POOL_CALLERS) ? caller : 0;
    best->FreeCaller = 0;
    best->CoalesceCaller = 0;
    if (dbg & KE_DEBUG_POOL_CANARY)
        best->Flags |= BLOCK_FLAG_CANARY;

    KiHdrToFtr(best)->Free  = BLOCK_STATE_USED;
    KiHdrToFtr(best)->Magic = BLOCK_MAGIC;
    KiHdrToFtr(best)->Size  = best->Size;

    memset(KiBlockUserBase(best), 0, requestedSize);
    if (best->Size > requestedSize)
        memset(KiBlockUserBase(best) + requestedSize, 0, best->Size - requestedSize);
    KiPoolWriteTailCanary(best);

    KeReleaseSpinlock(&h->Lock, flags);
    return KiBlockUserBase(best);
}

// Releases a block back to the heap.
static void KiHeapFree(HEAP* h, void* ptr, uintptr_t caller) {
    if (!ptr) return;

    uintptr_t u = (uintptr_t)ptr;
    if (!h || u < h->Start + HDR_SIZE || u >= h->Brk)
        KiPoolPanic("invalid free pointer outside heap", h, (BLOCK_HDR*)0, ptr);

    BLOCK_HDR* blk = (BLOCK_HDR*)(u - HDR_SIZE);

    uintptr_t flags = KeAcquireSpinlock(&h->Lock);

    if (blk->Magic != BLOCK_MAGIC)
        KiPoolPanic("invalid free pointer/bad header magic", h, blk, ptr);
    KiPoolValidateFooterLocked(h, blk, "allocated footer corrupt on free");
    if (blk->Free == BLOCK_STATE_FREE)
        KiPoolPanic("double free", h, blk, ptr);
    if (blk->Free == BLOCK_STATE_QUARANTINE)
        KiPoolPanic("double free/quarantined UAF", h, blk, ptr);
    if (blk->Free != BLOCK_STATE_USED)
        KiPoolPanic("invalid block state on free", h, blk, ptr);

    KiPoolValidateTailCanaryLocked(h, blk, "tail canary corrupt on free");

    uint32_t dbg = KiPoolDebugFlags();
    blk->FreeCaller = (dbg & KE_DEBUG_POOL_CALLERS) ? caller : 0;
    if (dbg & KE_DEBUG_POOL_POISON)
        KiPoolPoisonFreedPayload(blk);

    if (dbg & KE_DEBUG_POOL_QUARANTINE) {
        KiHeapQuarantineBlockLocked(h, blk, caller);
    } else {
        KiHeapReleaseFreeBlockLocked(h, blk, caller);
    }

    KeReleaseSpinlock(&h->Lock, flags);
}

// Resizes a heap allocation while preserving its contents.
static void* KiHeapRealloc(HEAP* h, void* ptr, size_t newSize, uintptr_t caller) {
    // If ptr is NULL, realloc acts like malloc.
    if (!ptr) {
        return KiHeapAlloc(h, newSize, POOL_TAG_DEFAULT, caller);
    }

    // If size is 0, realloc acts like free.
    if (newSize == 0) {
        KiHeapFree(h, ptr, caller);
        return (void*)0;
    }

    uintptr_t u = (uintptr_t)ptr;
    if (!h || u < h->Start + HDR_SIZE || u >= h->Brk)
        KiPoolPanic("invalid realloc pointer outside heap", h, (BLOCK_HDR*)0, ptr);

    BLOCK_HDR* blk = (BLOCK_HDR*)(u - HDR_SIZE);

    uintptr_t flags = KeAcquireSpinlock(&h->Lock);
    if (blk->Magic != BLOCK_MAGIC)
        KiPoolPanic("invalid realloc pointer/bad header magic", h, blk, ptr);
    KiPoolValidateFooterLocked(h, blk, "allocated footer corrupt on realloc");
    if (blk->Free != BLOCK_STATE_USED)
        KiPoolPanic("realloc of non-allocated block", h, blk, ptr);
    KiPoolValidateTailCanaryLocked(h, blk, "tail canary corrupt on realloc");

    size_t oldRequested = blk->RequestedSize;
    size_t oldCopy = oldRequested < newSize ? oldRequested : newSize;
    if (newSize <= oldRequested) {
        blk->RequestedSize = newSize;
        if (blk->Flags & BLOCK_FLAG_CANARY)
            KiPoolWriteTailCanary(blk);
        KeReleaseSpinlock(&h->Lock, flags);
        return ptr;
    }
    KeReleaseSpinlock(&h->Lock, flags);

    void* newPtr = KiHeapAlloc(h, newSize, POOL_TAG_DEFAULT, caller);
    if (!newPtr)
        return (void*)0;
    memcpy(newPtr, ptr, oldCopy);
    KiHeapFree(h, ptr, caller);
    return newPtr;
}

// Kernel pool.
static HEAP gPoolHeap;

typedef struct KI_POOL_VALIDATE_STATS {
    uint64_t Blocks;
    uint64_t FreeBlocks;
    uint64_t FreeListBlocks;
    uint64_t QuarantineBlocks;
    size_t UsedBytes;
    size_t FreeBytes;
} KI_POOL_VALIDATE_STATS;

// Validates the heap boundary-tag walk and free-list topology while the heap is locked.
static KSTATUS KiPoolValidateHeapLocked(HEAP* h, const char* reason, KI_POOL_VALIDATE_STATS* out) {
    KI_POOL_VALIDATE_STATS stats;
    memset(&stats, 0, sizeof(stats));

    if (!h || h->Start == 0 || h->Limit == 0)
        return NTLSTATUS_NOT_SUPPORTED;
    if (h->Brk < h->Start || h->Brk > h->Limit)
        KiPoolPanic("heap break outside limits", h, (BLOCK_HDR*)0, (void*)0);
    if (h->Brk == h->Start) {
        if (out) *out = stats;
        return NTLSTATUS_SUCCESS;
    }

    uintptr_t cursor = h->Start;
    while (cursor < h->Brk) {
        BLOCK_HDR* blk = (BLOCK_HDR*)cursor;
        KiPoolValidateFooterLocked(h, blk, reason ? reason : "heap block corrupt");
        if (blk->Free > BLOCK_STATE_QUARANTINE)
            KiPoolPanic("heap block has invalid state", h, blk, KiBlockUserBase(blk));

        stats.Blocks++;
        if (blk->Free == BLOCK_STATE_USED) {
            stats.UsedBytes += blk->RequestedSize;
            KiPoolValidateTailCanaryLocked(h, blk, "heap used tail canary corrupt");
        } else if (blk->Free == BLOCK_STATE_FREE) {
            stats.FreeBlocks++;
            stats.FreeBytes += blk->Size;
            KiPoolValidateFreedPoisonLocked(h, blk, "heap free payload poison corrupt");
        } else {
            stats.QuarantineBlocks++;
            KiPoolValidateFreedPoisonLocked(h, blk, "heap quarantine payload poison corrupt");
        }

        BLOCK_HDR* next = KiNextBlock(blk);
        if ((uintptr_t)next <= cursor)
            KiPoolPanic("heap block walk did not advance", h, blk, KiBlockUserBase(blk));
        cursor = (uintptr_t)next;
    }

    if (cursor != h->Brk)
        KiPoolPanic("heap block walk missed break", h, (BLOCK_HDR*)0, (void*)cursor);

    BLOCK_HDR* prev = (BLOCK_HDR*)0;
    uint64_t guard = stats.Blocks + 1u;
    for (BLOCK_HDR* cur = h->FreeList; cur; cur = cur->NextFree) {
        if (stats.FreeListBlocks >= guard)
            KiPoolPanic("free list cycle detected", h, cur, KiBlockUserBase(cur));
        stats.FreeListBlocks++;
        KiFreeListValidateNodeLocked(h, cur, reason ? reason : "free list corrupt");
        if (cur->PrevFree != prev)
            KiPoolPanic("free list previous chain corrupt", h, cur, KiBlockUserBase(cur));
        prev = cur;
    }

    if (stats.FreeListBlocks != stats.FreeBlocks)
        KiPoolPanic("free-list/free-block count mismatch", h, h->FreeList, h->FreeList ? KiBlockUserBase(h->FreeList) : (void*)0);

    if (out) *out = stats;
    return NTLSTATUS_SUCCESS;
}

// Validates the kernel pool and prints a compact checkpoint line.
KSTATUS MmValidateKernelPool(const char* reason) {
    uintptr_t flags = KeAcquireSpinlock(&gPoolHeap.Lock);
    KI_POOL_VALIDATE_STATS stats;
    KSTATUS status = KiPoolValidateHeapLocked(&gPoolHeap, reason, &stats);
    uintptr_t brk = gPoolHeap.Brk;
    uintptr_t start = gPoolHeap.Start;
    KeReleaseSpinlock(&gPoolHeap.Lock, flags);

    if (KSUCCESS(status)) {
        DbgPrint("MM: pool checkpoint '%s' ok blocks=%llu free=%llu free-list=%llu quarantine=%llu used=%llu free-bytes=%llu brk=0x%016llX usage=%llu KiB\n",
                 reason ? reason : "pool",
                 (unsigned long long)stats.Blocks,
                 (unsigned long long)stats.FreeBlocks,
                 (unsigned long long)stats.FreeListBlocks,
                 (unsigned long long)stats.QuarantineBlocks,
                 (unsigned long long)stats.UsedBytes,
                 (unsigned long long)stats.FreeBytes,
                 (unsigned long long)brk,
                 (unsigned long long)((brk - start) / 1024));
    }

    return status;
}

// Initializes the kernel pool and optionally commits its initial extent.
void MmPoolInitialize(size_t poolSize) {
    gPoolHeap.FreeList = (BLOCK_HDR*)0;
    gPoolHeap.Start = POOL_BASE;
    gPoolHeap.Brk = POOL_BASE;
    gPoolHeap.Limit = POOL_LIMIT;
    gPoolHeap.MapFlags = PAGE_PRESENT | PAGE_WRITE | PAGE_NX;
    gPoolHeap.Lock = (SPINLOCK)SPINLOCK_INIT;
    gPoolHeap.DecommitTailOnFree = 0;
    gPoolHeap.NextSequence = 0;
    gPoolHeap.NextFreeListSequence = 0;
    gPoolHeap.QuarantineCursor = 0;
    gPoolHeap.QuarantineCount = 0;
    memset(gPoolHeap.Quarantine, 0, sizeof(gPoolHeap.Quarantine));

    if (poolSize) {
        poolSize = (poolSize + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
        (void)KiHeapGrow(&gPoolHeap, poolSize, 0);
    }
}

// Allocates zeroed kernel-pool memory from the non-user heap.
__declspec(dllexport) void* MmAllocatePool(size_t size) {
    return KiHeapAlloc(&gPoolHeap, size, POOL_TAG_DEFAULT, (uintptr_t)_ReturnAddress());
}

// Allocates zeroed kernel-pool memory with a diagnostic tag.
__declspec(dllexport) void* MmAllocatePoolTag(size_t size, uint32_t tag) {
    return KiHeapAlloc(&gPoolHeap, size, tag, (uintptr_t)_ReturnAddress());
}

// Releases kernel-pool memory back to the non-user heap.
__declspec(dllexport) void MmFreePool(void* ptr) {
    KiHeapFree(&gPoolHeap, ptr, (uintptr_t)_ReturnAddress());
}

// Reallocates kernel-pool memory while preserving existing contents.
__declspec(dllexport) void* MmReallocatePool(void* ptr, size_t newSize) {
    return KiHeapRealloc(&gPoolHeap, ptr, newSize, (uintptr_t)_ReturnAddress());
}

// Describes a virtual address that falls inside the kernel-pool aperture.
void MmPoolDebugDescribeAddress(uintptr_t addr) {
    DbgPrint("MM: kernel pool start=%016llx brk=%016llx limit=%016llx usage=%llu KiB\n",
             (unsigned long long)gPoolHeap.Start,
             (unsigned long long)gPoolHeap.Brk,
             (unsigned long long)gPoolHeap.Limit,
             (unsigned long long)((gPoolHeap.Brk - gPoolHeap.Start) / 1024));

    if (addr < gPoolHeap.Start || addr >= gPoolHeap.Limit) return;
    if (addr >= gPoolHeap.Brk) {
        DbgPrint("MM: address %016llx is inside pool VA range but above committed brk by +0x%llx\n",
                 (unsigned long long)addr,
                 (unsigned long long)(addr - gPoolHeap.Brk));
        return;
    }

    uintptr_t flags = KeAcquireSpinlock(&gPoolHeap.Lock);
    for (BLOCK_HDR* cur = (BLOCK_HDR*)gPoolHeap.Start;
         (uintptr_t)cur < gPoolHeap.Brk;
         cur = KiNextBlock(cur)) {
        if (cur->Magic != BLOCK_MAGIC) {
            DbgPrint("MM: pool walk corrupt at hdr=0x%016llX magic=%08x\n",
                     (unsigned long long)(uintptr_t)cur,
                     cur->Magic);
            break;
        }
        uintptr_t hdr = (uintptr_t)cur;
        uintptr_t data = hdr + HDR_SIZE;
        uintptr_t end = data + cur->Size;
        uintptr_t ftr = end;
        BLOCK_FTR* footer = KiHdrToFtr(cur);
        if ((uintptr_t)footer + FTR_SIZE > gPoolHeap.Brk ||
            footer->Magic != BLOCK_MAGIC || footer->Size != cur->Size || footer->Free != cur->Free) {
            DbgPrint("MM: pool walk corrupt footer hdr=0x%016llX state=%u size=%llx footer=0x%016llX fmagic=%08x fstate=%u fsize=%llx\n",
                     (unsigned long long)(uintptr_t)cur,
                     cur->Free,
                     (unsigned long long)cur->Size,
                     (unsigned long long)(uintptr_t)footer,
                     ((uintptr_t)footer + FTR_SIZE <= gPoolHeap.Brk) ? footer->Magic : 0,
                     ((uintptr_t)footer + FTR_SIZE <= gPoolHeap.Brk) ? footer->Free : 0,
                     ((uintptr_t)footer + FTR_SIZE <= gPoolHeap.Brk) ? (unsigned long long)footer->Size : 0ULL);
            break;
        }
        if (addr >= hdr && addr < ftr + FTR_SIZE) {
            DbgPrint("MM: address 0x%016llX in pool block hdr=0x%016llX data=0x%016llX-0x%016llX size=%llx req=%llx state=%u(%s) tag=%08x flags=%08x seq=%llu alloc=0x%016llX free=0x%016llX fl_ins=0x%016llX fl_rem=0x%016llX coal=0x%016llX fl_seq=%llu\n",
                     (unsigned long long)addr,
                     (unsigned long long)hdr,
                     (unsigned long long)data,
                     (unsigned long long)end,
                     (unsigned long long)cur->Size,
                     (unsigned long long)cur->RequestedSize,
                     cur->Free,
                     KiPoolStateName(cur->Free),
                     cur->Tag,
                     cur->Flags,
                     (unsigned long long)cur->Sequence,
                     (unsigned long long)cur->AllocCaller,
                     (unsigned long long)cur->FreeCaller,
                     (unsigned long long)cur->FreeListInsertCaller,
                     (unsigned long long)cur->FreeListRemoveCaller,
                     (unsigned long long)cur->CoalesceCaller,
                     (unsigned long long)cur->FreeListSequence);
            break;
        }
        if (KiNextBlock(cur) <= cur) {
            DbgPrint("MM: pool walk stopped: non-advancing block hdr=0x%016llX size=%llx\n",
                     (unsigned long long)(uintptr_t)cur,
                     (unsigned long long)cur->Size);
            break;
        }
    }
    KeReleaseSpinlock(&gPoolHeap.Lock, flags);
}

// Returns the currently committed kernel-pool bytes.
size_t MmGetKernelPoolUsage(void) {
    return (size_t)(gPoolHeap.Brk - gPoolHeap.Start);
}