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>
---
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 a87304161d45..32c13175bc65 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 */
+			emit_st_r0_imm64(&prog, func_meta, stack_size + 8);
+			func_meta -= (1 << BPF_TRAMP_M_COOKIE);
+		}
 		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_st_r0_imm64(&prog, 0, cookie_off - 8 * i);
+		/* clear the return value to make sure fentry always get 0 */
+		emit_st_r0_imm64(&prog, 0, 8);
+	}
+	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_st_r0_imm64(&prog, func_meta, nregs_off);
+
 	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

> diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
> index a87304161d45..32c13175bc65 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 */
> +			emit_st_r0_imm64(&prog, func_meta, stack_size + 8);
> +			func_meta -= (1 << BPF_TRAMP_M_COOKIE);
> +		}

This decrement of func_meta is problematic for fsession programs. When
invoke_bpf() is called for fentry, func_meta gets decremented for each
program with call_session_cookie. Later, when invoke_bpf() is called
for fexit, it receives this already-decremented func_meta value.

Since fsession programs appear in both fentry and fexit lists (same
program object), call_session_cookie is true for both. The fexit part
will see a different cookie offset than the fentry part saw.

For example, with one fsession program:
  - fentry stores func_meta with cookie offset N, then decrements to N-1
  - fexit receives N-1, stores it, and the BPF program reads offset N-1

Should the func_meta offset be reset for fexit so that each fsession
program's fentry and fexit parts access the same cookie slot?

[ ... ]

> @@ -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_st_r0_imm64(&prog, 0, cookie_off - 8 * i);
> +		/* clear the return value to make sure fentry always get 0 */
> +		emit_st_r0_imm64(&prog, 0, 8);
> +	}
> +	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_st_r0_imm64(&prog, func_meta, nregs_off);
> +
>  	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)) {

At this point func_meta has already been decremented during the fentry
invoke_bpf() call above. When invoke_bpf() processes the fexit list and
encounters an fsession program (which has call_session_cookie set since
it is the same program object), it will store this modified func_meta.

The cookie offset encoded in func_meta will be different from what the
fentry part stored, causing bpf_session_cookie() to return a pointer to
a different memory location for fexit than for fentry.

> +			ret = -EINVAL;
> +			goto cleanup;
> +		}


---
AI reviewed your patch. Please fix the bug or email reply why it's not a bug.
See: https://github.com/kernel-patches/vmtest/blob/master/ci/claude/README.md

CI run summary: https://github.com/kernel-patches/bpf/actions/runs/20773233136