// Image section objects. This file owns parsed PE image identity, sparse prototype pages, and image-section caching.

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>
#include <kernel/ldr_section.h>
#include <kernel/mm_virt.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_pte.h>
#include <kernel/ob_object.h>
#include <kernel/io_file.h>
#include <kernel/io_volume.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_debug.h>

#define PE_SIGNATURE                      0x00004550UL
#define IMAGE_FILE_MACHINE_AMD64          0x8664
#define IMAGE_FILE_MACHINE_I386           0x014c
#define IMAGE_NT_OPTIONAL_HDR32_MAGIC     0x10b
#define IMAGE_NT_OPTIONAL_HDR64_MAGIC     0x20b
#define IMAGE_SCN_CNT_CODE                0x00000020u
#define IMAGE_SCN_CNT_INITIALIZED_DATA    0x00000040u
#define IMAGE_SCN_CNT_UNINITIALIZED_DATA  0x00000080u
#define IMAGE_SCN_MEM_EXECUTE             0x20000000u
#define IMAGE_SCN_MEM_READ                0x40000000u
#define IMAGE_SCN_MEM_WRITE               0x80000000u

#define LDR_IMAGE_SECTION_CACHE_MAX       32u

#pragma pack(push, 1)
typedef struct LDR_OPT_HEADER64 {
    uint16_t Magic;
    uint8_t  MajorLinkerVersion;
    uint8_t  MinorLinkerVersion;
    uint32_t SizeOfCode;
    uint32_t SizeOfInitializedData;
    uint32_t SizeOfUninitializedData;
    uint32_t AddressOfEntryPoint;
    uint32_t BaseOfCode;
    uint64_t ImageBase;
    uint32_t SectionAlignment;
    uint32_t FileAlignment;
    uint16_t MajorOSVersion;
    uint16_t MinorOSVersion;
    uint16_t MajorImageVersion;
    uint16_t MinorImageVersion;
    uint16_t MajorSubsystemVersion;
    uint16_t MinorSubsystemVersion;
    uint32_t Win32VersionValue;
    uint32_t SizeOfImage;
    uint32_t SizeOfHeaders;
    uint32_t CheckSum;
    uint16_t Subsystem;
    uint16_t DllCharacteristics;
    uint64_t SizeOfStackReserve;
    uint64_t SizeOfStackCommit;
    uint64_t SizeOfHeapReserve;
    uint64_t SizeOfHeapCommit;
    uint32_t LoaderFlags;
    uint32_t NumberOfRvaAndSizes;
    struct { uint32_t VirtualAddress; uint32_t Size; } DataDirectory[16];
} LDR_OPT_HEADER64;

typedef struct LDR_PE_HEADER {
    uint32_t Signature;
    uint16_t Machine;
    uint16_t NumberOfSections;
    uint32_t TimeDateStamp;
    uint32_t PointerToSymbolTable;
    uint32_t NumberOfSymbols;
    uint16_t SizeOfOptionalHeader;
    uint16_t Characteristics;
    LDR_OPT_HEADER64 OptionalHeader;
} LDR_PE_HEADER;

typedef struct LDR_SECTION_HEADER {
    char     Name[8];
    uint32_t VirtualSize;
    uint32_t VirtualAddress;
    uint32_t SizeOfRawData;
    uint32_t PointerToRawData;
    uint32_t PointerToRelocations;
    uint32_t PointerToLinenumbers;
    uint16_t NumberOfRelocations;
    uint16_t NumberOfLinenumbers;
    uint32_t Characteristics;
} LDR_SECTION_HEADER;
#pragma pack(pop)

struct IMAGE_SECTION_OBJECT {
    OBJECT Header;
    IMAGE_SECTION_METADATA Metadata;
    IMAGE_SECTION_DESCRIPTOR* Sections;
    PHYSICAL_ADDRESS* Frames;
    FILE_OBJECT* BackingFile;
    size_t HeadersFileSize;
    SPINLOCK Lock;
};

static IMAGE_SECTION_OBJECT* gImageSectionCache[LDR_IMAGE_SECTION_CACHE_MAX];
static SPINLOCK gImageSectionCacheLock = SPINLOCK_INIT;

// Aligns a PE image size to the hardware page size.
static size_t KiImageSectionAlignSize(size_t value) {
    return (value + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
}

// Mixes filesystem identity into a compact cache key.
static uint64_t KiImageSectionIdentityKey(const IO_FILE_IDENTITY* identity) {
    if (!identity)
        return 0;

    uint64_t h = 1469598103934665603ULL;
    h ^= identity->VolumeId;
    h *= 1099511628211ULL;
    h ^= identity->FileId;
    h *= 1099511628211ULL;
    h ^= identity->Generation;
    h *= 1099511628211ULL;
    h ^= identity->Size;
    h *= 1099511628211ULL;
    return h ? h : 1;
}

// Converts PE section characteristics to the final user page protection union for that section.
static uint64_t KiImageSectionProtection(uint32_t characteristics) {
    uint64_t prot = PAGE_PRESENT | PAGE_USER | PAGE_NX;
    if (characteristics & IMAGE_SCN_MEM_WRITE)
        prot |= PAGE_WRITE;
    if (characteristics & IMAGE_SCN_MEM_EXECUTE)
        prot &= ~PAGE_NX;
    return prot;
}

// Computes the required image mapping size after validating raw section extents.
static KSTATUS KiImageSectionComputeSize(const LDR_PE_HEADER* pe, size_t fileSize, size_t* outSize) {
    if (!pe || !outSize) return NTLSTATUS_INVALID_PARAMETER;

    const LDR_OPT_HEADER64* opt = &pe->OptionalHeader;
    size_t required = KiImageSectionAlignSize(opt->SizeOfImage);
    size_t headers = KiImageSectionAlignSize(opt->SizeOfHeaders);
    if (headers > required) required = headers;

    const LDR_SECTION_HEADER* sec = (const LDR_SECTION_HEADER*)((const uint8_t*)&pe->OptionalHeader + pe->SizeOfOptionalHeader);
    for (uint32_t i = 0; i < pe->NumberOfSections; i++) {
        size_t va = (size_t)sec[i].VirtualAddress;
        size_t virt = (size_t)sec[i].VirtualSize;
        size_t raw = (size_t)sec[i].SizeOfRawData;
        size_t span = virt ? virt : raw;
        if (raw > span) span = raw;
        if (!span) continue;

        if (va > (size_t)-1 - span) return NTLSTATUS_NO_MEMORY;
        size_t end = KiImageSectionAlignSize(va + span);
        if (end > required) required = end;

        if (raw) {
            size_t rawBase = (size_t)sec[i].PointerToRawData;
            if (rawBase > fileSize || raw > fileSize - rawBase)
                return NTLSTATUS_FAILURE;
        }
    }

    if (!required) return NTLSTATUS_FAILURE;
    *outSize = required;
    return NTLSTATUS_SUCCESS;
}

// Parses and validates the PE header that backs an image section object.
static KSTATUS KiImageSectionParseHeader(const void* buffer,
                                         size_t bufferSize,
                                         size_t fileSize,
                                         const LDR_PE_HEADER** peOut,
                                         const LDR_SECTION_HEADER** sectionsOut,
                                         size_t* imageSizeOut) {
    if (!buffer || !bufferSize || !peOut || !sectionsOut || !imageSizeOut)
        return NTLSTATUS_INVALID_PARAMETER;
    if (bufferSize < 0x40) return NTLSTATUS_FAILURE;

    const uint8_t* bytes = (const uint8_t*)buffer;
    uint32_t peOffset = *(const uint32_t*)(bytes + 0x3C);
    if ((size_t)peOffset > bufferSize || bufferSize - (size_t)peOffset < 24)
        return NTLSTATUS_FAILURE;

    const LDR_PE_HEADER* pe = (const LDR_PE_HEADER*)(bytes + peOffset);
    if (pe->Signature != PE_SIGNATURE)
        return NTLSTATUS_FAILURE;

    uint16_t optionalSize = pe->SizeOfOptionalHeader;
    if (optionalSize < sizeof(LDR_OPT_HEADER64) ||
        (size_t)peOffset + 24 + optionalSize > bufferSize) {
        return NTLSTATUS_FAILURE;
    }

    uint16_t optionalMagic = pe->OptionalHeader.Magic;
    if (pe->Machine == IMAGE_FILE_MACHINE_I386 || optionalMagic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
        DbgPrint("LdrImageSection: 32-bit PE images are not supported (machine=0x%X magic=0x%X)\n",
                 pe->Machine, optionalMagic);
        return NTLSTATUS_NOT_SUPPORTED;
    }

    if (pe->Machine != IMAGE_FILE_MACHINE_AMD64 || optionalMagic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
        DbgPrint("LdrImageSection: unsupported PE image (machine=0x%X magic=0x%X)\n",
                 pe->Machine, optionalMagic);
        return NTLSTATUS_NOT_SUPPORTED;
    }

    if (!pe->NumberOfSections)
        return NTLSTATUS_FAILURE;

    const LDR_SECTION_HEADER* sections = (const LDR_SECTION_HEADER*)((const uint8_t*)&pe->OptionalHeader + pe->SizeOfOptionalHeader);
    size_t sectionTableOffset = (size_t)peOffset + 24u + (size_t)pe->SizeOfOptionalHeader;
    size_t sectionTableBytes = (size_t)pe->NumberOfSections * sizeof(LDR_SECTION_HEADER);
    if (sectionTableOffset > bufferSize || sectionTableBytes > bufferSize - sectionTableOffset)
        return NTLSTATUS_FAILURE;

    size_t imageSize = 0;
    KSTATUS s = KiImageSectionComputeSize(pe, fileSize, &imageSize);
    if (!KSUCCESS(s)) return s;

    *peOut = pe;
    *sectionsOut = sections;
    *imageSizeOut = imageSize;
    return NTLSTATUS_SUCCESS;
}

// Releases all resources owned by an image section object.
static void KiDestroyImageSectionObject(OBJECT* object) {
    IMAGE_SECTION_OBJECT* section = (IMAGE_SECTION_OBJECT*)object;
    if (!section) return;

    if (section->Frames) {
        for (size_t i = 0; i < section->Metadata.PageCount; ++i) {
            if (section->Frames[i])
                (void)MmPfnFree(section->Frames[i]);
        }
        MmFreePool(section->Frames);
        section->Frames = (PHYSICAL_ADDRESS*)0;
    }

    if (section->Sections) {
        MmFreePool(section->Sections);
        section->Sections = (IMAGE_SECTION_DESCRIPTOR*)0;
    }

    if (section->BackingFile) {
        IoCloseFile(section->BackingFile);
        section->BackingFile = (FILE_OBJECT*)0;
    }
}

// Finds a cached image section while the image-section cache lock is held.
static IMAGE_SECTION_OBJECT* KiFindImageSectionLocked(const IO_FILE_IDENTITY* identity, uint64_t key) {
    if (!identity || !key)
        return (IMAGE_SECTION_OBJECT*)0;

    for (uint32_t i = 0; i < LDR_IMAGE_SECTION_CACHE_MAX; ++i) {
        IMAGE_SECTION_OBJECT* section = gImageSectionCache[i];
        if (section &&
            section->Metadata.ContentKey == key &&
            section->Metadata.VolumeId == identity->VolumeId &&
            section->Metadata.FileId == identity->FileId &&
            section->Metadata.FileGeneration == identity->Generation &&
            section->Metadata.RawSize == identity->Size) {
            return section;
        }
    }
    return (IMAGE_SECTION_OBJECT*)0;
}

// Verifies that a cached image section is still backed by the same file generation.
static KSTATUS KiValidateImageSectionBackingFile(IMAGE_SECTION_OBJECT* section) {
    if (!section || !section->BackingFile)
        return NTLSTATUS_INVALID_PARAMETER;

    IO_FILE_IDENTITY identity;
    KSTATUS s = IoQueryFileIdentity(section->BackingFile, &identity);
    if (!KSUCCESS(s))
        return s;

    if (identity.VolumeId != section->Metadata.VolumeId ||
        identity.FileId != section->Metadata.FileId ||
        identity.Generation != section->Metadata.FileGeneration ||
        identity.Size != section->Metadata.RawSize) {
        return NTLSTATUS_CHECKSUM_MISMATCH;
    }

    return NTLSTATUS_SUCCESS;
}

// Returns true when a virtual image page has initialized bytes in the backing file.
static bool KiImageSectionPageHasFileBytes(const IMAGE_SECTION_OBJECT* section, size_t pageIndex) {
    if (!section || pageIndex >= section->Metadata.PageCount)
        return false;

    size_t pageStart = pageIndex * PAGE_SIZE;
    size_t pageEnd = pageStart + PAGE_SIZE;
    if (pageStart < section->HeadersFileSize && pageEnd > 0)
        return true;

    for (uint32_t i = 0; i < section->Metadata.SectionCount; ++i) {
        const IMAGE_SECTION_DESCRIPTOR* desc = &section->Sections[i];
        size_t raw = desc->RawDataSize;
        if (desc->VirtualSize && desc->VirtualSize < raw)
            raw = desc->VirtualSize;
        if (!raw)
            continue;

        size_t start = desc->VirtualAddress;
        size_t end = start + raw;
        if (end < start)
            continue;
        if (pageStart < end && pageEnd > start)
            return true;
    }

    return false;
}

// Copies a file range into an image prototype page, leaving absent bytes zero-filled.
static KSTATUS KiReadImageSectionFileBytes(IMAGE_SECTION_OBJECT* section,
                                           uint64_t fileOffset,
                                           void* dst,
                                           size_t size) {
    if (!section || !section->BackingFile || (!dst && size))
        return NTLSTATUS_INVALID_PARAMETER;
    if (!size)
        return NTLSTATUS_SUCCESS;

    // Image prototype pages are the canonical cache for executable image bytes.
    // Read the backing filesystem directly so the fault path has one cache owner.
    if (!IoVolumeIsMounted(section->BackingFile->Volume))
        return NTLSTATUS_OBJECT_DEAD;

    size_t got = 0;
    KSTATUS s = IoVolumeRead(section->BackingFile->Volume,
                             section->BackingFile->FileCtx,
                             fileOffset,
                             dst,
                             size,
                             &got);
    if (!KSUCCESS(s))
        return s;
    if (got != size)
        return NTLSTATUS_FAILURE;
    return NTLSTATUS_SUCCESS;
}

// Fills one prototype frame from the image backing file.
static KSTATUS KiFillImageSectionFrameFromFile(IMAGE_SECTION_OBJECT* section,
                                               size_t pageIndex,
                                               PHYSICAL_ADDRESS frame) {
    if (!section || !frame || pageIndex >= section->Metadata.PageCount)
        return NTLSTATUS_INVALID_PARAMETER;

    uint8_t* page = (uint8_t*)MmPfnToVirtual(frame);
    if (!page)
        return NTLSTATUS_FAILURE;
    memset(page, 0, PAGE_SIZE);

    size_t pageStart = pageIndex * PAGE_SIZE;
    size_t pageEnd = pageStart + PAGE_SIZE;

    if (pageStart < section->HeadersFileSize) {
        size_t copyStart = pageStart;
        size_t copyEnd = pageEnd;
        if (copyEnd > section->HeadersFileSize)
            copyEnd = section->HeadersFileSize;
        KSTATUS s = KiReadImageSectionFileBytes(section,
                                                (uint64_t)copyStart,
                                                page,
                                                copyEnd - copyStart);
        if (!KSUCCESS(s))
            return s;
    }

    for (uint32_t i = 0; i < section->Metadata.SectionCount; ++i) {
        const IMAGE_SECTION_DESCRIPTOR* desc = &section->Sections[i];
        size_t raw = desc->RawDataSize;
        if (desc->VirtualSize && desc->VirtualSize < raw)
            raw = desc->VirtualSize;
        if (!raw)
            continue;

        size_t sectionStart = desc->VirtualAddress;
        size_t sectionEnd = sectionStart + raw;
        if (sectionEnd < sectionStart)
            return NTLSTATUS_FAILURE;
        if (pageStart >= sectionEnd || pageEnd <= sectionStart)
            continue;

        size_t copyStart = pageStart > sectionStart ? pageStart : sectionStart;
        size_t copyEnd = pageEnd < sectionEnd ? pageEnd : sectionEnd;
        size_t dstOffset = copyStart - pageStart;
        uint64_t srcOffset = (uint64_t)desc->RawDataOffset + (uint64_t)(copyStart - sectionStart);

        KSTATUS s = KiReadImageSectionFileBytes(section,
                                                srcOffset,
                                                page + dstOffset,
                                                copyEnd - copyStart);
        if (!KSUCCESS(s))
            return s;
    }

    return NTLSTATUS_SUCCESS;
}

// Ensures a prototype page exists for initialized image bytes.
static KSTATUS KiEnsureImageSectionFrame(IMAGE_SECTION_OBJECT* section, size_t pageIndex) {
    if (!section || pageIndex >= section->Metadata.PageCount)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&section->Lock);
    PHYSICAL_ADDRESS existing = section->Frames[pageIndex];
    KeReleaseSpinlock(&section->Lock, flags);
    if (existing)
        return NTLSTATUS_SUCCESS;

    if (section->BackingFile && !KiImageSectionPageHasFileBytes(section, pageIndex))
        return NTLSTATUS_SUCCESS;

    PHYSICAL_ADDRESS frame = MmPfnAllocate(MmPfnUseImage);
    if (!frame)
        return NTLSTATUS_NO_MEMORY;

    void* va = MmPfnToVirtual(frame);
    if (!va) {
        MmPfnFree(frame);
        return NTLSTATUS_FAILURE;
    }

    KSTATUS s = NTLSTATUS_SUCCESS;
    if (section->BackingFile) {
        s = KiValidateImageSectionBackingFile(section);
        if (KSUCCESS(s))
            s = KiFillImageSectionFrameFromFile(section, pageIndex, frame);
    } else {
        memset(va, 0, PAGE_SIZE);
    }
    if (!KSUCCESS(s)) {
        MmPfnFree(frame);
        return s;
    }

    flags = KeAcquireSpinlock(&section->Lock);
    if (section->Frames[pageIndex]) {
        KeReleaseSpinlock(&section->Lock, flags);
        MmPfnFree(frame);
        return NTLSTATUS_SUCCESS;
    }
    section->Frames[pageIndex] = frame;
    KeReleaseSpinlock(&section->Lock, flags);
    return NTLSTATUS_SUCCESS;
}

// Builds an uncached image section object backed directly by an open file.
static KSTATUS KiBuildImageSectionObjectFromFile(FILE_OBJECT* file,
                                                 const void* headerBuffer,
                                                 size_t headerBufferSize,
                                                 const IO_FILE_IDENTITY* identity,
                                                 uint64_t key,
                                                 IMAGE_SECTION_OBJECT** sectionOut) {
    if (!file || !headerBuffer || !headerBufferSize || !identity || !sectionOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *sectionOut = (IMAGE_SECTION_OBJECT*)0;

    size_t fileSize = (size_t)identity->Size;
    const LDR_PE_HEADER* pe = (const LDR_PE_HEADER*)0;
    const LDR_SECTION_HEADER* rawSections = (const LDR_SECTION_HEADER*)0;
    size_t imageSize = 0;
    KSTATUS s = KiImageSectionParseHeader(headerBuffer,
                                          headerBufferSize,
                                          fileSize,
                                          &pe,
                                          &rawSections,
                                          &imageSize);
    if (!KSUCCESS(s))
        return s;

    size_t pageCount = (imageSize + PAGE_SIZE - 1) / PAGE_SIZE;
    if (!pageCount || pageCount > ((size_t)-1 / sizeof(PHYSICAL_ADDRESS)))
        return NTLSTATUS_FAILURE;

    IMAGE_SECTION_OBJECT* section = (IMAGE_SECTION_OBJECT*)ObAllocateObject(ObjectTypeImageSection,
                                                                            sizeof(IMAGE_SECTION_OBJECT));
    if (!section)
        return NTLSTATUS_NO_MEMORY;
    section->Header.Destroy = KiDestroyImageSectionObject;

    PHYSICAL_ADDRESS* frames = (PHYSICAL_ADDRESS*)MmAllocatePoolTag(pageCount * sizeof(PHYSICAL_ADDRESS),
                                                                    MM_POOL_TAG('i', 's', 'f', 'r'));
    if (!frames) {
        ObDereferenceObject(&section->Header);
        return NTLSTATUS_NO_MEMORY;
    }
    memset(frames, 0, pageCount * sizeof(PHYSICAL_ADDRESS));

    IMAGE_SECTION_DESCRIPTOR* descriptors = (IMAGE_SECTION_DESCRIPTOR*)MmAllocatePoolTag(
        (size_t)pe->NumberOfSections * sizeof(IMAGE_SECTION_DESCRIPTOR),
        MM_POOL_TAG('i', 's', 'e', 'c'));
    if (!descriptors) {
        MmFreePool(frames);
        ObDereferenceObject(&section->Header);
        return NTLSTATUS_NO_MEMORY;
    }
    memset(descriptors, 0, (size_t)pe->NumberOfSections * sizeof(IMAGE_SECTION_DESCRIPTOR));

    section->Frames = frames;
    section->Sections = descriptors;
    section->BackingFile = file;
    section->HeadersFileSize = pe->OptionalHeader.SizeOfHeaders;
    section->Lock = (SPINLOCK)SPINLOCK_INIT;
    ObReferenceObject(&file->Header);

    section->Metadata.ContentKey = key;
    section->Metadata.VolumeId = identity->VolumeId;
    section->Metadata.FileId = identity->FileId;
    section->Metadata.FileGeneration = identity->Generation;
    section->Metadata.PreferredBase = pe->OptionalHeader.ImageBase;
    section->Metadata.EntryRva = pe->OptionalHeader.AddressOfEntryPoint;
    section->Metadata.HeaderSize = pe->OptionalHeader.SizeOfHeaders;
    section->Metadata.RawSize = fileSize;
    section->Metadata.ImageSize = imageSize;
    section->Metadata.PageCount = pageCount;
    section->Metadata.SectionCount = pe->NumberOfSections;

    for (uint32_t i = 0; i < pe->NumberOfSections; ++i) {
        IMAGE_SECTION_DESCRIPTOR* dst = &section->Sections[i];
        memcpy(dst->Name, rawSections[i].Name, sizeof(rawSections[i].Name));
        dst->Name[8] = '\0';
        dst->VirtualAddress = rawSections[i].VirtualAddress;
        dst->VirtualSize = rawSections[i].VirtualSize;
        dst->RawDataOffset = rawSections[i].PointerToRawData;
        dst->RawDataSize = rawSections[i].SizeOfRawData;
        dst->Characteristics = rawSections[i].Characteristics;
        dst->PageProtection = KiImageSectionProtection(rawSections[i].Characteristics);
    }

    *sectionOut = section;
    return NTLSTATUS_SUCCESS;
}

// Creates or references an image section object backed by an open file.
KSTATUS LdrCreateImageSectionFromFile(FILE_OBJECT* file,
                                      IMAGE_SECTION_OBJECT** sectionOut) {
    if (!file || !sectionOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *sectionOut = (IMAGE_SECTION_OBJECT*)0;

    IO_FILE_IDENTITY identity;
    KSTATUS s = IoQueryFileIdentity(file, &identity);
    if (!KSUCCESS(s))
        return s;

    size_t fileSize = (size_t)identity.Size;
    if (!fileSize)
        return NTLSTATUS_FAILURE;

    size_t firstRead = fileSize < 4096u ? fileSize : 4096u;
    uint8_t stackHeader[4096];
    size_t got = 0;
    s = IoReadFileAt(file, 0, stackHeader, firstRead, &got);
    if (!KSUCCESS(s) || got != firstRead)
        return KSUCCESS(s) ? NTLSTATUS_FAILURE : s;
    if (firstRead < 0x40)
        return NTLSTATUS_FAILURE;

    uint32_t peOffset = *(const uint32_t*)(stackHeader + 0x3C);
    if ((size_t)peOffset > fileSize || fileSize - (size_t)peOffset < 24)
        return NTLSTATUS_FAILURE;

    size_t minimumHeader = (size_t)peOffset + 24u;
    if (minimumHeader > firstRead) {
        if (minimumHeader > 65536u || minimumHeader > fileSize)
            return NTLSTATUS_FAILURE;
        void* tmp = MmAllocatePoolTag(minimumHeader, MM_POOL_TAG('i', 's', 'h', 'd'));
        if (!tmp)
            return NTLSTATUS_NO_MEMORY;
        s = IoReadFileAt(file, 0, tmp, minimumHeader, &got);
        if (!KSUCCESS(s) || got != minimumHeader) {
            MmFreePool(tmp);
            return KSUCCESS(s) ? NTLSTATUS_FAILURE : s;
        }
        uint16_t optSize = *(const uint16_t*)((const uint8_t*)tmp + peOffset + 20u);
        size_t fullHeader = (size_t)peOffset + 24u + (size_t)optSize;
        MmFreePool(tmp);
        if (fullHeader > firstRead) {
            firstRead = fullHeader;
        }
    } else {
        uint16_t optSize = *(const uint16_t*)(stackHeader + peOffset + 20u);
        size_t fullHeader = (size_t)peOffset + 24u + (size_t)optSize;
        if (fullHeader > firstRead)
            firstRead = fullHeader;
    }

    if (firstRead > fileSize || firstRead > 65536u)
        return NTLSTATUS_FAILURE;

    void* header = stackHeader;
    bool heapHeader = false;
    if (firstRead > sizeof(stackHeader)) {
        header = MmAllocatePoolTag(firstRead, MM_POOL_TAG('i', 's', 'h', 'd'));
        if (!header)
            return NTLSTATUS_NO_MEMORY;
        heapHeader = true;
        s = IoReadFileAt(file, 0, header, firstRead, &got);
        if (!KSUCCESS(s) || got != firstRead) {
            MmFreePool(header);
            return KSUCCESS(s) ? NTLSTATUS_FAILURE : s;
        }
    }

    const LDR_PE_HEADER* pe = (const LDR_PE_HEADER*)((const uint8_t*)header + peOffset);
    size_t tableBytes = (size_t)pe->NumberOfSections * sizeof(LDR_SECTION_HEADER);
    size_t tableEnd = (size_t)peOffset + 24u + (size_t)pe->SizeOfOptionalHeader + tableBytes;
    if (tableEnd > firstRead) {
        if (tableEnd > fileSize || tableEnd > 65536u) {
            if (heapHeader) MmFreePool(header);
            return NTLSTATUS_FAILURE;
        }
        void* bigger = MmAllocatePoolTag(tableEnd, MM_POOL_TAG('i', 's', 'h', 'd'));
        if (!bigger) {
            if (heapHeader) MmFreePool(header);
            return NTLSTATUS_NO_MEMORY;
        }
        s = IoReadFileAt(file, 0, bigger, tableEnd, &got);
        if (heapHeader) MmFreePool(header);
        if (!KSUCCESS(s) || got != tableEnd) {
            MmFreePool(bigger);
            return KSUCCESS(s) ? NTLSTATUS_FAILURE : s;
        }
        header = bigger;
        firstRead = tableEnd;
        heapHeader = true;
    }

    uint64_t key = KiImageSectionIdentityKey(&identity);

    uintptr_t flags = KeAcquireSpinlock(&gImageSectionCacheLock);
    IMAGE_SECTION_OBJECT* cached = KiFindImageSectionLocked(&identity, key);
    if (cached) {
        ObReferenceObject(&cached->Header);
        KeReleaseSpinlock(&gImageSectionCacheLock, flags);
        if (heapHeader) MmFreePool(header);
        *sectionOut = cached;
        return NTLSTATUS_SUCCESS;
    }
    KeReleaseSpinlock(&gImageSectionCacheLock, flags);

    IMAGE_SECTION_OBJECT* built = (IMAGE_SECTION_OBJECT*)0;
    s = KiBuildImageSectionObjectFromFile(file, header, firstRead, &identity, key, &built);
    if (heapHeader) MmFreePool(header);
    if (!KSUCCESS(s))
        return s;

    flags = KeAcquireSpinlock(&gImageSectionCacheLock);
    cached = KiFindImageSectionLocked(&identity, key);
    if (cached) {
        ObReferenceObject(&cached->Header);
        KeReleaseSpinlock(&gImageSectionCacheLock, flags);
        ObDereferenceObject(&built->Header);
        *sectionOut = cached;
        return NTLSTATUS_SUCCESS;
    }

    for (uint32_t i = 0; i < LDR_IMAGE_SECTION_CACHE_MAX; ++i) {
        if (!gImageSectionCache[i]) {
            ObReferenceObject(&built->Header);
            gImageSectionCache[i] = built;
            break;
        }
    }
    KeReleaseSpinlock(&gImageSectionCacheLock, flags);

    *sectionOut = built;
    return NTLSTATUS_SUCCESS;
}

// Adds a reference to an existing image section object.
void LdrReferenceImageSection(IMAGE_SECTION_OBJECT* section) {
    if (!section) return;
    ObReferenceObject(&section->Header);
}

// Releases a reference returned by LdrCreateImageSectionFromFile.
void LdrDereferenceImageSection(IMAGE_SECTION_OBJECT* section) {
    if (!section) return;
    ObDereferenceObject(&section->Header);
}

// Returns immutable metadata for an image section object.
const IMAGE_SECTION_METADATA* LdrImageSectionMetadata(const IMAGE_SECTION_OBJECT* section) {
    if (!section) return (const IMAGE_SECTION_METADATA*)0;
    return &section->Metadata;
}

// Returns the immutable PE section table owned by the image section object.
const IMAGE_SECTION_DESCRIPTOR* LdrImageSectionDescriptors(const IMAGE_SECTION_OBJECT* section,
                                                          uint32_t* countOut) {
    if (countOut) *countOut = section ? section->Metadata.SectionCount : 0;
    if (!section) return (const IMAGE_SECTION_DESCRIPTOR*)0;
    return section->Sections;
}

// Retains one prototype page frame owned by an image section object.
KSTATUS LdrRetainImageSectionPage(IMAGE_SECTION_OBJECT* section,
                                  size_t pageIndex,
                                  PHYSICAL_ADDRESS* frameOut) {
    if (!section || !frameOut) return NTLSTATUS_INVALID_PARAMETER;
    *frameOut = 0;
    if (pageIndex >= section->Metadata.PageCount)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS es = KiEnsureImageSectionFrame(section, pageIndex);
    if (!KSUCCESS(es))
        return es;

    uintptr_t flags = KeAcquireSpinlock(&section->Lock);
    PHYSICAL_ADDRESS frame = section->Frames[pageIndex];
    KeReleaseSpinlock(&section->Lock, flags);
    if (!frame)
        return NTLSTATUS_SUCCESS;

    KSTATUS s = MmPfnRetain(frame);
    if (!KSUCCESS(s))
        return s;

    *frameOut = frame;
    return NTLSTATUS_SUCCESS;
}

// Maps an image section into a process as an image-backed VAD.
KSTATUS LdrMapImageSection(PROCESS* process,
                           IMAGE_SECTION_OBJECT* section,
                           void** baseOut) {
    if (!process || !section || !baseOut) return NTLSTATUS_INVALID_PARAMETER;
    *baseOut = (void*)0;

    return MmVadCreateImageSectionMapping(process,
                                          section,
                                          PAGE_PRESENT | PAGE_WRITE | PAGE_USER | PAGE_NX,
                                          MM_VAD_ZERO | MM_VAD_GUARD_LOW | MM_VAD_GUARD_HIGH,
                                          baseOut);
}