Add kprobe_addr_iter_start() and kprobe_addr_iter_next() to look up
probes at nondecreasing addresses while retaining the next candidate
in the range iterator.

Provide for_each_addr_kprobe() for callers that advance the address
in the loop body. Return NULL for addresses without a registered probe
and require kprobe_mutex to remain held during iteration.

Signed-off-by: Jiri Olsa <jolsa@kernel.org>
---
 include/linux/kprobes.h | 10 ++++++++++
 kernel/kprobes.c        | 25 +++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/include/linux/kprobes.h b/include/linux/kprobes.h
index bdba893cc0b2..30d0de29ceb5 100644
--- a/include/linux/kprobes.h
+++ b/include/linux/kprobes.h
@@ -394,11 +394,21 @@ struct kprobe *get_kprobe(void *addr);
 struct kprobe *kprobe_iter_start(struct kprobe_iter *iter,
 				 unsigned long start, unsigned long end);
 struct kprobe *kprobe_iter_next(struct kprobe_iter *iter);
+struct kprobe *kprobe_addr_iter_start(struct kprobe_iter *iter,
+				      unsigned long start, unsigned long end);
+struct kprobe *kprobe_addr_iter_next(struct kprobe_iter *iter,
+				     unsigned long addr);
 
 #define for_each_kprobe(iter, p, start, end)			\
 	for ((p) = kprobe_iter_start(&(iter), (start), (end));	\
 	     (p); (p) = kprobe_iter_next(&(iter)))
 
+/* The caller holds kprobe_mutex and initializes and advances addr. */
+#define for_each_addr_kprobe(iter, kp, addr, end)			\
+	for ((kp) = kprobe_addr_iter_start(&(iter), (addr), (end));	\
+	     (addr) <= (end);						\
+	     (kp) = kprobe_addr_iter_next(&(iter), (addr)))
+
 /* kprobe_running() will just return the current_kprobe on this CPU */
 static inline struct kprobe *kprobe_running(void)
 {
diff --git a/kernel/kprobes.c b/kernel/kprobes.c
index bbd51e0fab7e..6655c24bb66f 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -477,6 +477,31 @@ struct kprobe *kprobe_iter_next(struct kprobe_iter *iter)
 	return kp;
 }
 
+/* Address iteration retains the candidate in iter->next until addr passes it. */
+struct kprobe *kprobe_addr_iter_start(struct kprobe_iter *iter,
+				      unsigned long start, unsigned long end)
+{
+	iter->next = kprobe_iter_start(iter, start, end);
+	return kprobe_addr_iter_next(iter, start);
+}
+
+/* Addresses must be nondecreasing; return NULL at addresses without a probe. */
+struct kprobe *kprobe_addr_iter_next(struct kprobe_iter *iter,
+				     unsigned long addr)
+{
+	struct kprobe *kp = iter->next;
+
+	lockdep_assert_held(&kprobe_mutex);
+	while (kp && (unsigned long)kp->addr < addr) {
+		kprobe_iter_next(iter);
+		kp = iter->next;
+	}
+	if (!kp || (unsigned long)kp->addr > iter->end)
+		return NULL;
+
+	return (unsigned long)kp->addr == addr ? kp : NULL;
+}
+
 static int aggr_pre_handler(struct kprobe *p, struct pt_regs *regs);
 
 /* Return true if 'p' is an aggregator */
-- 
2.43.0