Re: [PATCH RFCv2 bpf-next 1/4] bpf: Add support for kprobe multi wrapper attach

[Jiri Olsa <olsajiri@gmail.com>]

On Wed, Feb 28, 2024 at 05:23:05PM -0800, Andrii Nakryiko wrote:
> On Wed, Feb 28, 2024 at 1:03 AM Jiri Olsa <jolsa@kernel.org> wrote:
> >
> > Adding support to attach bpf program for entry and return probe
> > of the same function. This is common usecase and at the moment
> > it requires to create two kprobe multi links.
> >
> > Adding new attr.link_create.kprobe_multi.flags value that instructs
> > kernel to attach link program to both entry and exit probe.
> >
> > It's possible to control execution of the bpf program on return
> > probe simply by returning zero or non zero from the entry bpf
> > program execution to execute or not respectively the bpf program
> > on return probe.
> >
> > Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> > ---
> >  include/uapi/linux/bpf.h       |  3 ++-
> >  kernel/trace/bpf_trace.c       | 24 ++++++++++++++++++------
> >  tools/include/uapi/linux/bpf.h |  3 ++-
> >  3 files changed, 22 insertions(+), 8 deletions(-)
> >
> > diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
> > index d2e6c5fcec01..a430855c5bcd 100644
> > --- a/include/uapi/linux/bpf.h
> > +++ b/include/uapi/linux/bpf.h
> > @@ -1247,7 +1247,8 @@ enum bpf_perf_event_type {
> >   * BPF_TRACE_KPROBE_MULTI attach type to create return probe.
> >   */
> >  enum {
> > -       BPF_F_KPROBE_MULTI_RETURN = (1U << 0)
> > +       BPF_F_KPROBE_MULTI_RETURN  = (1U << 0),
> > +       BPF_F_KPROBE_MULTI_WRAPPER = (1U << 1),
> >  };
> >
> >  /* link_create.uprobe_multi.flags used in LINK_CREATE command for
> > diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
> > index 241ddf5e3895..726a8c71f0da 100644
> > --- a/kernel/trace/bpf_trace.c
> > +++ b/kernel/trace/bpf_trace.c
> > @@ -2587,6 +2587,7 @@ struct bpf_kprobe_multi_link {
> >         u32 mods_cnt;
> >         struct module **mods;
> >         u32 flags;
> > +       bool is_wrapper;
> 
> flags should be sufficient for this, why storing redundant bool field?

true

> 
> >  };
> >
> >  struct bpf_kprobe_multi_run_ctx {
> > @@ -2826,10 +2827,11 @@ kprobe_multi_link_handler(struct fprobe *fp, unsigned long fentry_ip,
> >                           void *data)
> >  {
> >         struct bpf_kprobe_multi_link *link;
> > +       int err;
> >
> >         link = container_of(fp, struct bpf_kprobe_multi_link, fp);
> > -       kprobe_multi_link_prog_run(link, get_entry_ip(fentry_ip), regs);
> > -       return 0;
> > +       err = kprobe_multi_link_prog_run(link, get_entry_ip(fentry_ip), regs);
> > +       return link->is_wrapper ? err : 0;
> >  }
> >
> >  static void
> > @@ -2967,6 +2969,7 @@ int bpf_kprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
> >         void __user *uaddrs;
> >         u64 *cookies = NULL;
> >         void __user *usyms;
> > +       bool is_wrapper;
> >         int err;
> >
> >         /* no support for 32bit archs yet */
> > @@ -2977,9 +2980,12 @@ int bpf_kprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
> >                 return -EINVAL;
> >
> >         flags = attr->link_create.kprobe_multi.flags;
> > -       if (flags & ~BPF_F_KPROBE_MULTI_RETURN)
> > +       if (flags & ~(BPF_F_KPROBE_MULTI_RETURN|
> > +                     BPF_F_KPROBE_MULTI_WRAPPER))
> 
> nit: spaces around | are missing, also keep on a single line?

ok

> 
> >                 return -EINVAL;
> >
> > +       is_wrapper = flags & BPF_F_KPROBE_MULTI_WRAPPER;
> > +
> >         uaddrs = u64_to_user_ptr(attr->link_create.kprobe_multi.addrs);
> >         usyms = u64_to_user_ptr(attr->link_create.kprobe_multi.syms);
> >         if (!!uaddrs == !!usyms)
> > @@ -3054,15 +3060,21 @@ int bpf_kprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
> >         if (err)
> >                 goto error;
> >
> > -       if (flags & BPF_F_KPROBE_MULTI_RETURN)
> > -               link->fp.exit_handler = kprobe_multi_link_exit_handler;
> > -       else
> > +       if (is_wrapper) {
> >                 link->fp.entry_handler = kprobe_multi_link_handler;
> > +               link->fp.exit_handler = kprobe_multi_link_exit_handler;
> > +       } else {
> > +               if (flags & BPF_F_KPROBE_MULTI_RETURN)
> > +                       link->fp.exit_handler = kprobe_multi_link_exit_handler;
> > +               else
> > +                       link->fp.entry_handler = kprobe_multi_link_handler;
> > +       }
> >
> 
> how about:
> 
> if (!(flags & BPF_F_KPROBE_MULTI_RETURN))
>     link->fp.entry_handler = kprobe_multi_link_handler;
> if (flags & (BPF_F_KPROBE_MULTI_RETURN | BPF_F_KPROBE_MULTI_WRAPPER))
>     link->fp.exit_handler = kprobe_multi_link_exit_handler;

I have another changes on top of that, which add more handlers,
so I guess I wanted to keep it more explicit, but your suggestion
is simpler, will change

thanks,
jirka

> 
> 
> >         link->addrs = addrs;
> >         link->cookies = cookies;
> >         link->cnt = cnt;
> >         link->flags = flags;
> > +       link->is_wrapper = is_wrapper;
> >
> >         if (cookies) {
> >                 /*
> > diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
> > index d2e6c5fcec01..a430855c5bcd 100644
> > --- a/tools/include/uapi/linux/bpf.h
> > +++ b/tools/include/uapi/linux/bpf.h
> > @@ -1247,7 +1247,8 @@ enum bpf_perf_event_type {
> >   * BPF_TRACE_KPROBE_MULTI attach type to create return probe.
> >   */
> >  enum {
> > -       BPF_F_KPROBE_MULTI_RETURN = (1U << 0)
> > +       BPF_F_KPROBE_MULTI_RETURN  = (1U << 0),
> > +       BPF_F_KPROBE_MULTI_WRAPPER = (1U << 1),
> >  };
> >
> >  /* link_create.uprobe_multi.flags used in LINK_CREATE command for
> > --
> > 2.43.2
> >