// SPDX-License-Identifier: GPL-2.0
/*
 * regmap_race.c - dynamic verification for HV-1 of the drivers/base
 * race-condition audit (regmap async completion vs map lifetime).
 *
 * Drives the regmap core through its exported API only
 * (__regmap_init/regmap_exit/regmap_write_async/regmap_async_complete/
 * regmap_async_complete_cb are all EXPORT_SYMBOL_GPL).  The kernel source
 * is NOT modified.
 *
 * The bus is a fake regmap_bus whose async_write defers the completion
 * callback (regmap_async_complete_cb) onto a system work with a
 * configurable delay, exactly mirroring what a real SPI controller does:
 * the callback runs in bus/controller context with no reference on map.
 *
 *   echo hv1b <ms> > action : HV-1 form 2 (wide window, devm-style):
 *              create map -> regmap_write_async -> regmap_exit IMMEDIATELY
 *              (no drain, identical to devm_regmap_release) -> wait for the
 *              late bus completion to call regmap_async_complete_cb on the
 *              freed map.  Expected: KASAN UAF at spin_lock(&map->async_lock).
 *
 *   echo hv1a        > action : HV-1 form 1 (narrow tail window,
 *              controlled window via kprobe): a kprobe is placed on the
 *              "jmp __wake_up" tail inside regmap_async_complete_cb
 *              (offset wake_off, computed from the running vmlinux) and
 *              stalls the callback AFTER it released map->async_lock and
 *              moved the node off async_list, but BEFORE wake_up().  The
 *              test thread then runs regmap_async_complete() (returns via
 *              the wait_event fast path while the callback is stalled),
 *              regmap_exit() (frees map + the async node), and finally the
 *              stall releases -> wake_up() on freed map.
 *              Expected: KASAN UAF inside __wake_up path on the freed map.
 *
 *   echo status      > action : stats.
 */
#define pr_fmt(fmt) "REGMAPRACE: " fmt

#include <linux/module.h>
#include <linux/init.h>
#include <linux/kernel.h>
#include <linux/kprobes.h>
#include <linux/delay.h>
#include <linux/workqueue.h>
#include <linux/slab.h>
#include <linux/regmap.h>
#include "../../linux/drivers/base/regmap/internal.h"
#include <linux/ktime.h>
#include <linux/kthread.h>
#include <linux/sched.h>
#include <linux/string.h>
#include <linux/atomic.h>

static int delay_ms = 300;
module_param(delay_ms, int, 0600);
MODULE_PARM_DESC(delay_ms, "fake-bus async completion delay in ms");

static int stall_us = 500000;
module_param(stall_us, int, 0600);
MODULE_PARM_DESC(stall_us, "hv1a controlled-window stall in us");

static ulong wake_off = 0x195;
module_param(wake_off, ulong, 0600);
MODULE_PARM_DESC(wake_off, "offset of the 'jmp __wake_up' tail inside regmap_async_complete_cb (per-build, from objdump)");

/* ------------------------------------------------------------------ */
/* fake regmap bus: async completion is deferred onto system work      */

struct fb_async {
	struct regmap_async core;
	struct delayed_work dwork;
};

static atomic_long_t stat_async_submit = ATOMIC_LONG_INIT(0);
static atomic_long_t stat_async_fire = ATOMIC_LONG_INIT(0);

static void fb_complete(struct work_struct *w)
{
	struct fb_async *fa = container_of(w, struct fb_async, dwork.work);

	atomic_long_inc(&stat_async_fire);
	pr_info("FBUS: completion FIRING async=%px map=%px (bus context)\n",
		&fa->core, fa->core.map);

	regmap_async_complete_cb(&fa->core, 0);

	pr_info("FBUS: completion RETURNED async=%px\n", &fa->core);
}

static struct regmap_async *fb_async_alloc(void)
{
	struct fb_async *fa;

	fa = kzalloc(sizeof(*fa), GFP_KERNEL);
	if (!fa)
		return NULL;
	INIT_DELAYED_WORK(&fa->dwork, fb_complete);
	return &fa->core;
}

static int use_kthread_complete;	/* hv1a: complete from a dedicated
					 * kernel thread (bus "controller"
					 * context we fully control) */
static struct task_struct *fb_thread;
static struct fb_async *pending_fa;
static DECLARE_COMPLETION(fb_go);
static DECLARE_COMPLETION(fb_thread_up);

static int fb_thread_fn(void *unused)
{
	complete(&fb_thread_up);
	while (!kthread_should_stop()) {
		wait_for_completion(&fb_go);
		if (kthread_should_stop())
			break;
		atomic_long_inc(&stat_async_fire);
		pr_info("FBUS: completion FIRING async=%px map=%px (bus kthread context)\n",
			&pending_fa->core, pending_fa->core.map);
		regmap_async_complete_cb(&pending_fa->core, 0);
		pr_info("FBUS: completion RETURNED async=%px\n", &pending_fa->core);
	}
	return 0;
}

static int fb_async_write(void *context, const void *reg, size_t reg_len,
			  const void *val, size_t val_len,
			  struct regmap_async *a)
{
	struct fb_async *fa = container_of(a, struct fb_async, core);

	atomic_long_inc(&stat_async_submit);
	if (use_kthread_complete) {
		pr_info("FBUS: async_write queued async=%px (kthread completion)\n",
			&fa->core);
		pending_fa = fa;
		complete(&fb_go);
		return 0;
	}
	pr_info("FBUS: async_write queued async=%px delay=%dms\n",
		&fa->core, delay_ms);
	schedule_delayed_work(&fa->dwork, msecs_to_jiffies(delay_ms));
	return 0;
}

static u8 fb_scratch[256];

static int fb_write(void *context, const void *data, size_t count)
{
	memcpy(fb_scratch, data, min(count, sizeof(fb_scratch)));
	return 0;
}

static int fb_read(void *context, const void *reg, size_t reg_size,
		   void *val, size_t val_size)
{
	memset(val, 0x5a, val_size);
	return 0;
}

static const struct regmap_bus fb_bus = {
	.write = fb_write,
	.read = fb_read,
	.async_write = fb_async_write,
	.async_alloc = fb_async_alloc,
};

static const struct regmap_config fb_cfg = {
	.reg_bits = 8,
	.val_bits = 8,
	.cache_type = REGCACHE_NONE,
};

static struct lock_class_key fb_lock_key;

/* ------------------------------------------------------------------ */
/* HV-1 form 2: exit with in-flight async (devm-style teardown)        */

static void do_hv1b(void)
{
	struct regmap *map;
	int ret;

	map = __regmap_init(NULL, &fb_bus, NULL, &fb_cfg,
			    &fb_lock_key, "fbmap");
	if (IS_ERR(map)) {
		pr_err("hv1b: __regmap_init failed %ld\n", PTR_ERR(map));
		return;
	}
	pr_info("HV1B: map=%px created; submitting async write then regmap_exit IMMEDIATELY (no drain, devm-style)\n",
		map);

	ret = regmap_write_async(map, 0x10, 0xAA);
	pr_info("HV1B: regmap_write_async ret=%d; calling regmap_exit NOW\n", ret);

	regmap_exit(map);
	pr_info("HV1B: regmap_exit returned (map freed); late bus completion fires in ~%dms ...\n",
		delay_ms);

	msleep(delay_ms + 2500);
	pr_info("HV1B: done\n");
}

/* ------------------------------------------------------------------ */
/* HV-1 form 1: controlled stall at the wake_up tail                   */

static atomic_t wake_state = ATOMIC_INIT(0); /* 0=idle 1=stalled 2=released */
static int probe_armed;

static int wake_stall_pre(struct kprobe *p, struct pt_regs *regs)
{
	/* Only the bus-completion kthread is stalled, at __wake_up() entry:
	 * i.e. AFTER the callback released map->async_lock and moved the node
	 * off async_list (predicate true), BEFORE wake_up() touches the
	 * waitqueue head.  This is exactly the audited tail window; no
	 * instruction-offset math is needed, so it is KASLR-robust. */
	if (fb_thread && current == fb_thread &&
	    atomic_xchg(&wake_state, 1) == 0) {
		pr_info("HV1A: callback STALLED at __wake_up entry (tail window) for %d us\n",
			stall_us);
		udelay(stall_us);
		atomic_set(&wake_state, 2);
		pr_info("HV1A: stall released, callback proceeds into wake_up on the waitqueue\n");
	}
	return 0;
}

static struct kprobe wake_kp = {
	.symbol_name = "__wake_up",
	.pre_handler = wake_stall_pre,
};

__maybe_unused static unsigned long sym_addr(const char *name)
{
	struct kprobe probe = { .symbol_name = name, };
	unsigned long addr = 0;

	if (register_kprobe(&probe) == 0) {
		addr = (unsigned long)probe.addr;
		unregister_kprobe(&probe);
	}
	return addr;
}

static void do_hv1a(void)
{
	struct regmap *map;
	u64 t0, t1;
	int ret, i;

	if (register_kprobe(&wake_kp)) {
		pr_err("hv1a: cannot register kprobe on __wake_up\n");
		return;
	}
	probe_armed = 1;
	atomic_set(&wake_state, 0);

	fb_thread = kthread_run(fb_thread_fn, NULL, "fb_cb");
	if (IS_ERR(fb_thread)) {
		pr_err("hv1a: kthread failed\n");
		fb_thread = NULL;
		unregister_kprobe(&wake_kp);
		probe_armed = 0;
		return;
	}
	wait_for_completion(&fb_thread_up);

	use_kthread_complete = 1;
	map = __regmap_init(NULL, &fb_bus, NULL, &fb_cfg,
			    &fb_lock_key, "fbmap1");
	if (IS_ERR(map)) {
		pr_err("hv1a: __regmap_init failed %ld\n", PTR_ERR(map));
		goto out;
	}
	pr_info("HV1A: map=%px created; submitting async write (kthread completion, stall=%d us)\n",
		map, stall_us);

	ret = regmap_write_async(map, 0x20, 0xBB);
	if (ret)
		pr_err("hv1a: write_async ret=%d\n", ret);

	/* wait until the bus-completion kthread is stalled inside the
	 * tail window (right before wake_up) */
	for (i = 0; i < 5000 && atomic_read(&wake_state) != 1; i++)
		msleep(1);
	if (atomic_read(&wake_state) != 1) {
		pr_err("HV1A: callback never reached the wake tail\n");
		regmap_async_complete(map);
		regmap_exit(map);
		goto out;
	}

	/*
	 * Node is off async_list, async_lock released, bus callback stalled
	 * before wake_up(): regmap_async_complete() must return via the
	 * wait_event fast path while the callback is still running.
	 */
	t0 = ktime_get_ns();
	ret = regmap_async_complete(map);
	t1 = ktime_get_ns();
	pr_info("HV1A: regmap_async_complete RETURNED ret=%d in %llu ns WHILE bus callback stalled in tail window\n",
		ret, t1 - t0);

	regmap_exit(map);
	pr_info("HV1A: regmap_exit returned (map freed, async node freed); stall releases in <=%d us -> wake_up on freed map\n",
		stall_us);

	/* let the stall expire and the (now freed) wake_up fire */
	for (i = 0; i < 3000 && atomic_read(&wake_state) != 2; i++)
		msleep(1);
	msleep(500);
	pr_info("HV1A: done\n");

out:
	use_kthread_complete = 0;
	if (fb_thread) {
		kthread_stop(fb_thread);
		fb_thread = NULL;
	}
	if (probe_armed) {
		unregister_kprobe(&wake_kp);
		probe_armed = 0;
	}
}

/* ------------------------------------------------------------------ */
/* action plumbing (same pattern as pmu_race.c)                        */

static char action_cmd[64];

static void action_workfn(struct work_struct *w)
{
	char cmd[64];
	char *p, *cp;

	strscpy(cmd, strstrip(action_cmd), sizeof(cmd));
	cp = cmd;
	p = strsep(&cp, " ");

	if (!strcmp(p, "hv1b"))
		do_hv1b();
	else if (!strcmp(p, "hv1a"))
		do_hv1a();
	else if (!strcmp(p, "status"))
		pr_info("status: async_submit=%ld async_fire=%ld probe_armed=%d wake_state=%d delay_ms=%d stall_us=%d\n",
			atomic_long_read(&stat_async_submit),
			atomic_long_read(&stat_async_fire), probe_armed,
			atomic_read(&wake_state), delay_ms, stall_us);
	else
		pr_info("unknown action '%s'\n", p);

	pr_info("ACTION_COMPLETE\n");
}

static DECLARE_WORK(action_work, action_workfn);

static int action_set(const char *val, const struct kernel_param *kp)
{
	int ret = param_set_copystring(val, kp);

	if (ret == 0)
		schedule_work(&action_work);
	return ret;
}

static const struct kernel_param_ops action_ops = {
	.set = action_set,
	.get = param_get_string,
};
static struct kparam_string action_param = {
	.maxlen = sizeof(action_cmd),
	.string = action_cmd,
};
module_param_cb(action, &action_ops, &action_param, 0200);
MODULE_PARM_DESC(action, "hv1b <ms> | hv1a | status");

static int __init regmap_race_init(void)
{
	pr_info("loaded (delay_ms=%d stall_us=%d wake_off=%#lx)\n",
		delay_ms, stall_us, wake_off);
	return 0;
}

static void __exit regmap_race_exit(void)
{
	if (probe_armed)
		unregister_kprobe(&wake_kp);
	pr_info("removed\n");
}

module_init(regmap_race_init);
module_exit(regmap_race_exit);
MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("HV-1 verification: regmap async completion vs map lifetime");