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

Andrii Nakryiko
2022-03-04T23:11:36Z

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
>