// User-mode smoke test for socket IRP dispatch, public async sockets, cancellation, and statistics.

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

#define EVENT_AUTO_RESET 0u
#define NET_IRP_SMOKE_UDP_PORT 40219u
#define NET_IRP_SMOKE_HTTP_PORT 80u

// 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 to the kernel console.
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("net_irp_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("net_irp_smoke: ");
    print(what);
    print(" status=");
    print_u64(status);
    print("\n");
    exit_now(1);
}

// Reads the public network IRP statistics snapshot.
static SYS_NET_IRP_STATS query_stats(void) {
    SYS_NET_IRP_STATS stats;
    memset(&stats, 0, sizeof(stats));
    uint32_t status = SysQueryKernelStats(SYS_KERNEL_STATS_NET_IRP, &stats, sizeof(stats), 0);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("query net irp stats", status);
    if (stats.Size != sizeof(stats))
        fail("net irp stats size mismatch");
    if (stats.Version != SYS_STATS_VERSION)
        fail("net irp stats version mismatch");
    return stats;
}

// Builds an IPv4 socket address.
static SYS_SOCKET_ADDRESS make_ipv4(uint8_t a, uint8_t b, uint8_t c, uint8_t d, uint16_t port) {
    SYS_SOCKET_ADDRESS address;
    memset(&address, 0, sizeof(address));
    address.Size = sizeof(address);
    address.Family = SYS_AF_INET;
    address.Port = port;
    address.Address[0] = a;
    address.Address[1] = b;
    address.Address[2] = c;
    address.Address[3] = d;
    return address;
}

// Sets one uint32 socket option.
static void set_u32_option(uint32_t socket, uint32_t option, uint32_t value) {
    SYS_SOCKET_OPTION_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Option = option;
    req.Value = &value;
    req.ValueSize = sizeof(value);
    uint32_t status = SysSetSocketOption(socket, &req);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("set socket option", status);
}

// Waits for an async socket request to publish a terminal IOSB status.
static void wait_for_async(uint32_t event, SYS_IO_STATUS_BLOCK* iosb, const char* what) {
    for (uint32_t i = 0; i < 200u && iosb->Status == NTLSTATUS_PENDING; ++i) {
        uint32_t status = SysWait(event, 0);
        if (status != NTLSTATUS_SUCCESS && status != NTLSTATUS_ALERTED &&
            status != NTLSTATUS_KERNEL_APC && status != NTLSTATUS_USER_APC) {
            fail_status(what, status);
        }
    }
    if (iosb->Status == NTLSTATUS_PENDING)
        fail("async request stayed pending");
}

// Verifies nonblocking UDP receive returns would-block through the IRP path.
static void check_nonblocking_receive(void) {
    uint32_t udp = 0;
    uint32_t status = SysCreateSocket(SYS_AF_INET, SYS_SOCK_DGRAM, SYS_IPPROTO_UDP, &udp);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("create udp socket", status);

    SYS_SOCKET_ADDRESS bindAddress = make_ipv4(0, 0, 0, 0, NET_IRP_SMOKE_UDP_PORT);
    status = SysBindSocket(udp, &bindAddress, sizeof(bindAddress));
    if (status != NTLSTATUS_SUCCESS) {
        SysCloseHandle(udp);
        fail_status("bind udp socket", status);
    }

    set_u32_option(udp, SYS_SOCKET_OPT_NONBLOCKING, 1);

    char buffer[16];
    SYS_SOCKET_ADDRESS remote;
    size_t got = 0;
    size_t remoteSize = sizeof(remote);
    memset(buffer, 0, sizeof(buffer));
    memset(&remote, 0, sizeof(remote));

    SYS_SOCKET_RECV_FROM_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Buffer = buffer;
    req.Length = sizeof(buffer);
    req.RemoteAddressOut = &remote;
    req.RemoteAddressSizeInOut = &remoteSize;
    req.ReceivedOut = &got;

    status = SysRecvFromSocket(udp, &req);
    SysCloseHandle(udp);
    if (status != NTLSTATUS_WOULD_BLOCK)
        fail_status("nonblocking udp recvfrom", status);
}

// Verifies a pending async UDP receive can be cancelled and waited through its request handle.
static void check_async_receive_cancel(void) {
    uint32_t udp = 0;
    uint32_t status = SysCreateSocket(SYS_AF_INET, SYS_SOCK_DGRAM, SYS_IPPROTO_UDP, &udp);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("create cancel udp socket", status);

    SYS_SOCKET_ADDRESS bindAddress = make_ipv4(0, 0, 0, 0, NET_IRP_SMOKE_UDP_PORT + 1u);
    status = SysBindSocket(udp, &bindAddress, sizeof(bindAddress));
    if (status != NTLSTATUS_SUCCESS) {
        SysCloseHandle(udp);
        fail_status("bind cancel udp socket", status);
    }

    uint32_t event = 0;
    status = SysCreateEvent(EVENT_AUTO_RESET, 0, &event);
    if (status != NTLSTATUS_SUCCESS) {
        SysCloseHandle(udp);
        fail_status("create event", status);
    }

    char buffer[16];
    SYS_SOCKET_ADDRESS remote;
    size_t remoteSize = sizeof(remote);
    uint32_t requestHandle = 0;
    SYS_IO_STATUS_BLOCK iosb;
    memset(buffer, 0, sizeof(buffer));
    memset(&remote, 0, sizeof(remote));
    iosb.Status = 0;
    iosb.Information = 0;

    SYS_ASYNC_SOCKET_REQUEST req;
    memset(&req, 0, sizeof(req));
    req.Size = sizeof(req);
    req.SocketHandle = udp;
    req.EventHandle = event;
    req.Operation = SYS_ASYNC_SOCKET_OP_RECV_FROM;
    req.Buffer = buffer;
    req.Length = sizeof(buffer);
    req.AddressOut = &remote;
    req.AddressSizeInOut = &remoteSize;
    req.Iosb = &iosb;
    req.RequestHandleOut = &requestHandle;

    SYS_NET_IRP_STATS beforeSubmit = query_stats();
    status = SysSocketIoAsync(&req);
    if (status != NTLSTATUS_PENDING || requestHandle == 0) {
        SysCloseHandle(event);
        SysCloseHandle(udp);
        fail_status("submit async recvfrom", status);
    }

    uint32_t enteredIrp = 0;
    for (uint32_t i = 0; i < 40u; ++i) {
        SYS_NET_IRP_STATS now = query_stats();
        if (now.ByOperation[SYS_NET_IRP_OP_RECEIVE_FROM] >
            beforeSubmit.ByOperation[SYS_NET_IRP_OP_RECEIVE_FROM]) {
            enteredIrp = 1;
            break;
        }
        SysSleep(10);
    }
    if (!enteredIrp)
        fail("async recvfrom did not enter the IRP path");

    uint32_t cancelStatus = SysCancelIo(requestHandle);
    if (cancelStatus != NTLSTATUS_SUCCESS && cancelStatus != NTLSTATUS_PENDING)
        fail_status("cancel async recvfrom", cancelStatus);

    wait_for_async(event, &iosb, "wait async recvfrom cancel");
    if (iosb.Status != NTLSTATUS_CANCELLED)
        fail_status("async recvfrom iosb", iosb.Status);
    if (iosb.Information != 0)
        fail("async recvfrom cancel transferred bytes");

    status = SysWaitSingle(requestHandle, 0, 0);
    if (status != NTLSTATUS_SUCCESS)
        fail_status("wait cancelled request handle", status);

    SysCloseHandle(requestHandle);
    SysCloseHandle(event);
    SysCloseHandle(udp);
}

// Tries a host-backed HTTP transfer to prove async send and receive success paths.
static void try_http_async_success(void) {
    uint32_t tcp = 0;
    uint32_t status = SysCreateSocket(SYS_AF_INET, SYS_SOCK_STREAM, SYS_IPPROTO_TCP, &tcp);
    if (status != NTLSTATUS_SUCCESS) {
        print("net_irp_smoke: http success path skipped create status=");
        print_u64(status);
        print("\n");
        return;
    }

    SYS_SOCKET_ADDRESS address = make_ipv4(10, 0, 2, 2, NET_IRP_SMOKE_HTTP_PORT);
    status = SysConnectSocket(tcp, &address, sizeof(address));
    if (status != NTLSTATUS_SUCCESS) {
        print("net_irp_smoke: http success path skipped connect status=");
        print_u64(status);
        print("\n");
        SysCloseHandle(tcp);
        return;
    }

    uint32_t event = 0;
    status = SysCreateEvent(EVENT_AUTO_RESET, 0, &event);
    if (status != NTLSTATUS_SUCCESS) {
        SysCloseHandle(tcp);
        fail_status("create http event", status);
    }

    static const char requestText[] = "GET / HTTP/1.0\r\nHost: 10.0.2.2\r\n\r\n";
    SYS_IO_STATUS_BLOCK sendIosb;
    uint32_t sendRequest = 0;
    memset(&sendIosb, 0, sizeof(sendIosb));

    SYS_ASYNC_SOCKET_REQUEST sendReq;
    memset(&sendReq, 0, sizeof(sendReq));
    sendReq.Size = sizeof(sendReq);
    sendReq.SocketHandle = tcp;
    sendReq.EventHandle = event;
    sendReq.Operation = SYS_ASYNC_SOCKET_OP_SEND;
    sendReq.Buffer = (void*)requestText;
    sendReq.Length = sizeof(requestText) - 1u;
    sendReq.Iosb = &sendIosb;
    sendReq.RequestHandleOut = &sendRequest;

    status = SysSocketIoAsync(&sendReq);
    if (status != NTLSTATUS_PENDING || sendRequest == 0)
        fail_status("submit async send", status);
    wait_for_async(event, &sendIosb, "wait async send");
    if (sendIosb.Status != NTLSTATUS_SUCCESS)
        fail_status("async send iosb", sendIosb.Status);
    if (sendIosb.Information != sizeof(requestText) - 1u)
        fail("async send byte count mismatch");
    SysCloseHandle(sendRequest);

    char response[128];
    SYS_IO_STATUS_BLOCK recvIosb;
    uint32_t recvRequest = 0;
    memset(response, 0, sizeof(response));
    memset(&recvIosb, 0, sizeof(recvIosb));

    SYS_ASYNC_SOCKET_REQUEST recvReq;
    memset(&recvReq, 0, sizeof(recvReq));
    recvReq.Size = sizeof(recvReq);
    recvReq.SocketHandle = tcp;
    recvReq.EventHandle = event;
    recvReq.Operation = SYS_ASYNC_SOCKET_OP_RECV;
    recvReq.Buffer = response;
    recvReq.Length = sizeof(response) - 1u;
    recvReq.Iosb = &recvIosb;
    recvReq.RequestHandleOut = &recvRequest;

    status = SysSocketIoAsync(&recvReq);
    if (status != NTLSTATUS_PENDING || recvRequest == 0)
        fail_status("submit async recv", status);
    wait_for_async(event, &recvIosb, "wait async recv");
    if (recvIosb.Status != NTLSTATUS_SUCCESS)
        fail_status("async recv iosb", recvIosb.Status);
    if (recvIosb.Information == 0)
        fail("async recv returned no bytes");
    SysCloseHandle(recvRequest);

    SysCloseHandle(event);
    SysCloseHandle(tcp);
}

// Verifies counters prove socket requests entered and completed through the IRP path.
static void check_stats_delta(const SYS_NET_IRP_STATS* before) {
    SYS_NET_IRP_STATS after = query_stats();
    if (after.Submitted <= before->Submitted)
        fail("net irp submitted counter did not advance");
    if (after.Completed <= before->Completed)
        fail("net irp completed counter did not advance");
    if (after.WouldBlock <= before->WouldBlock)
        fail("net irp would-block counter did not advance");
    if (after.CancelRequested <= before->CancelRequested)
        fail("net irp cancel-request counter did not advance");
    if (after.Cancelled <= before->Cancelled)
        fail("net irp cancelled counter did not advance");
    if (after.ByOperation[SYS_NET_IRP_OP_RECEIVE_FROM] <= before->ByOperation[SYS_NET_IRP_OP_RECEIVE_FROM])
        fail("net irp recvfrom operation counter did not advance");
}

// Runs the network IRP smoke test.
int main(void) {
    print("net_irp_smoke: start\n");
    SYS_NET_IRP_STATS before = query_stats();
    check_nonblocking_receive();
    check_async_receive_cancel();
    try_http_async_success();
    check_stats_delta(&before);
    print("net_irp_smoke: PASS\n");
    exit_now(0);
    return 0;
}