[PATCHv5 bpf-next 04/13] bpf: Add support for uprobe multi session context

[Jiri Olsa]

Placing bpf_session_run_ctx layer in between bpf_run_ctx and
bpf_uprobe_multi_run_ctx, so the session data can be retrieved
from uprobe_multi link.

Plus granting session kfuncs access to uprobe session programs.

Acked-by: Andrii Nakryiko <and...@kernel.org>
Signed-off-by: Jiri Olsa <jo...@kernel.org>
---
 kernel/trace/bpf_trace.c | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index 98e940ec184d..41f83d504bf6 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -3098,7 +3098,7 @@ struct bpf_uprobe_multi_link {
 };
 
 struct bpf_uprobe_multi_run_ctx {
-       struct bpf_run_ctx run_ctx;
+       struct bpf_session_run_ctx session_ctx;
        unsigned long entry_ip;
        struct bpf_uprobe *uprobe;
 };
@@ -3211,10 +3211,15 @@ static const struct bpf_link_ops 
bpf_uprobe_multi_link_lops = {
 
 static int uprobe_prog_run(struct bpf_uprobe *uprobe,
                           unsigned long entry_ip,
-                          struct pt_regs *regs)
+                          struct pt_regs *regs,
+                          bool is_return, void *data)
 {
        struct bpf_uprobe_multi_link *link = uprobe->link;
        struct bpf_uprobe_multi_run_ctx run_ctx = {
+               .session_ctx = {
+                       .is_return = is_return,
+                       .data = data,
+               },
                .entry_ip = entry_ip,
                .uprobe = uprobe,
        };
@@ -3233,7 +3238,7 @@ static int uprobe_prog_run(struct bpf_uprobe *uprobe,
 
        migrate_disable();
 
-       old_run_ctx = bpf_set_run_ctx(&run_ctx.run_ctx);
+       old_run_ctx = bpf_set_run_ctx(&run_ctx.session_ctx.run_ctx);
        err = bpf_prog_run(link->link.prog, regs);
        bpf_reset_run_ctx(old_run_ctx);
 
@@ -3263,7 +3268,7 @@ uprobe_multi_link_handler(struct uprobe_consumer *con, 
struct pt_regs *regs,
        int ret;
 
        uprobe = container_of(con, struct bpf_uprobe, consumer);
-       ret = uprobe_prog_run(uprobe, instruction_pointer(regs), regs);
+       ret = uprobe_prog_run(uprobe, instruction_pointer(regs), regs, false, 
data);
        if (uprobe->session)
                return ret ? UPROBE_HANDLER_IGNORE : 0;
        return ret;
@@ -3282,14 +3287,15 @@ uprobe_multi_link_ret_handler(struct uprobe_consumer 
*con, unsigned long func, s
         */
        if (uprobe->session && !data)
                return 0;
-       return uprobe_prog_run(uprobe, func, regs);
+       return uprobe_prog_run(uprobe, func, regs, true, data);
 }
 
 static u64 bpf_uprobe_multi_entry_ip(struct bpf_run_ctx *ctx)
 {
        struct bpf_uprobe_multi_run_ctx *run_ctx;
 
-       run_ctx = container_of(current->bpf_ctx, struct 
bpf_uprobe_multi_run_ctx, run_ctx);
+       run_ctx = container_of(current->bpf_ctx, struct 
bpf_uprobe_multi_run_ctx,
+                              session_ctx.run_ctx);
        return run_ctx->entry_ip;
 }
 
@@ -3297,7 +3303,8 @@ static u64 bpf_uprobe_multi_cookie(struct bpf_run_ctx 
*ctx)
 {
        struct bpf_uprobe_multi_run_ctx *run_ctx;
 
-       run_ctx = container_of(current->bpf_ctx, struct 
bpf_uprobe_multi_run_ctx, run_ctx);
+       run_ctx = container_of(current->bpf_ctx, struct 
bpf_uprobe_multi_run_ctx,
+                              session_ctx.run_ctx);
        return run_ctx->uprobe->cookie;
 }
 
@@ -3491,7 +3498,7 @@ static int bpf_kprobe_multi_filter(const struct bpf_prog 
*prog, u32 kfunc_id)
        if (!btf_id_set8_contains(&kprobe_multi_kfunc_set_ids, kfunc_id))
                return 0;
 
-       if (!is_kprobe_session(prog))
+       if (!is_kprobe_session(prog) && !is_uprobe_session(prog))
                return -EACCES;
 
        return 0;
-- 
2.46.1