// SPDX-License-Identifier: GPL-2.0

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/mutex.h>
#include <linux/ww_mutex.h>
#include <linux/spinlock.h>
#include <linux/lockdep.h>
#include <linux/rhashtable.h>
#include <linux/slab.h>

/*
 * ww_mutex machinery used to give a held lock a non-zero ->references.
 * All ww_mutexes of one ww_class share a single lockdep class, and
 * ww_mutex_lock() passes nest_lock = &ctx->dep_map, so acquiring the second
 * one sets ->references on the first one's held_lock entry.
 */
static DEFINE_WW_CLASS(repro_ww_class);
static struct ww_mutex ww_a;
static struct ww_mutex ww_b;

struct repro_obj {
	int			value;
	struct rhash_head	node;
};

static struct rhashtable ht;
static const struct rhashtable_params repro_params = {
	.key_len	= sizeof(int),
	.key_offset	= offsetof(struct repro_obj, value),
	.head_offset	= offsetof(struct repro_obj, node),
};

/*
 * Run fn() while holding two same-class ww_mutexes under one acquire context,
 * so that the first held_lock entry has ->references != 0.
 */
static int with_referenced_held_lock(void (*fn)(void))
{
	struct ww_acquire_ctx ctx;
	int ret;

	ww_mutex_init(&ww_a, &repro_ww_class);
	ww_mutex_init(&ww_b, &repro_ww_class);

	ww_acquire_init(&ctx, &repro_ww_class);

	ret = ww_mutex_lock(&ww_a, &ctx);
	if (ret) {
		pr_err("ww_mutex_lock(ww_a) failed: %d\n", ret);
		goto out_fini;
	}
	/* same class as ww_a plus a nest_lock -> bumps ->references */
	ret = ww_mutex_lock(&ww_b, &ctx);
	if (ret) {
		pr_err("ww_mutex_lock(ww_b) failed: %d\n", ret);
		ww_mutex_unlock(&ww_a);
		goto out_fini;
	}

	fn();

	ww_mutex_unlock(&ww_b);
	ww_mutex_unlock(&ww_a);
out_fini:
	ww_acquire_fini(&ctx);
	return ret;
}

static void lookup_rhashtable(void)
{
	int wanted = 42;

	pr_info("doing rhashtable_lookup_fast()\n");
	pr_info("rhashtable_lookup_fast() = %px\n",
		rhashtable_lookup_fast(&ht, &wanted, repro_params));
}

static int __init repro_rhashtable(void)
{
	struct rhashtable_iter iter;
	int ret;

	ret = rhashtable_init(&ht, &repro_params);
	if (ret) {
		pr_err("rhashtable_init failed: %d\n", ret);
		return ret;
	}

	/*
	 * rhashtable_walk_enter() takes spin_lock(&ht->lock), which is what
	 * registers subclass 0 of this instance's key under the name "key".
	 */
	rhashtable_walk_enter(&ht, &iter);
	rhashtable_walk_exit(&iter);

	ret = with_referenced_held_lock(lookup_rhashtable);

	rhashtable_destroy(&ht);
	return ret;
}

static int __init repro_init(void)
{
	if (!IS_ENABLED(CONFIG_PROVE_LOCKING)) {
		pr_err("CONFIG_PROVE_LOCKING is required\n");
		return -EOPNOTSUPP;
	}
	if (!debug_locks) {
		pr_err("lockdep is already disabled, cannot test\n");
		return -EOPNOTSUPP;
	}

	return repro_rhashtable();
}

static void __exit repro_exit(void)
{
	pr_info("rht_lockdep_repro: unloaded\n");
}

module_init(repro_init);
module_exit(repro_exit);

MODULE_DESCRIPTION("Reproducer for the rhashtable lockdep subclass collision");
MODULE_LICENSE("GPL");