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

[Jiri Olsa <olsajiri@gmail.com>]

On Fri, Mar 04, 2022 at 03:11:08PM -0800, Andrii Nakryiko wrote:
> 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).

hum, I don't see it at the moment.. I'll check on that and add the
comment or come up with more questions ;-)

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

the swap_words_64 has CONFIG_64BIT ifdef with some tweaks for 32bit,
so I wanted to use that.. but I agree with your other comment below
wrt performace, so will change

> 
> ?
> 
> > +}
> > +
> > +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.

ok, will change

> 
> >         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?

right, I did not realize this.. I'll add to cookie code directly

> 
> >  {
> >         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.

I was following your previous comment:
  https://lore.kernel.org/bpf/CAEf4BzbPeQbURZOD93TgPudOk3JD4odsZ9uwriNkrphes9V4dg@mail.gmail.com/

I like the idea that syms/addrs/cookies stay together,
because they are all related to cnt.. but yes, it's
'breaking' KABI in between these patches

jirka

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