// SPDX-License-Identifier: GPL-2.0
//
// Copyright (c) 2021 MediaTek Inc.

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/dma-buf.h>
#include <linux/types.h>
#include <linux/fs.h>
#include <linux/string.h>
#include <linux/proc_fs.h>
#include <linux/seq_file.h>
#include <linux/sizes.h>
#include <linux/notifier.h>
#include <linux/version.h>
#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 14, 0)
#include <linux/panic_notifier.h>
#endif

#include "mtk_iommu_dbg.h"
#include "mtk_iommu_trace.h"

#undef pr_fmt
#define pr_fmt(fmt) "IOMMUDBG: " fmt

static int db_vumap_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
#ifdef CONFIG_ARM64
	pr_debug("<%s> pre_handler: pc = 0x%lx, argu = 0x%lx\n",
#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 11, 0)
		ri->rp->kp.symbol_name,
#else
		ri->rph->rp->kp.symbol_name,
#endif
		(long)regs->pc,
		(long)regs_get_register(regs, 0));
#endif
	// regs[0] should be pointer of dma_buf
	add_vunmap_record_list((void *)regs_get_register(regs, 0));

	return 0;
}

static int db_vmap_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
#ifdef CONFIG_ARM64
	pr_debug("<%s> pre_handler: pc = 0x%lx, argu = 0x%lx\n",
#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 11, 0)
		ri->rp->kp.symbol_name,
#else
		ri->rph->rp->kp.symbol_name,
#endif
		(long)regs->pc,
		(long)regs_get_register(regs, 0));
#endif
	// regs[0] should be pointer of dma_buf
	add_vmap_record_list((void *)regs_get_register(regs, 0));
	return 0;
}

static int db_get_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
#ifdef CONFIG_ARM64
	pr_debug("<%s> pre_handler: pc = 0x%lx, argu = 0x%lx\n",
#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 11, 0)
		ri->rp->kp.symbol_name,
#else
		ri->rph->rp->kp.symbol_name,
#endif
		(long)regs->pc,
		(long)regs_get_register(regs, 0));
#endif
	// regs[0] should be fd
	might_add_fget_list((int)regs_get_register(regs, 0));
	return 0;
}

static int db_put_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
#ifdef CONFIG_ARM64
	pr_debug("<%s> pre_handler: pc = 0x%lx, argu = 0x%lx\n",
#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 11, 0)
		ri->rp->kp.symbol_name,
#else
		ri->rph->rp->kp.symbol_name,
#endif
		(long)regs->pc,
		(long)regs_get_register(regs, 0));
#endif
	// regs[0] should be pointer of dma_buf
	might_add_fput_list((void *)regs_get_register(regs, 0));
	return 0;
}

static int db_export_ret(struct kretprobe_instance *ri, struct pt_regs *regs)
{
#ifdef CONFIG_ARM64
	pr_debug("<%s> post_handler: pc = 0x%lx, ret = 0x%lx\n",
#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 11, 0)
		ri->rp->kp.symbol_name,
#else
		ri->rph->rp->kp.symbol_name,
#endif
		(long)regs->pc,
		(long)regs_return_value(regs));
#endif
	// regs[0] should be pointer of dma_buf
	might_add_fget_list_dma((void *)regs_return_value(regs));
	return 0;
}

static struct kretprobe kretp[] = {
	{
		.kp.symbol_name	= "dma_buf_vmap",
		.entry_handler	= db_vmap_entry,
		.handler   = NULL,
	},
	{
		.kp.symbol_name	= "dma_buf_vunmap",
		.entry_handler	= db_vumap_entry,
		.handler	= NULL,
	},
	{
		.kp.symbol_name	= "dma_buf_get",
		.entry_handler	= db_get_entry,
		.handler	= NULL,
	},
	{
		.kp.symbol_name	= "dma_buf_put",
		.entry_handler	= db_put_entry,
		.handler	= NULL,
	},
	{
		.kp.symbol_name	= "dma_buf_export",
		.entry_handler	= NULL,
		.handler	= db_export_ret,
	},
};

static int file_record_size_open(struct inode *inode, struct file *file)
{
	return single_open(file, fcount_record_size_show, PDE_DATA(inode));
}

static ssize_t file_record_size_write(struct file *filp,
		const char __user *buf, size_t count, loff_t *offp)
{
	int ret;
	unsigned long sz;

	ret = kstrtoul_from_user(buf, count, 0, &sz);
	if (ret) {
		pr_crit("kstrtoul_from_user failed with %d\n", ret);
		return count;
	}

	if (!IS_ALIGNED(sz, SZ_4K)) {
		pr_crit("size 0x%lx is not ailgned 4K\n", sz);
		return count;
	}

	fcount_record_size_store((size_t)sz);
	return count;
}

#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 6, 0)
static const struct file_operations trace_sz_proc = {
	.owner		= THIS_MODULE,
	.open		= file_record_size_open,
	.read		= seq_read,
	.write		= file_record_size_write,
	.release	= single_release,
};
#else
static const struct proc_ops trace_sz_proc = {
	.proc_open		= file_record_size_open,
	.proc_read		= seq_read,
	.proc_write		= file_record_size_write,
	.proc_release		= single_release,
};
#endif

static int file_record_inode_open(struct inode *inode, struct file *file)
{
	return single_open(file, fcount_record_inode_show, PDE_DATA(inode));
}

static ssize_t file_record_inode_write(struct file *filp,
		const char __user *buf, size_t count, loff_t *offp)
{
	int ret;
	unsigned long ino;

	ret = kstrtoul_from_user(buf, count, 0, &ino);
	if (ret) {
		pr_crit("kstrtoul_from_user failed with %d\n", ret);
		return ret;
	}
	fcount_record_inode_store(ino);

	return count;
}

#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 6, 0)
static const struct file_operations ino_proc = {
	.owner		= THIS_MODULE,
	.open		= file_record_inode_open,
	.read		= seq_read,
	.write		= file_record_inode_write,
	.release	= single_release,
};
#else
static const struct proc_ops ino_proc = {
	.proc_open		= file_record_inode_open,
	.proc_read		= seq_read,
	.proc_write		= file_record_inode_write,
	.proc_release		= single_release,
};
#endif
static int record_alive_open(struct inode *inode, struct file *file)
{
	return single_open(file, record_alive_show, PDE_DATA(inode));
}

static ssize_t record_alive_write(struct file *filp,
		const char __user *buf, size_t count, loff_t *offp)
{
	int ret;
	bool enable;

	ret = kstrtobool_from_user(buf, count, &enable);
	if (ret) {
		pr_crit("kstrtobool_from_user failed with %d\n", ret);
		return ret;
	}
	record_alive_store(enable);

	return count;
}

#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 6, 0)
static const struct file_operations alive_proc = {
	.owner		= THIS_MODULE,
	.open		= record_alive_open,
	.read		= seq_read,
	.write		= record_alive_write,
	.release	= single_release,
};
#else
static const struct proc_ops alive_proc = {
	.proc_open		= record_alive_open,
	.proc_read		= seq_read,
	.proc_write		= record_alive_write,
	.proc_release		= single_release,
};
#endif

static int proc_help_show(struct seq_file *m, void *v)
{
	seq_printf(m, "help help help help help\n");

	return 0;
}

static ssize_t leaker_write(struct file *filp,
		const char __user *buf, size_t count, loff_t *offp)
{
	int rc;
	bool en = false;

	rc = kstrtobool_from_user(buf, count, &en);
	if (rc)
		return rc;

	if (en)
		rc = leaker_start();
	else
		rc = leaker_stop();

	if (rc)
		return rc;

	return count;
}

int leaker_show(struct seq_file *m, void *v)
{
	struct dma_buf *db = get_leak_dmabuf();

	if (!db) {
		seq_printf(m, "leak dma buf not found!\n");
		return 0;
	}

	seq_printf(m, "dma(0x%lx) inode(%lu) fcount(%d)\n",
		(unsigned long)db, file_inode(db->file)->i_ino,
		(int)file_count(db->file));

	return 0;
}

static int leaker_open(struct inode *inode, struct file *file)
{
	return single_open(file, leaker_show, PDE_DATA(inode));
}

#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 6, 0)
static const struct file_operations leaker_proc = {
	.owner		= THIS_MODULE,
	.open		= leaker_open,
	.read		= seq_read,
	.write		= leaker_write,
	.release	= single_release,
};
#else
static const struct proc_ops leaker_proc = {
	.proc_open		= leaker_open,
	.proc_read		= seq_read,
	.proc_write		= leaker_write,
	.proc_release		= single_release,
};
#endif

static ssize_t af_write(struct file *filp,
		const char __user *buf, size_t count, loff_t *offp)
{
	int rc;
	unsigned long en;

	rc = kstrtoul_from_user(buf, count, 0, &en);
	if (rc)
		return rc;

	if (en)
		rc = alloc_free_start(en);
	else
		rc = alloc_free_stop();

	if (rc)
		return rc;

	return count;
}

int af_show(struct seq_file *m, void *v)
{
	seq_printf(m, "run count %ld * 512 times\n", get_alloc_free_cnt());
	return 0;
}

static int af_open(struct inode *inode, struct file *file)
{
	return single_open(file, af_show, PDE_DATA(inode));
}

#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 6, 0)
static const struct file_operations alloc_free_proc = {
	.owner		= THIS_MODULE,
	.open		= af_open,
	.read		= seq_read,
	.write		= af_write,
	.release	= single_release,
};
#else
static const struct proc_ops alloc_free_proc = {
	.proc_open		= af_open,
	.proc_read		= seq_read,
	.proc_write		= af_write,
	.proc_release		= single_release,
};
#endif

static ssize_t smc_info_write(struct file *filp,
		const char __user *buf, size_t count, loff_t *offp)
{
	int ret;
	unsigned long addr = ULONG_MAX;

	ret = kstrtoul_from_user(buf, count, 0, &addr);
	pr_err("ret:%d register_symbol addr: 0x%lx\n", ret, addr);
	if (ret)
		return ret;

	pr_err("register_symbol addr: 0x%lx\n", addr);
	ret = register_symbol(addr);
	if (ret)
		return ret;

	return count;
}

static int smc_info_open(struct inode *inode, struct file *file)
{
	return single_open(file, smc_info_dump, PDE_DATA(inode));
}

#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 6, 0)
static const struct file_operations smc_info_proc = {
	.owner		= THIS_MODULE,
	.open		= smc_info_open,
	.read		= seq_read,
	.write		= smc_info_write,
	.release	= single_release,
};
#else
static const struct proc_ops smc_info_proc = {
	.proc_open		= smc_info_open,
	.proc_read		= seq_read,
	.proc_write		= smc_info_write,
	.proc_release	= single_release,
};
#endif

static int hello_panic(struct notifier_block *self, unsigned long val, void *data)
{
	//TODO: print your debug message
	return 0;
}

static struct notifier_block mtk_helper_panic_nb = {
	.notifier_call = hello_panic,
};

static void
show_buf_handle(void *data, const char *action,
		struct mtk_iommu_buf_handle *buf_handle)
{
	pr_info("[%s] action: %s, addr: %llx, size: %zx, comm:%s, pid:%d\n",
			__func__, action,
			(unsigned long long)buf_handle->addr,
			(size_t)buf_handle->length,
			buf_handle->comm,
			buf_handle->pid);
}

static struct proc_dir_entry *dir;
static struct proc_dir_entry *dir_tort;
static struct proc_dir_entry *dir_smcinfo;

static int __init iommu_dbg_init(void)
{
	int i, ret;
	struct proc_dir_entry *d;

	for (i = 0; i < ARRAY_SIZE(kretp); ++i) {
		ret = register_kretprobe(&kretp[i]);
		if (ret < 0) {
			pr_err("register kretprobe for [%s] failed with %d\n",
				kretp[i].kp.symbol_name, ret);
			return ret;
		}
		pr_info("register kretprobe for [%s] success\n", kretp[i].kp.symbol_name);
	}

	dir = proc_mkdir("iommu_dbg", NULL);
	if (IS_ERR_OR_NULL(dir)) {
		pr_crit("create iommu debug folder in procfs failed!\n");
		return -EIO;
	}
	pr_info("register proc-dir [iommu_dbg] success!\n");

	dir_tort = proc_mkdir("iommu_torture", NULL);
	if (IS_ERR_OR_NULL(dir_tort)) {
		pr_err("Create iommu_torture folder in procfs failed\n");
		return -EIO;
	}
	pr_info("register proc-dir [iommu_torture] success!\n");

	d = proc_create_single("help", 0, dir, proc_help_show);
	if (IS_ERR_OR_NULL(d)) {
		pr_crit("Create vmap_record procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [help] success!\n");

	d = proc_create_single("vmap_record_dump", 0, dir, proc_vmap_show);
	if (IS_ERR_OR_NULL(d)) {
		pr_crit("Create vmap_record procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [vmap_record] success!\n");

	d = proc_create_single("file_record_dump", 0, dir, fcount_record_show_all);
	if (IS_ERR_OR_NULL(d)) {
		pr_crit("Create vmap_record procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [file_record_dump] success!\n");

	d = proc_create("trace_size", S_IRUGO | S_IWUGO, dir, &trace_sz_proc);
	if (IS_ERR_OR_NULL(d)) {
		pr_crit("Create trace_size procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [trace_size] success!\n");

	d = proc_create("trace_inode", S_IRUGO | S_IWUGO, dir, &ino_proc);
	if (IS_ERR_OR_NULL(d)) {
		pr_crit("Create trace_ino procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [trace_ino] success!\n");

	d = proc_create("record_alive", S_IRUGO | S_IWUGO, dir, &alive_proc);
	if (IS_ERR_OR_NULL(d)) {
		pr_crit("Create record_alive procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [record_alive] success!\n");

	d = proc_create("alloc_free", S_IRUGO | S_IWUGO, dir_tort, &alloc_free_proc);
	if (IS_ERR_OR_NULL(d)) {
		pr_err("Create alloc_free procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [alloc_free] success!\n");

	d = proc_create("leaker", S_IRUGO | S_IWUGO, dir_tort, &leaker_proc);
	if (IS_ERR_OR_NULL(d)) {
		pr_err("Create leaker procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [leaker] success!\n");

	dir_smcinfo = proc_mkdir("smc_debug", NULL);
	if (IS_ERR_OR_NULL(dir_smcinfo)) {
		pr_err("Create smc_debug folder in procfs failed\n");
		return -EIO;
	}
	pr_info("register proc [SMC DBG] success!\n");

	d = proc_create("smc_info", S_IRUGO | S_IWUGO, dir_smcinfo, &smc_info_proc);
	if (IS_ERR_OR_NULL(d)) {
		pr_err("Create smc_info procfs failed\n");
		return -ENOMEM;
	}
	pr_info("register proc [SMC INFO] success!\n");

	atomic_notifier_chain_register(&panic_notifier_list,
			&mtk_helper_panic_nb);
	if (register_trace_mtk_buf_handle(show_buf_handle, NULL))
		return -EINVAL;
	return 0;
}

static void __exit iommu_dbg_exit(void)
{
	int i;

	unregister_trace_mtk_buf_handle(show_buf_handle, NULL);
	atomic_notifier_chain_unregister(&panic_notifier_list,
			&mtk_helper_panic_nb);
	for (i = 0; i < ARRAY_SIZE(kretp); ++i)
		unregister_kretprobe(&kretp[i]);

	remove_proc_entry("vmap_record_dump", dir);
	remove_proc_entry("file_record_dump", dir);
	remove_proc_entry("trace_size", dir);
	remove_proc_entry("trace_inode", dir);
	remove_proc_entry("record_alive", dir);
	remove_proc_entry("alloc_free", dir_tort);
	remove_proc_entry("leaker", dir_tort);
	remove_proc_entry("iommu_dbg", NULL);
	remove_proc_entry("iommu_torture", NULL);

	unregister_symbol();
	remove_proc_entry("smc_info", dir_smcinfo);
	remove_proc_entry("smc_debug", NULL);
}

module_init(iommu_dbg_init)
module_exit(iommu_dbg_exit)
MODULE_LICENSE("GPL");
