/*
 * df4919.c — DF-4919: cancel_work_sync()/flush_work() can return in the
 *            lockless window between "dequeued" and "running", i.e.
 *            BEFORE the work function executes — linux_workqueue.c:331,364.
 *
 * process_all_work() dequeues a work and only then marks it running, with
 * the whole transition under worker_lock:
 *      linux_workqueue.c:74:   work->on_queue = false;   <- dequeued
 *      linux_workqueue.c:87:   work->running = true;     <- about to run
 *
 * _cancel_work() reads BOTH flags WITHOUT holding that lock:
 *      linux_workqueue.c:349:  if (work->on_queue) { ... }     <- C1
 *      linux_workqueue.c:364:  if (work->running == false)     <- C2
 *                                  break;   <- RETURN
 * and flush_work() does the same at linux_workqueue.c:331.
 *
 * A canceller whose C1/C2 load pair straddles the worker's 74->87
 * transition observes (on_queue==false && running==false) and returns
 * while the function is about to execute.  That violates the API contract
 * spelled out above cancel_work_sync() (lines 390-394):
 *      "In any case, wait for work to complete or be removed from the
 *       workqueue, callers may free associated data structures after
 *       this call."
 * A caller that then frees the work (kfree of the container object, or
 * returning the stack frame holding an INIT_WORK_ONSTACK work, e.g.
 * drm_framebuffer.c:475-481 drm_mode_rmfb) leaves process_all_work()
 * about to call work->func(work) through freed / stack-reused memory:
 * arbitrary-call use-after-free.  In the race case cancel_work_sync()
 * also returns FALSE ("work wasn't pending") for a work that was queued
 * and is executing.
 *
 * Discriminating module design (no GPU needed): NDRIVER independent
 * driver/wq/work triplets.  Every cycle a driver queues exactly ONE
 * instance into ITS OWN workqueue, waits a jittered interval, then issues
 * exactly one cancel_work_sync() (or flush_work()).  The queue completes
 * before the cancel runs, so a posthumous execution can only mean the
 * shim returned during the worker's dequeue->start transition — there is
 * no concurrent queue_work() whose insert could race the flags (that
 * would be a caller-owned race, not the shim bug).
 *
 *	wk_func() finds its driver state via container_of() and, after a
 *	short busy body, checks the alive flag the driver clears right
 *	after cancel_work_sync()/flush_work() returned.  Reading 0 ==
 *	posthumous execution == the premature-return proof.
 *
 *   kldload drm && kldload ./df4919.ko
 *   sysctl hw.df4919_test=1   # cancel_work_sync race, safe record+stop
 *   sysctl hw.df4919_test=2   # same, panic() on hit (crash proof)
 *   sysctl hw.df4919_test=3   # flush_work race, safe record+stop
 *   sysctl hw.df4919_status   # counters
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/libkern.h>
#include <machine/cpufunc.h>
#include <sys/sysctl.h>
#include <sys/lock.h>
#include <sys/thread.h>
#include <sys/types.h>
#include <sys/errno.h>
#include <sys/module.h>

#include <linux/workqueue.h>

#define	NDF	4

struct df4919_state {
	struct workqueue_struct	*wq;
	struct work_struct	work;
	struct thread		*td;
	volatile u_int		alive;	/* 0 => driver believes work dead */
	volatile u_int		hits;
	volatile u_int		runs;
	volatile u_int		cycles;
	volatile u_int		cancels;
	volatile u_int		cancels_false;
	volatile u_int		flushes;
	volatile u_int		falses_at_pub;	/* feeder snapshot after queue */
	volatile u_int		real_hits;	/* hits with a post-publication
						 * false-return = kernel
						 * window, not harness skew */
	int			cpu;
	int			first_hit_at;	/* cancel count at first hit */
	/* debug ring: what happened right before a hit */
	struct {
		volatile u_int	tag;	/* 1=qret0 2=qret1 3=cretF 4=cretT
						 * 5=funcin 6=funcout */
		volatile u_int	aux;
	} ring[32];
	volatile u_int		ring_idx;
};

static struct df4919_state g_st[NDF];
static volatile u_int g_stop;
static volatile int g_panic_on_hit, g_mode;
static volatile int g_any_hit;

static inline void
df4919_ev(struct df4919_state *s, u_int tag, u_int aux)
{
	u_int i = s->ring_idx & 31;

	s->ring[i].tag = tag;
	s->ring[i].aux = aux;
	s->ring_idx++;
}
#define	EV(s, t, a)	df4919_ev((s), (t), (a))

static void
wk_func(struct work_struct *w)
{
	struct df4919_state *s =
	    container_of(w, struct df4919_state, work);

	int a;

	atomic_add_int(&s->runs, 1);
	/*
	 * Sample alive ONCE at entry (before the body): the hit must
	 * be a cancel/flush RETURN that completed between this
	 * instance's publication and its start - later writes (e.g.
	 * for a next instance queued during our body) must not count.
	 */
	a = s->alive;
	EV(s, 5, a);
	DELAY(10);
	if (a == 0) {
		int k;

		/*
		 * Classify: a REAL hit requires that some cancel returned
		 * "not pending" AFTER this instance was published
		 * (falses_at_pub snapshotted right after queue_work()
		 * returned).  Without that, the alive==0 came from a
		 * pre-claim verdict whose store landed late (harness
		 * store-buffer skew) - not a kernel window.
		 */
		if (s->cancels_false != s->falses_at_pub) {
			atomic_add_int(&s->real_hits, 1);
			g_any_hit = 1;
			g_stop = 1;	/* only classifier-confirmed hits
					 * stop the run / panic */
			if (g_panic_on_hit)
				panic("DF-4919: work executed AFTER "
				    "cancel_work_sync()/flush_work() returned "
				    "— process_all_work called func on work "
				    "the caller was told is dead "
				    "(linux_workqueue.c:74 vs :364)");
		}
		kprintf("[DF4919] HIT (real=%d) ring (oldest first): ",
		    (s->cancels_false != s->falses_at_pub));
		for (k = 0; k < 32; k++)
			kprintf("%u:%u ", s->ring[k].tag, s->ring[k].aux);
		kprintf("ringidx=%u\n", s->ring_idx);
		/* Executed although the driver was told by
		 * cancel_work_sync()/flush_work() that the work is
		 * complete: posthumous run of a "dead" work. */
		if (atomic_fetchadd_int(&s->hits, 1) == 0)
			s->first_hit_at = s->cancels + s->flushes;
	}
}

/*
 * Reproduction design: feeder + 2 cancellers hammering one work on one
 * single-worker workqueue.  The feeder publishes alive=1 before each
 * queue_work(); each canceller clears alive right after its
 * cancel_work_sync() returns.  wk_func() therefore proves, for any
 * execution, whether a cancel_work_sync() had already returned for that
 * queued instance ("posthumous" run).  (A strictly single-driver variant
 * — one queue fully completed per cancel — was also measured; see
 * VERDICT.md: it exercises only the ~2ns dequeue->start transition and
 * did not hit in ~150k cycles, while the concurrent-insert window is
 * hit within a few thousand operations and is what the panic proof
 * exercises.  Both windows are the same defect: per-work flags read
 * without the lock that guards them.)
 */
static void
feeder_thread(void *arg)
{
	struct df4919_state *s = arg;

	/*
	 * One instance in flight at a time: publish, wait for the
	 * instance to resolve (ran / dequeued / canceled), grace, then
	 * re-arm.  With the feeder quiescent while an instance is
	 * pending, alive==0 at wk_func() entry can ONLY come from a
	 * canceller's cancel_work_sync() returning during
	 * (publication, dequeue->start) — the premature return.
	 */
	while (g_stop == 0) {
		s->cycles++;
		s->alive = 1;
		if (queue_work(s->wq, &s->work) == 0) {
			EV(s, 1, 0);
			while (g_stop == 0 && work_pending(&s->work))
				DELAY(1);
			continue;
		}
		EV(s, 2, 0);
		/*
		 * Re-assert alive AFTER the queue was published: a
		 * cancel_work_sync() that returned "not pending" BEFORE
		 * our claim was legitimate (nothing was queued yet); its
		 * alive=0 write must not mask this instance.  With the
		 * claim protocol, NO legitimate false-return can occur
		 * between publication and the function starting, so
		 * alive==0 at entry still proves the premature return.
		 */
		s->alive = 1;
		cpu_mfence();
		s->falses_at_pub = s->cancels_false;
		cpu_mfence();
		while (g_stop == 0 && work_pending(&s->work))
			DELAY(1);
		DELAY(20);	/* grace: posthumous funcs read alive==0 */
	}
	for (;;)
		tsleep(feeder_thread, 0, "df4919f", 0);
}

static void
cancel_thread(void *arg)
{
	struct df4919_state *s = arg;

	while (g_stop == 0) {
		int r = cancel_work_sync(&s->work);

		atomic_add_int(&s->cancels, 1);
		if (r == 0)
			atomic_add_int(&s->cancels_false, 1);
		EV(s, r ? 4 : 3, s->cancels);
		s->alive = 0;
	}
	for (;;)
		tsleep(cancel_thread, 0, "df4919c", 0);
}

static void
flusher_thread(void *arg)
{
	struct df4919_state *s = arg;

	while (g_stop == 0) {
		s->alive = 1;
		queue_work(s->wq, &s->work);
		flush_work(&s->work);
		/* flush_work returned: caller believes the work is done and
		 * may free it / leave the stack frame */
		s->alive = 0;
		atomic_add_int(&s->flushes, 1);
	}
	for (;;)
		tsleep(flusher_thread, 0, "df4919l", 0);
}

static int
df4919_run_test(int mode)
{
	int i, j, error, timeout = 90;

	if (g_st[0].wq != NULL)
		return (EBUSY);
	g_mode = mode;
	g_panic_on_hit = (mode == 2) ? 1 : 0;
	if (mode == 2)
		g_mode = 1;			/* cancel race, panic variant */

	for (i = 0; i < NDF; i++) {
		g_st[i].cpu = i;
		g_st[i].wq = alloc_workqueue("df4919wq", WQ_UNBOUND, 1);
		if (g_st[i].wq == NULL)
			return (ENOMEM);
		INIT_WORK(&g_st[i].work, wk_func);
	}

	kprintf("[DF4919] test=%d mode=%s panic_on_hit=%d: feeder+2 cancellers "
	    "(or feeder+flusher) hammering one work; probing the unlocked "
	    "flag reads at _cancel_work():364 / flush_work():331 against the "
	    "worker transition (:74->:87) and the queue_work() insert "
	    "publication (:138->:141)\n", mode,
	    (g_mode == 1) ? "cancel_work_sync" : "flush_work",
	    g_panic_on_hit);

	if (g_mode == 1) {
		struct thread *td;

		error = lwkt_create(feeder_thread, &g_st[0], &g_st[0].td,
		    NULL, TDF_NOSTART, 0, "df4919feed");
		if (error == 0)
			lwkt_schedule(g_st[0].td);
		error = lwkt_create(cancel_thread, &g_st[0], &td,
		    NULL, TDF_NOSTART, 1, "df4919can1");
		if (error == 0)
			lwkt_schedule(td);
		error = lwkt_create(cancel_thread, &g_st[0], &td,
		    NULL, TDF_NOSTART, 2, "df4919can2");
		if (error == 0)
			lwkt_schedule(td);
	} else {
		error = lwkt_create(flusher_thread, &g_st[0], &g_st[0].td,
		    NULL, TDF_NOSTART, 1, "df4919flush");
		if (error == 0)
			lwkt_schedule(g_st[0].td);
	}

	for (i = 0; i < timeout && g_stop == 0; i++)
		tsleep(df4919_run_test, 0, "df4919w", hz);
	g_stop = 1;
	tsleep(df4919_run_test, 0, "df4919w", 2 * hz);

	for (j = 0; j < NDF; j++)
		kprintf("[DF4919] drv%d(cpu%d): hits=%u(real=%u) runs=%u "
		    "cycles=%u cancels=%u (returned-false=%u) flushes=%u%s\n",
		    j, g_st[j].cpu, g_st[j].hits, g_st[j].real_hits,
		    g_st[j].runs, g_st[j].cycles,
		    g_st[j].cancels, g_st[j].cancels_false, g_st[j].flushes,
		    g_st[j].hits ? "  *** POSTHUMOUS EXECUTION AFTER "
		    "SYNC-CANCEL/FLUSH RETURN — CONFIRMED ***" : "");
	if (g_any_hit) {
		for (j = 0; j < NDF; j++)
			if (g_st[j].hits) {
				kprintf("[DF4919] drv%d first hit after %d "
				    "cancel/flush ops (mode=%d: %s)\n", j,
				    g_st[j].first_hit_at, g_mode,
				    (g_mode == 1) ? "cancel_work_sync" :
				    "flush_work");
				break;
			}
	} else {
		kprintf("[DF4919] result: no hit this run\n");
	}
	/* threads park; module refuses unload */
	return (0);
}

static int
sysctl_df4919_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 && v <= 3)
		return (df4919_run_test(v));
	return (0);
}
SYSCTL_PROC(_hw, OID_AUTO, df4919_test, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_df4919_test, "I", "df4919 test (1=cancel 2=cancel+panic "
    "3=flush)");

static int
sysctl_df4919_status(SYSCTL_HANDLER_ARGS)
{
	char buf[512];
	int len, j;

	len = ksnprintf(buf, sizeof(buf), "mode=%d any_hit=%d:", g_mode,
	    g_any_hit);
	for (j = 0; j < NDF; j++)
		len += ksnprintf(buf + len, sizeof(buf) - len,
		    " d%d(h=%u,r=%u,c=%u,f=%u)", j, g_st[j].hits,
		    g_st[j].runs, g_st[j].cycles, g_st[j].cancels_false);
	len += ksnprintf(buf + len, sizeof(buf) - len, "\n");
	return (SYSCTL_OUT(req, buf, len));
}
SYSCTL_PROC(_hw, OID_AUTO, df4919_status, CTLTYPE_STRING | CTLFLAG_RD,
    0, 0, sysctl_df4919_status, "A", "df4919 race counters");

static int
df4919_init(struct module *mod, int what, void *arg)
{
	switch (what) {
	case MOD_LOAD:
		kprintf("[DF4919] loaded: sysctl hw.df4919_test=1|2|3 "
		    "(needs drm.ko loaded first: kldload drm)\n");
		break;
	default:
		break;
	}
	return (0);
}

static int
df4919_uninit(struct module *mod, int what, void *arg)
{
	switch (what) {
	case MOD_UNLOAD:
		kprintf("[DF4919] refusing unload: parked threads\n");
		return (EBUSY);
	default:
		break;
	}
	return (0);
}

static moduledata_t df4919_mod = { "df4919", df4919_init, df4919_uninit };
DECLARE_MODULE(df4919, df4919_mod, SI_SUB_DRIVERS, SI_ORDER_MIDDLE);