// Userspace PE32+ image loader. The kernel sees only generic process memory and execution state.

#include <ntlite/pe_loader.h>
#include <ntlite/syscall.h>
#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>

#define PE_DOS_MAGIC 0x5A4Du
#define PE_SIGNATURE 0x00004550u
#define PE_MACHINE_AMD64 0x8664u
#define PE_FILE_RELOCS_STRIPPED 0x0001u
#define PE_OPTIONAL_MAGIC_PE32_PLUS 0x020Bu
#define PE_OPTIONAL_MIN_BYTES 112u
#define PE_SECTION_BYTES 40u
#define PE_DIRECTORY_IMPORT 1u
#define PE_DIRECTORY_RELOC 5u
#define PE_DIRECTORY_DELAY_IMPORT 13u
#define PE_ORDINAL_FLAG64 0x8000000000000000ull
#define PE_RELOC_DIR64 10u
#define PE_PAGE_SIZE 4096u
#define PE_SCN_MEM_EXECUTE 0x20000000u
#define PE_SCN_MEM_READ    0x40000000u
#define PE_SCN_MEM_WRITE   0x80000000u

typedef struct PE_VIEW {
    const uint8_t* File;
    size_t FileSize;
    size_t OptionalOffset;
    size_t SectionOffset;
    uint16_t SectionCount;
    uint16_t OptionalSize;
    uint16_t FileCharacteristics;
    uint32_t DirectoryCount;
    uint64_t ImageBase;
    uint32_t ImageSize;
    uint32_t HeaderSize;
    uint32_t EntryRva;
} PE_VIEW;

static bool PeRange(size_t offset, size_t length, size_t size) {
    return offset <= size && length <= size - offset;
}

static uint16_t PeU16(const void* p) { uint16_t v; memcpy(&v, p, sizeof(v)); return v; }
static uint32_t PeU32(const void* p) { uint32_t v; memcpy(&v, p, sizeof(v)); return v; }
static uint64_t PeU64(const void* p) { uint64_t v; memcpy(&v, p, sizeof(v)); return v; }
static void PeW64(void* p, uint64_t v) { memcpy(p, &v, sizeof(v)); }

static bool PeImageRange(uint32_t rva, size_t length, size_t imageSize) {
    return (size_t)rva <= imageSize && length <= imageSize - (size_t)rva;
}

static bool PeImageString(const uint8_t* image, size_t imageSize, uint32_t rva, size_t* lengthOut) {
    if ((size_t)rva >= imageSize) return false;
    for (size_t i = rva; i < imageSize; ++i) {
        if (image[i] == 0) {
            if (lengthOut) *lengthOut = i - rva;
            return true;
        }
    }
    return false;
}

static KSTATUS PeParse(const void* file, size_t fileSize, PE_VIEW* out) {
    if (!file || !out || fileSize < 64u) return NTLSTATUS_FAILURE;
    const uint8_t* b = (const uint8_t*)file;
    if (PeU16(b) != PE_DOS_MAGIC) return NTLSTATUS_FAILURE;
    size_t nt = PeU32(b + 0x3cu);
    if (!PeRange(nt, 24u, fileSize) || PeU32(b + nt) != PE_SIGNATURE) return NTLSTATUS_FAILURE;
    if (PeU16(b + nt + 4u) != PE_MACHINE_AMD64) return NTLSTATUS_NOT_SUPPORTED;
    uint16_t sections = PeU16(b + nt + 6u);
    uint16_t optionalSize = PeU16(b + nt + 20u);
    uint16_t fileCharacteristics = PeU16(b + nt + 22u);
    size_t optional = nt + 24u;
    if (optionalSize < PE_OPTIONAL_MIN_BYTES || !PeRange(optional, optionalSize, fileSize)) return NTLSTATUS_FAILURE;
    if (PeU16(b + optional) != PE_OPTIONAL_MAGIC_PE32_PLUS) return NTLSTATUS_NOT_SUPPORTED;

    uint64_t imageBase = PeU64(b + optional + 24u);
    uint32_t imageSize = PeU32(b + optional + 56u);
    uint32_t headerSize = PeU32(b + optional + 60u);
    uint32_t entryRva = PeU32(b + optional + 16u);
    uint32_t directories = PeU32(b + optional + 108u);
    uint32_t available = (uint32_t)((optionalSize - PE_OPTIONAL_MIN_BYTES) / 8u);
    if (!imageBase || !imageSize || !headerSize || headerSize > imageSize || entryRva >= imageSize || directories > available)
        return NTLSTATUS_FAILURE;
    if (headerSize > fileSize) return NTLSTATUS_FAILURE;
    if (directories > PE_DIRECTORY_DELAY_IMPORT) {
        const uint8_t* delay = b + optional + PE_OPTIONAL_MIN_BYTES + PE_DIRECTORY_DELAY_IMPORT * 8u;
        if (PeU32(delay) || PeU32(delay + 4u)) return NTLSTATUS_NOT_SUPPORTED;
    }

    size_t sectionOffset = optional + optionalSize;
    size_t sectionBytes = (size_t)sections * PE_SECTION_BYTES;
    if (sections && sectionBytes / PE_SECTION_BYTES != sections) return NTLSTATUS_FAILURE;
    if (!PeRange(sectionOffset, sectionBytes, fileSize)) return NTLSTATUS_FAILURE;

    for (uint16_t i = 0; i < sections; ++i) {
        const uint8_t* sec = b + sectionOffset + (size_t)i * PE_SECTION_BYTES;
        uint32_t va = PeU32(sec + 12u);
        uint32_t rawSize = PeU32(sec + 16u);
        uint32_t raw = PeU32(sec + 20u);
        uint32_t virtualSize = PeU32(sec + 8u);
        uint32_t span = virtualSize > rawSize ? virtualSize : rawSize;
        if ((va & (PE_PAGE_SIZE - 1u)) != 0) return NTLSTATUS_NOT_SUPPORTED;
        if (!PeImageRange(va, span, imageSize)) return NTLSTATUS_FAILURE;
        if (rawSize && !PeRange(raw, rawSize, fileSize)) return NTLSTATUS_FAILURE;
    }

    memset(out, 0, sizeof(*out));
    out->File = b;
    out->FileSize = fileSize;
    out->OptionalOffset = optional;
    out->SectionOffset = sectionOffset;
    out->SectionCount = sections;
    out->OptionalSize = optionalSize;
    out->FileCharacteristics = fileCharacteristics;
    out->DirectoryCount = directories;
    out->ImageBase = imageBase;
    out->ImageSize = imageSize;
    out->HeaderSize = headerSize;
    out->EntryRva = entryRva;
    return NTLSTATUS_SUCCESS;
}

static bool PeDirectory(const PE_VIEW* view, uint32_t index, uint32_t* rva, uint32_t* size) {
    if (!view || index >= view->DirectoryCount) return false;
    const uint8_t* d = view->File + view->OptionalOffset + PE_OPTIONAL_MIN_BYTES + (size_t)index * 8u;
    if (rva) *rva = PeU32(d);
    if (size) *size = PeU32(d + 4u);
    return true;
}

static KSTATUS PeAllocCurrent(size_t bytes, void** out) {
    if (!out || !bytes) return NTLSTATUS_INVALID_PARAMETER;
    size_t pages = (bytes + PE_PAGE_SIZE - 1u) / PE_PAGE_SIZE;
    SYS_ALLOC_VM_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Size = sizeof(req);
    req.Protect = SYS_VM_READWRITE;
    req.Flags = SYS_VM_FLAG_NONE;
    req.PageCount = pages;
    req.BaseOut = out;
    return SysAllocVm(&req);
}

static KSTATUS PeReadFile(SYS_HANDLE file, void** bufferOut, size_t* sizeOut) {
    if (!bufferOut || !sizeOut) return NTLSTATUS_INVALID_PARAMETER;
    *bufferOut = 0; *sizeOut = 0;
    uint64_t fileSize64 = 0;
    KSTATUS s = SysGetFileSize(file, &fileSize64);
    if (!KSUCCESS(s)) return s;
    if (!fileSize64 || fileSize64 > (uint64_t)SIZE_MAX) return NTLSTATUS_BUFFER_TOO_SMALL;
    size_t fileSize = (size_t)fileSize64;
    void* buffer = 0;
    s = PeAllocCurrent(fileSize, &buffer);
    if (!KSUCCESS(s)) return s;
    s = SysSeekFile(file, 0);
    if (!KSUCCESS(s)) { SysFreeVm(buffer); return s; }
    size_t got = 0;
    SYS_FILE_IO_REQUEST read = SysMakeFileIoRequest(file, buffer, fileSize, &got);
    s = SysReadFile(&read);
    if (!KSUCCESS(s) || got != fileSize) {
        SysFreeVm(buffer);
        return KSUCCESS(s) ? NTLSTATUS_DEVICE_ERROR : s;
    }
    *bufferOut = buffer; *sizeOut = fileSize;
    return NTLSTATUS_SUCCESS;
}

static KSTATUS PeRelocate(uint8_t* image, size_t imageSize, const PE_VIEW* view, uintptr_t actualBase) {
    uintptr_t delta = actualBase - (uintptr_t)view->ImageBase;
    if (!delta) return NTLSTATUS_SUCCESS;
    uint32_t rva = 0, size = 0;
    if (!PeDirectory(view, PE_DIRECTORY_RELOC, &rva, &size) || !rva || !size) {
        // An image with no base fixups can still be position-independent. Only the
        // COFF RELOCS_STRIPPED bit says missing relocation data makes rebasing invalid.
        return (view->FileCharacteristics & PE_FILE_RELOCS_STRIPPED)
                   ? NTLSTATUS_NOT_SUPPORTED
                   : NTLSTATUS_SUCCESS;
    }
    if (!PeImageRange(rva, size, imageSize)) return NTLSTATUS_FAILURE;
    size_t offset = rva;
    size_t end = (size_t)rva + size;
    while (offset < end) {
        if (end - offset < 8u) return NTLSTATUS_FAILURE;
        uint32_t pageRva = PeU32(image + offset);
        uint32_t blockSize = PeU32(image + offset + 4u);
        if (blockSize < 8u || blockSize > end - offset || ((blockSize - 8u) & 1u)) return NTLSTATUS_FAILURE;
        uint32_t count = (blockSize - 8u) / 2u;
        for (uint32_t i = 0; i < count; ++i) {
            uint16_t e = PeU16(image + offset + 8u + (size_t)i * 2u);
            uint16_t type = e >> 12;
            uint16_t pageOffset = e & 0x0fffu;
            if (type == 0) continue;
            if (type != PE_RELOC_DIR64) return NTLSTATUS_NOT_SUPPORTED;
            uint32_t fixRva = pageRva + pageOffset;
            if (fixRva < pageRva || !PeImageRange(fixRva, 8u, imageSize)) return NTLSTATUS_FAILURE;
            PeW64(image + fixRva, PeU64(image + fixRva) + delta);
        }
        offset += blockSize;
    }
    return offset == end ? NTLSTATUS_SUCCESS : NTLSTATUS_FAILURE;
}

static KSTATUS PeBindImports(uint8_t* image, size_t imageSize, const PE_VIEW* view,
                             NTL_PE_IMPORT_RESOLVER resolver, void* context) {
    uint32_t rva = 0, size = 0;
    if (!PeDirectory(view, PE_DIRECTORY_IMPORT, &rva, &size) || !rva || !size)
        return NTLSTATUS_SUCCESS;
    if (!resolver) return NTLSTATUS_NOT_FOUND;
    if (!PeImageRange(rva, size, imageSize)) return NTLSTATUS_FAILURE;

    size_t offset = rva;
    size_t end = (size_t)rva + size;
    while (offset + 20u <= end) {
        uint8_t* d = image + offset;
        uint32_t originalThunk = PeU32(d + 0u);
        uint32_t nameRva = PeU32(d + 12u);
        uint32_t firstThunk = PeU32(d + 16u);
        if (!originalThunk && !nameRva && !firstThunk) return NTLSTATUS_SUCCESS;
        size_t moduleLen = 0;
        if (!nameRva || !PeImageString(image, imageSize, nameRva, &moduleLen)) return NTLSTATUS_FAILURE;
        uint32_t sourceThunk = originalThunk ? originalThunk : firstThunk;
        if (!sourceThunk || !firstThunk) return NTLSTATUS_FAILURE;

        for (uint32_t index = 0;; ++index) {
            size_t sourceRva = (size_t)sourceThunk + (size_t)index * 8u;
            size_t targetRva = (size_t)firstThunk + (size_t)index * 8u;
            if (sourceRva + 8u > imageSize || targetRva + 8u > imageSize) return NTLSTATUS_FAILURE;
            uint64_t thunk = PeU64(image + sourceRva);
            if (!thunk) break;
            if (thunk & PE_ORDINAL_FLAG64) return NTLSTATUS_NOT_SUPPORTED;
            uint32_t hintNameRva = (uint32_t)thunk;
            if (!PeImageRange(hintNameRva, 3u, imageSize)) return NTLSTATUS_FAILURE;
            size_t symbolLen = 0;
            if (!PeImageString(image, imageSize, hintNameRva + 2u, &symbolLen)) return NTLSTATUS_FAILURE;
            uintptr_t address = 0;
            KSTATUS s = resolver(context,
                                 SysSvFromParts((const char*)image + nameRva, moduleLen),
                                 SysSvFromParts((const char*)image + hintNameRva + 2u, symbolLen),
                                 &address);
            if (!KSUCCESS(s)) return s;
            if (!address) return NTLSTATUS_NOT_FOUND;
            PeW64(image + targetRva, address);
        }
        offset += 20u;
    }
    return NTLSTATUS_FAILURE;
}

static uint32_t PeSectionProtect(uint32_t characteristics) {
    bool r = (characteristics & PE_SCN_MEM_READ) != 0;
    bool w = (characteristics & PE_SCN_MEM_WRITE) != 0;
    bool x = (characteristics & PE_SCN_MEM_EXECUTE) != 0;
    if (x && w) return SYS_VM_EXECUTE_READWRITE;
    if (x && r) return SYS_VM_EXECUTE_READ;
    if (x) return SYS_VM_EXECUTE;
    if (w) return SYS_VM_READWRITE;
    if (r) return SYS_VM_READ;
    return SYS_VM_NOACCESS;
}

static KSTATUS PeProtectRange(SYS_HANDLE process, void* base, size_t pages, uint32_t protect) {
    if (process != SYS_HANDLE_INVALID)
        return SysProtectProcessVm(process, base, pages, protect);
    return SysProtectVm(base, pages, protect, (uint32_t*)0);
}

static KSTATUS PeProtectImage(SYS_HANDLE process, uint8_t* image, const PE_VIEW* view) {
    size_t pages = ((size_t)view->ImageSize + PE_PAGE_SIZE - 1u) / PE_PAGE_SIZE;
    KSTATUS s = PeProtectRange(process, image, pages, SYS_VM_READ);
    if (!KSUCCESS(s)) return s;
    for (uint16_t i = 0; i < view->SectionCount; ++i) {
        const uint8_t* sec = view->File + view->SectionOffset + (size_t)i * PE_SECTION_BYTES;
        uint32_t va = PeU32(sec + 12u);
        uint32_t virtualSize = PeU32(sec + 8u);
        uint32_t rawSize = PeU32(sec + 16u);
        uint32_t span = virtualSize > rawSize ? virtualSize : rawSize;
        if (!span) continue;
        size_t sectionPages = ((size_t)span + PE_PAGE_SIZE - 1u) / PE_PAGE_SIZE;
        s = PeProtectRange(process, image + va, sectionPages, PeSectionProtect(PeU32(sec + 36u)));
        if (!KSUCCESS(s)) return s;
    }
    return NTLSTATUS_SUCCESS;
}

static KSTATUS PeAllocateImage(SYS_HANDLE process, size_t pages, void** baseOut) {
    if (process == SYS_HANDLE_INVALID) {
        SYS_ALLOC_VM_REQUEST alloc;
        memset(&alloc, 0, sizeof(alloc));
        alloc.Size = sizeof(alloc);
        alloc.Protect = SYS_VM_READWRITE;
        alloc.Flags = SYS_VM_FLAG_NONE;
        alloc.PageCount = pages;
        alloc.BaseOut = baseOut;
        return SysAllocVm(&alloc);
    }

    SYS_ALLOC_PROCESS_VM_REQUEST alloc;
    memset(&alloc, 0, sizeof(alloc));
    alloc.Size = sizeof(alloc);
    alloc.ProcessHandle = process;
    alloc.Protect = SYS_VM_READWRITE;
    alloc.Flags = SYS_VM_FLAG_NONE;
    alloc.PageCount = pages;
    alloc.BaseOut = baseOut;
    return SysAllocProcessVm(&alloc);
}

KSTATUS NtlLoadPeImage(const NTL_PE_LOAD_REQUEST* request) {
    if (!request || request->Size != sizeof(*request) || request->Flags != NTL_PE_LOAD_FLAG_NONE ||
        request->FileHandle == SYS_HANDLE_INVALID || !request->ImageOut)
        return NTLSTATUS_INVALID_PARAMETER;

    void* file = 0;
    size_t fileSize = 0;
    KSTATUS s = PeReadFile(request->FileHandle, &file, &fileSize);
    if (!KSUCCESS(s)) return s;

    PE_VIEW view;
    s = PeParse(file, fileSize, &view);
    if (!KSUCCESS(s)) { SysFreeVm(file); return s; }

    size_t pages = ((size_t)view.ImageSize + PE_PAGE_SIZE - 1u) / PE_PAGE_SIZE;
    void* imageBase = 0;
    s = PeAllocateImage(request->ProcessHandle, pages, &imageBase);
    if (!KSUCCESS(s)) { SysFreeVm(file); return s; }

    uint8_t* image = (uint8_t*)imageBase;
    memset(image, 0, pages * PE_PAGE_SIZE);
    memcpy(image, file, view.HeaderSize);
    for (uint16_t i = 0; i < view.SectionCount; ++i) {
        const uint8_t* sec = view.File + view.SectionOffset + (size_t)i * PE_SECTION_BYTES;
        uint32_t va = PeU32(sec + 12u);
        uint32_t rawSize = PeU32(sec + 16u);
        uint32_t raw = PeU32(sec + 20u);
        if (rawSize) memcpy(image + va, (const uint8_t*)file + raw, rawSize);
    }

    s = PeRelocate(image, view.ImageSize, &view, (uintptr_t)image);
    if (KSUCCESS(s)) {
        request->ImageOut->Base = (uintptr_t)image;
        request->ImageOut->Size = view.ImageSize;
        request->ImageOut->EntryPoint = (uintptr_t)image + view.EntryRva;
        request->ImageOut->PreferredBase = view.ImageBase;
        if (request->Prepared)
            s = request->Prepared(request->PreparedContext, request->ImageOut);
    }
    if (KSUCCESS(s)) s = PeBindImports(image, view.ImageSize, &view, request->ResolveImport, request->ResolveContext);
    if (KSUCCESS(s)) s = PeProtectImage(request->ProcessHandle, image, &view);
    if (!KSUCCESS(s) && request->ProcessHandle == SYS_HANDLE_INVALID)
        (void)SysFreeVm(image);
    SysFreeVm(file);
    return s;
}

KSTATUS NtlGetPeExport(uintptr_t imageBase, SYS_STRING_VIEW name, uintptr_t* addressOut) {
    if (!imageBase || !addressOut || !SysSvIsValid(name) || name.Length == 0 ||
        SysSvContainsByte(name, '\0'))
        return NTLSTATUS_INVALID_PARAMETER;
    *addressOut = 0;

    const uint8_t* image = (const uint8_t*)imageBase;
    if (PeU16(image) != PE_DOS_MAGIC) return NTLSTATUS_FAILURE;
    uint32_t nt = PeU32(image + 0x3cu);
    if (PeU32(image + nt) != PE_SIGNATURE || PeU16(image + nt + 4u) != PE_MACHINE_AMD64)
        return NTLSTATUS_FAILURE;
    const uint8_t* optional = image + nt + 24u;
    if (PeU16(optional) != PE_OPTIONAL_MAGIC_PE32_PLUS) return NTLSTATUS_FAILURE;
    uint32_t imageSize = PeU32(optional + 56u);
    uint32_t directories = PeU32(optional + 108u);
    if (directories < 1u) return NTLSTATUS_NOT_FOUND;
    uint32_t exportRva = PeU32(optional + PE_OPTIONAL_MIN_BYTES);
    uint32_t exportSize = PeU32(optional + PE_OPTIONAL_MIN_BYTES + 4u);
    if (!exportRva || !exportSize || !PeImageRange(exportRva, exportSize, imageSize) ||
        !PeImageRange(exportRva, 40u, imageSize))
        return NTLSTATUS_NOT_FOUND;

    const uint8_t* exp = image + exportRva;
    uint32_t functionCount = PeU32(exp + 20u);
    uint32_t nameCount = PeU32(exp + 24u);
    uint32_t functionsRva = PeU32(exp + 28u);
    uint32_t namesRva = PeU32(exp + 32u);
    uint32_t ordinalsRva = PeU32(exp + 36u);
    if (!PeImageRange(functionsRva, (size_t)functionCount * 4u, imageSize) ||
        !PeImageRange(namesRva, (size_t)nameCount * 4u, imageSize) ||
        !PeImageRange(ordinalsRva, (size_t)nameCount * 2u, imageSize))
        return NTLSTATUS_FAILURE;

    for (uint32_t i = 0; i < nameCount; ++i) {
        uint32_t nameRva = PeU32(image + namesRva + (size_t)i * 4u);
        size_t exportedLength = 0;
        if (!PeImageString(image, imageSize, nameRva, &exportedLength)) return NTLSTATUS_FAILURE;
        if (exportedLength != name.Length || memcmp(image + nameRva, name.Data, name.Length) != 0)
            continue;
        uint16_t ordinal = PeU16(image + ordinalsRva + (size_t)i * 2u);
        if (ordinal >= functionCount) return NTLSTATUS_FAILURE;
        uint32_t functionRva = PeU32(image + functionsRva + (size_t)ordinal * 4u);
        if (!functionRva || !PeImageRange(functionRva, 1u, imageSize)) return NTLSTATUS_NOT_FOUND;
        if (functionRva >= exportRva && functionRva < exportRva + exportSize)
            return NTLSTATUS_NOT_SUPPORTED;
        *addressOut = imageBase + functionRva;
        return NTLSTATUS_SUCCESS;
    }
    return NTLSTATUS_NOT_FOUND;
}