/*
 * df4842.c — DF-4842: DragonFly's linux-kthread shim kthread_stop() frees a
 * struct thread it does not own (and cannot prove has exited) after a fixed
 * 1-second tsleep.
 *
 *   sys/dev/drm/linux_kthread.c:89-101:
 *
 *	int kthread_stop(struct task_struct *ts)
 *	{
 *		set_bit(KTHREAD_SHOULD_STOP, &ts->kt_flags);
 *		kthread_unpark(ts);
 *		wake_up_process(ts);
 *		tsleep(kthread_stop, 0, "kstop", hz);   <-- fixed 1s, proves nothing
 *		lwkt_free_thread(ts->dfly_td);         <-- frees a thread it does
 *		return ts->kt_exitvalue;                   not own
 *	}
 *
 * The DFly thread reclaims itself: when the kthread body returns,
 * cpu_set_thread_handler() runs kthread_exit() -> lwkt_exit() which parks
 * the td in gd->gd_freetd for the NEXT exiting thread on that cpu to
 * objcache_put() (lwkt_thread.c:1684-1744).  kthread_stop's extra
 * lwkt_free_thread() therefore either:
 *
 *   case A (body exited < 1s): the td is parked in gd_freetd[cpu] AND
 *          objcache_put() into thread_cache by kthread_stop -> the same td
 *          exists in two freelists; the next thread exit on that cpu puts
 *          it into the cache a SECOND time -> two live threads can be
 *          handed the same struct thread.
 *
 *   case B (body still running at +1s, e.g. drm_sched_main() busy inside
 *          run_job()): lwkt_free_thread() on a TDF_RUNNING thread ->
 *          KKASSERT "lwkt_free_thread" fatal panic (or, if merely blocked,
 *          a freed-but-scheduled live td).
 *
 * Reached in-tree from drm_sched_fini() (sys/dev/drm/scheduler/
 * sched_main.c:652-656) — every amdgpu ring teardown / driver unload — and
 * from intel_engine_fini_breadcrumbs() (i915).
 *
 * test=1: case B  — deterministic 'panic: lwkt_free_thread' ~1s after the
 *                   sysctl (guest wedges at db>).
 * test=2: case A  — staged double-ownership: shows gd_freetd[cpu] == tdS
 *                   after fini, drives the second put, then scans the
 *                   thread cache from several cpus for tdS reappearing in
 *                   allocations it must not be part of.
 *
 * Root-only module test on the audit guest (no GPU required):
 *   kldload amdgpu           # pulls in drm.ko; scheduler symbols
 *   kldload ./df4842.ko
 *   sysctl hw.df4842_test=1  # case B: panic
 *   sysctl hw.df4842_test=2  # case A: staged double-free evidence
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/sysctl.h>
#include <sys/lock.h>
#include <sys/thread.h>
#include <sys/thread2.h>
#include <sys/globaldata.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 DF4842_NSCAN	16

struct dfjob {
	struct drm_sched_job	base;
};

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

static struct drm_sched_entity	g_ent;
static struct drm_sched_rq	*g_rq_list[1];

static volatile int	g_inrun;	/* run_job entered (test 1) */
static int		g_runcount;	/* run_job calls */

/* case A staging */
static struct thread	*g_tds;		/* the sched thread's DFly td */
static int		g_cpus;		/* cpu the sched thread ran on */
static volatile int	g_phase;	/* handshake with pinned helpers */
static struct thread	*g_probe_td, *g_sacr_td, *g_scanA, *g_scanB, *g_scanC;
static int		g_cpuF0, g_cpuF1; /* sysctl cpu before/after fini */
static volatile struct thread *g_hit_td[8];
static volatile int	g_hits;
static int		g_t1_done;

/* ------------------------------------------------------------------ */
/* fake dma-fence (never signalled)                                      */

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

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);
	int t0 = ticks;

	lockinit(&f->lock, "df4842f", 0, LK_CANRECURSE);
	dma_fence_init(&f->base, &g_fence_ops, &f->lock,
	    0x48420000 + g_runcount, 1);
	g_runcount++;

	kprintf("[DF4842] run_job(job=%p): entered on td=%p cpu=%d - "
	    "now BUSY-SPINNING 4s (no blocking: td stays TDF_RUNNING)\n",
	    (void *)job, curthread, mycpu->gd_cpuid);
	g_inrun = 1;

	/* pure busy-wait: no tsleep, no yield, no preemption of our state.
	 * NOTE: ticks must be read through a volatile qualifier or -O2
	 * register-caches the load and the loop never terminates. */
	while (*(volatile int *)&ticks - t0 < 3 * hz)
		;

	kprintf("[DF4842] run_job(job=%p): spin done (too late for "
	    "kthread_stop's 1s window)\n", (void *)job);
	return (&f->base);
}

static void
op_free_job(struct drm_sched_job *job)
{
	kfree(job);
}

static void
op_timedout_job(struct drm_sched_job *job)
{
	kprintf("[DF4842] 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,
};

/* ------------------------------------------------------------------ */
/* case A helpers (pinned lwkt threads; they NEVER exit once started)    */

static void
df4842_probe(void *arg)
{
	struct thread *ftd = mycpu->gd_freetd;

	kprintf("[DF4842] probe: cpu=%d gd_freetd=%p tdS=%p -> %s\n",
	    mycpu->gd_cpuid, ftd, (void *)g_tds,
	    (ftd == g_tds) ?
	    "== tdS STILL PARKED in gd_freetd (while kthread_stop already "
	    "objcache_put() it into thread_cache: DOUBLE OWNERSHIP)" :
	    "not tdS (unexpected)");
	g_phase = 1;
	for (;;)
		tsleep(&df4842_probe, 0, "df4842p", 0);
}

static void
df4842_sacr(void *arg)
{
	/*
	 * Returning makes the DFly kthread wrapper run kthread_exit() ->
	 * lwkt_exit(), which on this cpu consumes gd_freetd == tdS and
	 * objcache_put()s it into thread_cache a SECOND time.
	 */
	kprintf("[DF4842] sacr: exiting on cpu=%d - lwkt_exit() will now "
	    "objcache_put(gd_freetd=tdS %p) AGAIN (2nd put)\n",
	    mycpu->gd_cpuid, (void *)g_tds);
}

static void
df4842_scan(void *arg)
{
	struct thread *td;
	int i, myhits = 0;

	for (i = 0; i < DF4842_NSCAN; i++) {
		/*
		 * Inert allocations (never scheduled): each pulls a struct
		 * thread out of freetd/cache exactly like any real
		 * kthread/ithread creation would.
		 */
		td = lwkt_alloc_thread(NULL, LWKT_THREAD_STACK,
		    mycpu->gd_cpuid, 0);
		if (td == NULL) {
			kprintf("[DF4842] scan(cpu=%d): alloc %d FAILED\n",
			    mycpu->gd_cpuid, i);
			break;
		}
		if (td == g_tds) {
			int n;

			kprintf("[DF4842] scan(cpu=%d) alloc %d: thread_cache "
			    "handed out tdS %p — SAME struct thread in a NEW "
			    "allocation while %s\n",
			    mycpu->gd_cpuid, i, (void *)td,
			    (g_phase >= 3) ?
			    "the 2nd put also sits in the cache" :
			    "gd_freetd also held it");
			n = atomic_fetchadd_int(&g_hits, 1);
			if (n < 8)
				g_hit_td[n] = td;
			myhits++;
		}
	}
	kprintf("[DF4842] scan(cpu=%d): done, %d tdS hits in %d allocs\n",
	    mycpu->gd_cpuid, myhits, DF4842_NSCAN);
	atomic_add_int(&g_phase, 10);
	for (;;)
		tsleep(&df4842_scan, 0, "df4842s", 0);
}

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

static int
df4842_run_test1(void)
{
	struct drm_gpu_scheduler *s;
	struct dfjob *j;
	int i;

	kprintf("[DF4842] test1 (case B): scheduler thread busy in run_job() "
	    "when drm_sched_fini() runs\n");

	s = kzalloc(sizeof(*s), GFP_KERNEL);
	drm_sched_init(s, &g_ops, 1, 0, MAX_SCHEDULE_TIMEOUT, "df4842b");
	g_rq_list[0] = &s->sched_rq[DRM_SCHED_PRIORITY_NORMAL];
	drm_sched_entity_init(&g_ent, g_rq_list, 1, NULL);

	kprintf("[DF4842] sched thread: task=%p dfly_td=%p cpu=%d\n",
	    (void *)s->thread, (void *)s->thread->dfly_td,
	    s->thread->dfly_td->td_gd->gd_cpuid);

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

	for (i = 0; i < 100 && g_inrun == 0; i++)
		tsleep(&g_inrun, 0, "df4842w", 1);
	if (g_inrun == 0) {
		kprintf("[DF4842] run_job never entered?\n");
		return (ETIMEDOUT);
	}

	kprintf("[DF4842] calling drm_sched_fini(): kthread_stop sets "
	    "SHOULD_STOP, tsleeps its FIXED 1s, then lwkt_free_thread() on "
	    "the still-TDF_RUNNING td %p -> expect KKASSERT panic\n",
	    (void *)s->thread->dfly_td);
	g_t1_done = 1;

	drm_sched_fini(s);	/* never returns on an unpatched kernel */

	kprintf("[DF4842] drm_sched_fini RETURNED (no panic?!)\n");
	kfree(s);
	g_t1_done = 2;
	return (0);
}

static int
df4842_run_test2(void)
{
	struct drm_gpu_scheduler *s;
	int i, rounds, hits;

	kprintf("[DF4842] test2 (case A): idle sched thread exits fast; "
	    "td parked in gd_freetd AND objcache_put by kthread_stop\n");

	for (rounds = 1; rounds <= 3; rounds++) {
		g_hits = 0;
		g_phase = 0;
		g_cpuF0 = mycpu->gd_cpuid;

		s = kzalloc(sizeof(*s), GFP_KERNEL);
		drm_sched_init(s, &g_ops, 1, 0, MAX_SCHEDULE_TIMEOUT,
		    "df4842a");
		g_tds = s->thread->dfly_td;
		g_cpus = g_tds->td_gd->gd_cpuid;
		kprintf("[DF4842] round %d: sched thread td=%p on cpu=%d; "
		    "calling drm_sched_fini() (kthread_stop: 1s tsleep + "
		    "lwkt_free_thread)\n",
		    rounds, (void *)g_tds, g_cpus);

		drm_sched_fini(s);	/* idle: body exits, td parked + put */
		g_cpuF1 = mycpu->gd_cpuid;
		kfree(s);
		kprintf("[DF4842] drm_sched_fini returned; sysctl cpu was "
		    "%d, now %d; staging gd_freetd proof on cpu %d\n",
		    g_cpuF0, g_cpuF1, g_cpus);

		/* 1: prove the td is still parked in gd_freetd[cpuS] */
		lwkt_create(df4842_probe, NULL, &g_probe_td, NULL,
		    TDF_NOSTART, g_cpus, "df4842probe");
		lwkt_schedule(g_probe_td);
		for (i = 0; i < 100 && g_phase == 0; i++)
			tsleep(&g_phase, 0, "df4842h", 1);
		if (g_phase == 0) {
			kprintf("[DF4842] probe never ran?!\n");
			return (ETIMEDOUT);
		}
		g_phase = 2;

		/* 2: force the SECOND objcache_put via a thread exit on cpuS */
		kprintf("[DF4842] launching sacrificial exit thread on "
		    "cpu=%d\n", g_cpus);
		lwkt_create(df4842_sacr, NULL, &g_sacr_td, NULL,
		    TDF_NOSTART, g_cpus, "df4842sacr");
		lwkt_schedule(g_sacr_td);
		tsleep(&df4842_sacr, 0, "df4842x", 2);
		g_phase = 3;

		/* 3: scan the cache from the involved cpus */
		lwkt_create(df4842_scan, NULL, &g_scanA, NULL,
		    TDF_NOSTART, g_cpus, "df4842scanA");
		lwkt_schedule(g_scanA);
		if (g_cpuF1 != g_cpus) {
			lwkt_create(df4842_scan, NULL, &g_scanB, NULL,
			    TDF_NOSTART, g_cpuF1, "df4842scanB");
			lwkt_schedule(g_scanB);
		}
		lwkt_create(df4842_scan, NULL, &g_scanC, NULL,
		    TDF_NOSTART, mycpu->gd_cpuid, "df4842scanC");
		lwkt_schedule(g_scanC);

		for (i = 0; i < 200 && g_phase < 32; i++)
			tsleep(&g_phase, 0, "df4842h", 1);

		hits = g_hits;
		kprintf("[DF4842] round %d result: %d allocation hit(s) for "
		    "tdS %p after the double put\n", rounds, hits,
		    (void *)g_tds);
		if (hits > 0)
			break;
		kprintf("[DF4842] (magazine locality miss; retrying with a "
		    "fresh scheduler)\n");
	}

	if (g_hits == 0) {
		kprintf("[DF4842] test2: double-put is proven by construction "
		    "(gd_freetd==tdS probe + code path); cache dispense not "
		    "observed in these rounds\n");
	} else if (g_hits >= 2) {
		kprintf("[DF4842] *** DOUBLE DISPENSE CONFIRMED: %d "
		    "allocations received the SAME struct thread %p — two "
		    "live threads would share one td ***\n", g_hits,
		    (void *)g_tds);
	}
	return (0);
}

static int
sysctl_df4842_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 (df4842_run_test1());
	if (v == 2)
		return (df4842_run_test2());
	return (0);
}
SYSCTL_PROC(_hw, OID_AUTO, df4842_test, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_df4842_test, "I", "df4842 test selector (1=caseB panic, "
    "2=caseA double-put)");

static int
df4842_init(struct module *mod, int what, void *arg)
{
	switch (what) {
	case MOD_LOAD:
		kprintf("[DF4842] module loaded: sysctl hw.df4842_test=1 "
		    "(case B panic) or =2 (case A double-put)\n");
		break;
	default:
		break;
	}
	return (0);
}

static int
df4842_uninit(struct module *mod, int what, void *arg)
{
	switch (what) {
	case MOD_UNLOAD:
	default:
		break;
	}
	return (0);
}

static moduledata_t df4842_mod = { "df4842", df4842_init, df4842_uninit };
DECLARE_MODULE(df4842, df4842_mod, SI_SUB_DRIVERS, SI_ORDER_MIDDLE);