The stack will look like this:

  return value  -> 8 bytes
  argN          -> 8 bytes
  ...
  arg1          -> 8 bytes
  nr_args       -> 8 bytes
  ip (optional) -> 8 bytes
  cookie2       -> 8 bytes
  cookie1       -> 8 bytes

Implement and inline the bpf_session_cookie() for the fsession in the
verifier.

Signed-off-by: Menglong Dong <redacted>
---
v9:
- remove the definition of bpf_fsession_cookie()

v7:
- reuse bpf_session_cookie() instead of introduce new kfunc

v5:
- remove "cookie_cnt" in struct bpf_trampoline

v4:
- limit the maximum of the cookie count to 4
- store the session cookies before nr_regs in stack
---
 include/linux/bpf.h     | 15 +++++++++++++++
 kernel/bpf/trampoline.c | 13 +++++++++++--
 kernel/bpf/verifier.c   | 22 +++++++++++++++++++++-
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 2640ec2157e1..a416050e0dd2 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h

@@ -1231,6 +1231,7 @@ enum {

 #define BPF_TRAMP_M_NR_ARGS    0
 #define BPF_TRAMP_M_IS_RETURN  8
+#define BPF_TRAMP_M_COOKIE     9

 struct bpf_tramp_links {
        struct bpf_tramp_link *links[BPF_MAX_TRAMP_LINKS];

@@ -1783,6 +1784,7 @@ struct bpf_prog {
                                enforce_expected_attach_type:1, /* Enforce expected_attach_type checking at attach time */
                                call_get_stack:1, /* Do we call bpf_get_stack() or bpf_get_stackid() */
                                call_get_func_ip:1, /* Do we call get_func_ip() */
+                               call_session_cookie:1, /* Do we call bpf_session_cookie() */
                                tstamp_type_access:1, /* Accessed __sk_buff->tstamp_type */
                                sleepable:1;    /* BPF program is sleepable */
        enum bpf_prog_type      type;           /* Type of BPF program */

@@ -2191,6 +2193,19 @@ static inline int bpf_fsession_cnt(struct bpf_tramp_links *links)
        return cnt;
 }

+static inline int bpf_fsession_cookie_cnt(struct bpf_tramp_links *links)
+{
+       struct bpf_tramp_links fentries = links[BPF_TRAMP_FENTRY];
+       int cnt = 0;
+
+       for (int i = 0; i < links[BPF_TRAMP_FENTRY].nr_links; i++) {
+               if (fentries.links[i]->link.prog->call_session_cookie)
+                       cnt++;
+       }
+
+       return cnt;
+}
+
 int bpf_prog_ctx_arg_info_init(struct bpf_prog *prog,
                               const struct bpf_ctx_arg_aux *info, u32 cnt);

diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index 11e043049d68..29b4e00d860c 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c

@@ -592,6 +592,8 @@ static int bpf_freplace_check_tgt_prog(struct bpf_prog *tgt_prog)
        return 0;
 }

+#define BPF_TRAMP_MAX_COOKIES 4
+
 static int __bpf_trampoline_link_prog(struct bpf_tramp_link *link,
                                      struct bpf_trampoline *tr,
                                      struct bpf_prog *tgt_prog)

@@ -600,7 +602,7 @@ static int __bpf_trampoline_link_prog(struct bpf_tramp_link *link,
        struct bpf_tramp_link *link_exiting;
        struct bpf_fsession_link *fslink;
        struct hlist_head *prog_list;
-       int err = 0;
+       int err = 0, cookie_cnt = 0;
        int cnt = 0, i;

        kind = bpf_attach_type_to_tramp(link->link.prog);

@@ -637,11 +639,18 @@ static int __bpf_trampoline_link_prog(struct bpf_tramp_link *link,
                /* prog already linked */
                return -EBUSY;
        hlist_for_each_entry(link_exiting, prog_list, tramp_hlist) {
-               if (link_exiting->link.prog != link->link.prog)
+               if (link_exiting->link.prog != link->link.prog) {
+                       if (kind == BPF_TRAMP_FSESSION &&
+                           link_exiting->link.prog->call_session_cookie)
+                               cookie_cnt++;
                        continue;
+               }
                /* prog already linked */
                return -EBUSY;
        }
+       if (link->link.prog->call_session_cookie &&
+           cookie_cnt >= BPF_TRAMP_MAX_COOKIES)
+               return -E2BIG;

        hlist_add_head(&link->tramp_hlist, prog_list);
        if (kind == BPF_TRAMP_FSESSION) {

diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 1b0292a03186..b91fd8af2393 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c

@@ -12508,7 +12508,8 @@ get_kfunc_ptr_arg_type(struct bpf_verifier_env *env,
        bool arg_mem_size = false;

        if (meta->func_id == special_kfunc_list[KF_bpf_cast_to_kern_ctx] ||
-           meta->func_id == special_kfunc_list[KF_bpf_session_is_return])
+           meta->func_id == special_kfunc_list[KF_bpf_session_is_return] ||
+           meta->func_id == special_kfunc_list[KF_bpf_session_cookie])
                return KF_ARG_PTR_TO_CTX;

        if (argno + 1 < nargs &&

@@ -14294,6 +14295,9 @@ static int check_kfunc_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
                        return err;
        }

+       if (meta.func_id == special_kfunc_list[KF_bpf_session_cookie])
+               env->prog->call_session_cookie = true;
+
        return 0;
 }

@@ -22571,6 +22575,22 @@ static int fixup_kfunc_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
                insn_buf[1] = BPF_ALU64_IMM(BPF_RSH, BPF_REG_0, BPF_TRAMP_M_IS_RETURN);
                insn_buf[2] = BPF_ALU64_IMM(BPF_AND, BPF_REG_0, 1);
                *cnt = 3;
+       } else if (desc->func_id == special_kfunc_list[KF_bpf_session_cookie] &&
+                  env->prog->expected_attach_type == BPF_TRACE_FSESSION) {
+               /* inline bpf_session_cookie() for fsession:
+                *   __u64 *bpf_session_cookie(void *ctx)
+                *   {
+                *       u64 off = (((u64 *)ctx)[-1] >> BPF_TRAMP_M_COOKIE) & 0xFF;
+                *       return &((u64 *)ctx)[-off];
+                *   }
+                */
+               insn_buf[0] = BPF_LDX_MEM(BPF_DW, BPF_REG_0, BPF_REG_1, -8);
+               insn_buf[1] = BPF_ALU64_IMM(BPF_RSH, BPF_REG_0, BPF_TRAMP_M_COOKIE);
+               insn_buf[2] = BPF_ALU64_IMM(BPF_AND, BPF_REG_0, 0xFF);
+               insn_buf[3] = BPF_ALU64_IMM(BPF_LSH, BPF_REG_0, 3);
+               insn_buf[4] = BPF_ALU64_REG(BPF_SUB, BPF_REG_0, BPF_REG_1);
+               insn_buf[5] = BPF_ALU64_IMM(BPF_NEG, BPF_REG_0, 0);
+               *cnt = 6;
        }

        if (env->insn_aux_data[insn_idx].arg_prog) {
--

2.52.0