Re: [PATCH bpf-next v2 2/6] bpf: Introduce css_task open-coded iterator kfuncs

[Andrii Nakryiko]

On Tue, Sep 12, 2023 at 12:02 AM Chuyi Zhou <zhouchuyi@xxxxxxxxxxxxx> wrote:
>
> This patch adds kfuncs bpf_iter_css_task_{new,next,destroy} which allow
> creation and manipulation of struct bpf_iter_css_task in open-coded
> iterator style. These kfuncs actually wrapps css_task_iter_{start,next,
> end}. BPF programs can use these kfuncs through bpf_for_each macro for
> iteration of all tasks under a css.
>
> css_task_iter_*() would try to get the global spin-lock *css_set_lock*, so
> the bpf side has to be careful in where it allows to use this iter.
> Currently we only allow it in bpf_lsm and bpf iter-s.
>
> Signed-off-by: Chuyi Zhou <zhouchuyi@xxxxxxxxxxxxx>
> ---
>  include/uapi/linux/bpf.h       |  4 +++
>  kernel/bpf/helpers.c           |  3 +++
>  kernel/bpf/task_iter.c         | 48 ++++++++++++++++++++++++++++++++++
>  kernel/bpf/verifier.c          | 23 ++++++++++++++++
>  tools/include/uapi/linux/bpf.h |  4 +++
>  tools/lib/bpf/bpf_helpers.h    |  7 +++++
>  6 files changed, 89 insertions(+)
>
> diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
> index 73b155e52204..de02c0971428 100644
> --- a/include/uapi/linux/bpf.h
> +++ b/include/uapi/linux/bpf.h
> @@ -7318,4 +7318,8 @@ struct bpf_iter_num {
>         __u64 __opaque[1];
>  } __attribute__((aligned(8)));
>
> +struct bpf_iter_css_task {
> +       __u64 __opaque[1];
> +} __attribute__((aligned(8)));
> +
>  #endif /* _UAPI__LINUX_BPF_H__ */
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index b0a9834f1051..d6a16becfbb9 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -2504,6 +2504,9 @@ BTF_ID_FLAGS(func, bpf_dynptr_slice_rdwr, KF_RET_NULL)
>  BTF_ID_FLAGS(func, bpf_iter_num_new, KF_ITER_NEW)
>  BTF_ID_FLAGS(func, bpf_iter_num_next, KF_ITER_NEXT | KF_RET_NULL)
>  BTF_ID_FLAGS(func, bpf_iter_num_destroy, KF_ITER_DESTROY)
> +BTF_ID_FLAGS(func, bpf_iter_css_task_new, KF_ITER_NEW)
> +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_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 7473068ed313..d8539cc05ffd 100644
> --- a/kernel/bpf/task_iter.c
> +++ b/kernel/bpf/task_iter.c
> @@ -803,6 +803,54 @@ const struct bpf_func_proto bpf_find_vma_proto = {
>         .arg5_type      = ARG_ANYTHING,
>  };
>
> +struct bpf_iter_css_task_kern {
> +       struct css_task_iter *css_it;
> +} __attribute__((aligned(8)));
> +
> +__bpf_kfunc int bpf_iter_css_task_new(struct bpf_iter_css_task *it,
> +               struct cgroup_subsys_state *css, unsigned int flags)
> +{
> +       struct bpf_iter_css_task_kern *kit = (void *)it;
> +
> +       BUILD_BUG_ON(sizeof(struct bpf_iter_css_task_kern) != sizeof(struct bpf_iter_css_task));
> +       BUILD_BUG_ON(__alignof__(struct bpf_iter_css_task_kern) !=
> +                                       __alignof__(struct bpf_iter_css_task));
> +
> +       switch (flags) {
> +       case CSS_TASK_ITER_PROCS | CSS_TASK_ITER_THREADED:
> +       case CSS_TASK_ITER_PROCS:
> +       case 0:
> +               break;
> +       default:
> +               return -EINVAL;
> +       }
> +
> +       kit->css_it = kzalloc(sizeof(struct css_task_iter), GFP_KERNEL);

Dave used bpf_mem_alloc() inside his iterator, any reason to not use it here?


> +       if (!kit->css_it)
> +               return -ENOMEM;
> +       css_task_iter_start(css, flags, kit->css_it);
> +       return 0;
> +}
> +
> +__bpf_kfunc struct task_struct *bpf_iter_css_task_next(struct bpf_iter_css_task *it)
> +{
> +       struct bpf_iter_css_task_kern *kit = (void *)it;
> +
> +       if (!kit->css_it)
> +               return NULL;
> +       return css_task_iter_next(kit->css_it);
> +}
> +
> +__bpf_kfunc void bpf_iter_css_task_destroy(struct bpf_iter_css_task *it)
> +{
> +       struct bpf_iter_css_task_kern *kit = (void *)it;
> +
> +       if (!kit->css_it)
> +               return;
> +       css_task_iter_end(kit->css_it);
> +       kfree(kit->css_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/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
> index dbba2b806017..2367483bf4c2 100644
> --- a/kernel/bpf/verifier.c
> +++ b/kernel/bpf/verifier.c
> @@ -10332,6 +10332,7 @@ enum special_kfunc_type {
>         KF_bpf_dynptr_clone,
>         KF_bpf_percpu_obj_new_impl,
>         KF_bpf_percpu_obj_drop_impl,
> +       KF_bpf_iter_css_task_new,
>  };
>
>  BTF_SET_START(special_kfunc_set)
> @@ -10354,6 +10355,7 @@ BTF_ID(func, bpf_dynptr_slice_rdwr)
>  BTF_ID(func, bpf_dynptr_clone)
>  BTF_ID(func, bpf_percpu_obj_new_impl)
>  BTF_ID(func, bpf_percpu_obj_drop_impl)
> +BTF_ID(func, bpf_iter_css_task_new)
>  BTF_SET_END(special_kfunc_set)
>
>  BTF_ID_LIST(special_kfunc_list)
> @@ -10378,6 +10380,7 @@ BTF_ID(func, bpf_dynptr_slice_rdwr)
>  BTF_ID(func, bpf_dynptr_clone)
>  BTF_ID(func, bpf_percpu_obj_new_impl)
>  BTF_ID(func, bpf_percpu_obj_drop_impl)
> +BTF_ID(func, bpf_iter_css_task_new)
>
>  static bool is_kfunc_ret_null(struct bpf_kfunc_call_arg_meta *meta)
>  {
> @@ -10902,6 +10905,20 @@ static int process_kf_arg_ptr_to_rbtree_node(struct bpf_verifier_env *env,
>                                                   &meta->arg_rbtree_root.field);
>  }
>
> +static bool check_css_task_iter_allowlist(struct bpf_verifier_env *env)
> +{
> +       enum bpf_prog_type prog_type = resolve_prog_type(env->prog);
> +
> +       switch (prog_type) {
> +       case BPF_PROG_TYPE_LSM:
> +               return true;
> +       case BPF_TRACE_ITER:
> +               return env->prog->aux->sleepable;
> +       default:
> +               return false;
> +       }
> +}
> +
>  static int check_kfunc_args(struct bpf_verifier_env *env, struct bpf_kfunc_call_arg_meta *meta,
>                             int insn_idx)
>  {
> @@ -11152,6 +11169,12 @@ static int check_kfunc_args(struct bpf_verifier_env *env, struct bpf_kfunc_call_
>                         break;
>                 }
>                 case KF_ARG_PTR_TO_ITER:
> +                       if (meta->func_id == special_kfunc_list[KF_bpf_iter_css_task_new]) {
> +                               if (!check_css_task_iter_allowlist(env)) {
> +                                       verbose(env, "css_task_iter is only allowed in bpf_lsm and bpf iter-s\n");
> +                                       return -EINVAL;
> +                               }
> +                       }
>                         ret = process_iter_arg(env, regno, insn_idx, meta);
>                         if (ret < 0)
>                                 return ret;
> diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
> index 73b155e52204..de02c0971428 100644
> --- a/tools/include/uapi/linux/bpf.h
> +++ b/tools/include/uapi/linux/bpf.h
> @@ -7318,4 +7318,8 @@ struct bpf_iter_num {
>         __u64 __opaque[1];
>  } __attribute__((aligned(8)));
>
> +struct bpf_iter_css_task {
> +       __u64 __opaque[1];
> +} __attribute__((aligned(8)));
> +
>  #endif /* _UAPI__LINUX_BPF_H__ */
> diff --git a/tools/lib/bpf/bpf_helpers.h b/tools/lib/bpf/bpf_helpers.h
> index 77ceea575dc7..f48723c6c593 100644
> --- a/tools/lib/bpf/bpf_helpers.h
> +++ b/tools/lib/bpf/bpf_helpers.h
> @@ -303,6 +303,13 @@ extern int bpf_iter_num_new(struct bpf_iter_num *it, int start, int end) __weak
>  extern int *bpf_iter_num_next(struct bpf_iter_num *it) __weak __ksym;
>  extern void bpf_iter_num_destroy(struct bpf_iter_num *it) __weak __ksym;
>
> +struct bpf_iter_css_task;
> +struct cgroup_subsys_state;
> +extern int bpf_iter_css_task_new(struct bpf_iter_css_task *it,
> +               struct cgroup_subsys_state *css, unsigned int flags) __weak __ksym;
> +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;

please move this into bpf_experimental.h under selftests, this
shouldn't be in libbpf's stable API headers


> +
>  #ifndef bpf_for_each
>  /* bpf_for_each(iter_type, cur_elem, args...) provides generic construct for
>   * using BPF open-coded iterators without having to write mundane explicit
> --
> 2.20.1
>