Re: [PATCH 05/10] bpf: Add cookie support to programs attached with kprobe multi link

[Andrii Nakryiko <andrii.nakryiko@gmail.com>]

On Tue, Feb 22, 2022 at 9:07 AM Jiri Olsa <jolsa@kernel.org> wrote:
>
> Adding support to call bpf_get_attach_cookie helper from
> kprobe programs attached with kprobe multi link.
>
> The cookie is provided by array of u64 values, where each
> value is paired with provided function address or symbol
> with the same array index.
>
> Suggested-by: Andrii Nakryiko <andrii@kernel.org>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  include/linux/sort.h           |   2 +
>  include/uapi/linux/bpf.h       |   1 +
>  kernel/trace/bpf_trace.c       | 103 ++++++++++++++++++++++++++++++++-
>  lib/sort.c                     |   2 +-
>  tools/include/uapi/linux/bpf.h |   1 +
>  5 files changed, 107 insertions(+), 2 deletions(-)
>

[...]

>  BPF_CALL_1(bpf_get_attach_cookie_trace, void *, ctx)
>  {
>         struct bpf_trace_run_ctx *run_ctx;
> @@ -1297,7 +1312,9 @@ kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
>                         &bpf_get_func_ip_proto_kprobe_multi :
>                         &bpf_get_func_ip_proto_kprobe;
>         case BPF_FUNC_get_attach_cookie:
> -               return &bpf_get_attach_cookie_proto_trace;
> +               return prog->expected_attach_type == BPF_TRACE_KPROBE_MULTI ?
> +                       &bpf_get_attach_cookie_proto_kmulti :
> +                       &bpf_get_attach_cookie_proto_trace;
>         default:
>                 return bpf_tracing_func_proto(func_id, prog);
>         }
> @@ -2203,6 +2220,9 @@ struct bpf_kprobe_multi_link {
>         struct bpf_link link;
>         struct fprobe fp;
>         unsigned long *addrs;
> +       struct bpf_run_ctx run_ctx;

clever, I like it! Keep in mind, though, that this trick can only be
used here because this run_ctx is read-only (I'd leave the comment
here about this, I didn't realize immediately that this approach can't
be used for run_ctx that needs to be modified).

> +       u64 *cookies;
> +       u32 cnt;
>  };
>
>  static void bpf_kprobe_multi_link_release(struct bpf_link *link)
> @@ -2219,6 +2239,7 @@ static void bpf_kprobe_multi_link_dealloc(struct bpf_link *link)
>
>         kmulti_link = container_of(link, struct bpf_kprobe_multi_link, link);
>         kvfree(kmulti_link->addrs);
> +       kvfree(kmulti_link->cookies);
>         kfree(kmulti_link);
>  }
>
> @@ -2227,10 +2248,57 @@ static const struct bpf_link_ops bpf_kprobe_multi_link_lops = {
>         .dealloc = bpf_kprobe_multi_link_dealloc,
>  };
>
> +static void bpf_kprobe_multi_cookie_swap(void *a, void *b, int size, const void *priv)
> +{
> +       const struct bpf_kprobe_multi_link *link = priv;
> +       unsigned long *addr_a = a, *addr_b = b;
> +       u64 *cookie_a, *cookie_b;
> +
> +       cookie_a = link->cookies + (addr_a - link->addrs);
> +       cookie_b = link->cookies + (addr_b - link->addrs);
> +
> +       swap_words_64(addr_a, addr_b, size);
> +       swap_words_64(cookie_a, cookie_b, size);

is it smart to call (now) non-inlined function just to swap two longs
and u64s?..

unsigned long tmp1;
u64 tmp2;

tmp1 = *addr_a; *addr_a = addr_b; *addr_b = tmp1;
tmp2 = *cookie_a; *cookie_a = cookie_b; *cookie_b = tmp2;

?

> +}
> +
> +static int __bpf_kprobe_multi_cookie_cmp(const void *a, const void *b)
> +{
> +       const unsigned long *addr_a = a, *addr_b = b;
> +
> +       if (*addr_a == *addr_b)
> +               return 0;
> +       return *addr_a < *addr_b ? -1 : 1;
> +}
> +

[...]

> @@ -2238,12 +2306,16 @@ kprobe_multi_link_prog_run(struct bpf_kprobe_multi_link *link,
>                 goto out;
>         }
>
> +       old_run_ctx = bpf_set_run_ctx(&link->run_ctx);
> +
>         rcu_read_lock();

so looking at other code, I see that we first migrate_disable() and
then rcu_read_lock(), so let's swap? We also normally set/reset
run_ctx inside migrate+rcu_lock region. I'm not sure that's necessary,
but also shouldn't hurt to stay consistent.

>         migrate_disable();
>         err = bpf_prog_run(link->link.prog, regs);
>         migrate_enable();
>         rcu_read_unlock();
>
> +       bpf_reset_run_ctx(old_run_ctx);
> +
>   out:
>         __this_cpu_dec(bpf_prog_active);
>         return err;

[...]

> diff --git a/lib/sort.c b/lib/sort.c
> index b399bf10d675..91f7ce701cf4 100644
> --- a/lib/sort.c
> +++ b/lib/sort.c
> @@ -80,7 +80,7 @@ static void swap_words_32(void *a, void *b, size_t n)
>   * but it's possible to have 64-bit loads without 64-bit pointers (e.g.
>   * x32 ABI).  Are there any cases the kernel needs to worry about?
>   */
> -static void swap_words_64(void *a, void *b, size_t n)
> +void swap_words_64(void *a, void *b, size_t n)

I'm worried that this might change performance unintentionally in
other places (making the function global might pessimize inlining, I
think). So let's not do that, just do a straightforward swap in cookie
support code?

>  {
>         do {
>  #ifdef CONFIG_64BIT
> diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
> index 6c66138c1b9b..d18996502aac 100644
> --- a/tools/include/uapi/linux/bpf.h
> +++ b/tools/include/uapi/linux/bpf.h
> @@ -1482,6 +1482,7 @@ union bpf_attr {
>                         struct {
>                                 __aligned_u64   syms;
>                                 __aligned_u64   addrs;
> +                               __aligned_u64   cookies;

looks a bit weird to change layout of UAPI. That's not really a
problem, because both patches will land at the same time. But if you
move flags and cnt to the front of the struct it would a bit better.


>                                 __u32           cnt;
>                                 __u32           flags;
>                         } kprobe_multi;
> --
> 2.35.1
>