/*kprobe_example.c*/
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/kallsyms.h>
#include <linux/sched.h>
#include <linux/hardirq.h>
#include <linux/time.h>
#include <linux/hash.h>

#define PID 3702
#define FD 1

#define HASH_MAX 1024

struct t {
	struct task_struct *task;
	s64 start;
};

struct t table[HASH_MAX];
DEFINE_SPINLOCK(table_lock);

/*For each probe you need to allocate a kprobe structure*/
static struct kprobe writep;
static struct kprobe write_outp;

static atomic_t write;
static atomic64_t time;

/*kprobe pre_handler: called just before the probed instruction is executed*/
int
handler_pre(struct kprobe *p, struct pt_regs *regs)
{
	struct timeval tv;
	struct t *t;
	u32 hash;
	int i;

	if (current->parent->pid == PID && regs->rdi == FD) {
		do_gettimeofday(&tv);
		hash = hash_ptr(current, 32);
		spin_lock(&table_lock);
		for (i = 0; i < 10; i++) {
			t = &table[(hash + (1 << i))  % HASH_MAX];
			if (t->task == NULL) {
				t->task = current;
				t->start = timeval_to_ns(&tv);
				break;
			}

		}
		spin_unlock(&table_lock);
		atomic_inc(&write);
	}
        return (0);
}

int
write_out_pre(struct kprobe *p, struct pt_regs *regs)
{
	struct timeval tv;
	struct t *t;
	u32 hash;
	s64 end;
	int i;

	if (current->parent->pid == PID) {
		do_gettimeofday(&tv);
		hash = hash_ptr(current, 32);
		spin_lock(&table_lock);
		for (i = 0; i < 10; i++) {
			t = &table[(hash + (1 << i))  % HASH_MAX];
			if (t->task == current) {
				end = timeval_to_ns(&tv);
				atomic64_add(end - t->start, &time);
				t->task = NULL;
				break;
			}
		}
		spin_unlock(&table_lock);
	}

	return (0);
}

int
install_probe(struct kprobe *k,
    int (*pre_handler)(struct kprobe *, struct pt_regs *),
    void (*post_handler)(struct kprobe *, struct pt_regs *, unsigned long),
    void *addr)
{
	int ret;

	k->pre_handler = pre_handler;
	k->post_handler = post_handler;
	k->addr = addr;
	if (!addr) {
		printk("Couldn't find %s to plant kprobe\n", "do_fork");
		return (-1);
	}
	if ((ret = register_kprobe(k) < 0)) {
		printk("register_kprobe failed, returned %d\n", ret);
		return -1;
	}
	return (0);
}

int init_module(void)
{
	/* cache_alloc_refill entry */
	install_probe(&writep, handler_pre, NULL,
	    (void *)kallsyms_lookup_name("sys_write"));
	install_probe(&write_outp, write_out_pre, NULL,
	    (void *)0xffffffff8028a1b1UL);
	
        printk("kprobe registered\n");
        return 0;
}

void cleanup_module(void)
{
        unregister_kprobe(&writep);
        unregister_kprobe(&write_outp);

        printk("sys_write %d total time %ld (ns)", atomic_read(&write),
            atomic64_read(&time));
        if (atomic_read(&write) != 0)
                printk(" avg %ld (ns)\n", atomic64_read(&time) /
                    atomic_read(&write));
        else
                printk("\n");
}

MODULE_LICENSE("GPL");