// SPDX-License-Identifier: GPL-2.0
/*
 * usbdev_stall.c - controlled-window instrument for the devio Case-3 race
 * (usbdev_release vs in-flight async_completed / cancel_bulk_urbs).
 *
 * The audited window is inside async_completed() (cancel_bulk_urbs inlined):
 *
 *   call _raw_spin_unlock   @ async_completed+0x478   (devio.c:614)
 *   call usb_unlink_urb     @ async_completed+0x480   (devio.c:615)  <- window
 *   call usb_free_urb       @ async_completed+0x488   (devio.c:616)
 *   call _raw_spin_lock     @ async_completed+0x491   (devio.c:617)  <- UAF access
 *
 * (offsets measured on verification/build-kasan-7.3/vmlinux, KASAN build;
 *  cancel_bulk_urbs is inlined into async_completed in this build)
 *
 * mode A (default): kprobe on usb_unlink_urb, pre_handler matches the return
 *   address == async_completed_base + ret_off (0x485) and busy-waits
 *   stall_us microseconds. This places the stall INSIDE the drop-lock window
 *   (after spin_unlock, before usb_unlink_urb executes), which is the exact
 *   window analysed in the audit report section 5.4 steps 6c-6e.
 *
 * mode B: kprobe on _raw_spin_lock, same ra filter with ret_off=0x496,
 *   stalling after unlink has executed but before the re-lock (report step
 *   6e). _raw_spin_lock is hot; the ra filter makes it cheap enough for the
 *   idle initramfs workload.
 *
 * Both modes print ktime timestamps on every hit. A kretprobe on
 * usbdev_release prints entry/return timestamps so the serial log yields a
 * single readable timeline: if release returns while a stall hit is in
 * flight, kfree(ps) happened inside the window and the next ps access from
 * the callback is the audited UAF.
 */
#define pr_fmt(fmt) "USBDEVSTALL: " fmt

#include <linux/module.h>
#include <linux/init.h>
#include <linux/kernel.h>
#include <linux/kprobes.h>
#include <linux/ktime.h>
#include <linux/delay.h>
#include <linux/atomic.h>

static int stall_us = 5000;
module_param(stall_us, int, 0600);
MODULE_PARM_DESC(stall_us, "busy-wait length at the matched call site");

static int mode = 1;
module_param(mode, int, 0600);
MODULE_PARM_DESC(mode, "1=stall at usb_unlink_urb call site (0x485), 2=stall at re-lock site (0x496)");

static int max_hits = 2000;
module_param(max_hits, int, 0600);
MODULE_PARM_DESC(max_hits, "stop stalling after this many hits (0=unlimited)");

static ulong ret_off = 0x485;
module_param(ret_off, ulong, 0600);
MODULE_PARM_DESC(ret_off, "return-address offset inside async_completed");

static ulong ret_off_b = 0x496;
module_param(ret_off_b, ulong, 0600);

static unsigned long ac_base;
static atomic_t hits = ATOMIC_INIT(0);
static atomic_t near = ATOMIC_INIT(0);

static int stall_pre(struct kprobe *p, struct pt_regs *regs)
{
	unsigned long ra = *(unsigned long *)regs->sp;
	unsigned long want = ac_base + ret_off - 4;	/* kprobe base skips ENDBR */
	int h;

	if (ra == want) {
		h = atomic_inc_return(&hits);
		if (max_hits && h > max_hits)
			return 0;
		pr_warn("hit#%d t=%llu stall_us=%d ra=%lx (base+0x%lx)\n",
			h, ktime_get_ns(), stall_us, ra, ra - ac_base);
		udelay(stall_us);
		return 0;
	}
	if (ra > ac_base - 16 && ra < ac_base + 0x900) {
		h = atomic_inc_return(&near);
		if (h <= 8)
			pr_warn("near#%d ra=base+0x%lx (want base+0x%lx)\n",
				h, ra - ac_base, ret_off - 4);
	}
	return 0;
}

/* second probe for mode 2: raw_spin_lock, filtered the same way */
static int stall_pre_b(struct kprobe *p, struct pt_regs *regs)
{
	unsigned long ra = *(unsigned long *)regs->sp;
	unsigned long want = ac_base + ret_off_b - 4;	/* kprobe base skips ENDBR */
	int h;

	if (ra != want)
		return 0;

	h = atomic_inc_return(&hits);
	if (max_hits && h > max_hits)
		return 0;

	pr_warn("hit#%d(relock) t=%llu stall_us=%d\n",
		h, ktime_get_ns(), stall_us);
	udelay(stall_us);
	return 0;
}

static struct kprobe kp_unlink = {
	.symbol_name = "usb_unlink_urb",
	.pre_handler = stall_pre,
};

static struct kprobe kp_lock = {
	.symbol_name = "_raw_spin_lock",
	.pre_handler = stall_pre_b,
};

static struct kretprobe kr_release = {
	.handler = NULL,	/* set below */
	.entry_handler = NULL,
};

static int rel_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	pr_info("usbdev_release enter t=%llu pid=%d\n",
		ktime_get_ns(), current->pid);
	return 0;
}

static int rel_ret(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	pr_info("usbdev_release ret  t=%llu pid=%d\n",
		ktime_get_ns(), current->pid);
	return 0;
}

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 int __init usbdev_stall_init(void)
{
	int ret;

	ac_base = sym_addr("async_completed");
	if (!ac_base) {
		pr_err("cannot resolve async_completed\n");
		return -EINVAL;
	}
	pr_info("async_completed base=%lx (+0x%lx unlink call, +0x%lx relock call)\n",
		ac_base, ret_off, ret_off_b);

	if (mode == 2) {
		ret = register_kprobe(&kp_lock);
		if (ret) {
			pr_err("register _raw_spin_lock probe failed: %d\n", ret);
			return ret;
		}
	} else {
		ret = register_kprobe(&kp_unlink);
		if (ret) {
			pr_err("register usb_unlink_urb probe failed: %d\n", ret);
			return ret;
		}
	}

	kr_release.kp.symbol_name = "usbdev_release";
	kr_release.entry_handler = rel_entry;
	kr_release.handler = rel_ret;
	ret = register_kretprobe(&kr_release);
	if (ret)
		pr_warn("kretprobe on usbdev_release failed: %d (timeline degraded)\n", ret);

	pr_info("armed mode=%d stall_us=%d ret_off=0x%lx max_hits=%d\n",
		mode, stall_us, mode == 2 ? ret_off_b : ret_off, max_hits);
	return 0;
}

static void __exit usbdev_stall_exit(void)
{
	if (mode == 2)
		unregister_kprobe(&kp_lock);
	else
		unregister_kprobe(&kp_unlink);
	unregister_kretprobe(&kr_release);
	pr_info("removed hits=%d near_miss=%d\n",
		atomic_read(&hits), atomic_read(&near));
}

module_init(usbdev_stall_init);
module_exit(usbdev_stall_exit);
MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("controlled-window kprobe stall for devio usbdev_release race verification");