// Native capability-handle smoke test.
// Verifies monotonic identities, typed rights, derivation, aliasing, and revocation.

#include <stdint.h>
#include <stdbool.h>
#include <ntlite/syscall.h>

#define EVENT_AUTO_RESET 0u
#define HANDLE_STRESS_COUNT 512u
#define DERIVATION_SUCCESS_COUNT 15u

typedef struct TEST_QUEUE_MESSAGE {
    uint32_t Type;
    uint32_t Flags;
    uint64_t Param[6];
} TEST_QUEUE_MESSAGE;

typedef char TEST_QUEUE_MESSAGE_SIZE_CHECK[
    (sizeof(TEST_QUEUE_MESSAGE) == SYS_QUEUE_MESSAGE_SIZE) ? 1 : -1
];

// Prints one NUL-terminated string to the kernel console.
static void print(const char* s) {
    SYS_WRITE_CONSOLE_REQUEST req = { sizeof(SYS_WRITE_CONSOLE_REQUEST), 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 = { sizeof(SYS_WRITE_CONSOLE_REQUEST), SysSvFromParts(&c, 1) };
        SysWriteConsole(&req);
    }
}

// Terminates the current test process.
static void exit_now(int code) {
    while (true) SysExitProcess(code);
}

// Fails the test with one plain message.
static void fail(const char* msg) {
    print("handle_dup_smoke: ");
    print(msg);
    print("\n");
    exit_now(1);
}

// Fails the test with one status value.
static void fail_status(const char* what, KSTATUS status) {
    print("handle_dup_smoke: ");
    print(what);
    print(" status=");
    print_u64(status);
    print("\n");
    exit_now(1);
}

// Requires one exact status value.
static void expect_status(const char* what, KSTATUS got, KSTATUS expected) {
    if (got == expected) return;
    print("handle_dup_smoke: ");
    print(what);
    print(" got=");
    print_u64(got);
    print(" expected=");
    print_u64(expected);
    print("\n");
    exit_now(1);
}

// Duplicates one handle into the current process.
static KSTATUS duplicate_handle(SYS_HANDLE source,
                                SYS_HANDLE process,
                                SYS_RIGHTS desiredRights,
                                uint32_t flags,
                                SYS_HANDLE* out) {
    SYS_DUPLICATE_HANDLE_REQUEST req = {
        sizeof(SYS_DUPLICATE_HANDLE_REQUEST),
        source,
        process,
        desiredRights,
        flags,
        out,
    };
    return SysDuplicateHandle(&req);
}

// Derives one child capability in the current process.
static KSTATUS derive_handle(SYS_HANDLE source, SYS_RIGHTS desiredRights, SYS_HANDLE* out) {
    SYS_DERIVE_HANDLE_REQUEST req = {
        sizeof(SYS_DERIVE_HANDLE_REQUEST),
        source,
        desiredRights,
        out,
    };
    return SysDeriveHandle(&req);
}

// Sends one queue message.
static KSTATUS send_message(SYS_HANDLE queue, uint32_t type, uint64_t value) {
    TEST_QUEUE_MESSAGE msg = {0};
    msg.Type = type;
    msg.Param[0] = value;
    SYS_QUEUE_SEND_REQUEST req = {
        sizeof(req),
        queue,
        SYS_QUEUE_FLAG_NONE,
        0,
        &msg,
    };
    return SysSendQueue(&req);
}

// Receives one queue message.
static KSTATUS recv_message(SYS_HANDLE queue, TEST_QUEUE_MESSAGE* msg) {
    SYS_QUEUE_RECV_REQUEST req = {
        sizeof(req),
        queue,
        SYS_QUEUE_FLAG_NONE,
        0,
        msg,
    };
    return SysRecvQueue(&req);
}

// Verifies one received queue message.
static void check_message(const TEST_QUEUE_MESSAGE* msg, uint32_t type, uint64_t value) {
    if (msg->Type != type || msg->Param[0] != value) fail("message payload mismatch");
}

// Closes one handle and requires success.
static void close_handle(const char* what, SYS_HANDLE handle) {
    KSTATUS status = SysCloseHandle(handle);
    if (!KSUCCESS(status)) fail_status(what, status);
}

// Verifies that closed identities are not reused and stale values stay invalid.
static void test_monotonic_identity(void) {
    SYS_HANDLE first = SYS_HANDLE_INVALID;
    KSTATUS status = SysCreateEvent(EVENT_AUTO_RESET, 0, &first);
    if (!KSUCCESS(status)) fail_status("create first monotonic event", status);
    close_handle("close first monotonic event", first);

    SYS_HANDLE second = SYS_HANDLE_INVALID;
    status = SysCreateEvent(EVENT_AUTO_RESET, 0, &second);
    if (!KSUCCESS(status)) fail_status("create second monotonic event", status);
    if (second <= first) fail("handle identity was reused or moved backwards");

    expect_status("close stale handle", SysCloseHandle(first), NTLSTATUS_INVALID_HANDLE);
    close_handle("close second monotonic event", second);
}

// Verifies that live-handle capacity is no longer tied to the old 256-slot table.
static void test_dynamic_handle_capacity(void) {
    SYS_HANDLE handles[HANDLE_STRESS_COUNT];
    SYS_HANDLE previous = SYS_HANDLE_INVALID;

    for (uint32_t i = 0; i < HANDLE_STRESS_COUNT; ++i) {
        handles[i] = SYS_HANDLE_INVALID;
        KSTATUS status = SysCreateEvent(EVENT_AUTO_RESET, 0, &handles[i]);
        if (!KSUCCESS(status)) fail_status("create stress event", status);
        if (previous != SYS_HANDLE_INVALID && handles[i] <= previous)
            fail("stress handles were not strictly monotonic");
        previous = handles[i];
    }

    for (uint32_t i = HANDLE_STRESS_COUNT; i != 0; --i)
        close_handle("close stress event", handles[i - 1]);
}

// Verifies explicit rights attenuation on duplicate aliases.
static void test_duplicate_rights(SYS_HANDLE process, SYS_HANDLE queue,
                                  SYS_HANDLE* sameRightsOut,
                                  SYS_HANDLE* receiveOnlyOut,
                                  SYS_HANDLE* sendOnlyOut) {
    SYS_HANDLE receiveOnly = SYS_HANDLE_INVALID;
    KSTATUS status = duplicate_handle(queue, process, SYS_QUEUE_RIGHT_RECEIVE,
                                      SYS_DUPLICATE_HANDLE_FLAG_NONE, &receiveOnly);
    if (!KSUCCESS(status)) fail_status("duplicate receive-only", status);

    SYS_HANDLE sendOnly = SYS_HANDLE_INVALID;
    status = duplicate_handle(queue, process, SYS_QUEUE_RIGHT_SEND,
                              SYS_DUPLICATE_HANDLE_FLAG_NONE, &sendOnly);
    if (!KSUCCESS(status)) fail_status("duplicate send-only", status);

    SYS_HANDLE sameRights = SYS_HANDLE_INVALID;
    status = duplicate_handle(queue, process, 0,
                              SYS_DUPLICATE_HANDLE_FLAG_SAME_RIGHTS, &sameRights);
    if (!KSUCCESS(status)) fail_status("duplicate same rights", status);

    status = send_message(receiveOnly, 10, 100);
    expect_status("receive-only send denied", status, NTLSTATUS_ACCESS_DENIED);

    TEST_QUEUE_MESSAGE msg;
    status = recv_message(sendOnly, &msg);
    expect_status("send-only receive denied", status, NTLSTATUS_ACCESS_DENIED);

    status = send_message(queue, 11, 101);
    if (!KSUCCESS(status)) fail_status("send original", status);
    status = recv_message(receiveOnly, &msg);
    if (!KSUCCESS(status)) fail_status("receive receive-only", status);
    check_message(&msg, 11, 101);

    status = send_message(sendOnly, 12, 102);
    if (!KSUCCESS(status)) fail_status("send send-only", status);
    status = recv_message(queue, &msg);
    if (!KSUCCESS(status)) fail_status("receive original", status);
    check_message(&msg, 12, 102);

    status = send_message(sameRights, 13, 103);
    if (!KSUCCESS(status)) fail_status("send same-rights", status);
    status = recv_message(sameRights, &msg);
    if (!KSUCCESS(status)) fail_status("receive same-rights", status);
    check_message(&msg, 13, 103);

    SYS_HANDLE ignored = SYS_HANDLE_INVALID;
    status = duplicate_handle(receiveOnly, process, SYS_QUEUE_RIGHT_RECEIVE,
                              SYS_DUPLICATE_HANDLE_FLAG_NONE, &ignored);
    expect_status("duplicate non-duplicable handle", status, NTLSTATUS_ACCESS_DENIED);

    status = duplicate_handle(queue, process, SYS_QUEUE_RIGHT_RECEIVE | 0x80000000u,
                              SYS_DUPLICATE_HANDLE_FLAG_NONE, &ignored);
    expect_status("duplicate unknown rights", status, NTLSTATUS_INVALID_PARAMETER);

    status = duplicate_handle(queue, process, SYS_QUEUE_RIGHT_RECEIVE,
                              SYS_DUPLICATE_HANDLE_FLAG_SAME_RIGHTS, &ignored);
    expect_status("duplicate same-rights with desired rights", status, NTLSTATUS_INVALID_PARAMETER);

    status = duplicate_handle(queue, process, 0, SYS_DUPLICATE_HANDLE_FLAG_NONE, &ignored);
    expect_status("duplicate zero rights", status, NTLSTATUS_INVALID_PARAMETER);

    SYS_DUPLICATE_HANDLE_REQUEST badOut = {
        sizeof(SYS_DUPLICATE_HANDLE_REQUEST),
        queue,
        process,
        SYS_QUEUE_RIGHT_RECEIVE,
        SYS_DUPLICATE_HANDLE_FLAG_NONE,
        0,
    };
    status = SysDuplicateHandle(&badOut);
    expect_status("duplicate null output", status, NTLSTATUS_INVALID_PARAMETER);

    *sameRightsOut = sameRights;
    *receiveOnlyOut = receiveOnly;
    *sendOnlyOut = sendOnly;
}

// Verifies child grants attenuate authority and can be revoked independently.
static void test_derived_revocation(SYS_HANDLE process, SYS_HANDLE queue) {
    const SYS_RIGHTS childRights = SYS_QUEUE_RIGHT_RECEIVE |
                                   SYS_RIGHT_DUPLICATE |
                                   SYS_RIGHT_DERIVE |
                                   SYS_RIGHT_REVOKE;

    SYS_HANDLE child = SYS_HANDLE_INVALID;
    KSTATUS status = derive_handle(queue, childRights, &child);
    if (!KSUCCESS(status)) fail_status("derive child", status);

    SYS_HANDLE childAlias = SYS_HANDLE_INVALID;
    status = duplicate_handle(child, process, 0,
                              SYS_DUPLICATE_HANDLE_FLAG_SAME_RIGHTS, &childAlias);
    if (!KSUCCESS(status)) fail_status("duplicate child alias", status);

    SYS_HANDLE grandchild = SYS_HANDLE_INVALID;
    status = derive_handle(child, SYS_QUEUE_RIGHT_RECEIVE, &grandchild);
    if (!KSUCCESS(status)) fail_status("derive grandchild", status);

    SYS_HANDLE ignored = SYS_HANDLE_INVALID;
    status = derive_handle(child, SYS_QUEUE_RIGHT_SEND, &ignored);
    expect_status("derive cannot amplify rights", status, NTLSTATUS_ACCESS_DENIED);

    status = SysRevokeHandle(grandchild);
    expect_status("revoke without revoke right", status, NTLSTATUS_ACCESS_DENIED);

    TEST_QUEUE_MESSAGE msg;
    status = send_message(queue, 20, 200);
    if (!KSUCCESS(status)) fail_status("send for grandchild", status);
    status = recv_message(grandchild, &msg);
    if (!KSUCCESS(status)) fail_status("receive grandchild", status);
    check_message(&msg, 20, 200);

    status = SysRevokeHandle(child);
    if (!KSUCCESS(status)) fail_status("revoke child", status);

    ignored = SYS_HANDLE_INVALID;
    status = duplicate_handle(childAlias, process, 0,
                              SYS_DUPLICATE_HANDLE_FLAG_SAME_RIGHTS, &ignored);
    expect_status("duplicate revoked alias", status, NTLSTATUS_CAPABILITY_REVOKED);

    status = derive_handle(childAlias, SYS_QUEUE_RIGHT_RECEIVE, &ignored);
    expect_status("derive revoked alias", status, NTLSTATUS_CAPABILITY_REVOKED);

    status = recv_message(grandchild, &msg);
    expect_status("revoked grandchild receive", status, NTLSTATUS_CAPABILITY_REVOKED);

    // Revoking a derived grant must not revoke its parent grant.
    status = send_message(queue, 21, 201);
    if (!KSUCCESS(status)) fail_status("parent send after child revoke", status);
    status = recv_message(queue, &msg);
    if (!KSUCCESS(status)) fail_status("parent receive after child revoke", status);
    check_message(&msg, 21, 201);

    close_handle("close grandchild", grandchild);
    close_handle("close child alias", childAlias);
    close_handle("close child", child);
}

// Verifies the bounded derivation chain used for revocation traversal.
static void test_derivation_depth(SYS_HANDLE queue) {
    const SYS_RIGHTS rights = SYS_QUEUE_RIGHT_RECEIVE |
                              SYS_RIGHT_DERIVE |
                              SYS_RIGHT_REVOKE;
    SYS_HANDLE chain[DERIVATION_SUCCESS_COUNT];
    SYS_HANDLE parent = queue;

    for (uint32_t i = 0; i < DERIVATION_SUCCESS_COUNT; ++i) {
        chain[i] = SYS_HANDLE_INVALID;
        KSTATUS status = derive_handle(parent, rights, &chain[i]);
        if (!KSUCCESS(status)) fail_status("derive bounded chain", status);
        parent = chain[i];
    }

    SYS_HANDLE tooDeep = SYS_HANDLE_INVALID;
    KSTATUS status = derive_handle(parent, rights, &tooDeep);
    expect_status("derive beyond depth bound", status, NTLSTATUS_LIMIT_REACHED);

    for (uint32_t i = DERIVATION_SUCCESS_COUNT; i != 0; --i)
        close_handle("close bounded derivation", chain[i - 1]);
}

// Runs the native capability-handle smoke test.
int main(void) {
    print("handle_dup_smoke: start\n");

    test_monotonic_identity();
    test_dynamic_handle_capacity();

    uint32_t currentPid = SysGetCurrentPid();
    SYS_HANDLE process = SYS_HANDLE_INVALID;
    KSTATUS status = SysOpenProcess(currentPid, SYS_PROCESS_RIGHT_INSERT_HANDLE | SYS_RIGHT_QUERY, &process);
    if (!KSUCCESS(status)) fail_status("open current process", status);

    SYS_HANDLE queue = SYS_HANDLE_INVALID;
    status = SysCreateQueue(&queue);
    if (!KSUCCESS(status)) fail_status("create queue", status);

    SYS_HANDLE sameRights = SYS_HANDLE_INVALID;
    SYS_HANDLE receiveOnly = SYS_HANDLE_INVALID;
    SYS_HANDLE sendOnly = SYS_HANDLE_INVALID;
    test_duplicate_rights(process, queue, &sameRights, &receiveOnly, &sendOnly);
    test_derived_revocation(process, queue);
    test_derivation_depth(queue);

    // A duplicate is an alias of the same grant, so revocation through either
    // alias revokes the original grant as well.
    status = SysRevokeHandle(sameRights);
    if (!KSUCCESS(status)) fail_status("revoke duplicate alias", status);

    SYS_HANDLE ignored = SYS_HANDLE_INVALID;
    status = duplicate_handle(queue, process, 0,
                              SYS_DUPLICATE_HANDLE_FLAG_SAME_RIGHTS, &ignored);
    expect_status("duplicate root after alias revoke", status, NTLSTATUS_CAPABILITY_REVOKED);

    close_handle("close same-rights", sameRights);
    close_handle("close send-only", sendOnly);
    close_handle("close receive-only", receiveOnly);
    close_handle("close revoked queue", queue);
    close_handle("close process", process);

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