// Tests extract the unchanged production bodies. Kernel I/O is simulated.
#define WIN32_NO_STATUS
#include <windows.h>
#undef WIN32_NO_STATUS
#include <ntstatus.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
#include <cctype>
#define NT_SUCCESS(s) (((NTSTATUS)(s)) >= 0)
#define MAXUSHORT 0xffff
#define Executive 0
#define KernelMode 0
#define NotificationEvent 0
#define OBJ_KERNEL_HANDLE 0x200
#define InitializeObjectAttributes(a, n, f, r, s) ((a)->Attributes = (f))

struct OBJECT_ATTRIBUTES
{
    ULONG Attributes;
};

struct ANSI_STRING
{
    USHORT Length;
    USHORT MaximumLength;
    PCHAR Buffer;
};

typedef ANSI_STRING* PANSI_STRING;

struct IO_STATUS_BLOCK
{
    NTSTATUS Status;
    ULONG_PTR Information;
};

static int objectType;
static void* objectTypeValue = &objectType;
static void** ExEventObjectType = &objectTypeValue;
static unsigned checks;

static void
Check(
    bool value,
    const char* label
)
{
    ++checks;

    if (!value)
    {
        std::fprintf(stderr, "FAIL: %s\n", label);
        std::exit(1);
    }
}

static LONG
RtlCompareString(
    PANSI_STRING a,
    PANSI_STRING b,
    BOOLEAN insensitive
)
{
    if (a->Length != b->Length)
    {
        return 1;
    }

    for (USHORT i = 0; i < a->Length; ++i)
    {
        unsigned char x = (unsigned char)a->Buffer[i], y = (unsigned char)b->Buffer[i];

        if (insensitive)
        {
            x = (unsigned char)std::tolower(x);
            y = (unsigned char)std::tolower(y);
        }

        if (x != y)
        {
            return 1;
        }
    }

    return 0;
}

static void* exportPointer;
static ULONG exportBytes;

static PVOID
Directory(
    PVOID,
    BOOLEAN mapped,
    USHORT index,
    PULONG bytes
)
{
    Check(mapped && index == IMAGE_DIRECTORY_ENTRY_EXPORT, "mapped export lookup");
    *bytes = exportBytes;
    return exportPointer;
}

struct
{
    PVOID (*RtlImageDirectoryEntryToData)(
        PVOID,
        BOOLEAN,
        USHORT,
        PULONG
    );
} G_Common = {Directory};

static int eventToken;
static int fileToken;
static unsigned created;
static unsigned referenced;
static unsigned closed;
static unsigned dereferenced;
static unsigned reads;
static unsigned waits;
static int ioMode;
static bool outstanding;
static IO_STATUS_BLOCK* pendingIosb;
static ULONG requested;
static LONGLONG expectedOffset;

static NTSTATUS
ZwCreateEvent(
    HANDLE* event,
    ACCESS_MASK access,
    OBJECT_ATTRIBUTES* attributes,
    int type,
    BOOLEAN signaled
)
{
    Check(
        access == EVENT_ALL_ACCESS && attributes->Attributes == OBJ_KERNEL_HANDLE && type == NotificationEvent &&
            !signaled,
        "private kernel event"
    );

    if (ioMode == 4)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    ++created;
    *event = &eventToken;
    return STATUS_SUCCESS;
}

static NTSTATUS
ObReferenceObjectByHandle(
    HANDLE event,
    ACCESS_MASK access,
    void* type,
    int mode,
    void** object,
    void*
)
{
    Check(
        event == &eventToken && access == SYNCHRONIZE && type == *ExEventObjectType && mode == KernelMode,
        "reference own event"
    );

    if (ioMode == 5)
    {
        return STATUS_ACCESS_DENIED;
    }

    ++referenced;
    *object = &eventToken;
    return STATUS_SUCCESS;
}

static void
ZwClose(
    HANDLE event
)
{
    Check(event == &eventToken && !outstanding, "close only after completion");
    ++closed;
}

static void
ObDereferenceObject(
    void* object
)
{
    Check(object == &eventToken && !outstanding, "dereference only after completion");
    ++dereferenced;
}

static NTSTATUS
ZwReadFile(
    HANDLE file,
    HANDLE event,
    void*,
    void*,
    IO_STATUS_BLOCK* iosb,
    void*,
    ULONG size,
    LARGE_INTEGER* offset,
    void*
)
{
    ++reads;
    Check(
        file == &fileToken && event == &eventToken && offset && offset->QuadPart == expectedOffset,
        "file and explicit offset"
    );
    Check(iosb->Status == 0 && iosb->Information == 0, "initialized IOSB");
    requested = size;

    if (ioMode == 6)
    {
        return STATUS_ACCESS_DENIED;
    }

    if (ioMode == 1 || ioMode == 2)
    {
        outstanding = true;
        pendingIosb = iosb;
        return STATUS_PENDING;
    }

    iosb->Status = STATUS_SUCCESS;
    iosb->Information = ioMode == 3 ? size - 1 : size;
    return STATUS_SUCCESS;
}

static NTSTATUS
KeWaitForSingleObject(
    void* object,
    int reason,
    int mode,
    BOOLEAN alertable,
    LARGE_INTEGER* timeout
)
{
    ++waits;
    Check(
        object == &eventToken && reason == Executive && mode == KernelMode && !alertable && !timeout && outstanding,
        "noninterruptible per-request wait"
    );
    outstanding = false;
    pendingIosb->Status = ioMode == 2 ? STATUS_END_OF_FILE : STATUS_SUCCESS;
    pendingIosb->Information = ioMode == 2 ? 0 : requested;
    return STATUS_SUCCESS;
}

#include "common-under-test.inc"

static std::vector<UCHAR>
Image()
{
    std::vector<UCHAR> image(4096);
    auto dos = (IMAGE_DOS_HEADER*)image.data();
    dos->e_magic = IMAGE_DOS_SIGNATURE;
    dos->e_lfanew = 128;
    auto nt = (IMAGE_NT_HEADERS64*)(image.data() + 128);
    nt->Signature = IMAGE_NT_SIGNATURE;
    nt->OptionalHeader.Magic = IMAGE_NT_OPTIONAL_HDR64_MAGIC;
    nt->OptionalHeader.SizeOfImage = (ULONG)image.size();
    auto exports = (IMAGE_EXPORT_DIRECTORY*)(image.data() + 512);
    exports->NumberOfFunctions = 3;
    exports->NumberOfNames = 2;
    exports->AddressOfFunctions = 801;
    exports->AddressOfNames = 901;
    exports->AddressOfNameOrdinals = 1001;
    const ULONG functions[] = {2000, 2100, 2200};
    const ULONG names[] = {1100, 1120};
    const USHORT ordinals[] = {2, 0};
    std::memcpy(image.data() + 801, functions, sizeof(functions));
    std::memcpy(image.data() + 901, names, sizeof(names));
    std::memcpy(image.data() + 1001, ordinals, sizeof(ordinals));
    std::memcpy(image.data() + 1100, "Beta", 5);
    std::memcpy(image.data() + 1120, "Alpha", 6);
    return image;
}

static NTSTATUS
Lookup(
    std::vector<UCHAR>& image,
    const char* name,
    PVOID* output
)
{
    ANSI_STRING string = {(USHORT)std::strlen(name), (USHORT)std::strlen(name), (PCHAR)name};
    exportPointer = image.data() + 512;
    exportBytes = 256;
    return DomitoFindExportedFunctionAddress(image.data(), &string, output);
}

int
main()
{
    auto image = Image();
    PVOID output = nullptr;
    Check(Lookup(image, "Beta", &output) == STATUS_SUCCESS && output == image.data() + 2200, "ordinal-indexed EAT");
    Check(Lookup(image, "aLpHa", &output) == STATUS_SUCCESS && output == image.data() + 2000, "case-insensitive name");
    Check(Lookup(image, "Alpha", nullptr) == STATUS_SUCCESS, "optional output discovery");
    Check(Lookup(image, "missing", &output) == STATUS_NOT_FOUND && !output, "notfound clears output");
    auto exports = (IMAGE_EXPORT_DIRECTORY*)(image.data() + 512);
    const USHORT invalidOrdinal = 3;
    std::memcpy(image.data() + 1001, &invalidOrdinal, 2);
    Check(Lookup(image, "Beta", &output) == STATUS_INVALID_IMAGE_FORMAT && !output, "ordinal bounded");
    image = Image();
    const ULONG forwarder = 600;
    std::memcpy(image.data() + 801 + 8, &forwarder, 4);
    Check(Lookup(image, "Beta", &output) == STATUS_NOT_SUPPORTED && !output, "forwarder rejected");
    image = Image();
    const ULONG zero = 0;
    std::memcpy(image.data() + 801 + 8, &zero, 4);
    Check(Lookup(image, "Beta", &output) == STATUS_NOT_FOUND, "empty EAT slot");
    image = Image();
    const ULONG outside = 4096;
    std::memcpy(image.data() + 801 + 8, &outside, 4);
    Check(Lookup(image, "Beta", &output) == STATUS_INVALID_IMAGE_FORMAT, "function RVA bounded");
    image = Image();
    exports = (IMAGE_EXPORT_DIRECTORY*)(image.data() + 512);
    exports->NumberOfNames = 0xffffffff;
    Check(Lookup(image, "Beta", &output) == STATUS_INVALID_IMAGE_FORMAT, "table multiplication bounded");
    image = Image();
    exports = (IMAGE_EXPORT_DIRECTORY*)(image.data() + 512);
    exports->AddressOfFunctions = 4090;
    Check(Lookup(image, "Beta", &output) == STATUS_INVALID_IMAGE_FORMAT, "EAT span bounded");
    image = Image();
    const ULONG lateName = 4094;
    std::memcpy(image.data() + 901, &lateName, 4);
    image[4094] = 'x';
    image[4095] = 'x';
    Check(Lookup(image, "Beta", &output) == STATUS_INVALID_IMAGE_FORMAT, "unterminated name bounded");
    image = Image();
    std::memcpy(image.data() + 901, &outside, 4);
    Check(Lookup(image, "Beta", &output) == STATUS_INVALID_IMAGE_FORMAT, "name RVA bounded");
    image = Image();
    ANSI_STRING name = {4, 4, (PCHAR) "Beta"};
    exportPointer = image.data() + 512;
    exportBytes = 4096;
    Check(
        DomitoFindExportedFunctionAddress(image.data(), &name, &output) == STATUS_INVALID_IMAGE_FORMAT,
        "export directory span bounded"
    );
    output = (void*)1;
    Check(
        DomitoFindExportedFunctionAddress(nullptr, &name, &output) == STATUS_INVALID_PARAMETER && !output,
        "invalid arguments clear output"
    );
    char buffer[32];

    for (ioMode = 0; ioMode <= 6; ++ioMode)
    {
        created = referenced = closed = dereferenced = reads = waits = 0;
        outstanding = false;
        expectedOffset = 123;
        LARGE_INTEGER offset;
        offset.QuadPart = expectedOffset;
        NTSTATUS status = DomitoReadFileAt(&fileToken, buffer, sizeof(buffer), &offset);
        const NTSTATUS expected[] = {
            STATUS_SUCCESS,
            STATUS_SUCCESS,
            STATUS_END_OF_FILE,
            STATUS_END_OF_FILE,
            STATUS_INSUFFICIENT_RESOURCES,
            STATUS_ACCESS_DENIED,
            STATUS_ACCESS_DENIED
        };
        Check(status == expected[ioMode], "read status and exact-byte contract");
        Check(created == closed && referenced == dereferenced && !outstanding, "event lifetime balanced");
        Check(waits == ((ioMode == 1 || ioMode == 2) ? 1u : 0u), "wait pending only");
    }

    ioMode = 0;
    expectedOffset = 0;
    Check(DomitoReadFile(&fileToken, buffer, sizeof(buffer)) == STATUS_SUCCESS, "zero-offset wrapper");
    LARGE_INTEGER offset = {};
    offset.QuadPart = -1;
    Check(
        DomitoReadFileAt(&fileToken, buffer, sizeof(buffer), &offset) == STATUS_INVALID_PARAMETER,
        "negative offset rejected"
    );
    Check(
        DomitoReadFileAt(&fileToken, buffer, sizeof(buffer), nullptr) == STATUS_INVALID_PARAMETER,
        "null offset rejected"
    );
    offset.QuadPart = 0;
    Check(DomitoReadFileAt(&fileToken, nullptr, 0, &offset) == STATUS_SUCCESS, "zero byte read");
    Check(DomitoReadFileAt(&fileToken, nullptr, 1, &offset) == STATUS_INVALID_PARAMETER, "null buffer rejected");
    std::printf("PASS: %u common-helper checks (mock kernel I/O)\n", checks);
}