/*
 * df5200.c - DF-5200: DragonFly's LinuxKPI kthread shim provides NO park
 * quiescence: kthread_park() returns immediately without waiting for the
 * target to park, and kthread_parkme() never actually deschedules the
 * thread (lwkt_deschedule_self() on a running curthread is a no-op -
 * TDF_RUNQ is already clear, sys/kern/lwkt_thread.c:153-163).
 *
 *   sys/dev/drm/linux_kthread.c:103-118,126-133:
 *
 *	int kthread_park(struct task_struct *ts)
 *	{
 *		set_bit(KTHREAD_SHOULD_PARK, &ts->kt_flags);
 *		wake_up_process(ts);
 *		return ts->kt_exitvalue;	   <- returns IMMEDIATELY:
 *	}					   no wait-for-park at all
 *
 *	void kthread_parkme(void)
 *	{
 *		if (test_bit(KTHREAD_SHOULD_PARK, &current->kt_flags) == 0)
 *			return;
 *		lwkt_deschedule_self(curthread);  <- NO-OP for a running thread;
 *	}					    no lwkt_switch(), bit never
 *					    cleared here either
 *
 * Linux semantics (kernel/kthread.c): kthread_park() blocks on
 * wait_for_completion(&self->parked) until the target has actually parked;
 * the whole park/unpark API exists to QUIESCE a worker.  In-tree callers
 * depend on that contract:
 *
 *   drm_sched_entity_fini()  sys/dev/drm/scheduler/sched_entity.c:273-279
 *       [comment in-tree: "Park the kernel for a moment to make sure it
 *        isn't processing our enity."]
 *       kthread_park(sched->thread);
 *       kthread_unpark(sched->thread);
 *       ... drm_sched_entity_kill_jobs(entity) + kfree(entity->rq_list)
 *   amdgpu ring stop/start across GPU reset
 *			      sys/dev/drm/amd/amdgpu/amdgpu_device.c:3412,3441
 *   amdgpu debugfs ring reset sys/dev/drm/amd/amdgpu/amdgpu_debugfs.c:868,884
 *
 * Because the shim's park returns while the worker is still anywhere in
 * its loop (including mid drm_sched_select_entity()/entity-pop on the
 * entity being freed), the fini-side frees race the worker's use of the
 * same structures -> use-after-free.
 *
 * This module drives the REAL in-tree shim functions (kthread_run /
 * kthread_should_park / kthread_parkme / kthread_park / kthread_unpark
 * resolved from drm.ko) with a drm_sched_main()-shaped worker loop.
 *
 * test=1: quiescence measurement - after kthread_park() returns, the
 *         worker MUST be frozen (Linux contract).  Count how many times
 *         the "parked" worker kept running during a 500 ms window.
 * test=2: the drm_sched_entity_fini() shape - park(); kfree(entity);
 *         kzalloc(victim) [memory reuse]; [worker must be unable to run
 *         here]; check victim integrity after the window.
 * test=3: one round of test=2, then a deliberate panic() when the victim
 *         corruption is observed (proof-of-primitive crash signature,
 *         same style as DF-4842's KKASSERT panic).
 *
 * Root-only module test on the audit guest (no GPU required):
 *   kldload amdgpu            # pulls in drm.ko; provides the shim
 *   kldload /tmp/df5200/df5200.ko
 *   sysctl hw.df5200_test=1   # quiescence violation counts (no crash)
 *   sysctl hw.df5200_test=2   # UAF read/write evidence (no crash)
 *   sysctl hw.df5200_test=3   # same, then panic on proven corruption
 *
 * NOTE: teardown NEVER calls kthread_stop() - that is the (known,
 * separately-filed) DF-4842 bug; the worker exits via its own loop flag
 * and the DragonFly kthread machinery reclaims its td.
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/sysctl.h>
#include <sys/types.h>
#include <sys/errno.h>
#include <sys/module.h>

#include <linux/slab.h>
#include <linux/bitops.h>
#include <linux/kthread.h>
#include <linux/sched.h>

#define DF5200_MAGIC	0x00c0ffee00c0ffeeUL
#define DF5200_VICTIM	0xdeadbeefbaadf00dUL
#define DF5200_ROUNDS	5
#define DF5200_SHOULD_PARK_BIT	2	/* KTHREAD_SHOULD_PARK, linux_kthread.c:80 */

/*
 * A drm-sched-entity-shaped object (512-byte class).  The worker reads
 * magic (entity state) and writes payload slots + visits (job processing)
 * every loop iteration, exactly like drm_sched_main() touches its entity
 * after select_entity() hands it out.
 */
struct df5200_obj {
	uint64_t	magic;
	uint64_t	canary;
	uint64_t	visits;
	uint64_t	payload[61];
} __packed;

static struct task_struct	*g_task;
static struct df5200_obj	*g_ent;		/* the live "entity" */
static volatile int		 g_stop;
static volatile int		 g_exited;
static volatile int		 g_writes;	/* entity touches */
static volatile int		 g_badmagic;	/* reads of stale/freed state */

/*
 * drm_sched_main()-shaped worker (sched_main.c:541-560):
 *   while (!kthread_should_stop()) {
 *     wait_event_interruptible(wq, (!drm_sched_blocked(sched) &&
 *                                  (entity = drm_sched_select_entity(sched))) ...)
 *     ... process entity ...
 *   }
 * drm_sched_blocked() -> kthread_should_park()/kthread_parkme() is our
 * park check; per-iteration entity capture is g_ent.
 */
static int
df5200_worker(void *arg __unused)
{
	struct df5200_obj *ent;
	int slot = 0;

	while (!g_stop) {
		if (kthread_should_park())
			kthread_parkme();

		ent = g_ent;
		if (ent != NULL) {
			if (ent->magic != DF5200_MAGIC)
				atomic_add_int(&g_badmagic, 1);	/* UAF read */
			ent->payload[slot] = 0x4141414141414141UL;
			ent->visits++;
			atomic_add_int(&g_writes, 1);		/* UAF write (post-free) */
			slot = (slot + 1) % 61;
		}

		tsleep(df5200_worker, 0, "df5200w", 1);
	}
	g_exited = 1;

	return (0);
}

static void
df5200_init_worker(void)
{
	g_ent = kzalloc(sizeof(*g_ent), GFP_KERNEL);
	g_ent->magic = DF5200_MAGIC;
	g_ent->canary = DF5200_MAGIC;

	g_stop = 0;
	g_exited = 0;
	g_writes = 0;
	g_badmagic = 0;

	g_task = kthread_run(df5200_worker, NULL, "df5200wkr");

	/* let the worker start touching the entity */
	tsleep(df5200_init_worker, 0, "df5200i", hz / 4);
}

static void
df5200_fini_worker(void)
{
	int i;

	g_stop = 1;
	for (i = 0; i < 40 && !g_exited; i++)	/* up to 4 s */
		tsleep(df5200_fini_worker, 0, "df5200f", hz / 10);
	/* kthread_stop() deliberately NOT used: that is DF-4842's bug.
	 * The DFly kthread machinery reclaims the td itself; the
	 * task_struct leak is the (known) DF-2196 issue. */
	g_task = NULL;
	kfree(g_ent);
	g_ent = NULL;
}

static int
df5200_test1(void)
{
	uint64_t v0, v1;
	int r, viol = 0, total_d = 0;

	df5200_init_worker();

	for (r = 1; r <= DF5200_ROUNDS; r++) {
		kthread_park(g_task);	/* Linux: must not return unparked */
		v0 = g_ent->visits;
		tsleep(df5200_test1, 0, "df5200q", hz / 2);
		v1 = g_ent->visits;
		kthread_unpark(g_task);

		kprintf("[DF5200] t1 r%d/%d: park() returned; worker visits "
			"during 500ms 'parked' window: %ju -> %ju (%s) "
			"SHOULD_PARK bit now=%d\n",
			r, DF5200_ROUNDS, (uintmax_t)v0, (uintmax_t)v1,
			(v1 > v0) ? "RAN WHILE PARKED" : "frozen (correct)",
			test_bit(DF5200_SHOULD_PARK_BIT,
				 &g_task->kt_flags) ? 1 : 0);
		if (v1 > v0) {
			viol++;
			total_d += (int)(v1 - v0);
		}
		tsleep(df5200_test1, 0, "df5200q2", hz / 10);
	}

	kprintf("[DF5200] t1 RESULT: %d/%d rounds violated the park "
		"quiescence contract; %d post-park visits total\n",
		viol, DF5200_ROUNDS, total_d);

	df5200_fini_worker();
	return (0);
}

static int
df5200_test2(int dopanic)
{
	struct df5200_obj *old, *victim;
	int r, corruptions = 0, overlaps = 0, bad0;

	df5200_init_worker();
	bad0 = g_badmagic;

	for (r = 1; r <= DF5200_ROUNDS; r++) {
		kthread_park(g_task);		/* returns immediately (bug) */

		old = g_ent;
		kfree(old);			/* drm_sched_entity_fini() frees
						 * the entity right here */
		victim = kzalloc(sizeof(*victim), GFP_KERNEL); /* reuse */
		victim->magic = DF5200_VICTIM;
		victim->canary = DF5200_VICTIM;
		victim->visits = 0;
		if ((void *)victim == (void *)old)
			overlaps++;
		/*
		 * The racing worker still has `old` captured for this loop
		 * iteration (g_ent is still `old` - select_entity() had
		 * already handed it out when fini struck).  In Linux the
		 * worker CANNOT run here: it is parked.  In this shim it
		 * keeps looping: reads old->magic (freed memory / UAF read)
		 * and writes old->payload[*] + old->visits - hitting the
		 * reallocated victim when the addresses overlap.
		 */
		tsleep(df5200_test2, 0, "df5200u", hz / 5);

		if (victim->visits != 0 || victim->payload[0] != 0 ||
		    victim->payload[30] != 0) {
			corruptions++;
			kprintf("[DF5200] t2 r%d/%d: UAF WRITE PROVEN: live "
				"victim %p (freed entity %p, overlap=%s) "
				"corrupted by the 'parked' worker: "
				"visits=%ju payload[0]=%#jx payload[30]=%#jx\n",
				r, DF5200_ROUNDS, victim, old,
				((void *)victim == (void *)old) ? "yes" : "no",
				(uintmax_t)victim->visits,
				(uintmax_t)victim->payload[0],
				(uintmax_t)victim->payload[30]);
			if (dopanic) {
				panic("DF5200: kthread_park() quiescence "
					"broken - 'parked' worker UAF-wrote "
					"freed-then-reallocated memory "
					"(victim->visits=%ju payload[0]=%#jx)",
					(uintmax_t)victim->visits,
					(uintmax_t)victim->payload[0]);
			}
		}

		/*
		 * Re-arm a fresh entity BEFORE un-parking the worker: on a
		 * fixed kernel the worker is truly parked and never touches
		 * the freed memory at all; on the buggy kernel it already
		 * did during the window above.
		 */
		kfree(victim);
		g_ent = kzalloc(sizeof(*g_ent), GFP_KERNEL);
		g_ent->magic = DF5200_MAGIC;
		g_ent->canary = DF5200_MAGIC;

		kthread_unpark(g_task);
		tsleep(df5200_test2, 0, "df5200u2", hz / 10);
	}

	kprintf("[DF5200] t2 RESULT: %d/%d rounds corrupted the reallocated "
		"victim (%d same-address overlaps); %d UAF reads of freed "
		"memory (badmagic %d -> %d); total worker writes=%d\n",
		corruptions, DF5200_ROUNDS, overlaps,
		g_badmagic - bad0, bad0, g_badmagic, g_writes);

	df5200_fini_worker();
	return (0);
}

static int
sysctl_df5200_test(SYSCTL_HANDLER_ARGS)
{
	int test, error;

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

	switch (test) {
	case 1:
		return (df5200_test1());
	case 2:
		return (df5200_test2(0));
	case 3:
		return (df5200_test2(1));
	}
	return (EINVAL);
}
SYSCTL_NODE(_hw, OID_AUTO, df5200, CTLFLAG_RW, 0, "DF-5200 test");
SYSCTL_PROC(_hw_df5200, OID_AUTO, test, CTLTYPE_INT | CTLFLAG_RW,
	0, 0, sysctl_df5200_test, "I", "DF-5200 park-quiescence test selector");

static int
df5200_modev(module_t mod, int type, void *data)
{
	switch (type) {
	case MOD_LOAD:
		kprintf("[DF5200] loaded (kthread shim from drm.ko)\n");
		return (0);
	case MOD_UNLOAD:
		/* each test tears down its own worker before returning */
		return (0);
	}
	return (0);
}

static moduledata_t df5200_mod = { "df5200", df5200_modev, 0 };
DECLARE_MODULE(df5200, df5200_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);