bpf: Introduce task open coded iterator kfuncs

This patch adds kfuncs bpf_iter_task_{new,next,destroy} which allow
creation and manipulation of struct bpf_iter_task in open-coded iterator
style. BPF programs can use these kfuncs or through bpf_for_each macro to
iterate all processes in the system.

The API design keep consistent with SEC("iter/task"). bpf_iter_task_new()
accepts a specific task and iterating type which allows:

1. iterating all process in the system (BPF_TASK_ITER_ALL_PROCS)

2. iterating all threads in the system (BPF_TASK_ITER_ALL_THREADS)

3. iterating all threads of a specific task (BPF_TASK_ITER_PROC_THREADS)

Signed-off-by: Chuyi Zhou <zhouchuyi@bytedance.com>
Link: https://lore.kernel.org/r/20231018061746.111364-4-zhouchuyi@bytedance.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index c01441d..c259415 100644 (file)
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -2563,6 +2563,9 @@ BTF_ID_FLAGS(func, bpf_iter_task_vma_destroy, KF_ITER_DESTROY)
 BTF_ID_FLAGS(func, bpf_iter_css_task_new, KF_ITER_NEW | KF_TRUSTED_ARGS)
 BTF_ID_FLAGS(func, bpf_iter_css_task_next, KF_ITER_NEXT | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_iter_css_task_destroy, KF_ITER_DESTROY)
+BTF_ID_FLAGS(func, bpf_iter_task_new, KF_ITER_NEW | KF_TRUSTED_ARGS)
+BTF_ID_FLAGS(func, bpf_iter_task_next, KF_ITER_NEXT | KF_RET_NULL)
+BTF_ID_FLAGS(func, bpf_iter_task_destroy, KF_ITER_DESTROY)
 BTF_ID_FLAGS(func, bpf_dynptr_adjust)
 BTF_ID_FLAGS(func, bpf_dynptr_is_null)
 BTF_ID_FLAGS(func, bpf_dynptr_is_rdonly)

diff --git a/kernel/bpf/task_iter.c b/kernel/bpf/task_iter.c
index e412669..faa1712 100644 (file)
--- a/kernel/bpf/task_iter.c
+++ b/kernel/bpf/task_iter.c
@@ -952,6 +952,96 @@ __bpf_kfunc void bpf_iter_css_task_destroy(struct bpf_iter_css_task *it)
 
 __diag_pop();
 
+struct bpf_iter_task {
+       __u64 __opaque[3];
+} __attribute__((aligned(8)));
+
+struct bpf_iter_task_kern {
+       struct task_struct *task;
+       struct task_struct *pos;
+       unsigned int flags;
+} __attribute__((aligned(8)));
+
+enum {
+       /* all process in the system */
+       BPF_TASK_ITER_ALL_PROCS,
+       /* all threads in the system */
+       BPF_TASK_ITER_ALL_THREADS,
+       /* all threads of a specific process */
+       BPF_TASK_ITER_PROC_THREADS
+};
+
+__diag_push();
+__diag_ignore_all("-Wmissing-prototypes",
+                 "Global functions as their definitions will be in vmlinux BTF");
+
+__bpf_kfunc int bpf_iter_task_new(struct bpf_iter_task *it,
+               struct task_struct *task, unsigned int flags)
+{
+       struct bpf_iter_task_kern *kit = (void *)it;
+
+       BUILD_BUG_ON(sizeof(struct bpf_iter_task_kern) > sizeof(struct bpf_iter_task));
+       BUILD_BUG_ON(__alignof__(struct bpf_iter_task_kern) !=
+                                       __alignof__(struct bpf_iter_task));
+
+       kit->task = kit->pos = NULL;
+       switch (flags) {
+       case BPF_TASK_ITER_ALL_THREADS:
+       case BPF_TASK_ITER_ALL_PROCS:
+       case BPF_TASK_ITER_PROC_THREADS:
+               break;
+       default:
+               return -EINVAL;
+       }
+
+       if (flags == BPF_TASK_ITER_PROC_THREADS)
+               kit->task = task;
+       else
+               kit->task = &init_task;
+       kit->pos = kit->task;
+       kit->flags = flags;
+       return 0;
+}
+
+__bpf_kfunc struct task_struct *bpf_iter_task_next(struct bpf_iter_task *it)
+{
+       struct bpf_iter_task_kern *kit = (void *)it;
+       struct task_struct *pos;
+       unsigned int flags;
+
+       flags = kit->flags;
+       pos = kit->pos;
+
+       if (!pos)
+               return pos;
+
+       if (flags == BPF_TASK_ITER_ALL_PROCS)
+               goto get_next_task;
+
+       kit->pos = next_thread(kit->pos);
+       if (kit->pos == kit->task) {
+               if (flags == BPF_TASK_ITER_PROC_THREADS) {
+                       kit->pos = NULL;
+                       return pos;
+               }
+       } else
+               return pos;
+
+get_next_task:
+       kit->pos = next_task(kit->pos);
+       kit->task = kit->pos;
+       if (kit->pos == &init_task)
+               kit->pos = NULL;
+
+       return pos;
+}
+
+__bpf_kfunc void bpf_iter_task_destroy(struct bpf_iter_task *it)
+{
+}
+
+__diag_pop();
+
 DEFINE_PER_CPU(struct mmap_unlock_irq_work, mmap_unlock_work);
 
 static void do_mmap_read_unlock(struct irq_work *entry)

diff --git a/tools/testing/selftests/bpf/bpf_experimental.h b/tools/testing/selftests/bpf/bpf_experimental.h
index 6792ed2..2f6c747 100644 (file)
--- a/tools/testing/selftests/bpf/bpf_experimental.h
+++ b/tools/testing/selftests/bpf/bpf_experimental.h
@@ -465,5 +465,10 @@ extern int bpf_iter_css_task_new(struct bpf_iter_css_task *it,
 extern struct task_struct *bpf_iter_css_task_next(struct bpf_iter_css_task *it) __weak __ksym;
 extern void bpf_iter_css_task_destroy(struct bpf_iter_css_task *it) __weak __ksym;
 
+struct bpf_iter_task;
+extern int bpf_iter_task_new(struct bpf_iter_task *it,
+               struct task_struct *task, unsigned int flags) __weak __ksym;
+extern struct task_struct *bpf_iter_task_next(struct bpf_iter_task *it) __weak __ksym;
+extern void bpf_iter_task_destroy(struct bpf_iter_task *it) __weak __ksym;
 
 #endif