On Tue, Feb 22, 2022 at 9:07 AM Jiri Olsa <jolsa@xxxxxxxxxx> 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@xxxxxxxxxx> > Signed-off-by: Jiri Olsa <jolsa@xxxxxxxxxx> > --- > 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 >