bpf,x86: add fsession support for x86_64

Add BPF_TRACE_FSESSION supporting to x86_64, including:

1. clear the return value in the stack before fentry to make the fentry
   of the fsession can only get 0 with bpf_get_func_ret().

2. clear all the session cookies' value in the stack.

2. store the index of the cookie to ctx[-1] before the calling to fsession

3. store the "is_return" flag to ctx[-1] before the calling to fexit of
   the fsession.

Signed-off-by: Menglong Dong <dongml2@chinatelecom.cn>
Co-developed-by: Leon Hwang <leon.hwang@linux.dev>
Signed-off-by: Leon Hwang <leon.hwang@linux.dev>
Link: https://lore.kernel.org/r/20260124062008.8657-8-dongml2@chinatelecom.cn
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 2f31331955b5642aae71b3fdebe262f98b64d812..5a075e06cf45afa0235ca5ad2e99a0602e31a380 100644 (file)
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -3094,13 +3094,19 @@ static int emit_cond_near_jump(u8 **pprog, void *func, void *ip, u8 jmp_cond)
 
 static int invoke_bpf(const struct btf_func_model *m, u8 **pprog,
                      struct bpf_tramp_links *tl, int stack_size,
-                     int run_ctx_off, bool save_ret,
-                     void *image, void *rw_image)
+                     int run_ctx_off, int func_meta_off, bool save_ret,
+                     void *image, void *rw_image, u64 func_meta,
+                     int cookie_off)
 {
-       int i;
+       int i, cur_cookie = (cookie_off - stack_size) / 8;
        u8 *prog = *pprog;
 
        for (i = 0; i < tl->nr_links; i++) {
+               if (tl->links[i]->link.prog->call_session_cookie) {
+                       emit_store_stack_imm64(&prog, BPF_REG_0, -func_meta_off,
+                               func_meta | (cur_cookie << BPF_TRAMP_COOKIE_INDEX_SHIFT));
+                       cur_cookie--;
+               }
                if (invoke_bpf_prog(m, &prog, tl->links[i], stack_size,
                                    run_ctx_off, save_ret, image, rw_image))
                        return -EINVAL;
@@ -3218,12 +3224,14 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
                                         void *func_addr)
 {
        int i, ret, nr_regs = m->nr_args, stack_size = 0;
-       int regs_off, nregs_off, ip_off, run_ctx_off, arg_stack_off, rbx_off;
+       int regs_off, func_meta_off, ip_off, run_ctx_off, arg_stack_off, rbx_off;
        struct bpf_tramp_links *fentry = &tlinks[BPF_TRAMP_FENTRY];
        struct bpf_tramp_links *fexit = &tlinks[BPF_TRAMP_FEXIT];
        struct bpf_tramp_links *fmod_ret = &tlinks[BPF_TRAMP_MODIFY_RETURN];
        void *orig_call = func_addr;
+       int cookie_off, cookie_cnt;
        u8 **branches = NULL;
+       u64 func_meta;
        u8 *prog;
        bool save_ret;
 
@@ -3259,7 +3267,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
         *                 [ ...             ]
         * RBP - regs_off  [ reg_arg1        ]  program's ctx pointer
         *
-        * RBP - nregs_off [ regs count      ]  always
+        * RBP - func_meta_off [ regs count, etc ]  always
         *
         * RBP - ip_off    [ traced function ]  BPF_TRAMP_F_IP_ARG flag
         *
@@ -3282,15 +3290,20 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
        stack_size += nr_regs * 8;
        regs_off = stack_size;
 
-       /* regs count  */
+       /* function matedata, such as regs count  */
        stack_size += 8;
-       nregs_off = stack_size;
+       func_meta_off = stack_size;
 
        if (flags & BPF_TRAMP_F_IP_ARG)
                stack_size += 8; /* room for IP address argument */
 
        ip_off = stack_size;
 
+       cookie_cnt = bpf_fsession_cookie_cnt(tlinks);
+       /* room for session cookies */
+       stack_size += cookie_cnt * 8;
+       cookie_off = stack_size;
+
        stack_size += 8;
        rbx_off = stack_size;
 
@@ -3358,8 +3371,9 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
        /* mov QWORD PTR [rbp - rbx_off], rbx */
        emit_stx(&prog, BPF_DW, BPF_REG_FP, BPF_REG_6, -rbx_off);
 
+       func_meta = nr_regs;
        /* Store number of argument registers of the traced function */
-       emit_store_stack_imm64(&prog, BPF_REG_0, -nregs_off, nr_regs);
+       emit_store_stack_imm64(&prog, BPF_REG_0, -func_meta_off, func_meta);
 
        if (flags & BPF_TRAMP_F_IP_ARG) {
                /* Store IP address of the traced function */
@@ -3378,9 +3392,18 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
                }
        }
 
+       if (bpf_fsession_cnt(tlinks)) {
+               /* clear all the session cookies' value */
+               for (int i = 0; i < cookie_cnt; i++)
+                       emit_store_stack_imm64(&prog, BPF_REG_0, -cookie_off + 8 * i, 0);
+               /* clear the return value to make sure fentry always get 0 */
+               emit_store_stack_imm64(&prog, BPF_REG_0, -8, 0);
+       }
+
        if (fentry->nr_links) {
-               if (invoke_bpf(m, &prog, fentry, regs_off, run_ctx_off,
-                              flags & BPF_TRAMP_F_RET_FENTRY_RET, image, rw_image))
+               if (invoke_bpf(m, &prog, fentry, regs_off, run_ctx_off, func_meta_off,
+                              flags & BPF_TRAMP_F_RET_FENTRY_RET, image, rw_image,
+                              func_meta, cookie_off))
                        return -EINVAL;
        }
 
@@ -3440,9 +3463,14 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
                }
        }
 
+       /* set the "is_return" flag for fsession */
+       func_meta |= (1ULL << BPF_TRAMP_IS_RETURN_SHIFT);
+       if (bpf_fsession_cnt(tlinks))
+               emit_store_stack_imm64(&prog, BPF_REG_0, -func_meta_off, func_meta);
+
        if (fexit->nr_links) {
-               if (invoke_bpf(m, &prog, fexit, regs_off, run_ctx_off,
-                              false, image, rw_image)) {
+               if (invoke_bpf(m, &prog, fexit, regs_off, run_ctx_off, func_meta_off,
+                              false, image, rw_image, func_meta, cookie_off)) {
                        ret = -EINVAL;
                        goto cleanup;
                }