/*
 * df5201.c - DF-5201: linux_task_alloc() (sys/dev/drm/linux_sched.c:58-91)
 * installs td_linux_task on ANY thread evaluating `current`, including
 * pure DragonFly kernel threads (td_lwp == NULL).  Those threads exit via
 * lwkt_exit() (sys/kern/lwkt_thread.c:1684-1744), NEVER via lwp_exit(),
 * so the linux_task_drop_callback (sys/kern/kern_exit.c:708-709) never
 * runs for them:
 *
 *   1. the task_struct (kt_spin included) leaks permanently, and
 *   2. td->td_linux_task stays set when the exiting td is parked in
 *      gd->gd_freetd (lwkt_thread.c:1740-1742) and handed out again by
 *      lwkt_alloc_thread() (lwkt_thread.c:361-387), which does NOT clear
 *      it - the recycled thread inherits `current` == the DEAD thread's
 *      task_struct: task->dfly_td points at a different context,
 *      kt_spin/state become shared across live threads, and
 *      signal_pending()/fatal_signal_pending() consult the wrong lwp.
 *
 * This is the linux_task_alloc() site - distinct from DF-2196 (the
 * kthread_run() task), reachable whenever ANY non-drm kernel thread
 * (syncer, bufdaemon, irq thread, callout thread - anything executing
 * drm code that evaluates current, e.g. dma_fence_wait() ->
 * prepare_to_wait() -> set_current_state()) touches drm.
 *
 * Mechanism (deterministic on one CPU):
 *
 *   W1 = kthread_create()   -> fresh td (linux_task == NULL)
 *        W1 evaluates current -> linux_task_alloc installs T
 *        W1 body returns -> kthread_exit -> lwkt_exit:
 *            gd_freetd = td_W1  (td_linux_task == T still set)
 *   W2 = kthread_create()   -> lwkt_alloc_thread takes gd_freetd == td_W1
 *        W2's curthread->td_linux_task != NULL on entry: STALE T from the
 *        dead W1.  W2's `current` returns the DEAD W1's task_struct.
 *
 * test=1 runs that pair for ROUNDS rounds and prints every hit.
 *
 * Root-only module test on the audit guest (no GPU required):
 *   kldload amdgpu            # pulls in drm.ko; provides linux_task_alloc
 *   kldload /tmp/df5201/df5201.ko
 *   sysctl hw.df5201_test=1
 */
#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 <sys/thread.h>
#include <sys/thread2.h>
#include <sys/kthread.h>

#include <linux/sched.h>	/* struct task_struct */
#include <asm/current.h>	/* the `current` macro -> linux_task_alloc */

#define DF5201_ROUNDS	8

struct df5201_rec {
	int			round;
	struct thread		*td_w1;
	struct thread		*td_w2;
	struct task_struct	*task_w1;
	void			*stale_in;	/* W1 entry (expect NULL) */
	void			*stale_out;	/* W2 entry (hit if != NULL) */
	void			*parked_stale;	/* W1's parked td after exit:
						 * non-NULL => task leaked AND
						 * pointer left dangling */
	volatile int		done;
};

static int	g_hits;
static int	g_recycled;	/* fix-validation: W2 task == W1 task addr */
static int	g_parked_stale;	/* primary metric: W1's parked td keeps
				 * td_linux_task after exit (the leak) */
static int	g_rounds;

static void
df5201_w1(void *arg)
{
	struct df5201_rec *r = arg;

	r->td_w1 = curthread;
	r->stale_in = curthread->td_linux_task;	/* expect NULL (fresh td) */

	/* THE touch: what any kernel thread executing drm code does
	 * (dma_fence_wait/prepare_to_wait/set_current_state/... all
	 * evaluate `current`). */
	r->task_w1 = current;

	kprintf("[DF5201] r%d W1: td=%p stale_on_entry=%p; current -> "
		"linux_task_alloc installed task=%p (dfly_td=%p, pid=%d)\n",
		r->round, r->td_w1, r->stale_in, r->task_w1,
		r->task_w1->dfly_td, r->task_w1->pid);

	r->done = 1;
	wakeup(r);
	/* body returns -> kthread_exit() -> lwkt_exit(): td parked in
	 * gd_freetd with td_linux_task == task_w1 still installed, and
	 * NOTHING will ever drop it (linux_task_drop only runs from
	 * lwp_exit, which pure kthreads never execute). */
}

static void
df5201_w2(void *arg)
{
	struct df5201_rec *r = arg;
	struct task_struct *t;

	r->td_w2 = curthread;
	r->stale_out = curthread->td_linux_task;

	if (r->stale_out != NULL) {
		g_hits++;
		t = current;	/* returns the STALE task (macro sees != NULL) */
		kprintf("[DF5201] r%d W2: *** HIT *** recycled td=%p "
			"(== W1's td: %s) carries STALE td_linux_task=%p - "
			"W2's `current` is the DEAD W1's task_struct "
			"(t=%p dfly_td=%p same-as-w2-td:%s)\n",
			r->round, r->td_w2,
			(r->td_w2 == r->td_w1) ? "yes" : "no",
			r->stale_out, t,
			t->dfly_td,
			(t->dfly_td == curthread) ? "yes" : "no");
	} else {
		/*
		 * Fresh alloc: if it lands on W1's task address, W1's task
		 * was FREED at exit (fix working); on the unfixed kernel W1's
		 * task leaks and the addresses never collide.
		 */
		t = current;
		if (t == r->task_w1)
			g_recycled++;
		kprintf("[DF5201] r%d W2: td=%p clean (no stale task); "
			"fresh task=%p (%s)\n",
			r->round, r->td_w2, t,
			(t == r->task_w1) ?
			"RECYCLED W1's task - freed at exit" :
			"new address - W1's task still leaked");
	}

	r->done = 2;
	wakeup(r);
}

static int
df5201_test1(void)
{
	static struct df5201_rec rec;
	int i, w, err;

	g_hits = 0;
	g_rounds = DF5201_ROUNDS;

	for (i = 1; i <= DF5201_ROUNDS; i++) {
		bzero(&rec, sizeof(rec));
		rec.round = i;

		err = kthread_create(df5201_w1, &rec, NULL, "df5201a%d", i);
		if (err) {
			kprintf("[DF5201] kthread_create(W1) failed: %d\n", err);
			return (err);
		}
		for (w = 0; w < 20 && rec.done < 1; w++)
			tsleep(&rec, 0, "df5201a", hz / 20);

		/*
		 * W1 has exited here: its td sits in gd_freetd[cpu].
		 * Primary metric: does the parked td still carry the
		 * task_struct pointer?  (It is dead-but-cached memory;
		 * reading it is safe in this controlled test.)  With the
		 * lwkt_exit() fix the pointer is NULL and the task was
		 * freed; without it the task leaked and the pointer dangles.
		 */
		rec.parked_stale = rec.td_w1->td_linux_task;
		if (rec.parked_stale != NULL) {
			g_parked_stale++;
			kprintf("[DF5201] r%d PARKED: exited td=%p still "
				"carries td_linux_task=%p -> task LEAKED "
				"(no drop path on kthread exit)\n",
				rec.round, rec.td_w1, rec.parked_stale);
		} else {
			kprintf("[DF5201] r%d PARKED: exited td=%p is clean "
				"(task dropped at lwkt_exit)\n",
				rec.round, rec.td_w1);
		}

		err = kthread_create(df5201_w2, &rec, NULL, "df5201b%d", i);
		if (err) {
			kprintf("[DF5201] kthread_create(W2) failed: %d\n", err);
			return (err);
		}
		for (w = 0; w < 20 && rec.done < 2; w++)
			tsleep(&rec, 0, "df5201b", hz / 20);

		/* let W2's td also cycle back before the next round */
		tsleep(df5201_test1, 0, "df5201c", hz / 10);
	}

	kprintf("[DF5201] RESULT: %d/%d exited kthreads left a leaked, "
		"never-dropped task_struct on their parked td (THE LEAK); "
		"%d/%d rounds handed a RECYCLED td with stale td_linux_task "
		"to a live thread (bzero prevents this); %d/%d rounds "
		"recycled W1's task address\n",
		g_parked_stale, g_rounds, g_hits, g_rounds,
		g_recycled, g_rounds);

	return (0);
}

static int
sysctl_df5201_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 (df5201_test1());
	}
	return (EINVAL);
}
SYSCTL_NODE(_hw, OID_AUTO, df5201, CTLFLAG_RW, 0, "DF-5201 test");
SYSCTL_PROC(_hw_df5201, OID_AUTO, test, CTLTYPE_INT | CTLFLAG_RW,
	0, 0, sysctl_df5201_test, "I", "DF-5201 stale-td_linux_task test");

static int
df5201_modev(module_t mod, int type, void *data)
{
	switch (type) {
	case MOD_LOAD:
		kprintf("[DF5201] loaded (linux_task_alloc from drm.ko)\n");
		return (0);
	case MOD_UNLOAD:
		return (0);
	}
	return (0);
}

static moduledata_t df5201_mod = { "df5201", df5201_modev, 0 };
DECLARE_MODULE(df5201, df5201_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);