can_optimize() already iterates over kprobes in the function. Pass the
iterator's kprobe to recover_probed_instruction() so it doesn't need to
look up the kprobe again for each address. Keep the lookup for callers
that pass NULL.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 arch/x86/kernel/kprobes/common.h |  4 ++--
 arch/x86/kernel/kprobes/core.c   | 16 ++++++++--------
 arch/x86/kernel/kprobes/opt.c    |  9 ++++++---
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/arch/x86/kernel/kprobes/common.h b/arch/x86/kernel/kprobes/common.h
index e772276f5aa9..062746e20c35 100644
--- a/arch/x86/kernel/kprobes/common.h
+++ b/arch/x86/kernel/kprobes/common.h
@@ -79,9 +79,9 @@
 
 /* Ensure if the instruction can be boostable */
 extern bool can_boost(struct insn *insn, void *orig_addr);
-/* Recover instruction if given address is probed */
+/* Recover instruction; @kp is NULL if the caller needs a kprobe lookup. */
 extern unsigned long recover_probed_instruction(kprobe_opcode_t *buf,
-					 unsigned long addr);
+					 unsigned long addr, struct kprobe **kp);
 /*
  * Copy an instruction and adjust the displacement if the instruction
  * uses the %rip-relative addressing mode.
diff --git a/arch/x86/kernel/kprobes/core.c b/arch/x86/kernel/kprobes/core.c
index 133ff20caccd..6ad8b89a9ba1 100644
--- a/arch/x86/kernel/kprobes/core.c
+++ b/arch/x86/kernel/kprobes/core.c
@@ -201,12 +201,11 @@ bool can_boost(struct insn *insn, void *addr)
 }
 
 static unsigned long
-__recover_probed_insn(kprobe_opcode_t *buf, unsigned long addr)
+__recover_probed_insn(kprobe_opcode_t *buf, unsigned long addr, struct kprobe **kpp)
 {
-	struct kprobe *kp;
+	struct kprobe *kp = kpp ? *kpp : get_kprobe((void *)addr);
 	bool faddr;
 
-	kp = get_kprobe((void *)addr);
 	faddr = ftrace_location(addr) == addr;
 	/*
 	 * Use the current code if it is not modified by Kprobe
@@ -251,7 +250,8 @@ __recover_probed_insn(kprobe_opcode_t *buf, unsigned long addr)
  * for preventing to release referencing kprobes.
  * Returns zero if the instruction can not get recovered (or access failed).
  */
-unsigned long recover_probed_instruction(kprobe_opcode_t *buf, unsigned long addr)
+unsigned long recover_probed_instruction(kprobe_opcode_t *buf, unsigned long addr,
+					 struct kprobe **kp)
 {
 	unsigned long __addr;
 
@@ -259,7 +259,7 @@ unsigned long recover_probed_instruction(kprobe_opcode_t *buf, unsigned long add
 	if (__addr != addr)
 		return __addr;
 
-	return __recover_probed_insn(buf, addr);
+	return __recover_probed_insn(buf, addr, kp);
 }
 
 /* Check if insn is INT or UD */
@@ -304,7 +304,7 @@ static bool can_probe(unsigned long paddr)
 		 * relative-jump. Since the relative-jump itself is
 		 * normally used, we just go through if there is no kprobe.
 		 */
-		__addr = recover_probed_instruction(buf, addr);
+		__addr = recover_probed_instruction(buf, addr, NULL);
 		if (!__addr)
 			return false;
 
@@ -327,7 +327,7 @@ static bool can_probe(unsigned long paddr)
 	if (addr != paddr)
 		return false;
 
-	__addr = recover_probed_instruction(buf, addr);
+	__addr = recover_probed_instruction(buf, addr, NULL);
 	if (!__addr)
 		return false;
 
@@ -394,7 +394,7 @@ kprobe_opcode_t *arch_adjust_kprobe_addr(unsigned long addr, unsigned long offse
 int __copy_instruction(u8 *dest, u8 *src, u8 *real, struct insn *insn)
 {
 	kprobe_opcode_t buf[MAX_INSN_SIZE];
-	unsigned long recovered_insn = recover_probed_instruction(buf, (unsigned long)src);
+	unsigned long recovered_insn = recover_probed_instruction(buf, (unsigned long)src, NULL);
 	int ret;
 
 	if (!recovered_insn || !insn)
diff --git a/arch/x86/kernel/kprobes/opt.c b/arch/x86/kernel/kprobes/opt.c
index cdf69c939631..9f0af24155ba 100644
--- a/arch/x86/kernel/kprobes/opt.c
+++ b/arch/x86/kernel/kprobes/opt.c
@@ -257,7 +257,9 @@ static int insn_jump_into_range(struct insn *insn, unsigned long start, int len)
 /* Decode whole function to ensure any instructions don't jump into target */
 static int can_optimize(unsigned long paddr)
 {
-	unsigned long addr, size = 0, offset = 0;
+	unsigned long addr, end, size = 0, offset = 0;
+	struct kprobe_iter iter;
+	struct kprobe *kp;
 	struct insn insn;
 	kprobe_opcode_t buf[MAX_INSN_SIZE];
 
@@ -279,7 +281,8 @@ static int can_optimize(unsigned long paddr)
 
 	/* Decode instructions */
 	addr = paddr - offset;
-	while (addr < paddr - offset + size) { /* Decode until function end */
+	end = paddr - offset + size - 1;
+	for_each_addr_kprobe(iter, kp, addr, end) { /* Decode until function end */
 		unsigned long recovered_insn;
 		int ret;
 
@@ -289,7 +292,7 @@ static int can_optimize(unsigned long paddr)
 			 * we can't optimize kprobe in this function.
 			 */
 			return 0;
-		recovered_insn = recover_probed_instruction(buf, addr);
+		recovered_insn = recover_probed_instruction(buf, addr, &kp);
 		if (!recovered_insn)
 			return 0;
 
-- 
2.43.0