/* test_memwrite.c — write-signaled seam (MEMWRITE): atomic guest-memory write under an
 * exclusive lease. Stub mode (no VM): proves the full path cap -> grant -> lease-gate ->
 * route -> adapter -> ACT_ACK without actuation, plus the default-deny and fencing
 * invariants. The adapter never sees a control name (SISC).
 *
 *   1) happy path: CAP_MEMWRITE + a MEMWRITE lease -> CMD_MEMWRITE -> ACT_ACK{ok=1};
 *   2) extent default-deny: len > VMSIG_MEMWRITE_MAX and a missing SRC flag -> ACK{ok=0};
 *   3) lease gate: CMD_MEMWRITE WITHOUT an acquired lease -> dropped at the gate (no ACK);
 *   4) cap gate: a control WITHOUT CAP_MEMWRITE cannot acquire the lease (DENIED{NOCAP});
 *   5) in-flight fence: A holds the lease, queues a write, B preempts SYNCHRONOUSLY -> A's
 *      queued write is dropped by the fence (no ACK for A's corr), B's write actuates.
 * In-proc, under ASAN. */
#include "vmsig.h"
#include "memctx.h"        /* VMSIG_MEMWRITE_MAX: the adapter's extent bound (private) */
#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stdlib.h>

static int g_fail = 0;
#define CHECK(cond, msg) do { \
    if (!(cond)) { printf("  FAIL: %s\n", (msg)); g_fail = 1; } \
} while (0)

/* ---- in-proc control: records lease replies + MEMWRITE ACKs ---- */
typedef struct {
    void*    core;
    int      granted, denied, last_deny_reason;
    int      ack_ok[64];      /* ok flag per ACK in arrival order */
    uint32_t ack_corr[64];    /* corr per ACK                     */
    int      nack;
    int      stop_replies, replies;   /* stop the loop after N lease replies (0=off) */
    int      stop_acks;               /* stop the loop after N acks (0=off)          */
} cstate;

typedef struct { cstate* s; } cref;
static cref* g_refs[16]; static int g_nrefs = 0;
static cref* cref_new(cstate* s) {
    cref* r = calloc(1, sizeof *r); r->s = s;
    if (g_nrefs < 16) g_refs[g_nrefs++] = r;
    return r;
}
static void cref_free_all(void) { for (int i = 0; i < g_nrefs; i++) free(g_refs[i]); g_nrefs = 0; }

static int on_ev(void* user, const vmsig_event* ev) {
    cref* r = user; cstate* s = r->s;
    switch (ev->kind) {
    case VMSIG_EV_LEASE_GRANTED: s->granted++; s->replies++; break;
    case VMSIG_EV_LEASE_DENIED:
        s->denied++;
        s->last_deny_reason = (int)((const vmsig_lease_req*)ev->inln)->reason;
        s->replies++;
        break;
    case VMSIG_EV_ACT_ACK:
        if (s->nack < 64) {
            /* inln layout from mc_memwrite_ack: {int ok; uint32_t corr; uint32_t origin}. */
            int ok; memcpy(&ok, ev->inln, sizeof ok);
            s->ack_ok[s->nack]   = ok;
            s->ack_corr[s->nack] = ev->corr;
            s->nack++;
        }
        break;
    default: break;
    }
    if (s->stop_replies && s->replies >= s->stop_replies) vmsig_core_stop(s->core);
    if (s->stop_acks && s->nack >= s->stop_acks)          vmsig_core_stop(s->core);
    return 0;
}

static void* add_ctl(vmsig_core* core, cstate* s, uint32_t cap, uint32_t arb_prio) {
    cref* r = cref_new(s);
    vmsig_inproc_cfg cfg; memset(&cfg, 0, sizeof cfg);
    cfg.on_event = on_ev; cfg.user = r;
    cfg.sub.source_mask = 0xFFFFFFFFu; cfg.sub.prio_min = VMSIG_PRIO_BULK;
    void* ctl = vmsig_inproc_control_new(&cfg);
    vmsig_grant g; memset(&g, 0, sizeof g);
    g.endpoint_mask = 1ull << 0; g.source_mask = 0xFFFFFFFFu;
    g.cap_mask = cap | VMSIG_CAP_OBSERVE; g.arb_prio = arb_prio;
    vmsig_core_add_control(core, vmsig_inproc_control_ops(), ctl, &g);
    return ctl;
}

/* ---- DOWN send helpers ---- */
static int acquire_mw(void* ctl) {
    vmsig_event d; memset(&d, 0, sizeof d);
    d.kind = VMSIG_EV_CMD_ACQUIRE; d.source = VMSIG_SRC_MEMCTX; d.dir = VMSIG_DIR_DOWN;
    d.endpoint = 0; d.prio = VMSIG_PRIO_HIGH;
    vmsig_lease_req lr = { VMSIG_LEASE_MEMWRITE, 0 };
    memcpy(d.inln, &lr, sizeof lr);
    return vmsig_inproc_send(ctl, &d);
}

/* CMD_MEMWRITE with inline SRC; corr for tracking. flags: VMSIG_MW_SRC_* (0 => no SRC).
 * cr3 selects the target AS: 0 => kernel default (System DTB), non-zero => a process AS. */
static int send_write(void* ctl, uint64_t cr3, uint64_t gva, uint32_t len, uint32_t flags,
                      const void* src, uint32_t corr) {
    vmsig_event d; memset(&d, 0, sizeof d);
    d.kind = VMSIG_EV_CMD_MEMWRITE; d.source = VMSIG_SRC_MEMCTX; d.dir = VMSIG_DIR_DOWN;
    d.endpoint = 0; d.prio = VMSIG_PRIO_HIGH; d.corr = corr;
    vmsig_memwrite mw = { cr3, gva, len, flags };
    memcpy(d.inln, &mw, sizeof mw);
    if ((flags & VMSIG_MW_SRC_INLINE) && src && len <= VMSIG_MEMWRITE_INLINE)
        memcpy(d.inln + sizeof mw, src, len);
    return vmsig_inproc_send(ctl, &d);
}

/* Run the loop until N acks (used after queuing actuated writes). */
static void run_until_acks(cstate* s, int n) {
    vmsig_core* c = (vmsig_core*)s->core;
    s->stop_acks = n; s->stop_replies = 0;
    vmsig_core_run(c);
    s->stop_acks = 0;
}

/* ---- 1+2+3: happy path, extent default-deny, lease gate -------------------- */
static void test_path_and_deny(void) {
    printf("test_path_and_deny\n");
    vmsig_ctx*  ctx  = vmsig_ctx_new();
    vmsig_core* core = vmsig_core_new(ctx);
    cstate s; memset(&s, 0, sizeof s); s.core = core;

    void* A = add_ctl(core, &s, VMSIG_CAP_MEMWRITE, 10);
    CHECK(vmsig_core_add_adapter(core, vmsig_memctx_ops(), NULL, 0) >= 0, "add memctx");

    /* 3) lease gate: without ACQUIRE the write is dropped at the gate (-1, no actuation). */
    uint8_t pat[8] = { 0xDE, 0xAD, 0xBE, 0xEF, 1, 2, 3, 4 };
    CHECK(send_write(A, 0, 0x1000, 8, VMSIG_MW_SRC_INLINE, pat, 99) == -1,
          "3: CMD_MEMWRITE without a lease is dropped by the gate");

    /* acquire the MEMWRITE lease (synchronous intercept; UP reply paced by ctx). */
    CHECK(acquire_mw(A) == 0, "acquire submitted");

    /* 1) happy path: inline write -> queued -> ACT_ACK{ok=1}. Also drains the GRANTED reply. */
    CHECK(send_write(A, 0, 0x1000, 8, VMSIG_MW_SRC_INLINE, pat, 11) == 0,
          "1: owner's CMD_MEMWRITE passes the gate");

    /* 2) extent: len > MAX -> ACK{ok=0}, NOT actuated (queued ack on the loop thread). */
    CHECK(send_write(A, 0, 0x2000, VMSIG_MEMWRITE_MAX + 1, VMSIG_MW_SRC_INLINE, pat, 22) == 0,
          "2: over-extent write is accepted by the gate (denied inside the adapter)");
    /* 2b) missing SRC flag -> ACK{ok=0}. */
    CHECK(send_write(A, 0, 0x3000, 4, 0u, NULL, 33) == 0,
          "2b: no-SRC-flag write is accepted by the gate (denied inside the adapter)");

    /* expect 3 ACKs (corr 11/22/33) + the GRANTED reply. */
    run_until_acks(&s, 3);

    CHECK(s.granted == 1, "lease GRANTED once");
    int saw11_ok = -1, saw22_ok = -1, saw33_ok = -1, saw99 = 0;
    for (int i = 0; i < s.nack; i++) {
        if (s.ack_corr[i] == 11) saw11_ok = s.ack_ok[i];
        if (s.ack_corr[i] == 22) saw22_ok = s.ack_ok[i];
        if (s.ack_corr[i] == 33) saw33_ok = s.ack_ok[i];
        if (s.ack_corr[i] == 99) saw99 = 1;
    }
    CHECK(saw11_ok == 1, "1: happy-path write ACKs ok=1 (stub)");
    CHECK(saw22_ok == 0, "2: over-extent write ACKs ok=0 (default-deny)");
    CHECK(saw33_ok == 0, "2b: no-SRC-flag write ACKs ok=0 (default-deny)");
    CHECK(!saw99,        "3: the gate-dropped write produced no ACK");

    vmsig_core_free(core);
    vmsig_ctx_free(ctx);
}

/* ---- 1b: cr3 default + explicit — the cr3 field threads the whole seam -------
 * Stub does not actuate, so this asserts CONTRACT/ROUTE only: both a kernel-default
 * (cr3==0) and an arbitrary process cr3 (cr3!=0) pass cap -> grant -> lease-gate ->
 * route -> adapter -> ACT_ACK{ok=1}. A non-zero cr3 does NOT bypass any gate (it is
 * read by the adapter only after the core admitted the command). */
static void test_cr3_default_and_explicit(void) {
    printf("test_cr3_default_and_explicit\n");
    vmsig_ctx*  ctx  = vmsig_ctx_new();
    vmsig_core* core = vmsig_core_new(ctx);
    cstate s; memset(&s, 0, sizeof s); s.core = core;

    void* A = add_ctl(core, &s, VMSIG_CAP_MEMWRITE, 10);
    CHECK(vmsig_core_add_adapter(core, vmsig_memctx_ops(), NULL, 0) >= 0, "add memctx");

    CHECK(acquire_mw(A) == 0, "acquire submitted");

    uint8_t pat[8] = { 0xAA, 0xBB, 0xCC, 0xDD, 5, 6, 7, 8 };
    /* (a) kernel default: cr3 == 0 -> resolves to the adapter's System DTB on the worker. */
    CHECK(send_write(A, 0, 0x4000, 8, VMSIG_MW_SRC_INLINE, pat, 71) == 0,
          "1b-a: cr3==0 (kernel default) passes the gate");
    /* (b) explicit process AS: an arbitrary non-zero cr3 is carried through unchanged. */
    CHECK(send_write(A, 0xDEADBEEF000ull, 0x5000, 8, VMSIG_MW_SRC_INLINE, pat, 72) == 0,
          "1b-b: cr3!=0 (process AS) passes the gate");

    run_until_acks(&s, 2);

    CHECK(s.granted == 1, "1b: lease GRANTED once");
    int saw71_ok = -1, saw72_ok = -1;
    for (int i = 0; i < s.nack; i++) {
        if (s.ack_corr[i] == 71) saw71_ok = s.ack_ok[i];
        if (s.ack_corr[i] == 72) saw72_ok = s.ack_ok[i];
    }
    CHECK(saw71_ok == 1, "1b-a: cr3==0 write ACKs ok=1 (kernel default, stub)");
    CHECK(saw72_ok == 1, "1b-b: cr3!=0 write ACKs ok=1 (route proven, stub does not actuate)");

    vmsig_core_free(core);
    vmsig_ctx_free(ctx);
}

/* ---- 4: cap gate — no CAP_MEMWRITE cannot acquire the lease ----------------- */
static void test_cap_gate(void) {
    printf("test_cap_gate\n");
    vmsig_ctx*  ctx  = vmsig_ctx_new();
    vmsig_core* core = vmsig_core_new(ctx);
    cstate s; memset(&s, 0, sizeof s); s.core = core;

    void* NC = add_ctl(core, &s, 0u /* no MEMWRITE */, 10);
    CHECK(vmsig_core_add_adapter(core, vmsig_memctx_ops(), NULL, 0) >= 0, "add memctx");

    CHECK(acquire_mw(NC) == 0, "acquire submitted");
    s.stop_replies = 1; vmsig_core_run(core); s.stop_replies = 0;

    CHECK(s.denied == 1, "4: acquire without CAP_MEMWRITE -> DENIED");
    CHECK(s.last_deny_reason == VMSIG_LEASE_DENY_NOCAP, "4: reason=NOCAP");
    CHECK(s.granted == 0, "4: not granted");

    vmsig_core_free(core);
    vmsig_ctx_free(ctx);
}

/* ---- 5: in-flight fence — losing the lease before pump_down drops the write -- */
static void test_inflight_fence(void) {
    printf("test_inflight_fence\n");
    vmsig_ctx*  ctx  = vmsig_ctx_new();
    vmsig_core* core = vmsig_core_new(ctx);
    cstate s; memset(&s, 0, sizeof s); s.core = core;

    void* A = add_ctl(core, &s, VMSIG_CAP_MEMWRITE, 10);
    void* B = add_ctl(core, &s, VMSIG_CAP_MEMWRITE, 100);   /* higher prio: preempts */
    CHECK(vmsig_core_add_adapter(core, vmsig_memctx_ops(), NULL, 0) >= 0, "add memctx");

    uint8_t pat[4] = { 1, 2, 3, 4 };
    CHECK(acquire_mw(A) == 0, "A acquires");
    /* A queues a write (corr=55): passes the gate (A owns), lands in the DOWN queue. */
    CHECK(send_write(A, 0, 0x1000, 4, VMSIG_MW_SRC_INLINE, pat, 55) == 0, "A queues write 55");
    /* B preempts SYNCHRONOUSLY (acquire does not go through ctx). */
    CHECK(acquire_mw(B) == 0, "B preempts");
    /* B's own write (corr=66) — should actuate. */
    CHECK(send_write(B, 0, 0x2000, 4, VMSIG_MW_SRC_INLINE, pat, 66) == 0, "B queues write 66");

    run_until_acks(&s, 1);   /* B's 66 acks; A's 55 must be fenced (no ack) */

    int saw55 = 0, saw66 = 0;
    for (int i = 0; i < s.nack; i++) {
        if (s.ack_corr[i] == 55) saw55 = 1;
        if (s.ack_corr[i] == 66) saw66 = 1;
    }
    CHECK(!saw55, "5: ex-owner A's in-flight write is dropped by the fence");
    CHECK(saw66,  "5: new owner B's write actuates after preemption");
    CHECK(s.granted == 2, "5: A and B each got GRANTED");

    vmsig_core_free(core);
    vmsig_ctx_free(ctx);
}

int main(void) {
    printf("test_memwrite\n");
    /* POD layout: cr3+gva (2*u64) + len+flags (2*u32) = 24, +24 inline tail = inln[48]. */
    CHECK(sizeof(vmsig_memwrite) == 24, "vmsig_memwrite header is 24 bytes");
    CHECK(VMSIG_MEMWRITE_INLINE == 24u, "VMSIG_MEMWRITE_INLINE is 24");
    test_path_and_deny();
    test_cr3_default_and_explicit();
    test_cap_gate();
    test_inflight_fence();
    cref_free_all();
    printf("memwrite tests: %s\n", g_fail ? "FAIL" : "PASS");
    return g_fail ? 1 : 0;
}