// Win32 vectored exception smoke test. This file verifies Win32 VEH policy over the native exception runtime.

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <intrin.h>
#include <stdint.h>
#include <stddef.h>
#include <ntlite/syscall.h>

static volatile uint32_t gFirstHits;

static volatile uint32_t gTailHits;
static volatile uint32_t gContinueHits;
static volatile uint32_t gOnlyHits;
static volatile uint32_t gRaisedHits;

// 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("win32_exception_smoke: ");
    print(message);
    print("\n");
    exit_now(1);
}

// Fails the test with one integer value.
static void fail_value(const char* message, uint64_t value) {
    print("win32_exception_smoke: ");
    print(message);
    print(" value=");
    print_u64(value);
    print("\n");
    exit_now(1);
}

// Verifies the Win32 exception pointers contain a breakpoint and editable context.
static void verify_breakpoint(EXCEPTION_POINTERS* info) {
    if (!info || !info->ExceptionRecord || !info->ContextRecord)
        fail("handler received null exception pointers");
    if (info->ExceptionRecord->ExceptionCode != EXCEPTION_BREAKPOINT)
        fail_value("handler received unexpected exception code", info->ExceptionRecord->ExceptionCode);
    if (!info->ExceptionRecord->ExceptionAddress)
        fail("handler received null exception address");
    if (!info->ContextRecord->Rip || !info->ContextRecord->Rsp)
        fail("handler received empty context");
}

// Handles the first breakpoint before later exception handlers can observe it.
static LONG first_continue_handler(EXCEPTION_POINTERS* info) {
    verify_breakpoint(info);
    gFirstHits++;
    return EXCEPTION_CONTINUE_EXECUTION;
}

// Records an exception pass and asks Win32 policy to keep searching.
static LONG tail_search_handler(EXCEPTION_POINTERS* info) {
    verify_breakpoint(info);
    gTailHits++;
    return EXCEPTION_CONTINUE_SEARCH;
}

// Continues a breakpoint after the exception-handler list declines it.
static LONG continue_handler(EXCEPTION_POINTERS* info) {
    verify_breakpoint(info);
    gContinueHits++;
    return EXCEPTION_CONTINUE_EXECUTION;
}

// Handles a breakpoint after earlier handlers have been removed.
static LONG only_continue_handler(EXCEPTION_POINTERS* info) {
    verify_breakpoint(info);
    gOnlyHits++;
    return EXCEPTION_CONTINUE_EXECUTION;
}


// Handles a user-raised Win32 exception and verifies parameter translation.
static LONG raised_continue_handler(EXCEPTION_POINTERS* info) {
    if (!info || !info->ExceptionRecord || !info->ContextRecord)
        fail("raised handler received null exception pointers");
    if (info->ExceptionRecord->ExceptionCode != 0xE1111002u)
        fail_value("raised handler received unexpected exception code", info->ExceptionRecord->ExceptionCode);
    if (info->ExceptionRecord->ExceptionFlags & EXCEPTION_NONCONTINUABLE)
        fail("raised handler received unexpected noncontinuable flag");
    if (info->ExceptionRecord->NumberParameters != 2)
        fail_value("raised handler received wrong parameter count", info->ExceptionRecord->NumberParameters);
    if (info->ExceptionRecord->ExceptionInformation[0] != 0x1122u ||
        info->ExceptionRecord->ExceptionInformation[1] != 0x3344u)
        fail("raised handler received wrong parameters");
    gRaisedHits++;
    return EXCEPTION_CONTINUE_EXECUTION;
}

// Verifies one successful remove operation.
static void expect_remove(const char* what, unsigned long removed) {
    if (!removed)
        fail(what);
}

// Runs the Win32 vectored exception smoke test.
int main(void) {
    print("win32_exception_smoke: start\n");

    if (AddVectoredExceptionHandler(0, 0))
        fail("null exception handler was accepted");
    if (GetLastError() != 87)
        fail_value("null handler did not set ERROR_INVALID_PARAMETER", GetLastError());

    void* tail = AddVectoredExceptionHandler(0, tail_search_handler);
    void* first = AddVectoredExceptionHandler(1, first_continue_handler);
    if (!tail || !first)
        fail("exception handler registration failed");

    __debugbreak();
    if (gFirstHits != 1 || gTailHits != 0)
        fail("first vectored handler did not run before tail handler");

    expect_remove("remove first failed", RemoveVectoredExceptionHandler(first));

    void* cont = AddVectoredContinueHandler(0, continue_handler);
    if (!cont)
        fail("continue handler registration failed");

    __debugbreak();
    if (gTailHits != 1 || gContinueHits != 1)
        fail("continue handler did not run after exception search");

    expect_remove("remove tail failed", RemoveVectoredExceptionHandler(tail));
    expect_remove("remove continue failed", RemoveVectoredContinueHandler(cont));
    if (RemoveVectoredContinueHandler(cont))
        fail("second continue-handler removal succeeded unexpectedly");
    if (GetLastError() != 6)
        fail_value("second removal did not set ERROR_INVALID_HANDLE", GetLastError());

    void* only = AddVectoredExceptionHandler(0, only_continue_handler);
    if (!only)
        fail("single handler registration failed");

    __debugbreak();
    if (gOnlyHits != 1)
        fail("single vectored exception handler did not continue execution");

    expect_remove("remove single handler failed", RemoveVectoredExceptionHandler(only));

    void* raised = AddVectoredExceptionHandler(1, raised_continue_handler);
    if (!raised)
        fail("raised handler registration failed");
    ULONG_PTR args[2] = { 0x1122u, 0x3344u };
    RaiseException(0xE1111002u, 0, 2, args);
    if (gRaisedHits != 1)
        fail("raised vectored handler did not continue execution");
    expect_remove("remove raised handler failed", RemoveVectoredExceptionHandler(raised));

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