// User-mode smoke test for native exception dispatcher registration.

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <intrin.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("exception_dispatcher_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("exception_dispatcher_smoke: ");
    print(what);
    print(" status=");
    print_u64(status);
    print("\n");
    exit_now(1);
}

// Dummy dispatcher used only as a validated user-mode code pointer.
static void smoke_dispatcher(SYS_EXCEPTION_FRAME* frame) {
    (void)frame;
}

// Dispatcher hit count used by the delivery/continue stage.
static volatile uint32_t gBreakpointDispatchCount;

// Continues from a breakpoint exception using the kernel-provided context.
static void smoke_continue_dispatcher(SYS_EXCEPTION_FRAME* frame) {
    if (!frame)
        fail("dispatcher received null frame");
    if (frame->Size != sizeof(*frame) || frame->Version != SYS_EXCEPTION_ABI_VERSION)
        fail("dispatcher frame ABI mismatch");
    if (frame->ExceptionRecord.Code != SYS_EXCEPTION_CODE_BREAKPOINT)
        fail("dispatcher received unexpected exception code");
    if (frame->DispatcherContext != (void*)(uintptr_t)0xBEEF)
        fail("dispatcher context mismatch during delivery");

    gBreakpointDispatchCount++;

    SYS_CONTINUE_EXCEPTION_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Disposition = SYS_EXCEPTION_DISPOSITION_CONTINUE_EXECUTION;
    request.Context = &frame->ContextRecord;

    uint32_t status = SysContinueException(&request);
    fail_status("continue exception returned", status);
}

// Requires one syscall to return the expected status.
static void expect_status(const char* what, uint32_t got, uint32_t expected) {
    if (got != expected) {
        print("exception_dispatcher_smoke: ");
        print(what);
        print(" expected=");
        print_u64(expected);
        print(" got=");
        print_u64(got);
        print("\n");
        exit_now(1);
    }
}

// Reads the current dispatcher registration.
static SYS_EXCEPTION_DISPATCHER_REQUEST get_registration(void) {
    SYS_EXCEPTION_DISPATCHER_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    uint32_t status = SysGetExceptionDispatcher(&request);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("get dispatcher", status);
    if (request.Size != sizeof(request))
        fail("get dispatcher size mismatch");
    return request;
}

// Runs the native dispatcher registration smoke test.
int main(void) {
    print("exception_dispatcher_smoke: start\n");

    expect_status("null get", SysGetExceptionDispatcher(0), NTLSTATUS_INVALID_PARAMETER);
    expect_status("null set", SysSetExceptionDispatcher(0), NTLSTATUS_INVALID_PARAMETER);

    SYS_EXCEPTION_DISPATCHER_REQUEST initial = get_registration();
    SYS_EXCEPTION_DISPATCHER_REQUEST request;

    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request) - 1;
    request.DispatcherRoutine = smoke_dispatcher;
    expect_status("bad size", SysSetExceptionDispatcher(&request), NTLSTATUS_INVALID_PARAMETER);

    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Flags = 1;
    request.DispatcherRoutine = smoke_dispatcher;
    expect_status("bad flags", SysSetExceptionDispatcher(&request), NTLSTATUS_INVALID_PARAMETER);

    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.DispatcherRoutine = (SYS_USER_EXCEPTION_DISPATCHER)(uintptr_t)0xFFFF800000000000ull;
    expect_status("kernel dispatcher", SysSetExceptionDispatcher(&request), NTLSTATUS_ACCESS_DENIED);

    char tinyStack[32];
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.DispatcherRoutine = smoke_dispatcher;
    request.ExceptionStackBase = tinyStack;
    request.ExceptionStackSize = sizeof(tinyStack);
    expect_status("tiny exception stack", SysSetExceptionDispatcher(&request), NTLSTATUS_BUFFER_TOO_SMALL);

    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.DispatcherRoutine = smoke_dispatcher;
    request.DispatcherContext = (void*)(uintptr_t)0x1234;
    uint32_t status = SysSetExceptionDispatcher(&request);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("set dispatcher", status);

    SYS_EXCEPTION_DISPATCHER_REQUEST snapshot = get_registration();
    if (snapshot.DispatcherRoutine != smoke_dispatcher)
        fail("dispatcher routine mismatch");
    if (snapshot.DispatcherContext != (void*)(uintptr_t)0x1234)
        fail("dispatcher context mismatch");
    if (snapshot.ExceptionStackBase || snapshot.ExceptionStackSize)
        fail("unexpected exception stack registration");

    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.DispatcherRoutine = smoke_continue_dispatcher;
    request.DispatcherContext = (void*)(uintptr_t)0xBEEF;
    status = SysSetExceptionDispatcher(&request);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("set continue dispatcher", status);

    __debugbreak();
    if (gBreakpointDispatchCount != 1)
        fail("breakpoint was not delivered exactly once");

    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    status = SysSetExceptionDispatcher(&request);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("clear dispatcher", status);

    snapshot = get_registration();
    if (snapshot.DispatcherRoutine || snapshot.DispatcherContext ||
        snapshot.ExceptionStackBase || snapshot.ExceptionStackSize)
        fail("dispatcher did not clear");

    if (initial.DispatcherRoutine) {
        status = SysSetExceptionDispatcher(&initial);
        if (status != NTLSTATUS_SUCCESS)
            fail_status("restore initial dispatcher", status);
    }

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