// User-mode smoke test for public dynamic module loading.

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

// Prints one NUL-terminated string to the kernel console.
static void print(const char* s) {
    SYS_WRITE_CONSOLE_REQUEST req = { SysSvFromCstr(s) };
    SysWriteConsole(&req);
}

// Prints one unsigned integer.
static void print_u64(uint64_t v) {
    char buf[32];
    int i = 0;
    if (v == 0) {
        print("0");
        return;
    }
    while (v && i < (int)sizeof(buf)) {
        buf[i++] = (char)('0' + (v % 10));
        v /= 10;
    }
    while (i--) {
        char c = buf[i];
        SYS_WRITE_CONSOLE_REQUEST req = { SysSvFromParts(&c, 1) };
        SysWriteConsole(&req);
    }
}

// Terminates the current process.
static void exit_now(int32_t code) {
    while (true) SysExitProcess(code);
}

// Fails the test with one message.
static void fail(const char* message) {
    print("module_load_smoke: ");
    print(message);
    print("\n");
    exit_now(1);
}

// Fails the test with one status value.
static void fail_status(const char* what, uint32_t status) {
    print("module_load_smoke: ");
    print(what);
    print(" status=");
    print_u64(status);
    print("\n");
    exit_now(1);
}

// Checks whether one public module snapshot has the expected shape.
static void check_module_info(const SYS_MODULE_INFO* info) {
    if (!info) fail("null module info");
    if (info->Size != sizeof(*info)) fail("module info size mismatch");
    if (info->Version != SYS_MODULE_INFO_VERSION) fail("module info version mismatch");
    if (info->Base == 0 || info->ImageSize == 0) fail("module has empty mapping");
    if ((info->Flags & SYS_MODULE_FLAG_PUBLIC_REF) == 0) fail("module missing public reference flag");
    if (info->PublicRefCount == 0 || info->TotalRefCount == 0) fail("module reference counts missing");
    if (strcmp(info->Name, "posix.dll") != 0) fail("module name mismatch");
}

// Returns whether the current module list contains the specified base address.
static int module_list_contains(uintptr_t base) {
    size_t count = 0;
    uint32_t status = SysGetModuleList(0, 0, &count);
    if (status != NTLSTATUS_SUCCESS) fail_status("query module count", status);
    if (count == 0) return 0;
    if (count > 32) fail("module count unexpectedly large");

    SYS_MODULE_INFO modules[32];
    memset(modules, 0, sizeof(modules));
    status = SysGetModuleList(modules, count, &count);
    if (status != NTLSTATUS_SUCCESS) fail_status("query module list", status);

    for (size_t i = 0; i < count; i++) {
        if (modules[i].Base == base)
            return 1;
    }
    return 0;
}


// Verifies export lookup through a module handle.
static void check_export_lookup(void) {
    uint32_t handle = 0;
    SYS_LOAD_MODULE_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Size = sizeof(req);
    req.Path = SysSvFromCstr("syscall.dll");
    req.ModuleHandleOut = &handle;

    uint32_t status = SysLoadModule(&req);
    if (status != NTLSTATUS_SUCCESS) fail_status("load syscall.dll", status);
    if (handle == 0) fail("syscall load returned null handle");

    uintptr_t address = 0;
    SYS_GET_MODULE_EXPORT_REQUEST exportReq;
    memset(&exportReq, 0, sizeof(exportReq));
    exportReq.Size = sizeof(exportReq);
    exportReq.ModuleHandle = handle;
    exportReq.Name = SysSvFromCstr("SysGetCurrentPid");
    exportReq.AddressOut = &address;

    status = SysGetModuleExport(&exportReq);
    if (status != NTLSTATUS_SUCCESS) fail_status("lookup SysGetCurrentPid", status);
    if (address == 0) fail("export address is zero");

    status = SysUnloadModule(handle);
    if (status != NTLSTATUS_SUCCESS) fail_status("unload syscall.dll handle", status);
}

// Runs the public dynamic module loader smoke test.
int main(void) {
    print("module_load_smoke: start\n");
    check_export_lookup();

    uint32_t handle = 0;
    SYS_MODULE_INFO info;
    memset(&info, 0, sizeof(info));

    SYS_LOAD_MODULE_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Size = sizeof(req);
    req.Path = SysSvFromCstr("posix.dll");
    req.ModuleHandleOut = &handle;
    req.InfoOut = &info;

    uint32_t status = SysLoadModule(&req);
    if (status != NTLSTATUS_SUCCESS) fail_status("load posix.dll", status);
    if (handle == 0) fail("load returned null handle");
    check_module_info(&info);
    if (!module_list_contains(info.Base)) fail("loaded module absent from list");

    SYS_MODULE_INFO byHandle;
    memset(&byHandle, 0, sizeof(byHandle));
    status = SysGetModuleInfo(handle, &byHandle);
    if (status != NTLSTATUS_SUCCESS) fail_status("get module info", status);
    check_module_info(&byHandle);
    if (byHandle.Base != info.Base) fail("module handle base mismatch");

    status = SysUnloadModule(handle);
    if (status != NTLSTATUS_SUCCESS) fail_status("unload module", status);

    memset(&byHandle, 0, sizeof(byHandle));
    status = SysGetModuleInfo(handle, &byHandle);
    if (status != NTLSTATUS_INVALID_HANDLE) fail_status("query unloaded handle", status);

    print("module_load_smoke: PASS\n");
    exit_now(0);
    return 0;
}