/* uretprobe_example.c, adjusted to add syscall tracing */
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/uprobes.h>
#include <linux/ptrace.h>
#include <linux/utrace.h>
#include <linux/err.h>

/*
 * Usage:
 * insmod argetek1.ko pid=<pid> func=<addr>
 * where <pid> identifies the probed process, and <addr> is the virtual
 * address of the probed function.
 *
 * Traces calls and returns from the specified function, plus calls
 * and returns of system calls made by the function and its callees.
 */

static int pid = 0;
module_param(pid, int, 0);
MODULE_PARM_DESC(pid, "pid");

static long func = 0;
module_param(func, long, 0);
MODULE_PARM_DESC(func, "func");

static struct uprobe usp;
static struct uretprobe rp;
struct task_struct *tracee;
struct utrace_attached_engine *engine;

static u32 report_death(struct utrace_attached_engine *e,
	struct task_struct *tsk)
{
	BUG_ON(tsk != tracee);
	engine = NULL;
	return (UTRACE_ACTION_NEWSTATE | UTRACE_ACTION_DETACH);
}

static u32 report_syscall_entry(struct utrace_attached_engine *engine,
	struct task_struct *tsk, struct pt_regs *regs)
{
	printk("%d: syscall #%ld ", current->pid, regs->orig_eax);
	return UTRACE_ACTION_RESUME;
}

static u32 report_syscall_exit(struct utrace_attached_engine *engine,
	struct task_struct *tsk, struct pt_regs *regs)
{
	printk("returns %#lx\n", regs_return_value(regs));
	return UTRACE_ACTION_RESUME;
}

static const struct utrace_engine_ops utrace_ops =
{
	.report_syscall_entry = report_syscall_entry,
	.report_syscall_exit = report_syscall_exit,
	.report_death = report_death
};

static int attach_engine(void)
{
	unsigned long flags;
	tracee = current;
	engine = utrace_attach(tracee, UTRACE_ATTACH_CREATE, &utrace_ops, NULL);
	if (IS_ERR(engine)) {
		long err = PTR_ERR(engine);
		printk(KERN_ERR "utrace_attach() failed; returned %ld\n", err);
		return (int) err;
	}
	flags = engine->flags;
	flags |= UTRACE_EVENT(DEATH);
	utrace_set_flags(tracee, engine, flags);
	return 0;
}

static void enable_syscall_trace(void)
{
	unsigned long flags;
	if (!engine)
		(void) attach_engine();
	if (IS_ERR(engine))
		return;
	flags = engine->flags;
	flags |= UTRACE_EVENT(SYSCALL_ENTRY) | UTRACE_EVENT(SYSCALL_EXIT);
	utrace_set_flags(current, engine, flags);
}

static void disable_syscall_trace(void)
{
	unsigned long flags;
	if (!engine || IS_ERR(engine))
		return;
	flags = engine->flags;
	flags &= ~(UTRACE_EVENT(SYSCALL_ENTRY) | UTRACE_EVENT(SYSCALL_EXIT));
	utrace_set_flags(current, engine, flags);
}

static void uprobe_handler(struct uprobe *u, struct pt_regs *regs)
{
	printk(KERN_INFO "Function at %#lx called\n", u->vaddr);
	enable_syscall_trace();
}

static void uretprobe_handler(struct uretprobe_instance *ri,
	struct pt_regs *regs)
{
	printk(KERN_INFO "Function at %#lx returns %#lx\n\n",
		ri->rp->u.vaddr, regs_return_value(regs));
	disable_syscall_trace();
}

int __init init_module(void)
{
	int ret;

	/* Register the entry probe. */
	usp.pid = pid;
	usp.vaddr = func;
	usp.handler = uprobe_handler;
	printk(KERN_INFO "Registering uprobe on pid %d, vaddr %#lx\n",
		usp.pid, usp.vaddr);
	ret = register_uprobe(&usp);
	if (ret != 0) {
		printk(KERN_ERR "register_uprobe() failed, returned %d\n", ret);
		return -1;
	}

	/* Register the return probe. */
	rp.u.pid = pid;
	rp.u.vaddr = func;
	rp.handler = uretprobe_handler;
	printk(KERN_INFO "Registering return probe on pid %d, vaddr %#lx\n",
		rp.u.pid, rp.u.vaddr);
	ret = register_uretprobe(&rp);
	if (ret != 0) {
		printk(KERN_ERR "register_uretprobe() failed, returned %d\n",
			ret);
		unregister_uprobe(&usp);
		return -1;
	}
	return 0;
}

void __exit cleanup_module(void)
{
	printk(KERN_INFO "Unregistering probes on pid %d, vaddr %#lx\n",
		usp.pid, usp.vaddr);
	unregister_uprobe(&usp);
	unregister_uretprobe(&rp);
	if (engine && !IS_ERR(engine))
		utrace_detach(tracee, engine);
}
MODULE_LICENSE("GPL");