Add open coded BPF iterators for wakeup_sources, which opens up more
options for BPF programs that need to traverse through wakeup_sources.

Signed-off-by: Samuel Wu <wusamuel@google.com>
---
 kernel/bpf/helpers.c            |  3 +++
 kernel/bpf/wakeup_source_iter.c | 34 +++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 9eaa4185e0a7..ca34d7614c3a 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -4518,6 +4518,9 @@ BTF_ID_FLAGS(func, bpf_iter_dmabuf_new, KF_ITER_NEW | KF_SLEEPABLE)
 BTF_ID_FLAGS(func, bpf_iter_dmabuf_next, KF_ITER_NEXT | KF_RET_NULL | KF_SLEEPABLE)
 BTF_ID_FLAGS(func, bpf_iter_dmabuf_destroy, KF_ITER_DESTROY | KF_SLEEPABLE)
 #endif
+BTF_ID_FLAGS(func, bpf_iter_wakeup_source_new, KF_ITER_NEW | KF_SLEEPABLE)
+BTF_ID_FLAGS(func, bpf_iter_wakeup_source_next, KF_ITER_NEXT | KF_RET_NULL | KF_SLEEPABLE)
+BTF_ID_FLAGS(func, bpf_iter_wakeup_source_destroy, KF_ITER_DESTROY | KF_SLEEPABLE)
 BTF_ID_FLAGS(func, __bpf_trap)
 BTF_ID_FLAGS(func, bpf_strcmp);
 BTF_ID_FLAGS(func, bpf_strcasecmp);
diff --git a/kernel/bpf/wakeup_source_iter.c b/kernel/bpf/wakeup_source_iter.c
index ab83d212a1f9..149baecfe436 100644
--- a/kernel/bpf/wakeup_source_iter.c
+++ b/kernel/bpf/wakeup_source_iter.c
@@ -90,6 +90,40 @@ static struct bpf_iter_reg bpf_wakeup_source_reg_info = {
 	.seq_info		= &wakeup_source_iter_seq_info,
 };
 
+struct bpf_iter_wakeup_source {
+	struct wakeup_source *ws;
+	int srcuidx;
+};
+
+__bpf_kfunc_start_defs();
+
+__bpf_kfunc int bpf_iter_wakeup_source_new(struct bpf_iter_wakeup_source *it)
+{
+	it->srcuidx = wakeup_sources_read_lock();
+	it->ws = wakeup_sources_walk_start();
+
+	return 0;
+}
+
+__bpf_kfunc struct wakeup_source *bpf_iter_wakeup_source_next(struct bpf_iter_wakeup_source *it)
+{
+	struct wakeup_source *prev = it->ws;
+
+	if (!prev)
+		return NULL;
+
+	it->ws = wakeup_sources_walk_next(it->ws);
+
+	return prev;
+}
+
+__bpf_kfunc void bpf_iter_wakeup_source_destroy(struct bpf_iter_wakeup_source *it)
+{
+	wakeup_sources_read_unlock(it->srcuidx);
+}
+
+__bpf_kfunc_end_defs();
+
 DEFINE_BPF_ITER_FUNC(wakeup_source, struct bpf_iter_meta *meta,
 		     struct wakeup_source *wakeup_source)
 BTF_ID_LIST_SINGLE(bpf_wakeup_source_btf_id, struct, wakeup_source)
-- 
2.52.0.457.g6b5491de43-goog

> diff --git a/kernel/bpf/wakeup_source_iter.c b/kernel/bpf/wakeup_source_iter.c
> index ab83d212a1f9..149baecfe436 100644
> --- a/kernel/bpf/wakeup_source_iter.c
> +++ b/kernel/bpf/wakeup_source_iter.c

[ ... ]

> +__bpf_kfunc int bpf_iter_wakeup_source_new(struct bpf_iter_wakeup_source *it)
> +{
> +	it->srcuidx = wakeup_sources_read_lock();
> +	it->ws = wakeup_sources_walk_start();
> +
> +	return 0;
> +}
> +
> +__bpf_kfunc struct wakeup_source *bpf_iter_wakeup_source_next(struct bpf_iter_wakeup_source *it)
> +{
> +	struct wakeup_source *prev = it->ws;
> +
> +	if (!prev)
> +		return NULL;
> +
> +	it->ws = wakeup_sources_walk_next(it->ws);
> +
> +	return prev;
> +}

When the wakeup_sources list is empty, wakeup_sources_walk_start() returns
a bogus pointer via list_entry_rcu()/container_of() on the list head itself,
rather than NULL.  The check "if (!prev)" at line 112 will not catch this
bogus pointer since it is non-NULL.

Does this return a bogus pointer to the BPF program when the wakeup_sources
list is empty?  If so, the BPF program could then access invalid memory when
it tries to read fields of the returned wakeup_source.

For reference, wakeup_sources_walk_start() in drivers/base/power/wakeup.c:

    struct wakeup_source *wakeup_sources_walk_start(void)
    {
        struct list_head *ws_head = &wakeup_sources;
        return list_entry_rcu(ws_head->next, struct wakeup_source, entry);
    }

When wakeup_sources is empty, ws_head->next == ws_head, so this returns
container_of(ws_head, struct wakeup_source, entry) which points to memory
before the wakeup_sources variable rather than a valid wakeup_source struct.

Other BPF open-coded iterators like bpf_iter_dmabuf handle this differently
by initializing to NULL and getting the first element lazily in _next().
Would a similar approach work here?


---
AI reviewed your patch. Please fix the bug or email reply why it's not a bug.
See: https://github.com/kernel-patches/vmtest/blob/master/ci/claude/README.md

CI run summary: https://github.com/kernel-patches/bpf/actions/runs/20835022582