On Wed, Oct 11, 2023 at 5:09 AM Chuyi Zhou <zhouchuyi@xxxxxxxxxxxxx> wrote:
>
> 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@xxxxxxxxxxxxx>
> ---
> kernel/bpf/helpers.c | 3 +
> kernel/bpf/task_iter.c | 82 +++++++++++++++++++
> .../testing/selftests/bpf/bpf_experimental.h | 5 ++
> 3 files changed, 90 insertions(+)
>
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index cb24c4a916df..690763751f6e 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -2555,6 +2555,9 @@ BTF_ID_FLAGS(func, bpf_iter_num_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 2cfcb4dd8a37..caeddad3d2f1 100644
> --- a/kernel/bpf/task_iter.c
> +++ b/kernel/bpf/task_iter.c
> @@ -856,6 +856,88 @@ __bpf_kfunc void bpf_iter_css_task_destroy(struct bpf_iter_css_task *it)
> bpf_mem_free(&bpf_global_ma, kit->css_it);
> }
>
> +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 {
> + BPF_TASK_ITER_ALL_PROCS,
> + BPF_TASK_ITER_ALL_THREADS,
> + BPF_TASK_ITER_PROC_THREADS
> +};
> +
> +__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)
> + goto out;
> +
> + 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;
> + goto out;
> + }
> + } else
> + goto out;
nit: this should have {} around it to match the other if branch
but actually, why goto out instead of return pos? same above, return
pos instead of goto out?
> +
> +get_next_task:
> + kit->pos = next_task(kit->pos);
> + kit->task = kit->pos;
> + if (kit->pos == &init_task)
> + kit->pos = NULL;
> +
> +out:
> + return pos;
> +}
> +
> +__bpf_kfunc void bpf_iter_task_destroy(struct bpf_iter_task *it)
> +{
> +}
> +
> 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 8b53537e0f27..1ec82997cce7 100644
> --- a/tools/testing/selftests/bpf/bpf_experimental.h
> +++ b/tools/testing/selftests/bpf/bpf_experimental.h
> @@ -457,5 +457,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
> --
> 2.20.1
>
