On Mon, Jul 1, 2024 at 9:42 AM Jiri Olsa <jolsa@xxxxxxxxxx> 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@xxxxxxxxxx>
> ---
> 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@xxxxxxxxxx>
> 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
>
