// SPDX-License-Identifier: GPL-2.0
/*
 * fprobe - Simple ftrace probe wrapper for function entry.
 */
#define pr_fmt(fmt) "fprobe: " fmt

#include <linux/err.h>
#include <linux/fprobe.h>
#include <linux/kallsyms.h>
#include <linux/kprobes.h>
#include <linux/slab.h>
#include <linux/sort.h>

static void fprobe_handler(unsigned long ip, unsigned long parent_ip,
			   struct ftrace_ops *ops, struct ftrace_regs *fregs)
{
	struct fprobe *fp;
	int bit;

	fp = container_of(ops, struct fprobe, ops);
	if (fprobe_disabled(fp))
		return;

	bit = ftrace_test_recursion_trylock(ip, parent_ip);
	if (bit < 0) {
		fp->nmissed++;
		return;
	}

	if (fp->entry_handler)
		fp->entry_handler(fp, ip, ftrace_get_regs(fregs));

	ftrace_test_recursion_unlock(bit);
}
NOKPROBE_SYMBOL(fprobe_handler);

/* Convert ftrace location address from symbols */
static unsigned long *get_ftrace_locations(const char **syms, int num)
{
	unsigned long addr, size;
	unsigned long *addrs;
	int i;

	/* Convert symbols to symbol address */
	addrs = kcalloc(num, sizeof(*addrs), GFP_KERNEL);
	if (!addrs)
		return ERR_PTR(-ENOMEM);

	for (i = 0; i < num; i++) {
		addr = kallsyms_lookup_name(syms[i]);
		if (!addr)	/* Maybe wrong symbol */
			goto error;

		/* Convert symbol address to ftrace location. */
		if (!kallsyms_lookup_size_offset(addr, &size, NULL) || !size)
			goto error;

		addr = ftrace_location_range(addr, addr + size - 1);
		if (!addr) /* No dynamic ftrace there. */
			goto error;

		addrs[i] = addr;
	}

	return addrs;

error:
	kfree(addrs);

	return ERR_PTR(-ENOENT);
}

static void fprobe_init(struct fprobe *fp)
{
	fp->nmissed = 0;
	fp->ops.func = fprobe_handler;
	fp->ops.flags |= FTRACE_OPS_FL_SAVE_REGS;
}

/**
 * register_fprobe() - Register fprobe to ftrace by pattern.
 * @fp: A fprobe data structure to be registered.
 * @filter: A wildcard pattern of probed symbols.
 * @notfilter: A wildcard pattern of NOT probed symbols.
 *
 * Register @fp to ftrace for enabling the probe on the symbols matched to @filter.
 * If @notfilter is not NULL, the symbols matched the @notfilter are not probed.
 *
 * Return 0 if @fp is registered successfully, -errno if not.
 */
int register_fprobe(struct fprobe *fp, const char *filter, const char *notfilter)
{
	unsigned char *str;
	int ret, len;

	if (!fp || !filter)
		return -EINVAL;

	fprobe_init(fp);

	len = strlen(filter);
	str = kstrdup(filter, GFP_KERNEL);
	ret = ftrace_set_filter(&fp->ops, str, len, 0);
	kfree(str);
	if (ret)
		return ret;

	if (notfilter) {
		len = strlen(notfilter);
		str = kstrdup(notfilter, GFP_KERNEL);
		ret = ftrace_set_notrace(&fp->ops, str, len, 0);
		kfree(str);
		if (ret)
			goto out;
	}

	ret = register_ftrace_function(&fp->ops);
out:
	if (ret)
		ftrace_free_filter(&fp->ops);
	return ret;
}
EXPORT_SYMBOL_GPL(register_fprobe);

/**
 * register_fprobe_ips() - Register fprobe to ftrace by address.
 * @fp: A fprobe data structure to be registered.
 * @addrs: An array of target ftrace location addresses.
 * @num: The number of entries of @addrs.
 *
 * Register @fp to ftrace for enabling the probe on the address given by @addrs.
 * The @addrs must be the addresses of ftrace location address, which may be
 * the symbol address + arch-dependent offset.
 * If you unsure what this mean, please use other registration functions.
 *
 * Return 0 if @fp is registered successfully, -errno if not.
 */
int register_fprobe_ips(struct fprobe *fp, unsigned long *addrs, int num)
{
	int ret;

	if (!fp || !addrs || num <= 0)
		return -EINVAL;

	fprobe_init(fp);

	ret = ftrace_set_filter_ips(&fp->ops, addrs, num, 0, 0);
	if (!ret)
		ret = register_ftrace_function(&fp->ops);

	if (ret)
		ftrace_free_filter(&fp->ops);

	return ret;
}
EXPORT_SYMBOL_GPL(register_fprobe_ips);

/**
 * register_fprobe_syms() - Register fprobe to ftrace by symbols.
 * @fp: A fprobe data structure to be registered.
 * @syms: An array of target symbols.
 * @num: The number of entries of @syms.
 *
 * Register @fp to the symbols given by @syms array. This will be useful if
 * you are sure the symbols exist in the kernel.
 *
 * Return 0 if @fp is registered successfully, -errno if not.
 */
int register_fprobe_syms(struct fprobe *fp, const char **syms, int num)
{
	unsigned long *addrs;
	int ret;

	if (!fp || !syms || num <= 0)
		return -EINVAL;

	addrs = get_ftrace_locations(syms, num);
	if (IS_ERR(addrs))
		return PTR_ERR(addrs);

	ret = register_fprobe_ips(fp, addrs, num);

	kfree(addrs);

	return ret;
}
EXPORT_SYMBOL_GPL(register_fprobe_syms);

/**
 * unregister_fprobe() - Unregister fprobe from ftrace
 * @fp: A fprobe data structure to be unregistered.
 *
 * Unregister fprobe (and remove ftrace hooks from the function entries).
 *
 * Return 0 if @fp is unregistered successfully, -errno if not.
 */
int unregister_fprobe(struct fprobe *fp)
{
	int ret;

	if (!fp || fp->ops.func != fprobe_handler)
		return -EINVAL;

	ret = unregister_ftrace_function(&fp->ops);

	if (!ret)
		ftrace_free_filter(&fp->ops);

	return ret;
}
EXPORT_SYMBOL_GPL(unregister_fprobe);