// PE image loader. This file owns image mapping, relocation, import resolution, and export lookup.

#include <stdint.h>
#include <string.h>
#include <kernel/ldr_image.h>
#include <kernel/ldr_section.h>
#include <ntlite/boot_info.h>
#include <kernel/ke_process.h>
#include <kernel/io_file.h>
#include <kernel/mm_virt.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_system_root.h>
#include <ntlite/syscall.h>

#ifndef LDR_DEBUG_VERBOSE
#define LDR_DEBUG_VERBOSE 0
#endif

#if LDR_DEBUG_VERBOSE
#define LdrDbg(...) DbgPrint(__VA_ARGS__)
#else
#define LdrDbg(...) ((void)0)
#endif

int _stricmp(const char* a, const char* b);

// PE32+ structures (just what we need).
#define PE_SIGNATURE        0x00004550UL   // 'PE\0\0'
#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

#pragma pack(push, 1)
typedef struct {
    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];
} OPT_HEADER64;

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

typedef struct {
    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;
} SECTION_HEADER;

// Implements the align Pe Size operation.
static size_t KiAlignPeSize(size_t value) {
    return (value + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
}

// Implements the compute Required Image Size operation.
static KSTATUS KiComputeRequiredImageSize(const PE_HEADER* pe, size_t fileSize, size_t* outSize) {
    if (!pe || !outSize) return STATUS_INVALID_PARAMETER;

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

    const SECTION_HEADER* sec = (const SECTION_HEADER*)((const uint8_t*)&pe->OptionalHeader + pe->SizeOfOptionalHeader);
    for (int 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 STATUS_NO_MEMORY;
        size_t end = KiAlignPeSize(va + span);
        if (end > required) required = end;

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

    if (!required) return STATUS_FAILURE;
    *outSize = required;
    return STATUS_SUCCESS;
}

typedef struct {
    uint32_t Characteristics;
    uint32_t TimeDateStamp;
    uint16_t MajorVersion;
    uint16_t MinorVersion;
    uint32_t Name;
    uint32_t Base;
    uint32_t NumberOfFunctions;
    uint32_t NumberOfNames;
    uint32_t AddressOfFunctions;
    uint32_t AddressOfNames;
    uint32_t AddressOfNameOrdinals;
} EXPORT_DIR;

typedef struct {
    uint32_t OriginalFirstThunk;
    uint32_t TimeDateStamp;
    uint32_t ForwarderChain;
    uint32_t Name;
    uint32_t FirstThunk;
} IMPORT_DESC;

typedef struct {
    uint32_t VirtualAddress;
    uint32_t SizeOfBlock;
} BASE_RELOC;
#pragma pack(pop)

#define IMAGE_REL_BASED_DIR64 10
#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

// Module registry.
#define MAX_MODULES  32

typedef struct {
    char      Name[32];
    uintptr_t Base;
    size_t    Size;
} MODULE_ENTRY;

static MODULE_ENTRY gModules[MAX_MODULES];
static uint32_t     gModuleCount = 0;

// Ensures a user image page is privately writable before the kernel patches it.
static KSTATUS KiEnsureUserImageWrite(PROCESS* process, bool user, uintptr_t address, size_t size) {
    if (!user) return STATUS_SUCCESS;
    return MmVadMakeImageRangePrivate(process, (void*)address, size);
}

// Implements the canonical Module Name operation.
static const char* KiCanonicalModuleName(const char* name) {
    if (!name) return name;

    if (!_stricmp(name, "kernel32.dll") ||
        !_stricmp(name, "user32.dll")   ||
        !_stricmp(name, "gdi32.dll")    ||
        !_stricmp(name, "ole32.dll")    ||
        !_stricmp(name, "winmm.dll")) {
        return "win32.dll";
    }

    if (!_stricmp(name, "msvcrt.dll")) {
        return "posix.dll";
    }

    return name;
}

// Returns module Base In Table.
static uintptr_t KiFindModuleBaseInTable(const MODULE_ENTRY* modules, uint32_t count, const char* name) {
    for (uint32_t i = 0; i < count; i++) {
        if (_stricmp(modules[i].Name, name) == 0) return modules[i].Base;
    }
    return 0;
}

// Registers module Entry.
static void KiRegisterModuleEntry(MODULE_ENTRY* entry, const char* name, uintptr_t base, size_t size) {
    strncpy(entry->Name, name, 31);
    entry->Name[31] = '\0';
    entry->Base = base;
    entry->Size = size;
}

// Implements the ldr Register Modules From Boot Info operation.
void LdrRegisterModulesFromBootInfo(const BOOT_INFO* bootInfo) {
    for (uint32_t i = 0; i < bootInfo->ModuleCount && gModuleCount < MAX_MODULES; i++) {
        // Only register modules that have a virtual base (Kernel, HAL).
        if (bootInfo->Modules[i].VirtBase != 0) {
            KiRegisterModuleEntry(&gModules[gModuleCount], 
                                bootInfo->Modules[i].Name,
                                bootInfo->Modules[i].VirtBase, 
                                bootInfo->Modules[i].ImageSize);
            gModuleCount++;
        }
    }
}

// Implements the ldr Register Module operation.
void LdrRegisterModule(const char* name, uintptr_t base) {
    if (gModuleCount < MAX_MODULES) {
        size_t size = base ? LdrGetImageSize(base) : 0;
        KiRegisterModuleEntry(&gModules[gModuleCount], name, base, size);
        gModuleCount++;
    }
}

// Finds the loaded kernel module that owns an address.
const char* LdrFindKernelModuleByAddress(uintptr_t address, uintptr_t* baseOut, size_t* sizeOut) {
    if (baseOut) *baseOut = 0;
    if (sizeOut) *sizeOut = 0;

    for (uint32_t i = 0; i < gModuleCount; i++) {
        uintptr_t base = gModules[i].Base;
        size_t size = gModules[i].Size;
        if (!base || !size)
            continue;
        if (address >= base && address < base + size) {
            if (baseOut) *baseOut = base;
            if (sizeOut) *sizeOut = size;
            return gModules[i].Name;
        }
    }
    return "unknown";
}

// Returns global Module Base.
static uintptr_t KiFindGlobalModuleBase(const char* name) {
    return KiFindModuleBaseInTable(gModules, gModuleCount, KiCanonicalModuleName(name));
}

static uintptr_t KiFindProcessModuleBaseInTable(
    const PROCESS_MODULE* modules,
    uint32_t count,
    const char* name
) {
    for (uint32_t i = 0; i < count; i++) {
        if (_stricmp(modules[i].Name, KiCanonicalModuleName(name)) == 0) {
            return modules[i].Base;
        }
    }
    return 0;
}

// Returns process Or Global Module Base.
static uintptr_t KiFindProcessOrGlobalModuleBase(PROCESS* process, const char* name) {
    if (process) {
        uintptr_t base = KiFindProcessModuleBaseInTable(process->Modules, process->ModuleCount, name);
        if (base) return base;
    }
    return KiFindGlobalModuleBase(name);
}

// Returns export At Base.
static void* KiGetExportAtBase(uintptr_t base, const char* functionName) {
    if (!base) return (void *)0;

    uint32_t peOffset = *(uint32_t *)(base + 0x3C);
    PE_HEADER* pe = (PE_HEADER *)(base + peOffset);
    if (pe->Signature != PE_SIGNATURE) return (void *)0;

    OPT_HEADER64* opt = &pe->OptionalHeader;
    if (opt->NumberOfRvaAndSizes < 1) return (void *)0;

    uint32_t expRva = opt->DataDirectory[0].VirtualAddress;
    if (!expRva) return (void *)0;

    EXPORT_DIR* exp = (EXPORT_DIR *)(base + expRva);
    uint32_t* names     = (uint32_t *)(base + exp->AddressOfNames);
    uint16_t* ordinals  = (uint16_t *)(base + exp->AddressOfNameOrdinals);
    uint32_t* functions = (uint32_t *)(base + exp->AddressOfFunctions);

    for (uint32_t i = 0; i < exp->NumberOfNames; i++) {
        const char* name = (const char *)(base + names[i]);
        if (strcmp(name, functionName) == 0) {
            uint16_t ord = ordinals[i];
            return (void *)(base + functions[ord]);
        }
    }
    return (void *)0;
}

// Export lookup.
void* LdrGetExport(const char* moduleName, const char* functionName) {
    uintptr_t base = KiFindGlobalModuleBase(KiCanonicalModuleName(moduleName));
    return KiGetExportAtBase(base, functionName);
}

// Implements the ldr Get Export For Process operation.
void* LdrGetExportForProcess(PROCESS* process, const char* moduleName, const char* functionName) {
    uintptr_t base = KiFindProcessOrGlobalModuleBase(process, KiCanonicalModuleName(moduleName));
    return KiGetExportAtBase(base, functionName);
}

// Looks up an export in a specific process module mapping.
void* LdrGetModuleExportForProcess(PROCESS* process, uintptr_t imageBase, const char* functionName) {
    if (!process || !imageBase || !functionName) return (void*)0;

    uintptr_t flags = KeAcquireSpinlock(&process->ModuleLock);
    int found = 0;
    for (uint32_t i = 0; i < process->ModuleCount; i++) {
        if (process->Modules[i].Base == imageBase) {
            found = 1;
            break;
        }
    }
    KeReleaseSpinlock(&process->ModuleLock, flags);
    if (!found) return (void*)0;

    return KiGetExportAtBase(imageBase, functionName);
}

// Implements the ldr Get Image Size operation.
size_t LdrGetImageSize(uintptr_t base) {
    if (!base) return 0;
    uint32_t peOffset = *(uint32_t *)(base + 0x3C);
    PE_HEADER* pe = (PE_HEADER *)(base + peOffset);
    if (pe->Signature != PE_SIGNATURE) return 0;
    return (size_t)pe->OptionalHeader.SizeOfImage;
}

// Implements the ldr Get Mapped Image Size operation.
size_t LdrGetMappedImageSize(PROCESS* process, uintptr_t base) {
    if (!base) return 0;
    if (process) {
        size_t regionSize = MmVadGetSize(process, (const void*)base);
        if (regionSize) return regionSize;
    }
    size_t globalRegionSize = MmVadGetSizeByBase((const void*)base);
    if (globalRegionSize) return globalRegionSize;
    return LdrGetImageSize(base);
}

// Releases loaded Image For Process.
static void KiFreeLoadedImageForProcess(PROCESS* process, uintptr_t imageBase, bool user) {
    if (!imageBase) return;
    size_t imageSize = user ? LdrGetMappedImageSize(process, imageBase) : LdrGetImageSize(imageBase);
    if (user) MmFreeUserImageForProcess(process, (void*)imageBase, imageSize);
    else MmFreePool((void*)imageBase);
}

// Returns the final page flags for a PE section.
static uint64_t KiSectionPageProtection(const SECTION_HEADER* section, bool user) {
    uint64_t prot = PAGE_PRESENT | (user ? PAGE_USER : 0);

    // Some old toolchains mark BSS-like sections as uninitialized data without
    // setting MEM_WRITE. Treating those pages as read-only breaks programs that
    // legitimately write their zero-fill storage.
    if ((section->Characteristics & IMAGE_SCN_MEM_WRITE) ||
        (section->Characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA)) {
        prot |= PAGE_WRITE;
    }

    if (!((section->Characteristics & IMAGE_SCN_MEM_EXECUTE) ||
          (section->Characteristics & IMAGE_SCN_CNT_CODE))) {
        prot |= PAGE_NX;
    }
    return prot;
}

// Applies final page protections to a mapped PE image.
KSTATUS LdrProtectImageForProcess(PROCESS* process, uintptr_t imageBase, bool user) {
    if (!imageBase) return STATUS_INVALID_PARAMETER;
    if (user && !process) return STATUS_INVALID_PARAMETER;

    size_t mappedSize = user ? LdrGetMappedImageSize(process, imageBase) : LdrGetImageSize(imageBase);
    if (!mappedSize) return STATUS_FAILURE;

    uint32_t peOffset = *(uint32_t *)(imageBase + 0x3C);
    if ((size_t)peOffset > mappedSize || mappedSize - (size_t)peOffset < sizeof(uint32_t) + 20)
        return STATUS_FAILURE;

    PE_HEADER* pe = (PE_HEADER *)(imageBase + peOffset);
    if (pe->Signature != PE_SIGNATURE) return STATUS_FAILURE;

    size_t sectionTableOffset = (size_t)peOffset + 24u + (size_t)pe->SizeOfOptionalHeader;
    size_t sectionTableSize = (size_t)pe->NumberOfSections * sizeof(SECTION_HEADER);
    if (sectionTableOffset > mappedSize || sectionTableSize > mappedSize - sectionTableOffset)
        return STATUS_FAILURE;

    size_t pageCount = (mappedSize + PAGE_SIZE - 1) / PAGE_SIZE;
    uint64_t* pageProt = (uint64_t*)MmAllocatePoolTag(pageCount * sizeof(uint64_t), MM_POOL_TAG('i', 'p', 'r', 't'));
    if (!pageProt) return STATUS_NO_MEMORY;

    uint64_t headerProt = PAGE_PRESENT | (user ? PAGE_USER : 0) | PAGE_NX;
    for (size_t i = 0; i < pageCount; ++i)
        pageProt[i] = headerProt;

    SECTION_HEADER* sec = (SECTION_HEADER *)(imageBase + sectionTableOffset);
    for (int i = 0; i < pe->NumberOfSections; i++) {
        size_t protectSize = sec[i].VirtualSize ? sec[i].VirtualSize : sec[i].SizeOfRawData;
        if (!protectSize) continue;

        uintptr_t secBase = imageBase + (uintptr_t)sec[i].VirtualAddress;
        uintptr_t secEnd = secBase + protectSize;
        uintptr_t mapEnd = imageBase + mappedSize;

        if (secBase < imageBase || secBase >= mapEnd) {
            LdrDbg("LDR: protect skipping out-of-range section %d base=%016llx size=%llx image=%016llx mapped=%llx\n",
                   i, (unsigned long long)secBase, (unsigned long long)protectSize,
                   (unsigned long long)imageBase, (unsigned long long)mappedSize);
            continue;
        }
        if (secEnd < secBase) {
            MmFreePool(pageProt);
            return STATUS_NO_MEMORY;
        }
        if (secEnd > mapEnd) {
            LdrDbg("LDR: protect clamping section %d end=%016llx to mapEnd=%016llx\n",
                   i, (unsigned long long)secEnd, (unsigned long long)mapEnd);
            secEnd = mapEnd;
        }

        size_t firstPage = (size_t)((secBase - imageBase) / PAGE_SIZE);
        size_t lastPage = (size_t)((secEnd - 1 - imageBase) / PAGE_SIZE);
        uint64_t prot = KiSectionPageProtection(&sec[i], user);
        for (size_t page = firstPage; page <= lastPage && page < pageCount; ++page) {
            // Protection is page-granular. If sections ever share a page, the
            // hardware page must allow the union of their access requirements.
            if (prot & PAGE_WRITE)
                pageProt[page] |= PAGE_WRITE;
            if ((prot & PAGE_NX) == 0)
                pageProt[page] &= ~PAGE_NX;
        }
    }

    size_t runStart = 0;
    KSTATUS s = STATUS_SUCCESS;
    while (runStart < pageCount) {
        uint64_t prot = pageProt[runStart];
        size_t runEnd = runStart + 1;
        while (runEnd < pageCount && pageProt[runEnd] == prot)
            runEnd++;

        uintptr_t va = imageBase + runStart * PAGE_SIZE;
        size_t bytes = (runEnd - runStart) * PAGE_SIZE;
        if (user)
            s = MmVadProtect(process, (void*)va, bytes, prot);
        else
            s = MmPteProtectRange(va, bytes, prot);

        if (!KSUCCESS(s)) {
            LdrDbg("LDR: protect failed va=%016llx size=%llx prot=%016llx st=%d image=%016llx mapped=%llx\n",
                   (unsigned long long)va, (unsigned long long)bytes,
                   (unsigned long long)prot, s,
                   (unsigned long long)imageBase, (unsigned long long)mappedSize);
            break;
        }
        runStart = runEnd;
    }

    MmFreePool(pageProt);
    return s;
}

// Entry point lookup.
uintptr_t LdrGetEntryPoint(uintptr_t base) {
    uint32_t pe_offset = *(uint32_t*)(base + 0x3C);
    uint32_t entry_rva = *(uint32_t*)(base + pe_offset + 0x28);
    return base + entry_rva;
}

// Base Relocation.
static KSTATUS LdrApplyRelocationsForProcess(PROCESS* process, uintptr_t base, uintptr_t delta, bool user) {
    uint32_t peOffset = *(uint32_t *)(base + 0x3C);
    PE_HEADER* pe = (PE_HEADER *)(base + peOffset);
    OPT_HEADER64* opt = &pe->OptionalHeader;

    if (opt->NumberOfRvaAndSizes < 6) return STATUS_SUCCESS;
    uint32_t relocRva = opt->DataDirectory[5].VirtualAddress;
    uint32_t relocSize = opt->DataDirectory[5].Size;
    if (!relocRva || !relocSize) return STATUS_SUCCESS;

    BASE_RELOC* block = (BASE_RELOC *)(base + relocRva);
    uintptr_t end = (uintptr_t)block + relocSize;

    while ((uintptr_t)block < end && block->SizeOfBlock > 0) {
        uint32_t count = (block->SizeOfBlock - sizeof(BASE_RELOC)) / sizeof(uint16_t);
        uint16_t* entries = (uint16_t *)(block + 1);

        for (uint32_t i = 0; i < count; i++) {
            int type = entries[i] >> 12;
            int offset = entries[i] & 0xFFF;

            // For 64-bit PEs, we mostly care about DIR64 (Type 10).
            if (type == IMAGE_REL_BASED_DIR64) {
                uint64_t* ptr = (uint64_t *)(base + block->VirtualAddress + offset);
                KSTATUS ws = KiEnsureUserImageWrite(process, user, (uintptr_t)ptr, sizeof(*ptr));
                if (!KSUCCESS(ws)) return ws;
                *ptr += delta;
            }
        }
        block = (BASE_RELOC *)((uint8_t *)block + block->SizeOfBlock);
    }
    return STATUS_SUCCESS;
}

// Import resolution.
typedef void* (*EXPORT_LOOKUP_FN)(void* ctx, const char* moduleName, const char* functionName);

// Implements the global Export Lookup operation.
static void* KiGlobalExportLookup(void* ctx, const char* moduleName, const char* functionName) {
    (void)ctx;
    return LdrGetExport(moduleName, functionName);
}

// Processes export Lookup.
static void* KiProcessExportLookup(void* ctx, const char* moduleName, const char* functionName) {
    return LdrGetExportForProcess((PROCESS*)ctx, moduleName, functionName);
}

static KSTATUS LdrResolveImportsCommon(
    PROCESS* process,
    uintptr_t imageBase,
    EXPORT_LOOKUP_FN lookup,
    void* ctx
) {
    uint32_t peOffset = *(uint32_t *)(imageBase + 0x3C);
    PE_HEADER* pe = (PE_HEADER *)(imageBase + peOffset);
    OPT_HEADER64* opt = &pe->OptionalHeader;

    if (opt->NumberOfRvaAndSizes < 2)
        return STATUS_SUCCESS;

    uint32_t impRva = opt->DataDirectory[1].VirtualAddress;
    if (!impRva)
        return STATUS_SUCCESS;

    IMPORT_DESC* imp = (IMPORT_DESC *)(imageBase + impRva);

    while (imp->Name) {
        const char* dllName = (const char *)(imageBase + imp->Name);

        //Compatibility remap.
        dllName = KiCanonicalModuleName(dllName);
        LdrDbg("LDR: resolve import dll=%s image=%016llx pid=%u\n", dllName, (unsigned long long)imageBase, process ? process->ProcessId : 0);

        // Ensure dependency is loaded into the process.

        if (process) {
            uintptr_t modBase =
                KiFindProcessOrGlobalModuleBase(process, dllName);

            if (!modBase) {
                KSTATUS s =
                    LdrLoadModuleForProcess(process, dllName, LDR_MODULE_LOAD_IMPORT, &modBase);

                if (!KSUCCESS(s)) {
                    DbgPrint(
                        "LdrResolveImports: Failed to load dependency %s (%d)\n",
                        dllName,
                        s
                    );
                    return s;
                }
            }
        }

        uint64_t* thunk =
            (uint64_t *)(imageBase + imp->FirstThunk);

        uint64_t* origThunk =
            (uint64_t *)(imageBase +
                (imp->OriginalFirstThunk
                    ? imp->OriginalFirstThunk
                    : imp->FirstThunk));

        while (*origThunk) {
            if (*origThunk & (1ULL << 63)) {
                DbgPrint(
                    "LdrResolveImports: Ordinal imports not supported! (%s)\n",
                    dllName
                );
                return STATUS_NOT_SUPPORTED;
            }

            const char* funcName =
                (const char *)(imageBase +
                    (*origThunk & 0xFFFFFFFF) + 2);

            void* addr = lookup(ctx, dllName, funcName);

            if (!addr) {
                DbgPrint(
                    "LdrResolveImports: Function not found! (%s!%s)\n",
                    dllName,
                    funcName
                );
                return STATUS_NOT_FOUND;
            }

            KSTATUS ws = KiEnsureUserImageWrite(process, process != (PROCESS*)0, (uintptr_t)thunk, sizeof(*thunk));
            if (!KSUCCESS(ws)) return ws;
            *thunk = (uint64_t)addr;

            thunk++;
            origThunk++;
        }

        imp++;
    }

    return STATUS_SUCCESS;
}

// Implements the ldr Resolve Imports operation.
KSTATUS LdrResolveImports(uintptr_t imageBase) {
    return LdrResolveImportsCommon(
        (PROCESS*)0,
        imageBase,
        KiGlobalExportLookup,
        (void*)0
    );
}

// Implements the ldr Resolve Imports For Process operation.
KSTATUS LdrResolveImportsForProcess(PROCESS* process, uintptr_t imageBase) {
    return LdrResolveImportsCommon(
        process,
        imageBase,
        KiProcessExportLookup,
        process
    );
}

// Loads a user image from an open file through the image-section object model.
KSTATUS LdrLoadImageFileForProcess(PROCESS* process, FILE_OBJECT* file, uintptr_t* baseOut) {
    if (!process || !file || !baseOut)
        return STATUS_INVALID_PARAMETER;
    *baseOut = 0;

    IMAGE_SECTION_OBJECT* section = (IMAGE_SECTION_OBJECT*)0;
    KSTATUS s = LdrCreateImageSectionFromFile(file, &section);
    if (!KSUCCESS(s))
        return s;

    void* view = (void*)0;
    s = LdrMapImageSection(process, section, &view);
    if (!KSUCCESS(s)) {
        LdrDereferenceImageSection(section);
        return s;
    }

    const IMAGE_SECTION_METADATA* metadata = LdrImageSectionMetadata(section);
    uintptr_t imageBase = (uintptr_t)view;
    uintptr_t delta = imageBase - metadata->PreferredBase;
    if (delta != 0) {
        s = LdrApplyRelocationsForProcess(process, imageBase, delta, true);
        if (!KSUCCESS(s)) {
            KiFreeLoadedImageForProcess(process, imageBase, true);
            LdrDereferenceImageSection(section);
            return s;
        }
    }

    LdrDereferenceImageSection(section);
    *baseOut = imageBase;
    return STATUS_SUCCESS;
}

// Converts a process module entry into the public module ABI.
static void KiFillProcessModuleInfo(const PROCESS_MODULE* mod, SYS_MODULE_INFO* info) {
    memset(info, 0, sizeof(*info));
    info->Size = sizeof(*info);
    info->Version = SYS_MODULE_INFO_VERSION;
    info->Flags = mod->Flags;
    if (mod->PublicRefCount)
        info->Flags |= SYS_MODULE_FLAG_PUBLIC_REF;
    if (mod->ImportRefCount)
        info->Flags |= SYS_MODULE_FLAG_IMPORT_PINNED;
    if (mod->PublicRefCount && mod->ImportRefCount == 0)
        info->Flags |= SYS_MODULE_FLAG_UNLOADABLE;
    info->Base = mod->Base;
    info->ImageSize = mod->Size;
    info->EntryPoint = mod->EntryPoint;
    info->PublicRefCount = mod->PublicRefCount;
    info->ImportRefCount = mod->ImportRefCount;
    info->TotalRefCount = mod->PublicRefCount + mod->ImportRefCount;
    strncpy(info->Name, mod->Name, sizeof(info->Name) - 1);
    info->Name[sizeof(info->Name) - 1] = '\0';
}

// Adds one reference of the requested class to an existing process module.
static void KiReferenceProcessModule(PROCESS_MODULE* mod, uint32_t loadFlags) {
    if (loadFlags & LDR_MODULE_LOAD_PUBLIC)
        mod->PublicRefCount++;
    if (loadFlags & LDR_MODULE_LOAD_IMPORT)
        mod->ImportRefCount++;
}

// Finds a process module by canonical name while the caller owns ModuleLock.
static PROCESS_MODULE* KiFindProcessModuleByNameLocked(PROCESS* process, const char* name) {
    for (uint32_t i = 0; i < process->ModuleCount; i++) {
        if (_stricmp(process->Modules[i].Name, name) == 0)
            return &process->Modules[i];
    }
    return (PROCESS_MODULE*)0;
}

// Finds a process module by image base while the caller owns ModuleLock.
static PROCESS_MODULE* KiFindProcessModuleByBaseLocked(PROCESS* process, uintptr_t base) {
    for (uint32_t i = 0; i < process->ModuleCount; i++) {
        if (process->Modules[i].Base == base)
            return &process->Modules[i];
    }
    return (PROCESS_MODULE*)0;
}

// Registers a mapped process module or records a reference to an existing one.
static KSTATUS LdrRegisterProcessModule(PROCESS* process, const char* name, uintptr_t base, uint32_t loadFlags, uintptr_t* registeredBaseOut) {
    if (!process || !name || !base || !loadFlags) return STATUS_INVALID_PARAMETER;
    if (registeredBaseOut) *registeredBaseOut = 0;

    uintptr_t flags = KeAcquireSpinlock(&process->ModuleLock);
    PROCESS_MODULE* existing = KiFindProcessModuleByNameLocked(process, name);
    if (existing) {
        KiReferenceProcessModule(existing, loadFlags);
        if (registeredBaseOut) *registeredBaseOut = existing->Base;
        KeReleaseSpinlock(&process->ModuleLock, flags);
        return STATUS_ALREADY_EXISTS;
    }

    if (process->ModuleCount >= PROCESS_MODULE_MAX) {
        KeReleaseSpinlock(&process->ModuleLock, flags);
        return STATUS_NO_MEMORY;
    }

    PROCESS_MODULE* mod = &process->Modules[process->ModuleCount++];
    memset(mod, 0, sizeof(*mod));
    strncpy(mod->Name, name, sizeof(mod->Name) - 1);
    mod->Name[sizeof(mod->Name) - 1] = '\0';
    mod->Base = base;
    mod->Size = LdrGetMappedImageSize(process, base);
    mod->EntryPoint = LdrGetEntryPoint(base);
    KiReferenceProcessModule(mod, loadFlags);
    if (registeredBaseOut) *registeredBaseOut = mod->Base;
    KeReleaseSpinlock(&process->ModuleLock, flags);
    return STATUS_SUCCESS;
}

// Loads a DLL-like user module into a process and records the requested reference.
KSTATUS LdrLoadModuleForProcess(PROCESS* process, const char* filename, uint32_t loadFlags, uintptr_t* baseOut) {
    if (!process || !filename || !loadFlags) return STATUS_INVALID_PARAMETER;

    filename = KiCanonicalModuleName(filename);

    uintptr_t flags = KeAcquireSpinlock(&process->ModuleLock);
    PROCESS_MODULE* existing = KiFindProcessModuleByNameLocked(process, filename);
    if (existing) {
        KiReferenceProcessModule(existing, loadFlags);
        if (baseOut) *baseOut = existing->Base;
        KeReleaseSpinlock(&process->ModuleLock, flags);
        return STATUS_SUCCESS;
    }
    KeReleaseSpinlock(&process->ModuleLock, flags);

    FILE_OBJECT* f;
    KSTATUS s = IoOpenFile(SysSvFromCstr(filename), &f);
    if (!KSUCCESS(s)) {
        char relative[260];
        size_t nameLen = strnlen(filename, sizeof(relative) - 8);
        if (nameLen + 8 >= sizeof(relative))
            return STATUS_BUFFER_TOO_SMALL;
        memcpy(relative, "system/", 7);
        memcpy(relative + 7, filename, nameLen);
        relative[7 + nameLen] = '\0';

        s = KeOpenSystemRootFile(relative, &f);
        if (!KSUCCESS(s)) {
            DbgPrint("LdrLoadModuleForProcess: Failed to open %s on system root (%d)\n", filename, s);
            return s;
        }
    }

    uint64_t fileSize = IoGetFileSize(f);
    LdrDbg("LDR: load module pid=%u file=%s size=%llu\n", process->ProcessId, filename, fileSize);

    uintptr_t imageBase = 0;
    s = LdrLoadImageFileForProcess(process, f, &imageBase);
    IoCloseFile(f);
    if (!KSUCCESS(s)) {
        DbgPrint("LDR: load image for module %s failed st=%d\n", filename, s);
        return s;
    }
    LdrDbg("LDR: module mapped pid=%u file=%s base=%016llx size=%llx\n",
           process->ProcessId, filename, (unsigned long long)imageBase, (unsigned long long)LdrGetMappedImageSize(process, imageBase));

    s = LdrResolveImportsForProcess(process, imageBase);
    if (!KSUCCESS(s)) {
        KiFreeLoadedImageForProcess(process, imageBase, true);
        return s;
    }

    s = LdrProtectImageForProcess(process, imageBase, true);
    if (!KSUCCESS(s)) {
        KiFreeLoadedImageForProcess(process, imageBase, true);
        return s;
    }

    uintptr_t registeredBase = 0;
    s = LdrRegisterProcessModule(process, filename, imageBase, loadFlags, &registeredBase);
    if (s == STATUS_ALREADY_EXISTS) {
        KiFreeLoadedImageForProcess(process, imageBase, true);
        imageBase = registeredBase;
        s = STATUS_SUCCESS;
    }
    if (!KSUCCESS(s)) {
        KiFreeLoadedImageForProcess(process, imageBase, true);
        return s;
    }

    LdrDbg("LDR: module registered pid=%u name=%s base=%016llx modules=%u\n",
           process->ProcessId, filename, (unsigned long long)imageBase, process->ModuleCount);
    if (baseOut) *baseOut = imageBase;
    return STATUS_SUCCESS;
}

// Releases one public module load reference and unmaps the image when no references remain.
KSTATUS LdrUnloadModuleForProcess(PROCESS* process, uintptr_t imageBase) {
    if (!process || !imageBase) return STATUS_INVALID_PARAMETER;

    char name[PROCESS_MODULE_NAME_MAX];
    uintptr_t base = 0;
    int shouldUnmap = 0;

    uintptr_t flags = KeAcquireSpinlock(&process->ModuleLock);
    PROCESS_MODULE* mod = KiFindProcessModuleByBaseLocked(process, imageBase);
    if (!mod) {
        KeReleaseSpinlock(&process->ModuleLock, flags);
        return STATUS_NOT_FOUND;
    }
    if (mod->PublicRefCount == 0) {
        KeReleaseSpinlock(&process->ModuleLock, flags);
        return STATUS_INVALID_PARAMETER;
    }

    mod->PublicRefCount--;
    if (mod->PublicRefCount == 0 && mod->ImportRefCount == 0) {
        base = mod->Base;
        strncpy(name, mod->Name, sizeof(name) - 1);
        name[sizeof(name) - 1] = '\0';
        uint32_t index = (uint32_t)(mod - process->Modules);
        if (index + 1 < process->ModuleCount) {
            memmove(&process->Modules[index],
                    &process->Modules[index + 1],
                    (process->ModuleCount - index - 1) * sizeof(process->Modules[0]));
        }
        process->ModuleCount--;
        shouldUnmap = 1;
    }
    KeReleaseSpinlock(&process->ModuleLock, flags);

    if (shouldUnmap) {
        LdrDbg("LDR: unload module pid=%u name=%s base=%016llx\n",
               process->ProcessId, name, (unsigned long long)base);
        KiFreeLoadedImageForProcess(process, base, true);
    }

    return STATUS_SUCCESS;
}

// Queries one loaded process module by image base.
KSTATUS LdrQueryModuleForProcess(PROCESS* process, uintptr_t imageBase, SYS_MODULE_INFO* infoOut) {
    if (!process || !imageBase || !infoOut) return STATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&process->ModuleLock);
    PROCESS_MODULE* mod = KiFindProcessModuleByBaseLocked(process, imageBase);
    if (!mod) {
        KeReleaseSpinlock(&process->ModuleLock, flags);
        return STATUS_NOT_FOUND;
    }
    KiFillProcessModuleInfo(mod, infoOut);
    KeReleaseSpinlock(&process->ModuleLock, flags);
    return STATUS_SUCCESS;
}

// Copies a bounded snapshot of the process module table.
KSTATUS LdrCopyProcessModuleInfo(PROCESS* process, SYS_MODULE_INFO* infoOut, size_t maxCount, size_t* countOut) {
    if (!process || !countOut) return STATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&process->ModuleLock);
    size_t count = process->ModuleCount;
    *countOut = count;
    size_t toCopy = count;
    if (toCopy > maxCount) toCopy = maxCount;
    if (toCopy && !infoOut) {
        KeReleaseSpinlock(&process->ModuleLock, flags);
        return STATUS_INVALID_PARAMETER;
    }
    for (size_t i = 0; i < toCopy; i++)
        KiFillProcessModuleInfo(&process->Modules[i], &infoOut[i]);
    KeReleaseSpinlock(&process->ModuleLock, flags);
    return (maxCount < count) ? STATUS_BUFFER_TOO_SMALL : STATUS_SUCCESS;
}

// Loads a boot/kernel PE image from a raw memory buffer into kernel space.
static KSTATUS KiLdrLoadKernelImageFromBuffer(const void* buffer, size_t size, uintptr_t* baseOut) {
    if (!buffer || !size || !baseOut) return STATUS_INVALID_PARAMETER;
    if (size < 0x40) return STATUS_FAILURE;

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

    uint32_t signature = *(const uint32_t *)(bytes + peOffset);
    if (signature != PE_SIGNATURE) return STATUS_FAILURE;

    uint16_t machine = *(const uint16_t *)(bytes + peOffset + 4);
    uint16_t optionalSize = *(const uint16_t *)(bytes + peOffset + 20);
    if (optionalSize < sizeof(OPT_HEADER64) ||
        (size_t)peOffset + 24 + optionalSize > size) {
        return STATUS_FAILURE;
    }

    uint16_t optionalMagic = *(const uint16_t *)(bytes + peOffset + 24);
    if (machine == IMAGE_FILE_MACHINE_I386 || optionalMagic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
        DbgPrint("LdrLoadKernelImage: hard error: 32-bit PE images are not supported (machine=0x%X magic=0x%X)\n",
                 machine, optionalMagic);
        return STATUS_NOT_SUPPORTED;
    }

    if (machine != IMAGE_FILE_MACHINE_AMD64 || optionalMagic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
        DbgPrint("LdrLoadKernelImage: unsupported PE image (machine=0x%X magic=0x%X)\n",
                 machine, optionalMagic);
        return STATUS_NOT_SUPPORTED;
    }

    PE_HEADER* pe = (PE_HEADER *)(bytes + peOffset);
    OPT_HEADER64* opt = &pe->OptionalHeader;
    size_t imageSize = 0;
    KSTATUS sizeStatus = KiComputeRequiredImageSize(pe, size, &imageSize);
    if (!KSUCCESS(sizeStatus)) return sizeStatus;

    LdrDbg("LDR: kernel image header sizeOfImage=%x required=%llx preferred=%016llx sections=%u\n",
           opt->SizeOfImage,
           (unsigned long long)imageSize,
           (unsigned long long)opt->ImageBase,
           pe->NumberOfSections);

    void* mem = MmAllocatePoolTag(imageSize, MM_POOL_TAG('i', 'm', 'g', ' '));
    if (!mem) {
        DbgPrint("LDR: kernel image allocation failed required=%llx\n",
                 (unsigned long long)imageSize);
        return STATUS_NO_MEMORY;
    }
    memset(mem, 0, imageSize);

    memcpy(mem, buffer, opt->SizeOfHeaders);
    SECTION_HEADER* sec = (SECTION_HEADER *)((uint8_t *)&pe->OptionalHeader + pe->SizeOfOptionalHeader);
    for (int i = 0; i < pe->NumberOfSections; i++) {
        void* dst = (uint8_t *)mem + sec[i].VirtualAddress;
        const void* src = (const uint8_t *)buffer + sec[i].PointerToRawData;
        if (sec[i].SizeOfRawData) {
            size_t copySize = sec[i].SizeOfRawData;
            if (sec[i].VirtualSize && sec[i].VirtualSize < copySize) copySize = sec[i].VirtualSize;
            memcpy(dst, src, copySize);
        }

        MmPteProtectRange((uintptr_t)mem + sec[i].VirtualAddress,
                         sec[i].VirtualSize ? sec[i].VirtualSize : sec[i].SizeOfRawData,
                         PAGE_PRESENT | PAGE_WRITE);
    }

    uintptr_t delta = (uintptr_t)mem - opt->ImageBase;
    if (delta != 0) {
        KSTATUS rs = LdrApplyRelocationsForProcess((PROCESS*)0, (uintptr_t)mem, delta, false);
        if (!KSUCCESS(rs)) {
            MmFreePool(mem);
            return rs;
        }
    }

    // Kernel drivers are still pool-backed images. Section protections cannot
    // safely make pages read-only yet because imports and driver-owned writable
    // image data are not separated from callback code by a section object. What
    // must be true is simpler: every page in the loaded image must be executable
    // enough for driver callbacks, including section-alignment slack near the
    // end of the image.
    KSTATUS xs = MmPteProtectRange((uintptr_t)mem, imageSize, PAGE_PRESENT | PAGE_WRITE);
    if (!KSUCCESS(xs)) {
        MmFreePool(mem);
        return xs;
    }

    *baseOut = (uintptr_t)mem;
    return STATUS_SUCCESS;
}

// Loads a kernel image from a bootloader or driver raw buffer.
KSTATUS LdrLoadKernelImage(const void* buffer, size_t size, uintptr_t* baseOut) {
    return KiLdrLoadKernelImageFromBuffer(buffer, size, baseOut);
}