bpf: Add support for kprobe session attach

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 kprobe multi links.

Adding new BPF_TRACE_KPROBE_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>
Signed-off-by: Andrii Nakryiko <andrii@kernel.org>
Link: https://lore.kernel.org/bpf/20240430112830.1184228-2-jolsa@kernel.org

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index d94a72593ead2788964a658336be4c313b097873..90706a47f6ffe974d346502c5511d9ac63a4f86e 100644 (file)
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -1115,6 +1115,7 @@ enum bpf_attach_type {
        BPF_CGROUP_UNIX_GETSOCKNAME,
        BPF_NETKIT_PRIMARY,
        BPF_NETKIT_PEER,
+       BPF_TRACE_KPROBE_SESSION,
        __MAX_BPF_ATTACH_TYPE
 };
 

diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index f655adf42e3960f8078f97d38eb1da0d1e71dc2d..13ad74ecf2cd4bcd95cc8f53ff7bbf632e5addc8 100644 (file)
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -4016,11 +4016,15 @@ static int bpf_prog_attach_check_attach_type(const struct bpf_prog *prog,
                if (prog->expected_attach_type == BPF_TRACE_KPROBE_MULTI &&
                    attach_type != BPF_TRACE_KPROBE_MULTI)
                        return -EINVAL;
+               if (prog->expected_attach_type == BPF_TRACE_KPROBE_SESSION &&
+                   attach_type != BPF_TRACE_KPROBE_SESSION)
+                       return -EINVAL;
                if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI &&
                    attach_type != BPF_TRACE_UPROBE_MULTI)
                        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)
                        return -EINVAL;
                return 0;
@@ -5281,7 +5285,8 @@ static int link_create(union bpf_attr *attr, bpfptr_t uattr)
        case BPF_PROG_TYPE_KPROBE:
                if (attr->link_create.attach_type == BPF_PERF_EVENT)
                        ret = bpf_perf_link_attach(attr, prog);
-               else if (attr->link_create.attach_type == BPF_TRACE_KPROBE_MULTI)
+               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)
                        ret = bpf_uprobe_multi_link_attach(attr, prog);

diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 0ba722b57af3ded2f9c87164e6395b1b4020b02e..06a9671834b6cb28f433f67ddc88735af02d4369 100644 (file)
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -1631,6 +1631,17 @@ bpf_tracing_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
        }
 }
 
+static bool is_kprobe_multi(const struct bpf_prog *prog)
+{
+       return prog->expected_attach_type == BPF_TRACE_KPROBE_MULTI ||
+              prog->expected_attach_type == BPF_TRACE_KPROBE_SESSION;
+}
+
+static inline bool is_kprobe_session(const struct bpf_prog *prog)
+{
+       return prog->expected_attach_type == BPF_TRACE_KPROBE_SESSION;
+}
+
 static const struct bpf_func_proto *
 kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 {
@@ -1646,13 +1657,13 @@ kprobe_prog_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
                return &bpf_override_return_proto;
 #endif
        case BPF_FUNC_get_func_ip:
-               if (prog->expected_attach_type == BPF_TRACE_KPROBE_MULTI)
+               if (is_kprobe_multi(prog))
                        return &bpf_get_func_ip_proto_kprobe_multi;
                if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI)
                        return &bpf_get_func_ip_proto_uprobe_multi;
                return &bpf_get_func_ip_proto_kprobe;
        case BPF_FUNC_get_attach_cookie:
-               if (prog->expected_attach_type == BPF_TRACE_KPROBE_MULTI)
+               if (is_kprobe_multi(prog))
                        return &bpf_get_attach_cookie_proto_kmulti;
                if (prog->expected_attach_type == BPF_TRACE_UPROBE_MULTI)
                        return &bpf_get_attach_cookie_proto_umulti;
@@ -2834,10 +2845,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 is_kprobe_session(link->link.prog) ? err : 0;
 }
 
 static void
@@ -2981,7 +2993,7 @@ int bpf_kprobe_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_KPROBE_MULTI)
+       if (!is_kprobe_multi(prog))
                return -EINVAL;
 
        flags = attr->link_create.kprobe_multi.flags;
@@ -3062,10 +3074,10 @@ 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 (!(flags & BPF_F_KPROBE_MULTI_RETURN))
                link->fp.entry_handler = kprobe_multi_link_handler;
+       if ((flags & BPF_F_KPROBE_MULTI_RETURN) || is_kprobe_session(prog))
+               link->fp.exit_handler = kprobe_multi_link_exit_handler;
 
        link->addrs = addrs;
        link->cookies = cookies;

diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index d94a72593ead2788964a658336be4c313b097873..90706a47f6ffe974d346502c5511d9ac63a4f86e 100644 (file)
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -1115,6 +1115,7 @@ enum bpf_attach_type {
        BPF_CGROUP_UNIX_GETSOCKNAME,
        BPF_NETKIT_PRIMARY,
        BPF_NETKIT_PEER,
+       BPF_TRACE_KPROBE_SESSION,
        __MAX_BPF_ATTACH_TYPE
 };