diff mbox series

[bpf-next,v5,3/8] bpf: Introduce task open coded iterator kfuncs

Message ID 20231011120857.251943-4-zhouchuyi@bytedance.com
State New
Headers
Series Add Open-coded task, css_task and css iters |

Commit Message

Chuyi Zhou Oct. 11, 2023, 12:08 p.m. UTC 
  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>
---
 kernel/bpf/helpers.c                          |  3 +
 kernel/bpf/task_iter.c                        | 82 +++++++++++++++++++
 .../testing/selftests/bpf/bpf_experimental.h  |  5 ++
 3 files changed, 90 insertions(+)

Comments

Andrii Nakryiko Oct. 13, 2023, 9:27 p.m. UTC | #1 
On Wed, Oct 11, 2023 at 5:09 AM Chuyi Zhou <zhouchuyi@bytedance.com> 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@bytedance.com>
> ---
>  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
>
Chuyi Zhou Oct. 14, 2023, 2:02 a.m. UTC | #2 
Hello,

在 2023/10/14 05:27, Andrii Nakryiko 写道:
> On Wed, Oct 11, 2023 at 5:09 AM Chuyi Zhou <zhouchuyi@bytedance.com> 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@bytedance.com>
>> ---
>>   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?
> 

Thanks for the review.


IIUC, do you mean:

diff --git a/kernel/bpf/task_iter.c b/kernel/bpf/task_iter.c
index 0772545568f1..b35debf19edb 100644
--- a/kernel/bpf/task_iter.c
+++ b/kernel/bpf/task_iter.c
@@ -913,7 +913,7 @@ __bpf_kfunc struct task_struct 
*bpf_iter_task_next(struct bpf_iter_task *it)
         pos = kit->pos;

         if (!pos)
-               goto out;
+               return pos;

         if (flags == BPF_TASK_ITER_ALL_PROCS)
                 goto get_next_task;
@@ -922,18 +922,22 @@ __bpf_kfunc struct task_struct 
*bpf_iter_task_next(struct bpf_iter_task *it)
         if (kit->pos == kit->task) {
                 if (flags == BPF_TASK_ITER_PROC_THREADS) {
                         kit->pos = NULL;
-                       goto out;
+                       return pos;
                 }
         } else
-               goto out;
+               return pos;

+       /*
+        * goto get_next_task means:
+        * case 1: flags == BPF_TASK_ITER_ALL_PROCS
+        * case 2: kit->pos == kit->task && flags == 
BPF_TASK_ITER_ALL_THREADS
+        */
  get_next_task:
         kit->pos = next_task(kit->pos);
         kit->task = kit->pos;
         if (kit->pos == &init_task)
                 kit->pos = NULL;

-out:
         return pos;



BTW, do you have some comments on patch-8 ? or I should send next 
version and pass all the CI first ?

Thanks.

> 
>> +
>> +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
>>
Chuyi Zhou Oct. 16, 2023, 8:40 a.m. UTC | #3 
在 2023/10/11 20:08, Chuyi Zhou 写道:
> 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>
> ---
>   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
> +};
> +

In next version, I would add the missing __diag_ignore_all for 
-Wmissing-prototypes in Patch2 ~ Patch4 to avoid kernel build warning.

Thanks.

> +__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;
> +
> +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);
>
Andrii Nakryiko Oct. 16, 2023, 8:18 p.m. UTC | #4 
On Fri, Oct 13, 2023 at 7:02 PM Chuyi Zhou <zhouchuyi@bytedance.com> wrote:
>
> Hello,
>
> 在 2023/10/14 05:27, Andrii Nakryiko 写道:
> > On Wed, Oct 11, 2023 at 5:09 AM Chuyi Zhou <zhouchuyi@bytedance.com> 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@bytedance.com>
> >> ---
> >>   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?
> >
>
> Thanks for the review.
>
>
> IIUC, do you mean:
>

yes, goto only makes sense when there is some common clean up or error
handling logic, in this case it's a plain return result, so no point.


> diff --git a/kernel/bpf/task_iter.c b/kernel/bpf/task_iter.c
> index 0772545568f1..b35debf19edb 100644
> --- a/kernel/bpf/task_iter.c
> +++ b/kernel/bpf/task_iter.c
> @@ -913,7 +913,7 @@ __bpf_kfunc struct task_struct
> *bpf_iter_task_next(struct bpf_iter_task *it)
>          pos = kit->pos;
>
>          if (!pos)
> -               goto out;
> +               return pos;
>
>          if (flags == BPF_TASK_ITER_ALL_PROCS)
>                  goto get_next_task;
> @@ -922,18 +922,22 @@ __bpf_kfunc struct task_struct
> *bpf_iter_task_next(struct bpf_iter_task *it)
>          if (kit->pos == kit->task) {
>                  if (flags == BPF_TASK_ITER_PROC_THREADS) {
>                          kit->pos = NULL;
> -                       goto out;
> +                       return pos;
>                  }
>          } else
> -               goto out;
> +               return pos;
>
> +       /*
> +        * goto get_next_task means:
> +        * case 1: flags == BPF_TASK_ITER_ALL_PROCS
> +        * case 2: kit->pos == kit->task && flags ==
> BPF_TASK_ITER_ALL_THREADS
> +        */
>   get_next_task:
>          kit->pos = next_task(kit->pos);
>          kit->task = kit->pos;
>          if (kit->pos == &init_task)
>                  kit->pos = NULL;
>
> -out:
>          return pos;
>
>
>
> BTW, do you have some comments on patch-8 ? or I should send next
> version and pass all the CI first ?
>

I didn't think too hard about changes you are proposing, but yes, CI
should be green on submission, of course

> Thanks.
>
> >
> >> +
> >> +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
> >>
diff mbox series

Patch

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;
+
+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