4 files changed, 19 insertions, 15 deletions

diff --git a/kprobe/dbi_kprobes.c b/kprobe/dbi_kprobes.c
index 964eb0d0..f389e595 100644
--- a/kprobe/dbi_kprobes.c
+++ b/kprobe/dbi_kprobes.c
@@ -638,7 +638,7 @@ static int pre_handler_kretprobe(struct kprobe *p, struct pt_regs *regs)
 		ri->task = current;
 
 		if (rp->entry_handler) {
-			rp->entry_handler(ri, regs, ri->rp->priv_arg);
+			rp->entry_handler(ri, regs);
 		}
 
 		arch_prepare_kretprobe(ri, regs);
@@ -703,7 +703,7 @@ int trampoline_probe_handler(struct kprobe *p, struct pt_regs *regs)
 		if (ri->rp && ri->rp->handler) {
 			__get_cpu_var(current_kprobe) = &ri->rp->kp;
 			get_kprobe_ctlblk()->kprobe_status = KPROBE_HIT_ACTIVE;
-			ri->rp->handler(ri, regs, ri->rp->priv_arg);
+			ri->rp->handler(ri, regs);
 			__get_cpu_var(current_kprobe) = NULL;
 		}
 
diff --git a/kprobe/dbi_kprobes.h b/kprobe/dbi_kprobes.h
index 7622740b..b0beec91 100644
--- a/kprobe/dbi_kprobes.h
+++ b/kprobe/dbi_kprobes.h
@@ -92,7 +92,7 @@ typedef int (*kprobe_pre_handler_t) (struct kprobe *, struct pt_regs *);
 typedef int (*kprobe_break_handler_t) (struct kprobe *, struct pt_regs *);
 typedef void (*kprobe_post_handler_t) (struct kprobe *, struct pt_regs *, unsigned long flags);
 typedef int (*kprobe_fault_handler_t) (struct kprobe *, struct pt_regs *, int trapnr);
-typedef int (*kretprobe_handler_t) (struct kretprobe_instance *, struct pt_regs *, void *);
+typedef int (*kretprobe_handler_t) (struct kretprobe_instance *, struct pt_regs *);
 
 struct kprobe
 {
@@ -195,7 +195,6 @@ struct kretprobe
 	struct kprobe kp;
 	kretprobe_handler_t handler;
 	kretprobe_handler_t entry_handler;
-	void *priv_arg;
 	int maxactive;
 	int nmissed;
 	size_t data_size;
diff --git a/ks_features/ks_features.c b/ks_features/ks_features.c
index 41123782..c061e08e 100644
--- a/ks_features/ks_features.c
+++ b/ks_features/ks_features.c
@@ -90,23 +90,30 @@ static void dec_counter(size_t id)
 }
 
 /* ========================= HANDLERS ========================= */
-static int entry_handler(struct kretprobe_instance *ri, struct pt_regs *regs, void *priv_arg)
+static int entry_handler(struct kretprobe_instance *ri, struct pt_regs *regs)
 {
-	struct ks_probe *ksp = (struct ks_probe *)priv_arg;
-	const char *fmt = ksp->args;
-	int sub_type = ksp->sub_type;
+	struct kretprobe *rp = ri->rp;
 
-	entry_event(fmt, regs, PT_KS, sub_type);
+	if (rp) {
+		struct ks_probe *ksp = container_of(rp, struct ks_probe, rp);
+		const char *fmt = ksp->args;
+		int sub_type = ksp->sub_type;
+
+		entry_event(fmt, regs, PT_KS, sub_type);
+	}
 
 	return 0;
 }
 
-static int ret_handler(struct kretprobe_instance *ri, struct pt_regs *regs, void *priv_arg)
+static int ret_handler(struct kretprobe_instance *ri, struct pt_regs *regs)
 {
-	struct ks_probe *ksp = (struct ks_probe *)priv_arg;
-	unsigned long func_addr = ri->rp ? ri->rp->kp.addr : 0;
+	struct kretprobe *rp = ri->rp;
 
-	exit_event(regs, func_addr);
+	if (rp) {
+		unsigned long func_addr = rp->kp.addr;
+
+		exit_event(regs, func_addr);
+	}
 
 	return 0;
 }
@@ -198,7 +205,6 @@ static int register_syscall(size_t id)
 
 	ksp[id].rp.entry_handler = entry_handler;
 	ksp[id].rp.handler = ret_handler;
-	ksp[id].rp.priv_arg = &ksp[id];
 
 	ret = dbi_register_kretprobe(&ksp[id].rp);
 
diff --git a/ks_manager/ks_manager.c b/ks_manager/ks_manager.c
index 400b26da..f9e7c0fb 100644
--- a/ks_manager/ks_manager.c
+++ b/ks_manager/ks_manager.c
@@ -44,7 +44,6 @@ static struct probe *create_probe(unsigned long addr, void *pre_handler,
 	p->p.jp.pre_entry = pre_handler;
 	p->p.jp.entry = jp_handler;
 	p->p.rp.handler = rp_handler;
-	p->p.jp.priv_arg = p->p.rp.priv_arg = (void *)&p->p;
 	INIT_HLIST_NODE(&p->hlist);
 
 	return p;