[PATCH 05/13] kprobes: Add an iterator for address ranges

[Jiri Olsa <jolsa@kernel.org>]

Add an iterator that walks registered probes in address order within
a specified range. Use the kprobe rbtree to find the first probe at
or after the start address, then advance through the tree.

Require kprobe_mutex to be held while iterating and provide a
for_each_kprobe() helper.

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

diff --git a/include/linux/kprobes.h b/include/linux/kprobes.h
index 6f16cc01a16c..bdba893cc0b2 100644
--- a/include/linux/kprobes.h
+++ b/include/linux/kprobes.h
@@ -95,6 +95,11 @@ struct kprobe {
 	u32 flags;
 };
 
+struct kprobe_iter {
+	struct kprobe *next;
+	unsigned long end;
+};
+
 /* Kprobe status flags */
 #define KPROBE_FLAG_GONE	1 /* breakpoint has already gone */
 #define KPROBE_FLAG_DISABLED	2 /* probe is temporarily disabled */
@@ -386,6 +391,14 @@ static inline void kprobe_ftrace_kill(void) {}
 /* Get the kprobe at this addr (if any) - called with preemption disabled */
 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);
+
+#define for_each_kprobe(iter, p, start, end)			\
+	for ((p) = kprobe_iter_start(&(iter), (start), (end));	\
+	     (p); (p) = kprobe_iter_next(&(iter)))
+
 /* 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 6be1430a4ebe..bbd51e0fab7e 100644
--- a/kernel/kprobes.c
+++ b/kernel/kprobes.c
@@ -435,6 +435,48 @@ static void kprobe_tree_replace(struct kprobe *old, struct kprobe *new)
 	RB_CLEAR_NODE(&old->rb_node);
 }
 
+struct kprobe *kprobe_iter_start(struct kprobe_iter *iter, unsigned long start,
+				 unsigned long end)
+{
+	struct rb_node *node = kprobe_tree.rb_node;
+	struct kprobe *kp, *next = NULL;
+
+	lockdep_assert_held(&kprobe_mutex);
+	iter->next = NULL;
+	iter->end = end;
+	if (start > end)
+		return NULL;
+
+	while (node) {
+		kp = rb_entry(node, struct kprobe, rb_node);
+		if ((unsigned long)kp->addr >= start) {
+			next = kp;
+			node = node->rb_left;
+		} else {
+			node = node->rb_right;
+		}
+	}
+
+	iter->next = next;
+	return kprobe_iter_next(iter);
+}
+
+struct kprobe *kprobe_iter_next(struct kprobe_iter *iter)
+{
+	struct kprobe *kp = iter->next;
+	struct rb_node *node;
+
+	lockdep_assert_held(&kprobe_mutex);
+	if (!kp || (unsigned long)kp->addr > iter->end) {
+		iter->next = NULL;
+		return NULL;
+	}
+
+	node = rb_next(&kp->rb_node);
+	iter->next = node ? rb_entry(node, struct kprobe, rb_node) : NULL;
+	return kp;
+}
+
 static int aggr_pre_handler(struct kprobe *p, struct pt_regs *regs);
 
 /* Return true if 'p' is an aggregator */
-- 
2.43.0