[PATCH bpf v2 2/4] rethook,fprobe,kprobes: Check a failure in the rethook_hook()

[Masami Hiramatsu <mhiramat@xxxxxxxxxx>]

Since there are possible to fail to hook the function return (depends on
archtecutre implememtation), rethook_hook() should return the error
in that case and caller must check it.

Signed-off-by: Masami Hiramatsu <mhiramat@xxxxxxxxxx>
---
 arch/x86/kernel/rethook.c |    4 +++-
 include/linux/rethook.h   |    4 ++--
 kernel/kprobes.c          |    8 +++++---
 kernel/trace/fprobe.c     |    5 ++++-
 kernel/trace/rethook.c    |   12 ++++++++++--
 5 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/arch/x86/kernel/rethook.c b/arch/x86/kernel/rethook.c
index 8a1c0111ae79..c92b4875e3b9 100644
--- a/arch/x86/kernel/rethook.c
+++ b/arch/x86/kernel/rethook.c
@@ -114,7 +114,7 @@ void arch_rethook_fixup_return(struct pt_regs *regs,
 }
 NOKPROBE_SYMBOL(arch_rethook_fixup_return);
 
-void arch_rethook_prepare(struct rethook_node *rh, struct pt_regs *regs, bool mcount)
+int arch_rethook_prepare(struct rethook_node *rh, struct pt_regs *regs, bool mcount)
 {
 	unsigned long *stack = (unsigned long *)regs->sp;
 
@@ -123,5 +123,7 @@ void arch_rethook_prepare(struct rethook_node *rh, struct pt_regs *regs, bool mc
 
 	/* Replace the return addr with trampoline addr */
 	stack[0] = (unsigned long) arch_rethook_trampoline;
+
+	return 0;
 }
 NOKPROBE_SYMBOL(arch_rethook_prepare);
diff --git a/include/linux/rethook.h b/include/linux/rethook.h
index c8ac1e5afcd1..07b9c6663b8e 100644
--- a/include/linux/rethook.h
+++ b/include/linux/rethook.h
@@ -63,12 +63,12 @@ void rethook_free(struct rethook *rh);
 void rethook_add_node(struct rethook *rh, struct rethook_node *node);
 struct rethook_node *rethook_try_get(struct rethook *rh);
 void rethook_recycle(struct rethook_node *node);
-void rethook_hook(struct rethook_node *node, struct pt_regs *regs, bool mcount);
+int rethook_hook(struct rethook_node *node, struct pt_regs *regs, bool mcount);
 unsigned long rethook_find_ret_addr(struct task_struct *tsk, unsigned long frame,
 				    struct llist_node **cur);
 
 /* Arch dependent code must implement arch_* and trampoline code */
-void arch_rethook_prepare(struct rethook_node *node, struct pt_regs *regs, bool mcount);
+int arch_rethook_prepare(struct rethook_node *node, struct pt_regs *regs, bool mcount);
 void arch_rethook_trampoline(void);
 
 /**
diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index dbe57df2e199..7fd7f1195bde 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -2109,10 +2109,12 @@ static int pre_handler_kretprobe(struct kprobe *p, struct pt_regs *regs)
 
 	ri = container_of(rhn, struct kretprobe_instance, node);
 
-	if (rp->entry_handler && rp->entry_handler(ri, regs))
+	if (rp->entry_handler && rp->entry_handler(ri, regs)) {
 		rethook_recycle(rhn);
-	else
-		rethook_hook(rhn, regs, kprobe_ftrace(p));
+	} else if (rethook_hook(rhn, regs, kprobe_ftrace(p)) < 0) {
+		rethook_recycle(rhn);
+		rp->nmissed++;
+	}
 
 	return 0;
 }
diff --git a/kernel/trace/fprobe.c b/kernel/trace/fprobe.c
index 89d9f994ebb0..d3b13294d545 100644
--- a/kernel/trace/fprobe.c
+++ b/kernel/trace/fprobe.c
@@ -48,7 +48,10 @@ static void fprobe_handler(unsigned long ip, unsigned long parent_ip,
 		}
 		fpr = container_of(rh, struct fprobe_rethook_node, node);
 		fpr->entry_ip = ip;
-		rethook_hook(rh, ftrace_get_regs(fregs), true);
+		if (rethook_hook(rh, ftrace_get_regs(fregs), true) < 0) {
+			rethook_recycle(rh);
+			fp->nmissed++;
+		}
 	}
 
 out:
diff --git a/kernel/trace/rethook.c b/kernel/trace/rethook.c
index b56833700d23..e7db83438e45 100644
--- a/kernel/trace/rethook.c
+++ b/kernel/trace/rethook.c
@@ -174,11 +174,19 @@ NOKPROBE_SYMBOL(rethook_try_get);
  * from ftrace (mcount) callback, @mcount must be set true. If this is called
  * from the real function entry (e.g. kprobes) @mcount must be set false.
  * This is because the way to hook the function return depends on the context.
+ * This returns 0 if succeeded to hook the function return, or -errno if
+ * failed.
  */
-void rethook_hook(struct rethook_node *node, struct pt_regs *regs, bool mcount)
+int rethook_hook(struct rethook_node *node, struct pt_regs *regs, bool mcount)
 {
-	arch_rethook_prepare(node, regs, mcount);
+	int ret;
+
+	ret = arch_rethook_prepare(node, regs, mcount);
+	if (ret < 0)
+		return ret;
+
 	__llist_add(&node->llist, &current->rethooks);
+	return 0;
 }
 NOKPROBE_SYMBOL(rethook_hook);