/* SPDX-License-Identifier: GPL-2.0-only */

#include <linux/init.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <asm/kvm_pkvm_module.h>

static unsigned long pkvm_module_token;
int kvm_nvhe_sym(test_lazy_pte_hyp_init)(const struct pkvm_module_ops *ops);
void kvm_nvhe_sym(test_lazy_pte_hyp_hvc)(struct user_pt_regs *regs);

static int __init test_lazy_pte_init(void)
{
	int ret = pkvm_load_el2_module(kvm_nvhe_sym(test_lazy_pte_hyp_init),
				       &pkvm_module_token);
	int hvc, cnt = 1 << 16;
	void *page;

	if (ret) {
		pr_err("Failed to load EL2 module: %d\n", ret);
		return ret;
	}

	ret = pkvm_register_el2_mod_call(kvm_nvhe_sym(test_lazy_pte_hyp_hvc), pkvm_module_token);
	if (ret < 0)
		return ret;

	hvc = ret;

	page = (void *)__get_free_page(GFP_KERNEL);
	if (!page)
		return -ENOMEM;

	while (cnt--) {
		ret = pkvm_el2_mod_call(hvc, virt_to_pfn(page), true);
		if (ret) {
			pr_err("Failed to enable dirty logging on page %px\n", page);
			break;
		}
		ret = pkvm_el2_mod_call(hvc, virt_to_pfn(page), false);
		if (ret) {
			pr_err("Failed to disable dirty logging on page %px\n", page);
			break;
		}
	}

	return 0;
}
module_init(test_lazy_pte_init);

MODULE_LICENSE("GPL v2");