[RFC PATCH 5/9] arm64/kprobes: Invoke pre/post handlers inside instrumentation

[Hongyan Xia]

From: Hongyan Xia <hongyan.xia at transsion.com>

The kprobe pre_handler() and post_handler() are arbitrary instrumentable
code and can themselves trace, fault, or hit other kprobes. They are
the only parts of the kprobe debug exception flow that legitimately
need instrumentation.

Signed-off-by: Hongyan Xia <hongyan.xia at transsion.com>
---
 arch/arm64/include/asm/kprobes.h   |  9 +++------
 arch/arm64/kernel/probes/kprobes.c | 22 ++++++++++++++++++----
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/arch/arm64/include/asm/kprobes.h b/arch/arm64/include/asm/kprobes.h
index 35ce2c94040e..a694f7d34f45 100644
--- a/arch/arm64/include/asm/kprobes.h
+++ b/arch/arm64/include/asm/kprobes.h
@@ -48,11 +48,8 @@ void __kprobes *trampoline_probe_handler(struct pt_regs *regs);
 
 #endif /* CONFIG_KPROBES */
 
-int __kprobes kprobe_brk_handler(struct pt_regs *regs,
-				 unsigned long esr);
-int __kprobes kprobe_ss_brk_handler(struct pt_regs *regs,
-				 unsigned long esr);
-int __kprobes kretprobe_brk_handler(struct pt_regs *regs,
-				 unsigned long esr);
+int noinstr kprobe_brk_handler(struct pt_regs *regs, unsigned long esr);
+int noinstr kprobe_ss_brk_handler(struct pt_regs *regs, unsigned long esr);
+int noinstr kretprobe_brk_handler(struct pt_regs *regs, unsigned long esr);
 
 #endif /* _ARM_KPROBES_H */
diff --git a/arch/arm64/kernel/probes/kprobes.c b/arch/arm64/kernel/probes/kprobes.c
index 1b12341b2af3..e9fa66fa4217 100644
--- a/arch/arm64/kernel/probes/kprobes.c
+++ b/arch/arm64/kernel/probes/kprobes.c
@@ -301,8 +301,11 @@ post_kprobe_handler(struct kprobe *cur, struct kprobe_ctlblk *kcb, struct pt_reg
 	}
 	/* call post handler */
 	kcb->kprobe_status = KPROBE_HIT_SSDONE;
+
+	instrumentation_begin();
 	if (cur->post_handler)
 		cur->post_handler(cur, regs, 0);
+	instrumentation_end();
 
 	reset_current_kprobe();
 }
@@ -359,24 +362,28 @@ int __kprobes kprobe_fault_handler(struct pt_regs *regs, unsigned int fsr)
 	return 0;
 }
 
-int __kprobes
+int noinstr
 kprobe_brk_handler(struct pt_regs *regs, unsigned long esr)
 {
 	struct kprobe *p, *cur_kprobe;
 	struct kprobe_ctlblk *kcb;
 	unsigned long addr = instruction_pointer(regs);
+	bool handled;
 
 	kcb = get_kprobe_ctlblk();
 	cur_kprobe = kprobe_running();
 
+	instrumentation_begin();
 	p = get_kprobe((kprobe_opcode_t *) addr);
 	if (WARN_ON_ONCE(!p)) {
+		instrumentation_end();
 		/*
 		 * Something went wrong. This BRK used an immediate reserved
 		 * for kprobes, but we couldn't find any corresponding probe.
 		 */
 		return DBG_HOOK_ERROR;
 	}
+	instrumentation_end();
 
 	if (cur_kprobe) {
 		/* Hit a kprobe inside another kprobe */
@@ -387,6 +394,10 @@ kprobe_brk_handler(struct pt_regs *regs, unsigned long esr)
 		set_current_kprobe(p);
 		kcb->kprobe_status = KPROBE_HIT_ACTIVE;
 
+		instrumentation_begin();
+		handled = p->pre_handler && p->pre_handler(p, regs);
+		instrumentation_end();
+
 		/*
 		 * If we have no pre-handler or it returned 0, we
 		 * continue with normal processing.  If we have a
@@ -394,7 +405,7 @@ kprobe_brk_handler(struct pt_regs *regs, unsigned long esr)
 		 * modify the execution path and not need to single-step
 		 * Let's just reset current kprobe and exit.
 		 */
-		if (!p->pre_handler || !p->pre_handler(p, regs))
+		if (!handled)
 			setup_singlestep(p, regs, kcb, 0);
 		else
 			reset_current_kprobe();
@@ -403,7 +414,7 @@ kprobe_brk_handler(struct pt_regs *regs, unsigned long esr)
 	return DBG_HOOK_HANDLED;
 }
 
-int __kprobes
+int noinstr
 kprobe_ss_brk_handler(struct pt_regs *regs, unsigned long esr)
 {
 	struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
@@ -422,13 +433,16 @@ kprobe_ss_brk_handler(struct pt_regs *regs, unsigned long esr)
 	return DBG_HOOK_ERROR;
 }
 
-int __kprobes
+int noinstr
 kretprobe_brk_handler(struct pt_regs *regs, unsigned long esr)
 {
 	if (regs->pc != (unsigned long)__kretprobe_trampoline)
 		return DBG_HOOK_ERROR;
 
+	instrumentation_begin();
 	regs->pc = kretprobe_trampoline_handler(regs, (void *)regs->regs[29]);
+	instrumentation_end();
+
 	return DBG_HOOK_HANDLED;
 }
 
-- 
2.47.3