// SPDX-License-Identifier: GPL-2.0
/*
 * typcrace.c - dynamic verification module for drivers/usb/typec audit
 * (docs/manual-review-drivers-usb-typec-2026-11-02.round3.md)
 *
 * Statement-faithful replicas of the audited driver state machines (anx7411,
 * stusb160x), driven via the exported typec class API on a fake typec_port,
 * with one-shot kprobe stalls making the audited interleavings deterministic.
 * Kernel source unmodified (build-kasan-typec-7.3).
 *
 * tc0  - audit Case 3 interleaving 1: anx7411 work uses amode[] pointer
 *        (check-then-use, anx7411.c:602-613) while remove frees the device.
 *        Work is parked by a kprobe at its use point
 *        (typec_altmode_update_active) AFTER it read amode[i]; the remove
 *        thread unregisters (frees) the altmode during the park; on resume
 *        the use dereferences freed memory -> KASAN UAF.
 * tc0b - audit Case 3 interleaving 2: double device_unregister. Remove is
 *        parked by a one-shot kprobe at typec_unregister_altmode BEFORE it
 *        executes; the queued work then unregisters the same altmode
 *        (device freed); on resume remove's unregister runs on the freed
 *        device -> UAF inside typec_unregister_altmode/class core.
 * tc1  - audit Case 4 (stusb160x form): irq thread parked by one-shot kprobe
 *        at typec_unregister_partner (having seen chip->partner != NULL);
 *        remove thread then unregisters the same partner and unregisters the
 *        port; irq resumes -> unregister on freed partner -> UAF.
 */
#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/delay.h>
#include <linux/ktime.h>
#include <linux/workqueue.h>
#include <linux/kthread.h>
#include <linux/completion.h>
#include <linux/mutex.h>
#include <linux/slab.h>
#include <linux/usb/typec.h>
#include <linux/usb/typec_altmode.h>

#define MAX_ALTMODE 4

struct tyc_typec {
	struct typec_port *port;
	struct typec_partner *partner;
	struct typec_altmode *amode[MAX_ALTMODE];
};

/* ---------------- anx7411 replica (Case 3) ---------------- */
struct tyc_anx {
	struct tyc_typec typec;
	struct mutex lock;		/* analog of ctx->lock: work side only */
	struct work_struct work;
	struct workqueue_struct *wq;
	struct completion work_in_use;	/* work parked at its use point */
};

static void anx_unregister_altmode_replica(struct tyc_anx *ctx)
{
	int i;

	/* anx7411.c:535-549 -- NO ctx->lock (audited) */
	for (i = 0; i < MAX_ALTMODE; i++)
		if (ctx->typec.amode[i]) {
			typec_unregister_altmode(ctx->typec.amode[i]);
			ctx->typec.amode[i] = NULL;
		}
}

static void anx_unregister_partner_replica(struct tyc_anx *ctx)
{
	/* anx7411.c:585-591 -- NO ctx->lock (audited) */
	if (ctx->typec.partner) {
		typec_unregister_partner(ctx->typec.partner);
		ctx->typec.partner = NULL;
	}
}

static bool anx_use_mode;	/* tc0: work parks at use point */
static struct tyc_anx *g_anx0;	/* tc0 target for kprobe lookup */

/*
 * tc0 kprobe target: typec_altmode_update_active is only reached from the
 * work's use point in this module (the remove thread never calls it).
 * Park here AFTER the work has read amode[i] (audited check-then-use gap).
 */
static int kp_use_pre(struct kprobe *p, struct pt_regs *regs)
{
	unsigned long long t0;
	struct tyc_anx *ctx = g_anx0;

	if (!anx_use_mode || !ctx)
		return 0;
	anx_use_mode = false;
	pr_info("TC0: work parked at use point (adev=%px), remove frees now\n",
		(void *)regs->di);
	complete(&ctx->work_in_use);
	t0 = ktime_get_mono_fast_ns();
	while (ktime_get_mono_fast_ns() - t0 < 400ULL * NSEC_PER_MSEC)
		cpu_relax();
	pr_info("TC0: work resumes use on (possibly freed) adev\n");
	return 0;
}

static struct kprobe kp_use = {
	.symbol_name = "typec_altmode_update_active",
	.pre_handler = kp_use_pre,
};

/* work replica of anx7411_work_func (anx7411.c:906-984) */
static void anx_work_func(struct work_struct *w)
{
	struct tyc_anx *ctx = container_of(w, struct tyc_anx, work);
	int i;

	mutex_lock(&ctx->lock);		/* :916 */

	/* "update_altmode" check-then-use loop (:602-613) */
	for (i = 0; i < MAX_ALTMODE; i++) {
		if (!ctx->typec.amode[i])
			continue;
		pr_info("TC0: work read amode[%d]=%px svid=%04x, calling update_active\n",
			i, ctx->typec.amode[i], ctx->typec.amode[i]->svid);
		typec_altmode_update_active(ctx->typec.amode[i], true);
	}
	/* POW_OFF branch (:943-946) */
	pr_info("TC0: work unregister altmode+partner\n");
	anx_unregister_altmode_replica(ctx);
	anx_unregister_partner_replica(ctx);

	mutex_unlock(&ctx->lock);
}

static int anx_remove_thread(void *data)
{
	struct tyc_anx *ctx = data;

	pr_info("TC0*: remove enters (audited order: unregister FIRST)\n");
	anx_unregister_altmode_replica(ctx);	/* :1566 */
	anx_unregister_partner_replica(ctx);	/* :1567 */
	pr_info("TC0*: remove stops work (too late, audited)\n");
	cancel_work_sync(&ctx->work);		/* :1571 */
	pr_info("TC0*: remove done\n");
	return 0;
}

/* one-shot stall at typec_unregister_altmode (tc0b): park the FIRST caller */
static bool oneshot_armed;
static struct completion tc_done_wait;

static int kp_unreg_pre(struct kprobe *p, struct pt_regs *regs)
{
	unsigned long long t0;

	if (!oneshot_armed)
		return 0;
	oneshot_armed = false;
	pr_info("TC0b: remove parked at typec_unregister_altmode (adev=%px)\n",
		(void *)regs->di);
	complete(&tc_done_wait);
	t0 = ktime_get_mono_fast_ns();
	while (ktime_get_mono_fast_ns() - t0 < 400ULL * NSEC_PER_MSEC)
		cpu_relax();
	pr_info("TC0b: parked unregister resumes (device likely freed)\n");
	return 0;
}

/* ---------------- stusb160x replica (Case 4) ---------------- */
struct tyc_stu {
	struct tyc_typec typec;
	struct task_struct *irq_thread;
	struct completion irq_parked;
};

static int stu_irq_thread(void *data)
{
	struct tyc_stu *ctx = data;

	while (!kthread_should_stop()) {
		set_current_state(TASK_INTERRUPTIBLE);
		schedule_timeout(msecs_to_jiffies(20));
		if (kthread_should_stop())
			break;
		/* stusb160x_detach(:388-389) check-then-use */
		if (ctx->typec.partner) {
			pr_info("TC1: irq saw partner=%px, unregistering\n",
				ctx->typec.partner);
			typec_unregister_partner(ctx->typec.partner);
			ctx->typec.partner = NULL;
		}
	}
	return 0;
}

static int stu_remove_thread(void *data)
{
	struct tyc_stu *ctx = data;

	pr_info("TC1: remove enters (devm-irq window open)\n");
	/* stusb160x.c:805-808 */
	if (ctx->typec.partner) {
		pr_info("TC1: remove saw partner=%px, unregistering\n",
			ctx->typec.partner);
		typec_unregister_partner(ctx->typec.partner);
		ctx->typec.partner = NULL;
	}
	pr_info("TC1: remove unregisters port\n");
	typec_unregister_port(ctx->typec.port);
	ctx->typec.port = NULL;
	pr_info("TC1: remove done\n");
	return 0;
}

/* ---------------- orchestration ---------------- */
static struct tyc_anx *g_anx;
static struct tyc_stu *g_stu;


static struct kprobe kp_unreg;

static int tc_setup_common(struct tyc_typec *t, struct device *parent)
{
	struct typec_capability cap = { };
	struct typec_partner_desc pdesc = { };
	struct typec_altmode_desc adesc = { };
	int i;

	cap.type = TYPEC_PORT_DRP;
	cap.data = TYPEC_PORT_DRD;
	cap.revision = USB_TYPEC_REV_1_2;
	cap.driver_data = t;
	t->port = typec_register_port(parent, &cap);
	if (IS_ERR(t->port)) {
		pr_err("typec_register_port failed: %pe\n", t->port);
		return PTR_ERR(t->port);
	}
	t->partner = typec_register_partner(t->port, &pdesc);
	if (IS_ERR(t->partner)) {
		pr_err("typec_register_partner failed: %pe\n", t->partner);
		typec_unregister_port(t->port);
		t->port = NULL;
		return PTR_ERR(t->partner);
	}
	for (i = 0; i < 2; i++) {
		adesc.svid = 0x1234 + i;
		adesc.mode = i + 1;
		adesc.vdo = 0xdeadbeef;
		adesc.roles = TYPEC_PORT_DRD;
		t->amode[i] = typec_partner_register_altmode(t->partner, &adesc);
		if (IS_ERR(t->amode[i])) {
			pr_err("partner_register_altmode[%d] failed: %pe\n",
			       i, t->amode[i]);
			t->amode[i] = NULL;
			return -EFAULT;
		}
	}
	return 0;
}

static int run_tc0(void)
{
	struct device *root_dev;
	struct task_struct *rm;
	int ret;

	root_dev = root_device_register("tycroot0");
	if (IS_ERR(root_dev))
		return PTR_ERR(root_dev);

	g_anx = kzalloc(sizeof(*g_anx), GFP_KERNEL);
	if (!g_anx) {
		root_device_unregister(root_dev);
		return -ENOMEM;
	}
	mutex_init(&g_anx->lock);
	init_completion(&g_anx->work_in_use);
	INIT_WORK(&g_anx->work, anx_work_func);
	g_anx->wq = create_singlethread_workqueue("tycanxwq");
	ret = tc_setup_common(&g_anx->typec, root_dev);
	if (ret)
		goto out;

	ret = register_kprobe(&kp_use);
	if (ret) {
		pr_err("TC0: kprobe(use) failed %d\n", ret);
		goto out;
	}
	g_anx0 = g_anx;
	anx_use_mode = true;

	/* 1. queue work; it parks inside kprobe at its first use, having
	 *    already read amode[0] (audited check-then-use gap). */
	queue_work(g_anx->wq, &g_anx->work);
	/* 2. wait until parked */
	if (!wait_for_completion_timeout(&g_anx->work_in_use, msecs_to_jiffies(2000))) {
		pr_err("TC0: work never parked; abort\n");
		goto out;
	}
	/* 3. remove thread frees altmodes+partner under the parked work */
	rm = kthread_run(anx_remove_thread, g_anx, "tyc-anxrm");
	if (!IS_ERR(rm))
		kthread_stop(rm);	/* joins: remove completes inside window */
	/* 4. work auto-resumes after 400ms -> UAF */
	pr_info("TC0: RESULT see KASAN reports above\n");
	msleep(300);
	anx_unregister_altmode_replica(g_anx);
	anx_unregister_partner_replica(g_anx);
out:
	unregister_kprobe(&kp_use);
	g_anx0 = NULL;
	if (g_anx) {
		if (g_anx->wq)
			destroy_workqueue(g_anx->wq);
		kfree(g_anx);
		g_anx = NULL;
	}
	root_device_unregister(root_dev);
	return 0;
}

static int run_tc0b(void)
{
	struct device *root_dev;
	struct task_struct *rm;
	int ret;

	root_dev = root_device_register("tycroot0b");
	if (IS_ERR(root_dev))
		return PTR_ERR(root_dev);

	g_anx = kzalloc(sizeof(*g_anx), GFP_KERNEL);
	if (!g_anx) {
		root_device_unregister(root_dev);
		return -ENOMEM;
	}
	mutex_init(&g_anx->lock);
	init_completion(&tc_done_wait);
	INIT_WORK(&g_anx->work, anx_work_func);
	g_anx->wq = create_singlethread_workqueue("tycanxwqb");
	ret = tc_setup_common(&g_anx->typec, root_dev);
	if (ret)
		goto out;

	kp_unreg = (struct kprobe){
		.symbol_name = "typec_unregister_altmode",
		.pre_handler = kp_unreg_pre,
	};
	ret = register_kprobe(&kp_unreg);
	if (ret) {
		pr_err("TC0b: kprobe failed %d\n", ret);
		goto out;
	}

	/* 1. remove thread parks inside kprobe BEFORE its unregister */
	oneshot_armed = true;
	rm = kthread_run(anx_remove_thread, g_anx, "tyc-anxrm0b");
	if (IS_ERR(rm)) {
		pr_err("TC0b: kthread failed\n");
		goto out;
	}
	/* 2. wait parked */
	if (!wait_for_completion_timeout(&tc_done_wait, msecs_to_jiffies(2000))) {
		pr_err("TC0b: remove never parked\n");
		kthread_stop(rm);
		goto out;
	}
	/* 3. now let the work run: it unregisters the same altmodes+partner
	 *    (device freed) while remove is parked mid-unregister */
	queue_work(g_anx->wq, &g_anx->work);
	msleep(100);
	/* 4. remove resumes -> unregister on freed device */
	pr_info("TC0b: remove resumes at ~300ms\n");
	msleep(700);			/* let everything settle */
	cancel_work_sync(&g_anx->work);
	pr_info("TC0b: RESULT see KASAN reports above\n");
out:
	unregister_kprobe(&kp_unreg);
	if (g_anx) {
		if (g_anx->wq)
			destroy_workqueue(g_anx->wq);
		kfree(g_anx);
		g_anx = NULL;
	}
	root_device_unregister(root_dev);
	return 0;
}

static int kp_park_pre(struct kprobe *p, struct pt_regs *regs)
{
	unsigned long long t0;

	if (!oneshot_armed)
		return 0;
	oneshot_armed = false;
	pr_info("TC1: irq thread parked at typec_unregister_partner (partner=%px)\n",
		(void *)regs->di);
	complete(&tc_done_wait);
	t0 = ktime_get_mono_fast_ns();
	while (ktime_get_mono_fast_ns() - t0 < 400ULL * NSEC_PER_MSEC)
		cpu_relax();
	pr_info("TC1: irq resumes unregister (partner likely freed)\n");
	return 0;
}

static int run_tc1(void)
{
	struct device *root_dev;
	int ret;

	root_dev = root_device_register("tycroot1");
	if (IS_ERR(root_dev))
		return PTR_ERR(root_dev);

	g_stu = kzalloc(sizeof(*g_stu), GFP_KERNEL);
	if (!g_stu) {
		root_device_unregister(root_dev);
		return -ENOMEM;
	}
	init_completion(&tc_done_wait);
	init_completion(&g_stu->irq_parked);
	ret = tc_setup_common(&g_stu->typec, root_dev);
	if (ret)
		goto out;

	kp_unreg = (struct kprobe){
		.symbol_name = "typec_unregister_partner",
		.pre_handler = kp_park_pre,
	};
	ret = register_kprobe(&kp_unreg);
	if (ret) {
		pr_err("TC1: kprobe failed %d\n", ret);
		goto out;
	}

	/* 1. irq thread parks at its unregister (partner still alive) */
	oneshot_armed = true;
	g_stu->irq_thread = kthread_run(stu_irq_thread, g_stu, "tyc-stuirq");
	if (IS_ERR(g_stu->irq_thread)) {
		g_stu->irq_thread = NULL;
		goto out;
	}
	if (!wait_for_completion_timeout(&tc_done_wait, msecs_to_jiffies(2000))) {
		pr_err("TC1: irq never parked\n");
		goto out;
	}
	/* 2. remove sequence inline on the insmod thread (partner still
	 *    alive -> unregister frees it, then port); parked irq thread is
	 *    on another CPU */
	stu_remove_thread(g_stu);
	/* 3. irq resumes -> unregister on freed partner */
	pr_info("TC1: RESULT see KASAN reports above\n");
	msleep(300);
out:
	unregister_kprobe(&kp_unreg);
	if (g_stu) {
		if (g_stu->irq_thread)
			kthread_stop(g_stu->irq_thread);
		kfree(g_stu);
		g_stu = NULL;
	}
	root_device_unregister(root_dev);
	return 0;
}

/* ---------------- tc2: Case 1-B deadlock structure replica ----------------
 * Faithful structure replica of the audited UCSI unbind deadlock:
 *   - unbind thread holds "con->lock" and calls cancel_delayed_work_sync()
 *     (ucsi.c:2003-2004 -> mode_selection.c:277);
 *   - the mode-selection work, inside its activate window, must take the
 *     same lock via the ucsi_con_mutex_lock() trylock loop whose exit
 *     condition is the STALE con->status CONNECTED snapshot
 *     (ucsi.c:2272-2277; ucsi.h:457 reads ->status, never refreshed on the
 *     unbind path).
 * Expected result: both sides park forever; the VM hangs until timeout.
 */
struct tyc_dead {
	struct mutex lock;		/* con->lock analog */
	bool connected_snapshot;	/* con->status CONNECTED bit (stale) */
	struct delayed_work work;
	struct completion in_loop;	/* handshake: work inside trylock loop */
	bool ctl;			/* ctl=true: lock first, queue later */
};

static void dead_work_fn(struct work_struct *w)
{
	struct tyc_dead *d = container_of(w, struct tyc_dead, work.work);
	bool locked = false;

	pr_info("TC2: work entered activate window, needs con->lock\n");
	/* first trylock fails: unbind already holds the lock */
	locked = mutex_trylock(&d->lock) != 0;
	if (!locked)
		complete(&d->in_loop);
	/* ucsi_con_mutex_lock() loop, ucsi.c:2272-2277, verbatim shape */
	while (d->connected_snapshot && !locked) {
		locked = mutex_trylock(&d->lock) != 0;
		if (d->connected_snapshot && !locked) {
			msleep(20);
		}
	}
	if (locked)
		mutex_unlock(&d->lock);
	pr_info("TC2: work finished\n");
}

/* audited interleaving: work enters the loop first, unbind then takes the
 * lock and cancels -- expected permanent hang (VM killed by timeout). */
static int run_tc2(void)
{
	struct tyc_dead *d;

	d = kzalloc(sizeof(*d), GFP_KERNEL);
	if (!d)
		return -ENOMEM;
	mutex_init(&d->lock);
	init_completion(&d->in_loop);
	d->connected_snapshot = true;	/* cable connected snapshot */
	d->ctl = false;
	INIT_DELAYED_WORK(&d->work, dead_work_fn);
	/* audited order: unbind holds con->lock FIRST, work starts and
	 * fails its trylocks while the lock stays held */
	mutex_lock(&d->lock);
	pr_info("TC2: unbind holds con->lock; scheduling work\n");
	schedule_delayed_work(&d->work, 0);
	if (!wait_for_completion_timeout(&d->in_loop, msecs_to_jiffies(2000))) {
		pr_err("TC2: work never entered loop\n");
		mutex_unlock(&d->lock);
		kfree(d);
		return -ETIMEDOUT;
	}
	pr_info("TC2: work looping on failed trylocks; unbind calls cancel (expect hang)\n");
	cancel_delayed_work_sync(&d->work);	/* audited hang */
	pr_info("TC2: cancel returned (should never print)\n");
	mutex_unlock(&d->lock);
	kfree(d);
	return 0;
}

/* control: unbind takes the lock BEFORE the work ever runs; cancel() kills
 * the pending timer -- no deadlock (audited escape-branch 1). */
static int run_tc2ctl(void)
{
	struct tyc_dead *d;

	d = kzalloc(sizeof(*d), GFP_KERNEL);
	if (!d)
		return -ENOMEM;
	mutex_init(&d->lock);
	init_completion(&d->in_loop);
	d->connected_snapshot = true;
	d->ctl = true;
	INIT_DELAYED_WORK(&d->work, dead_work_fn);

	mutex_lock(&d->lock);
	pr_info("TC2CTL: unbind holds lock first, queueing work\n");
	schedule_delayed_work(&d->work, msecs_to_jiffies(50));
	cancel_delayed_work_sync(&d->work);	/* pending -> killed */
	pr_info("TC2CTL: cancel returned (expected, pending branch)\n");
	mutex_unlock(&d->lock);
	msleep(200);
	kfree(d);
	return 0;
}

/* ---------------- tcnat: natural hammering (no probes) ---------------- */
static int tcnat_rounds = 1000;
module_param(tcnat_rounds, int, 0444);

static int run_tcnat(void)
{
	struct device *root_dev;
	int round;

	for (round = 0; round < tcnat_rounds; round++) {
		root_dev = root_device_register("tycrootnat");
		if (IS_ERR(root_dev))
			return PTR_ERR(root_dev);
		g_anx = kzalloc(sizeof(*g_anx), GFP_KERNEL);
		if (!g_anx) {
			root_device_unregister(root_dev);
			return -ENOMEM;
		}
		mutex_init(&g_anx->lock);
		INIT_WORK(&g_anx->work, anx_work_func);
		g_anx->wq = create_singlethread_workqueue("tycnatwq");
		if (tc_setup_common(&g_anx->typec, root_dev)) {
			destroy_workqueue(g_anx->wq);
			kfree(g_anx);
			g_anx = NULL;
			root_device_unregister(root_dev);
			return -EFAULT;
		}
		/* natural race: queue work, then immediately remove */
		queue_work(g_anx->wq, &g_anx->work);
		anx_remove_thread(g_anx);
		cancel_work_sync(&g_anx->work);
		destroy_workqueue(g_anx->wq);
		kfree(g_anx);
		g_anx = NULL;
		root_device_unregister(root_dev);
		if (!(round % 100))
			pr_info("TCNAT: round %d done\n", round);
	}
	pr_info("TCNAT: %d rounds done (check dmesg for KASAN reports)\n",
		tcnat_rounds);
	return 0;
}

/* ---------------- module params & entry ---------------- */
static char *action = "status";
module_param(action, charp, 0444);
MODULE_PARM_DESC(action, "tc0|tc0b|tc1|tc2|tcnat|status");

static int tyc_run_action(void)
{
	if (!strcmp(action, "tc0"))
		return run_tc0();
	if (!strcmp(action, "tc0b"))
		return run_tc0b();
	if (!strcmp(action, "tc1"))
		return run_tc1();
	if (!strcmp(action, "tc2"))
		return run_tc2();
	if (!strcmp(action, "tc2ctl"))
		return run_tc2ctl();
	if (!strcmp(action, "tcnat"))
		return run_tcnat();
	pr_info("no action, done\n");
	return 0;
}

static int __init tyc_init(void)
{
	pr_info("action=%s\n", action);
	return tyc_run_action();
}
module_init(tyc_init);

static void __exit tyc_exit(void)
{
	pr_info("exit\n");
}
module_exit(tyc_exit);
MODULE_LICENSE("GPL");