// SPDX-License-Identifier: GPL-2.0
/*
 * iommurace.c - dynamic verification driver for the IOMMU lifecycle race
 * audit (HV-3 / HV-4 in docs/manual-review-drivers-base-2026-10-20.md).
 *
 * Registers a FAKE IOMMU instance (struct iommu_device + struct iommu_ops)
 * through the real EXPORT_SYMBOL_GPL core APIs, and drives the REAL core
 * code paths of drivers/iommu/iommu.c with REAL platform-bus devices:
 *
 *   entry A (instance unregister traversal):
 *       iommu_device_unregister() -> bus_for_each_dev -> remove_iommu_group
 *   entry B (client removal):
 *       platform_device_unregister() -> device_del ->
 *       BUS_NOTIFY_REMOVED_DEVICE -> iommu_bus_notifier -> iommu_release_device
 *   entry C (client probe):
 *       platform_device_add() -> device_add -> BUS_NOTIFY_ADD_DEVICE ->
 *       iommu_bus_notifier -> iommu_probe_device -> iommu_init_device
 *
 * Clients are bound to the fake instance by calling iommu_fwspec_init()
 * (EXPORT_SYMBOL_GPL) between platform_device_alloc() and
 * platform_device_add(), so the ADD notifier's probe matches the fake ops
 * via iommu_fwspec_ops() on the shared software fwnode.
 *
 * actions (echo <cmd> > /sys/module/iommurace/parameters/action):
 *   setup           - carrier device + software fwnode + iommu_device_sysfs_add
 *                     + iommu_device_register (real bus_iommu_probe sweep)
 *   unreg           - one iommu_device_unregister (real traversal)
 *   reg             - one iommu_device_register
 *   regloop N       - N sequential register+unregister rounds (baseline)
 *   client N        - N clients added then removed sequentially (baseline;
 *                     expect zero splats with a stable registered instance)
 *   ab N            - HV-3 A x B: kthread A loops register+unregister,
 *                     kthread B loops client add+del; real paths on both sides
 *   ac N            - HV-4 A x C: kthread A loops register+unregister,
 *                     kthread C loops client add only (W1/W2 windows,
 *                     clients kept for post-mortem inspection)
 *   ac1 N           - same as ac but device_group = generic_single_device_group
 *                     (W3: unregister tail singleton_group race surface)
 *   cleanup         - best-effort teardown (unregister, drop clients, sysfs)
 *   status          - log stats
 *
 * Kernel source is NOT modified; everything is driven through exported
 * iommu core APIs, matching the methodology of pmu_race/mcrsim/regmap_race.
 *
 * NOTE: ab/ac/ac1 may leak module references when the audit's "W2 defective
 * state" hits (probe succeeded -> try_module_get, later removal skips
 * iommu_deinit_device -> no module_put). rmmod then fails with EBUSY; the
 * plans therefore do not rmmod after race actions - the leak itself is a
 * verified prediction of the report.
 */
#define pr_fmt(fmt) "IOMMURACE: " fmt

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/kthread.h>
#include <linux/delay.h>
#include <linux/err.h>
#include <linux/string.h>
#include <linux/slab.h>
#include <linux/completion.h>
#include <linux/atomic.h>
#include <linux/platform_device.h>
#include <linux/iommu.h>
#include <linux/property.h>

/* ------------------------------------------------------------------ */
/* fake IOMMU instance                                                 */

static struct iommu_device fom_iommu;
static struct iommu_ops fom_ops;
static struct platform_device *carrier;
static struct fwnode_handle *sw_node;
static atomic_t reg_state = ATOMIC_INIT(0); /* 1 = registered */
static int use_singleton;

/* stats */
static unsigned long stat_reg, stat_unreg, stat_reg_fail;
static unsigned long stat_add_ok, stat_add_fail, stat_fw_fail, stat_del;
static unsigned long stat_b_done, stat_c_done, stat_a_done;

/* fake static identity domain: attach is a no-op, free must not kfree */
static int fom_attach_dev(struct iommu_domain *dom, struct device *dev,
			  struct iommu_domain *old)
{
	return 0;
}

static void fom_dom_free(struct iommu_domain *dom)
{
	pr_info("fom_dom_free on static identity domain (ignored)\n");
}

static const struct iommu_domain_ops fom_dom_ops = {
	.attach_dev = fom_attach_dev,
	.free = fom_dom_free,
};

static struct iommu_domain fom_identity_dom = {
	.type = IOMMU_DOMAIN_IDENTITY,
	.ops = &fom_dom_ops,
	.owner = &fom_ops,
	.pgsize_bitmap = 0,
	.geometry = {
		.aperture_start = 0,
		.aperture_end = ~0UL,
		.force_aperture = false,
	},
};

static struct iommu_device *fom_probe_device(struct device *dev)
{
	return &fom_iommu;
}

static void fom_release_device(struct device *dev)
{
}

static int fom_def_domain_type(struct device *dev)
{
	return IOMMU_DOMAIN_IDENTITY;
}

static int fom_register(void)
{
	int ret;

	if (atomic_xchg(&reg_state, 1) != 0)
		return 0; /* already registered */

	fom_ops = (struct iommu_ops){
		.owner			= THIS_MODULE,
		.probe_device		= fom_probe_device,
		.release_device		= fom_release_device,
		.device_group		= use_singleton ?
					  generic_single_device_group :
					  generic_device_group,
		.def_domain_type	= fom_def_domain_type,
		.identity_domain	= &fom_identity_dom,
	};

	ret = iommu_device_register(&fom_iommu, &fom_ops, &carrier->dev);
	if (ret) {
		atomic_set(&reg_state, 0);
		stat_reg_fail++;
		return ret;
	}
	stat_reg++;
	return 0;
}

static void fom_unregister(void)
{
	if (atomic_xchg(&reg_state, 0) != 1)
		return; /* not registered */
	iommu_device_unregister(&fom_iommu);
	stat_unreg++;
}

/* ------------------------------------------------------------------ */
/* fake client devices (real platform bus devices)                     */

#define MAX_CLI 512
static struct platform_device *cli[MAX_CLI];
static int nr_cli;

static int fom_new_client(struct platform_device **out)
{
	struct platform_device *pdev;
	int ret;

	pdev = platform_device_alloc("fomcli", PLATFORM_DEVID_AUTO);
	if (!pdev)
		return -ENOMEM;

	/*
	 * Bind the client's fwspec to the fake instance BEFORE device_add so
	 * the ADD notifier's iommu_probe_device() matches our fake ops.
	 * Fails with -ENODEV/-EPROBE_DEFER when the instance is currently
	 * unregistered / mid-registration: the natural non-racy branch.
	 */
	ret = iommu_fwspec_init(&pdev->dev, sw_node);
	if (ret) {
		stat_fw_fail++;
		platform_device_put(pdev);
		return ret;
	}

	ret = platform_device_add(pdev);
	if (ret) {
		stat_add_fail++;
		platform_device_put(pdev);
		return ret;
	}
	stat_add_ok++;
	*out = pdev;
	return 0;
}

static void fom_del_client(struct platform_device *pdev)
{
	platform_device_unregister(pdev);
	stat_del++;
}

/* ------------------------------------------------------------------ */
/* race threads                                                        */

struct rargs {
	int rounds;
};

static struct rargs ra, rb;

static int thread_a(void *data) /* entry A: register/unregister rounds */
{
	struct rargs *a = data;
	int i;

	for (i = 0; i < a->rounds; i++) {
		fom_register();
		fom_unregister();
		if ((i % 500) == 499)
			pr_info("thread_a round %d\n", i + 1);
		/*
		 * No sleep: maximize overlap of the unregister traversal
		 * with concurrent client add/del/probe on other CPUs.
		 */
	}
	stat_a_done++;
	pr_info("thread_a finished %d rounds\n", a->rounds);
	return 0;
}

static int thread_b(void *data) /* entry B: client add+del churn */
{
	struct rargs *a = data;
	int i;
	struct platform_device *pdev;

	for (i = 0; i < a->rounds; i++) {
		if (fom_new_client(&pdev))
			continue;
		fom_del_client(pdev);
		if ((i % 500) == 499)
			pr_info("thread_b round %d\n", i + 1);
	}
	stat_b_done++;
	pr_info("thread_b finished %d rounds (add_ok=%lu fw_fail=%lu)\n",
		a->rounds, stat_add_ok, stat_fw_fail);
	return 0;
}

static int thread_c(void *data) /* entry C: client add only (HV-4) */
{
	struct rargs *a = data;
	int i;

	for (i = 0; i < a->rounds; i++) {
		struct platform_device *pdev;

		if (nr_cli < MAX_CLI) {
			if (fom_new_client(&pdev) == 0)
				cli[nr_cli++] = pdev;
		}
		if ((i % 500) == 499)
			pr_info("thread_c round %d\n", i + 1);
	}
	stat_c_done++;
	pr_info("thread_c finished %d rounds (added=%d)\n", a->rounds, nr_cli);
	return 0;
}

static void run_race(const char *kind, int rounds)
{
	struct task_struct *ta, *tb;

	ra.rounds = rounds;
	rb.rounds = rounds;

	if (strcmp(kind, "ab") == 0)
		tb = kthread_run(thread_b, &rb, "iommurace_b");
	else
		tb = kthread_run(thread_c, &rb, "iommurace_c");
	ta = kthread_run(thread_a, &ra, "iommurace_a");
	if (IS_ERR(ta) || IS_ERR(tb)) {
		pr_err("kthread creation failed\n");
		return;
	}
	/* block until both sides complete (bounded: a crashed thread never
	 * increments its done counter - bail out after 150s and let the
	 * plan continue; the KASAN report in dmesg is the evidence) */
	{
		int waited = 0;

		while (stat_a_done + stat_b_done + stat_c_done < 2 &&
		       waited < 150000) {
			msleep(200);
			waited += 200;
		}
	}
	pr_info("race %s rounds=%d complete\n", kind, rounds);
}

/* ------------------------------------------------------------------ */
/* action parameter                                                    */

static char action_buf[64];

static int action_set(const char *val, const struct kernel_param *kp)
{
	char cmd[16];
	long n = 0;
	int got;

	strscpy(action_buf, val, sizeof(action_buf));
	got = sscanf(val, "%15s %ld", cmd, &n);
	if (got < 1)
		return -EINVAL;

	if (strcmp(cmd, "setup") == 0) {
		int ret;

		if (carrier) {
			pr_info("setup: already done\n");
			return 0;
		}
		sw_node = fwnode_create_software_node(NULL, NULL);
		if (IS_ERR(sw_node)) {
			sw_node = NULL;
			return PTR_ERR(sw_node);
		}
		carrier = platform_device_alloc("fomhw", PLATFORM_DEVID_AUTO);
		if (!carrier) {
			fwnode_remove_software_node(sw_node);
			sw_node = NULL;
			return -ENOMEM;
		}
		device_set_node(&carrier->dev, sw_node);
		ret = platform_device_add(carrier);
		if (ret) {
			platform_device_put(carrier);
			carrier = NULL;
			fwnode_remove_software_node(sw_node);
			sw_node = NULL;
			return ret;
		}
		ret = iommu_device_sysfs_add(&fom_iommu, &carrier->dev, NULL,
					     "fom_iommu0");
		if (ret) {
			platform_device_unregister(carrier);
			carrier = NULL;
			fwnode_remove_software_node(sw_node);
			sw_node = NULL;
			return ret;
		}
		ret = fom_register();
		pr_info("setup done (register ret=%d)\n", ret);
		return ret;
	}

	if (strcmp(cmd, "unreg") == 0) {
		fom_unregister();
		pr_info("unreg done\n");
		return 0;
	}
	if (strcmp(cmd, "reg") == 0) {
		int ret = fom_register();

		pr_info("reg ret=%d\n", ret);
		return ret;
	}
	if (strcmp(cmd, "regloop") == 0 && got == 2) {
		long i;

		for (i = 0; i < n; i++) {
			fom_register();
			fom_unregister();
		}
		pr_info("regloop %ld done\n", n);
		return 0;
	}
	if (strcmp(cmd, "client") == 0 && got == 2) {
		long i;
		int err = 0;

		if (n > MAX_CLI)
			n = MAX_CLI;
		for (i = 0; i < n; i++) {
			if (fom_new_client(&cli[nr_cli])) {
				err = 1;
				break;
			}
			nr_cli++;
		}
		msleep(500); /* let probes settle */
		for (i = 0; i < nr_cli; i++)
			fom_del_client(cli[i]);
		nr_cli = 0;
		pr_info("client %ld done (err=%d)\n", n, err);
		return 0;
	}
	if ((strcmp(cmd, "ab") == 0 || strcmp(cmd, "ac") == 0 ||
	     strcmp(cmd, "ac1") == 0) && got == 2) {
		use_singleton = (cmd[2] == '1');
		stat_a_done = stat_b_done = stat_c_done = 0;
		run_race(cmd, (int)n);
		return 0;
	}
	if (strcmp(cmd, "cleanup") == 0) {
		int i;

		fom_unregister();
		for (i = 0; i < nr_cli; i++)
			if (cli[i])
				fom_del_client(cli[i]);
		nr_cli = 0;
		pr_info("cleanup done (module refcount leak => rmmod EBUSY is expected after races)\n");
		return 0;
	}
	if (strcmp(cmd, "status") == 0) {
		pr_info("stats: reg=%lu unreg=%lu reg_fail=%lu add_ok=%lu add_fail=%lu fw_fail=%lu del=%lu\n",
			stat_reg, stat_unreg, stat_reg_fail, stat_add_ok,
			stat_add_fail, stat_fw_fail, stat_del);
		pr_info("clients tracked=%d singleton=%d\n", nr_cli, use_singleton);
		return 0;
	}

	return -EINVAL;
}

static const struct kernel_param_ops action_ops = {
	.set = action_set,
	.get = param_get_string,
};
module_param_cb(action, &action_ops, &action_buf, 0600);
MODULE_PARM_DESC(action, "setup|unreg|reg|regloop N|client N|ab N|ac N|ac1 N|cleanup|status");

static int __init iommurace_init(void)
{
	pr_info("loaded (kernel has IOMMU_API)\n");
	return 0;
}

static void __exit iommurace_exit(void)
{
	int i;

	fom_unregister();
	for (i = 0; i < nr_cli; i++)
		if (cli[i])
			fom_del_client(cli[i]);
	nr_cli = 0;
	if (carrier) {
		iommu_device_sysfs_remove(&fom_iommu);
		platform_device_unregister(carrier);
		fwnode_remove_software_node(sw_node);
	}
	pr_info("unloaded\n");
}

module_init(iommurace_init);
module_exit(iommurace_exit);

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("IOMMU lifecycle race verification driver (HV-3/HV-4)");