Re: [PATCHv2 bpf-next 2/9] bpf: Add support for uprobe multi session attach

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

On Mon, Jul 1, 2024 at 9:42 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 use case which at the moment
> requires to create two uprobe multi links.
>
> Adding new BPF_TRACE_UPROBE_SESSION attach type that instructs
> kernel to attach single 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 the bpf program on return
> probe respectively.
>
> Signed-off-by: Jiri Olsa <jolsa@kernel.org>
> ---
>  include/uapi/linux/bpf.h       |  1 +
>  kernel/bpf/syscall.c           |  9 +++++++--
>  kernel/trace/bpf_trace.c       | 25 +++++++++++++++++++------
>  tools/include/uapi/linux/bpf.h |  1 +
>  4 files changed, 28 insertions(+), 8 deletions(-)
>

LGTM

Acked-by: Andrii Nakryiko <andrii@kernel.org>

> diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
> index 35bcf52dbc65..1d93cb014884 100644
> --- a/include/uapi/linux/bpf.h
> +++ b/include/uapi/linux/bpf.h
> @@ -1116,6 +1116,7 @@ enum bpf_attach_type {
>         BPF_NETKIT_PRIMARY,
>         BPF_NETKIT_PEER,
>         BPF_TRACE_KPROBE_SESSION,
> +       BPF_TRACE_UPROBE_SESSION,
>         __MAX_BPF_ATTACH_TYPE
>  };
>
> diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
> index 869265852d51..2a63a528fa3c 100644
> --- a/kernel/bpf/syscall.c
> +++ b/kernel/bpf/syscall.c
> @@ -4049,10 +4049,14 @@ static int bpf_prog_attach_check_attach_type(const struct bpf_prog *prog,
>                 if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI &&
>                     attach_type != BPF_TRACE_UPROBE_MULTI)
>                         return -EINVAL;
> +               if (prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION &&
> +                   attach_type != BPF_TRACE_UPROBE_SESSION)
> +                       return -EINVAL;
>                 if (attach_type != BPF_PERF_EVENT &&
>                     attach_type != BPF_TRACE_KPROBE_MULTI &&
>                     attach_type != BPF_TRACE_KPROBE_SESSION &&
> -                   attach_type != BPF_TRACE_UPROBE_MULTI)
> +                   attach_type != BPF_TRACE_UPROBE_MULTI &&
> +                   attach_type != BPF_TRACE_UPROBE_SESSION)
>                         return -EINVAL;
>                 return 0;
>         case BPF_PROG_TYPE_SCHED_CLS:
> @@ -5315,7 +5319,8 @@ static int link_create(union bpf_attr *attr, bpfptr_t uattr)
>                 else if (attr->link_create.attach_type == BPF_TRACE_KPROBE_MULTI ||
>                          attr->link_create.attach_type == BPF_TRACE_KPROBE_SESSION)
>                         ret = bpf_kprobe_multi_link_attach(attr, prog);
> -               else if (attr->link_create.attach_type == BPF_TRACE_UPROBE_MULTI)
> +               else if (attr->link_create.attach_type == BPF_TRACE_UPROBE_MULTI ||
> +                        attr->link_create.attach_type == BPF_TRACE_UPROBE_SESSION)
>                         ret = bpf_uprobe_multi_link_attach(attr, prog);
>                 break;
>         default:
> diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
> index 02d052639dfe..1b19c1cdb5e1 100644
> --- a/kernel/trace/bpf_trace.c
> +++ b/kernel/trace/bpf_trace.c
> @@ -1645,6 +1645,17 @@ static inline bool is_kprobe_session(const struct bpf_prog *prog)
>         return prog->expected_attach_type == BPF_TRACE_KPROBE_SESSION;
>  }
>
> +static inline bool is_uprobe_multi(const struct bpf_prog *prog)
> +{
> +       return prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI ||
> +              prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION;
> +}
> +
> +static inline bool is_uprobe_session(const struct bpf_prog *prog)
> +{
> +       return prog->expected_attach_type == BPF_TRACE_UPROBE_SESSION;
> +}
> +
>  static const struct bpf_func_proto *
>  kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
>  {
> @@ -1662,13 +1673,13 @@ kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
>         case BPF_FUNC_get_func_ip:
>                 if (is_kprobe_multi(prog))
>                         return &bpf_get_func_ip_proto_kprobe_multi;
> -               if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI)
> +               if (is_uprobe_multi(prog))
>                         return &bpf_get_func_ip_proto_uprobe_multi;
>                 return &bpf_get_func_ip_proto_kprobe;
>         case BPF_FUNC_get_attach_cookie:
>                 if (is_kprobe_multi(prog))
>                         return &bpf_get_attach_cookie_proto_kmulti;
> -               if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI)
> +               if (is_uprobe_multi(prog))
>                         return &bpf_get_attach_cookie_proto_umulti;
>                 return &bpf_get_attach_cookie_proto_trace;
>         default:
> @@ -3387,7 +3398,7 @@ int bpf_uprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
>         if (sizeof(u64) != sizeof(void *))
>                 return -EOPNOTSUPP;
>
> -       if (prog->expected_attach_type != BPF_TRACE_UPROBE_MULTI)
> +       if (!is_uprobe_multi(prog))
>                 return -EINVAL;
>
>         flags = attr->link_create.uprobe_multi.flags;
> @@ -3463,10 +3474,12 @@ int bpf_uprobe_multi_link_attach(const union bpf_attr *attr, struct bpf_prog *pr
>
>                 uprobes[i].link = link;
>
> -               if (flags & BPF_F_UPROBE_MULTI_RETURN)
> -                       uprobes[i].consumer.ret_handler = uprobe_multi_link_ret_handler;
> -               else
> +               if (!(flags & BPF_F_UPROBE_MULTI_RETURN))
>                         uprobes[i].consumer.handler = uprobe_multi_link_handler;
> +               if (flags & BPF_F_UPROBE_MULTI_RETURN || is_uprobe_session(prog))
> +                       uprobes[i].consumer.ret_handler = uprobe_multi_link_ret_handler;
> +               if (is_uprobe_session(prog))
> +                       uprobes[i].consumer.session = true;
>
>                 if (pid)
>                         uprobes[i].consumer.filter = uprobe_multi_link_filter;
> diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
> index 35bcf52dbc65..1d93cb014884 100644
> --- a/tools/include/uapi/linux/bpf.h
> +++ b/tools/include/uapi/linux/bpf.h
> @@ -1116,6 +1116,7 @@ enum bpf_attach_type {
>         BPF_NETKIT_PRIMARY,
>         BPF_NETKIT_PEER,
>         BPF_TRACE_KPROBE_SESSION,
> +       BPF_TRACE_UPROBE_SESSION,
>         __MAX_BPF_ATTACH_TYPE
>  };
>
> --
> 2.45.2
>