// Win32 runtime SEH smoke test. This file verifies explicit runtime frames before compiler SEH exists.

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

typedef LONG (*PRTL_EXCEPTION_FRAME_HANDLER)(EXCEPTION_POINTERS* ExceptionInfo, PVOID UserContext);

typedef struct _RTL_EXCEPTION_FRAME {
    DWORD Size;
    DWORD Flags;
    struct _RTL_EXCEPTION_FRAME* Previous;
    PRTL_EXCEPTION_FRAME_HANDLER Handler;
    PVOID UserContext;
} RTL_EXCEPTION_FRAME;

PVOID RtlPushExceptionFrame(RTL_EXCEPTION_FRAME* Frame,
                            PRTL_EXCEPTION_FRAME_HANDLER Handler,
                            PVOID UserContext);
BOOL RtlPopExceptionFrame(RTL_EXCEPTION_FRAME* Frame);

static volatile uint32_t gTopHits;

static volatile uint32_t gBottomHits;
static volatile uint32_t gVehHits;

// 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_seh_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_seh_smoke: ");
    print(message);
    print(" value=");
    print_u64(value);
    print("\n");
    exit_now(1);
}

// Verifies a breakpoint exception 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->ContextRecord->Rip || !info->ContextRecord->Rsp)
        fail("handler received empty context");
}

// Handles the first breakpoint from the newest runtime SEH frame.
static LONG top_frame_handler(EXCEPTION_POINTERS* info, void* userContext) {
    verify_breakpoint(info);
    if (userContext != (void*)0x1111u)
        fail("top frame context was not preserved");
    gTopHits++;
    return EXCEPTION_CONTINUE_EXECUTION;
}

// Handles a breakpoint after the top frame is removed.
static LONG bottom_frame_handler(EXCEPTION_POINTERS* info, void* userContext) {
    verify_breakpoint(info);
    if (userContext != (void*)0x2222u)
        fail("bottom frame context was not preserved");
    gBottomHits++;
    return EXCEPTION_CONTINUE_EXECUTION;
}

// Proves VEH still sees exceptions before the runtime SEH chain.
static LONG veh_search_handler(EXCEPTION_POINTERS* info) {
    verify_breakpoint(info);
    gVehHits++;
    return EXCEPTION_CONTINUE_SEARCH;
}

// Runs the runtime-registered Win32 SEH smoke test.
int main(void) {
    print("win32_seh_smoke: start\n");

    if (RtlPushExceptionFrame(NULL, top_frame_handler, NULL))
        fail("null frame was accepted");
    if (GetLastError() != 87)
        fail_value("null frame did not set ERROR_INVALID_PARAMETER", GetLastError());

    RTL_EXCEPTION_FRAME bottom;
    RTL_EXCEPTION_FRAME top;
    if (!RtlPushExceptionFrame(&bottom, bottom_frame_handler, (void*)0x2222u))
        fail("bottom SEH frame registration failed");
    if (!RtlPushExceptionFrame(&top, top_frame_handler, (void*)0x1111u))
        fail("top SEH frame registration failed");

    if (RtlPopExceptionFrame(&bottom))
        fail("non-top SEH frame pop succeeded unexpectedly");
    if (GetLastError() != 87)
        fail_value("non-top pop did not set ERROR_INVALID_PARAMETER", GetLastError());

    __debugbreak();
    if (gTopHits != 1 || gBottomHits != 0)
        fail("top SEH frame did not run before bottom frame");

    if (!RtlPopExceptionFrame(&top))
        fail("top SEH frame pop failed");

    void* veh = AddVectoredExceptionHandler(0, veh_search_handler);
    if (!veh)
        fail("VEH registration failed");

    __debugbreak();
    if (gVehHits != 1 || gBottomHits != 1)
        fail("VEH did not search before SEH handled the exception");

    if (!RemoveVectoredExceptionHandler(veh))
        fail("VEH removal failed");
    if (!RtlPopExceptionFrame(&bottom))
        fail("bottom SEH frame pop failed");

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