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 <redacted>
Co-developed-by: Leon Hwang <redacted>
Signed-off-by: Leon Hwang <redacted>
---
v5:
- add the variable "func_meta"
- define cookie_off in a new line

v4:
- some adjustment to the 1st patch, such as we get the fsession prog from
  fentry and fexit hlist
- remove the supporting of skipping fexit with fentry return non-zero

v2:
- add session cookie support
- add the session stuff after return value, instead of before nr_args
---
 arch/x86/net/bpf_jit_comp.c | 33 ++++++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index d94f7038c441..0671a434c00d 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c

@@ -3094,12 +3094,17 @@ 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)
+                     void *image, void *rw_image, u64 func_meta)
 {
        int i;
        u8 *prog = *pprog;

        for (i = 0; i < tl->nr_links; i++) {
+               if (tl->links[i]->link.prog->call_session_cookie) {
+                       /* 'stack_size + 8' is the offset of func_md in stack */

+               }
                if (invoke_bpf_prog(m, &prog, tl->links[i], stack_size,
                                    run_ctx_off, save_ret, image, rw_image))
                        return -EINVAL;

@@ -3222,7 +3227,9 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
        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;

@@ -3290,6 +3297,11 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im

        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;

@@ -3383,9 +3395,19 @@ 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, cookie_off - 8 * i, 0);
+               /* clear the return value to make sure fentry always get 0 */
+               emit_store_stack_imm64(&prog, 8, 0);
+       }
+       func_meta = nr_regs + (((cookie_off - regs_off) / 8) << BPF_TRAMP_M_COOKIE);

+
        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))
+                              flags & BPF_TRAMP_F_RET_FENTRY_RET, image, rw_image,
+                              func_meta))
                        return -EINVAL;
        }

@@ -3445,9 +3467,14 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
                }
        }

+       /* set the "is_return" flag for fsession */
+       func_meta += (1 << BPF_TRAMP_M_IS_RETURN);
+       if (bpf_fsession_cnt(tlinks))
+               emit_store_stack_imm64(&prog, nregs_off, func_meta);
+
        if (fexit->nr_links) {
                if (invoke_bpf(m, &prog, fexit, regs_off, run_ctx_off,
-                              false, image, rw_image)) {
+                              false, image, rw_image, func_meta)) {
                        ret = -EINVAL;
                        goto cleanup;
                }
--

2.52.0