// Native exception pipeline integration smoke test.

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

#define CHILD_DIVIDE_BY_ZERO "child-divide-by-zero"
#define CHILD_ACCESS_VIOLATION "child-access-violation"
#define CHILD_INVALID_RESUME "child-invalid-resume"
#define CHILD_NONCONTINUABLE "child-noncontinuable"
#define SMOKE_NONCONTINUABLE_CODE 0xE1111003u

// 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_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_smoke: ");
    print(what);
    print(" status=");
    print_u64(status);
    print("\n");
    exit_now(1);
}

// Returns whether a string view exactly matches a literal string.
static int sv_equals_cstr(SYS_STRING_VIEW sv, const char* literal) {
    size_t len = strlen(literal);
    return sv.Length == len && sv.Data && memcmp(sv.Data, literal, len) == 0;
}

// Returns the current process command line.
static SYS_STRING_VIEW current_command_line(void) {
    SYS_PEB* peb = SysGetPeb();
    if (!peb)
        return SysSvFromParts(NULL, 0);
    return peb->CommandLine;
}

// 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_smoke: ");
        print(what);
        print(" expected=");
        print_u64(expected);
        print(" got=");
        print_u64(got);
        print("\n");
        exit_now(1);
    }
}

// Runs this image as a child and requires the expected exit code.
static void run_child_expect_exit(const char* mode, uint32_t expectedExit) {
    SYS_PEB* peb = SysGetPeb();
    if (!peb || !peb->ImagePath.Data || peb->ImagePath.Length == 0)
        fail("missing image path");

    uint32_t child = 0;
    SYS_CREATE_PROCESS_REQUEST request = {
        peb->ImagePath,
        SysSvFromCstr(mode),
        &child,
    };
    uint32_t status = SysCreateProcess(&request);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("create child", status);

    status = SysWait(child, SYS_WAIT_FLAG_NONE);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("wait child", status);

    int32_t exitCode = 0;
    status = SysGetProcessExitCode(child, &exitCode);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("query child", status);
    SysCloseHandle(child);

    if ((uint32_t)exitCode != expectedExit) {
        print("exception_smoke: child ");
        print(mode);
        print(" expected exit=");
        print_u64(expectedExit);
        print(" got=");
        print_u64((uint32_t)exitCode);
        print("\n");
        exit_now(1);
    }
}

// Triggers a divide-by-zero with no registered native handlers.
static void child_divide_by_zero(void) {
    volatile int one = 1;
    volatile int zero = 0;
    volatile int sink = one / zero;
    (void)sink;
    exit_now(2);
}

// Observes an access violation and exits from inside the native handler.
static int access_violation_handler(SYS_EXCEPTION_RECORD* record,
                                    SYS_USER_CONTEXT* context,
                                    void* userContext) {
    (void)context;
    if (userContext != (void*)(uintptr_t)0xA11C)
        exit_now(10);
    if (!record || record->Code != SYS_EXCEPTION_CODE_ACCESS_VIOLATION)
        exit_now(11);
    if (record->FaultAddress != 0)
        exit_now(12);
    if (record->ParameterCount != 2 || record->Parameters[1] != 0)
        exit_now(13);
    exit_now(0);
    return NTL_EXCEPTION_TERMINATE_PROCESS;
}

// Triggers a null read so the handler can inspect the fault address.
static void child_access_violation(void) {
    void* handle = NtlAddExceptionHandler(1, access_violation_handler, (void*)(uintptr_t)0xA11C);
    if (!handle)
        exit_now(20);

    volatile uint8_t* p = (volatile uint8_t*)0;
    volatile uint8_t value = *p;
    (void)value;
    exit_now(21);
}

// Returns a resume context with an invalid user RIP.
static int invalid_resume_handler(SYS_EXCEPTION_RECORD* record,
                                  SYS_USER_CONTEXT* context,
                                  void* userContext) {
    (void)userContext;
    if (!record || !context || record->Code != SYS_EXCEPTION_CODE_BREAKPOINT)
        exit_now(30);
    context->Rip = 0;
    return NTL_EXCEPTION_CONTINUE_EXECUTION;
}

// Triggers a breakpoint whose handler asks the kernel to resume at RIP zero.
static void child_invalid_resume(void) {
    void* handle = NtlAddExceptionHandler(1, invalid_resume_handler, NULL);
    if (!handle)
        exit_now(31);
    __debugbreak();
    exit_now(32);
}

// Attempts to continue a noncontinuable user-raised exception.
static int noncontinuable_handler(SYS_EXCEPTION_RECORD* record,
                                  SYS_USER_CONTEXT* context,
                                  void* userContext) {
    (void)context;
    (void)userContext;
    if (!record || record->Code != SMOKE_NONCONTINUABLE_CODE)
        exit_now(40);
    if ((record->Flags & SYS_EXCEPTION_FLAG_NONCONTINUABLE) == 0)
        exit_now(41);
    return NTL_EXCEPTION_CONTINUE_EXECUTION;
}

// Raises a noncontinuable exception that a handler must not resume.
static void child_noncontinuable(void) {
    void* handle = NtlAddExceptionHandler(1, noncontinuable_handler, NULL);
    if (!handle)
        exit_now(42);

    SYS_RAISE_EXCEPTION_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Flags = SYS_RAISE_EXCEPTION_FLAG_NONCONTINUABLE;
    request.Code = SMOKE_NONCONTINUABLE_CODE;

    uint32_t status = SysRaiseException(&request);
    exit_now((int32_t)status);
}

// Dispatches the child process mode selected by the command line.
static void run_child_if_requested(void) {
    SYS_STRING_VIEW command = current_command_line();
    if (sv_equals_cstr(command, CHILD_DIVIDE_BY_ZERO))
        child_divide_by_zero();
    if (sv_equals_cstr(command, CHILD_ACCESS_VIOLATION))
        child_access_violation();
    if (sv_equals_cstr(command, CHILD_INVALID_RESUME))
        child_invalid_resume();
    if (sv_equals_cstr(command, CHILD_NONCONTINUABLE))
        child_noncontinuable();
}

// Runs validation that does not intentionally fault this process.
static void verify_continue_validation(void) {
    expect_status("null continue", SysContinueException(NULL), NTLSTATUS_INVALID_PARAMETER);

    SYS_CONTINUE_EXCEPTION_REQUEST request;
    memset(&request, 0, sizeof(request));
    request.Size = sizeof(request);
    request.Disposition = SYS_EXCEPTION_DISPOSITION_CONTINUE_EXECUTION;
    expect_status("inactive continue", SysContinueException(&request), NTLSTATUS_INVALID_PARAMETER);
}

// Runs the native exception integration smoke test.
int main(void) {
    run_child_if_requested();

    print("exception_smoke: start\n");

    verify_continue_validation();
    run_child_expect_exit(CHILD_DIVIDE_BY_ZERO, NTLSTATUS_PROCESS_DIVIDE_BY_ZERO);
    run_child_expect_exit(CHILD_ACCESS_VIOLATION, NTLSTATUS_SUCCESS);
    run_child_expect_exit(CHILD_INVALID_RESUME, NTLSTATUS_ACCESS_DENIED);
    run_child_expect_exit(CHILD_NONCONTINUABLE, NTLSTATUS_PROCESS_UNKNOWN_FAULT);

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