bpf: Support stack arguments for bpf functions

Currently BPF functions (subprogs) are limited to 5 register arguments.
With [1], the compiler can emit code that passes additional arguments
via a dedicated stack area through bpf register BPF_REG_PARAMS (r11),
introduced in an earlier patch ([2]).

The compiler uses positive r11 offsets for incoming (callee-side) args
and negative r11 offsets for outgoing (caller-side) args, following the
x86_64/arm64 calling convention direction. There is an 8-byte gap at
offset 0 separating two regions:
  Incoming (callee reads):   r11+8 (arg6), r11+16 (arg7), ...
  Outgoing (caller writes):  r11-8 (arg6), r11-16 (arg7), ...

The following is an example to show how stack arguments are saved
and transferred between caller and callee:

  int foo(int a1, int a2, int a3, int a4, int a5, int a6, int a7) {
    ...
    bar(a1, a2, a3, a4, a5, a6, a7, a8);
    ...
  }

  Caller (foo)                           Callee (bar)
  ============                           ============
  Incoming (positive offsets):           Incoming (positive offsets):

  r11+8:  [incoming arg 6]               r11+8:  [incoming arg 6] <-+
  r11+16: [incoming arg 7]               r11+16: [incoming arg 7] <-|+
                                         r11+24: [incoming arg 8] <-||+
  Outgoing (negative offsets):                                      |||
  r11-8:  [outgoing arg 6 to bar] -------->-------------------------+||
  r11-16: [outgoing arg 7 to bar] -------->--------------------------+|
  r11-24: [outgoing arg 8 to bar] -------->---------------------------+

If the bpf function has more than one call:

  int foo(int a1, int a2, int a3, int a4, int a5, int a6, int a7) {
    ...
    bar1(a1, a2, a3, a4, a5, a6, a7, a8);
    ...
    bar2(a1, a2, a3, a4, a5, a6, a7, a8, a9);
    ...
  }

  Caller (foo)                             Callee (bar2)
  ============                             ==============
  Incoming (positive offsets):             Incoming (positive offsets):

  r11+8:  [incoming arg 6]                 r11+8:  [incoming arg 6] <+
  r11+16: [incoming arg 7]                 r11+16: [incoming arg 7] <|+
                                           r11+24: [incoming arg 8] <||+
  Outgoing for bar2 (negative offsets):    r11+32: [incoming arg 9] <|||+
  r11-8:  [outgoing arg 6] ---->----------->-------------------------+|||
  r11-16: [outgoing arg 7] ---->----------->--------------------------+||
  r11-24: [outgoing arg 8] ---->----------->---------------------------+|
  r11-32: [outgoing arg 9] ---->----------->----------------------------+

The verifier tracks outgoing stack arguments in stack_arg_regs[] and
out_stack_arg_cnt in bpf_func_state, separately from the regular
r10 stack. The callee does not copy incoming args — it reads them
directly from the caller's outgoing slots at positive r11 offsets.
Similar to stacksafe(), introduce stack_arg_safe() to do pruning
check.

Outgoing stack arg slots are invalidated when the callee returns
(e.g. in prepare_func_exit), not at call time. This allows the callee to
read incoming args from the caller's outgoing slots during
verification. The following are a few examples.

Example 1:
  *(u64 *)(r11 - 8) = r6;
  *(u64 *)(r11 - 16) = r7;
  call bar1;                // arg6 = r6, arg7 = r7
  call bar2;                // expected with 2 stack arguments, failed

Example 2:
To fix the Example 1:
  *(u64 *)(r11 - 8) = r6;
  *(u64 *)(r11 - 16) = r7;
  call bar1;                // arg6 = r6, arg7 = r7
  *(u64 *)(r11 - 8) = r8;
  *(u64 *)(r11 - 16) = r9;
  call bar2;                // arg6 = r8, arg7 = r9

Example 3:
The compiler can hoist the shared stack arg stores above the branch:
  *(u64 *)(r11 - 16) = r7;
  if cond goto else;
    *(u64 *)(r11 - 8) = r8;
    call bar1;               // arg6 = r8, arg7 = r7
    goto end;
  else:
    *(u64 *)(r11 - 8) = r9;
    call bar2;               // arg6 = r9, arg7 = r7
  end:

Example 4:
Within a loop:
  loop:
    *(u64 *)(r11 - 8) = r6;  // arg6, before loop
    call bar;                // reuses arg6 each iteration
    if ... goto loop;

A separate max_out_stack_arg_cnt field in bpf_subprog_info tracks
the deepest outgoing slot actually written. This intends to
reject programs that write to slots beyond what any callee expects.
It is necessary for JIT.

Similar to typical compiler generated code, enforce the following
orderings:
  - all stack arg reads must be ahead of any stack arg write
  - all stack arg reads must be before any bpf func, kfunc and helpers
This is needed as JIT may emit 'mov' insns for read/write with
the same register and bpf function, kfunc and helper will invalidate
all arguments immediately after the call.

Callback functions with stack arguments need kernel setup parameter
types (including stack parameters) properly and then callback function
can retrieve such information for verification purpose.

Global subprogs and freplace with >5 args are not yet supported.

  [1] https://github.com/llvm/llvm-project/pull/189060
  [2] https://lore.kernel.org/bpf/20260423033506.2542005-1-yonghong.song@linux.dev/

Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
Link: https://lore.kernel.org/r/20260513045015.2385013-1-yonghong.song@linux.dev
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
index 203fb751eeaeb500fc914d5ee29cb02f4d7fdd0c..5398a02a12802adc66f4a8f7b08b57e74be6f528 100644 (file)
--- a/include/linux/bpf_verifier.h
+++ b/include/linux/bpf_verifier.h
@@ -402,6 +402,7 @@ struct bpf_func_state {
        bool in_callback_fn;
        bool in_async_callback_fn;
        bool in_exception_callback_fn;
+       bool no_stack_arg_load;
        /* For callback calling functions that limit number of possible
         * callback executions (e.g. bpf_loop) keeps track of current
         * simulated iteration number.
@@ -427,6 +428,9 @@ struct bpf_func_state {
         * `stack`. allocated_stack is always a multiple of BPF_REG_SIZE.
         */
        int allocated_stack;
+
+       u16 out_stack_arg_cnt; /* Number of outgoing on-stack argument slots */
+       struct bpf_reg_state *stack_arg_regs; /* Outgoing on-stack arguments */
 };
 
 #define MAX_CALL_FRAMES 8
@@ -465,8 +469,10 @@ struct bpf_jmp_history_entry {
        u64 linked_regs;
 };
 
-/* Maximum number of register states that can exist at once */
-#define BPF_ID_MAP_SIZE ((MAX_BPF_REG + MAX_BPF_STACK / BPF_REG_SIZE) * MAX_CALL_FRAMES)
+/* Maximum number of bpf_reg_state objects that can exist at once */
+#define MAX_STACK_ARG_SLOTS (MAX_BPF_FUNC_ARGS - MAX_BPF_FUNC_REG_ARGS)
+#define BPF_ID_MAP_SIZE ((MAX_BPF_REG + MAX_BPF_STACK / BPF_REG_SIZE + \
+                         MAX_STACK_ARG_SLOTS) * MAX_CALL_FRAMES)
 struct bpf_verifier_state {
        /* call stack tracking */
        struct bpf_func_state *frame[MAX_CALL_FRAMES];
@@ -561,12 +567,27 @@ bpf_get_spilled_reg(int slot, struct bpf_func_state *frame, u32 mask)
        return NULL;
 }
 
+static inline struct bpf_reg_state *
+bpf_get_spilled_stack_arg(int slot, struct bpf_func_state *frame)
+{
+       if (slot < frame->out_stack_arg_cnt &&
+           frame->stack_arg_regs[slot].type != NOT_INIT)
+               return &frame->stack_arg_regs[slot];
+       return NULL;
+}
+
 /* Iterate over 'frame', setting 'reg' to either NULL or a spilled register. */
 #define bpf_for_each_spilled_reg(iter, frame, reg, mask)                       \
        for (iter = 0, reg = bpf_get_spilled_reg(iter, frame, mask);            \
             iter < frame->allocated_stack / BPF_REG_SIZE;              \
             iter++, reg = bpf_get_spilled_reg(iter, frame, mask))
 
+/* Iterate over 'frame', setting 'reg' to either NULL or a spilled stack arg. */
+#define bpf_for_each_spilled_stack_arg(iter, frame, reg)               \
+       for (iter = 0, reg = bpf_get_spilled_stack_arg(iter, frame);   \
+            iter < frame->out_stack_arg_cnt;                          \
+            iter++, reg = bpf_get_spilled_stack_arg(iter, frame))
+
 #define bpf_for_each_reg_in_vstate_mask(__vst, __state, __reg, __mask, __expr)   \
        ({                                                               \
                struct bpf_verifier_state *___vstate = __vst;            \
@@ -584,6 +605,11 @@ bpf_get_spilled_reg(int slot, struct bpf_func_state *frame, u32 mask)
                                        continue;                        \
                                (void)(__expr);                          \
                        }                                                \
+                       bpf_for_each_spilled_stack_arg(___j, __state, __reg) { \
+                               if (!__reg)                              \
+                                       continue;                        \
+                               (void)(__expr);                          \
+                       }                                                \
                }                                                        \
        })
 
@@ -815,12 +841,21 @@ struct bpf_subprog_info {
        bool keep_fastcall_stack: 1;
        bool changes_pkt_data: 1;
        bool might_sleep: 1;
-       u8 arg_cnt:3;
+       u8 arg_cnt:4;
 
        enum priv_stack_mode priv_stack_mode;
-       struct bpf_subprog_arg_info args[MAX_BPF_FUNC_REG_ARGS];
+       struct bpf_subprog_arg_info args[MAX_BPF_FUNC_ARGS];
+       u16 stack_arg_cnt; /* incoming + max outgoing */
+       u16 max_out_stack_arg_cnt;
 };
 
+static inline u16 bpf_in_stack_arg_cnt(const struct bpf_subprog_info *sub)
+{
+       if (sub->arg_cnt > MAX_BPF_FUNC_REG_ARGS)
+               return sub->arg_cnt - MAX_BPF_FUNC_REG_ARGS;
+       return 0;
+}
+
 struct bpf_verifier_env;
 
 struct backtrack_state {

diff --git a/kernel/bpf/btf.c b/kernel/bpf/btf.c
index 099d7ca5a980001c365e9a21dc6fa138222f0f37..4fb8641546b8d24e37c3445cb6282beadf37ff65 100644 (file)
--- a/kernel/bpf/btf.c
+++ b/kernel/bpf/btf.c
@@ -7865,10 +7865,16 @@ int btf_prepare_func_args(struct bpf_verifier_env *env, int subprog)
        args = (const struct btf_param *)(t + 1);
        nargs = btf_type_vlen(t);
        sub->arg_cnt = nargs;
-       if (nargs > MAX_BPF_FUNC_REG_ARGS) {
-               if (!is_global)
-                       return -EINVAL;
-               bpf_log(log, "Global function %s() with %d > %d args. Buggy compiler.\n",
+       if (nargs > MAX_BPF_FUNC_ARGS) {
+               bpf_log(log, "kernel supports at most %d parameters, function %s has %d\n",
+                       MAX_BPF_FUNC_ARGS, tname, nargs);
+               return -EFAULT;
+       }
+       if (nargs > MAX_BPF_FUNC_REG_ARGS)
+               sub->stack_arg_cnt = nargs - MAX_BPF_FUNC_REG_ARGS;
+
+       if (is_global && nargs > MAX_BPF_FUNC_REG_ARGS) {
+               bpf_log(log, "global function %s has %d > %d args, stack args not supported\n",
                        tname, nargs, MAX_BPF_FUNC_REG_ARGS);
                return -EINVAL;
        }

diff --git a/kernel/bpf/fixups.c b/kernel/bpf/fixups.c
index fba9e8c008787547111ae60f4a508d19ea38500c..ba86039789fdf9ab30550a1b6fa0a86db7d0d156 100644 (file)
--- a/kernel/bpf/fixups.c
+++ b/kernel/bpf/fixups.c
@@ -1378,9 +1378,21 @@ int bpf_fixup_call_args(struct bpf_verifier_env *env)
        struct bpf_prog *prog = env->prog;
        struct bpf_insn *insn = prog->insnsi;
        bool has_kfunc_call = bpf_prog_has_kfunc_call(prog);
-       int i, depth;
+       int depth;
 #endif
-       int err = 0;
+       int i, err = 0;
+
+       for (i = 0; i < env->subprog_cnt; i++) {
+               struct bpf_subprog_info *subprog = &env->subprog_info[i];
+               u16 outgoing = subprog->stack_arg_cnt - bpf_in_stack_arg_cnt(subprog);
+
+               if (subprog->max_out_stack_arg_cnt > outgoing) {
+                       verbose(env,
+                               "func#%d writes %u stack arg slots, but calls only require %u\n",
+                               i, subprog->max_out_stack_arg_cnt, outgoing);
+                       return -EINVAL;
+               }
+       }
 
        if (env->prog->jit_requested &&
            !bpf_prog_is_offloaded(env->prog->aux)) {

diff --git a/kernel/bpf/states.c b/kernel/bpf/states.c
index bd9c229450503a3604d021c555752139429278e5..3ce6d2652b27c22a0f68bdcfe9c6dd856de9d6cc 100644 (file)
--- a/kernel/bpf/states.c
+++ b/kernel/bpf/states.c
@@ -833,6 +833,32 @@ static bool stacksafe(struct bpf_verifier_env *env, struct bpf_func_state *old,
        return true;
 }
 
+/*
+ * Compare stack arg slots between old and current states.
+ * Outgoing stack args are path-local state and must agree for pruning.
+ */
+static bool stack_arg_safe(struct bpf_verifier_env *env, struct bpf_func_state *old,
+                          struct bpf_func_state *cur, struct bpf_idmap *idmap,
+                          enum exact_level exact)
+{
+       int i, nslots;
+
+       nslots = max(old->out_stack_arg_cnt, cur->out_stack_arg_cnt);
+       for (i = 0; i < nslots; i++) {
+               struct bpf_reg_state *old_arg, *cur_arg;
+               struct bpf_reg_state not_init = { .type = NOT_INIT };
+
+               old_arg = i < old->out_stack_arg_cnt ?
+                         &old->stack_arg_regs[i] : &not_init;
+               cur_arg = i < cur->out_stack_arg_cnt ?
+                         &cur->stack_arg_regs[i] : &not_init;
+               if (!regsafe(env, old_arg, cur_arg, idmap, exact))
+                       return false;
+       }
+
+       return true;
+}
+
 static bool refsafe(struct bpf_verifier_state *old, struct bpf_verifier_state *cur,
                    struct bpf_idmap *idmap)
 {
@@ -915,6 +941,9 @@ static bool func_states_equal(struct bpf_verifier_env *env, struct bpf_func_stat
        if (old->callback_depth > cur->callback_depth)
                return false;
 
+       if (!old->no_stack_arg_load && cur->no_stack_arg_load)
+               return false;
+
        for (i = 0; i < MAX_BPF_REG; i++)
                if (((1 << i) & live_regs) &&
                    !regsafe(env, &old->regs[i], &cur->regs[i],
@@ -924,6 +953,9 @@ static bool func_states_equal(struct bpf_verifier_env *env, struct bpf_func_stat
        if (!stacksafe(env, old, cur, &env->idmap_scratch, exact))
                return false;
 
+       if (!stack_arg_safe(env, old, cur, &env->idmap_scratch, exact))
+               return false;
+
        return true;
 }
 

diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 1dd9736c2a1393685d040bf6a12c364cf8237c1b..a29b3003cbeca5d59f661146cbb3e766bfb6bed4 100644 (file)
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -1362,6 +1362,18 @@ static int copy_stack_state(struct bpf_func_state *dst, const struct bpf_func_st
                return -ENOMEM;
 
        dst->allocated_stack = src->allocated_stack;
+
+       /* copy stack args state */
+       n = src->out_stack_arg_cnt;
+       if (n) {
+               dst->stack_arg_regs = copy_array(dst->stack_arg_regs, src->stack_arg_regs, n,
+                                                sizeof(struct bpf_reg_state),
+                                                GFP_KERNEL_ACCOUNT);
+               if (!dst->stack_arg_regs)
+                       return -ENOMEM;
+       }
+
+       dst->out_stack_arg_cnt = src->out_stack_arg_cnt;
        return 0;
 }
 
@@ -1403,6 +1415,23 @@ static int grow_stack_state(struct bpf_verifier_env *env, struct bpf_func_state
        return 0;
 }
 
+static int grow_stack_arg_slots(struct bpf_verifier_env *env,
+                               struct bpf_func_state *state, int cnt)
+{
+       size_t old_n = state->out_stack_arg_cnt;
+
+       if (old_n >= cnt)
+               return 0;
+
+       state->stack_arg_regs = realloc_array(state->stack_arg_regs, old_n, cnt,
+                                             sizeof(struct bpf_reg_state));
+       if (!state->stack_arg_regs)
+               return -ENOMEM;
+
+       state->out_stack_arg_cnt = cnt;
+       return 0;
+}
+
 /* Acquire a pointer id from the env and update the state->refs to include
  * this new pointer reference.
  * On success, returns a valid pointer id to associate with the register
@@ -1565,6 +1594,7 @@ static void free_func_state(struct bpf_func_state *state)
 {
        if (!state)
                return;
+       kfree(state->stack_arg_regs);
        kfree(state->stack);
        kfree(state);
 }
@@ -4050,6 +4080,103 @@ static int check_stack_write(struct bpf_verifier_env *env,
        return err;
 }
 
+/*
+ * Write a value to the outgoing stack arg area.
+ * off is a negative offset from r11 (e.g. -8 for arg6, -16 for arg7).
+ */
+static int check_stack_arg_write(struct bpf_verifier_env *env, struct bpf_func_state *state,
+                                int off, struct bpf_reg_state *value_reg)
+{
+       int max_stack_arg_regs = MAX_BPF_FUNC_ARGS - MAX_BPF_FUNC_REG_ARGS;
+       struct bpf_subprog_info *subprog = &env->subprog_info[state->subprogno];
+       int spi = -off / BPF_REG_SIZE - 1;
+       struct bpf_reg_state *arg;
+       int err;
+
+       if (spi >= max_stack_arg_regs) {
+               verbose(env, "stack arg write offset %d exceeds max %d stack args\n",
+                       off, max_stack_arg_regs);
+               return -EINVAL;
+       }
+
+       err = grow_stack_arg_slots(env, state, spi + 1);
+       if (err)
+               return err;
+
+       /* Track the max outgoing stack arg slot count. */
+       if (spi + 1 > subprog->max_out_stack_arg_cnt)
+               subprog->max_out_stack_arg_cnt = spi + 1;
+
+       if (value_reg) {
+               state->stack_arg_regs[spi] = *value_reg;
+       } else {
+               /* BPF_ST: store immediate, treat as scalar */
+               arg = &state->stack_arg_regs[spi];
+               arg->type = SCALAR_VALUE;
+               __mark_reg_known(arg, env->prog->insnsi[env->insn_idx].imm);
+       }
+       state->no_stack_arg_load = true;
+       return 0;
+}
+
+/*
+ * Read a value from the incoming stack arg area.
+ * off is a positive offset from r11 (e.g. +8 for arg6, +16 for arg7).
+ */
+static int check_stack_arg_read(struct bpf_verifier_env *env, struct bpf_func_state *state,
+                               int off, int dst_regno)
+{
+       struct bpf_subprog_info *subprog = &env->subprog_info[state->subprogno];
+       struct bpf_verifier_state *vstate = env->cur_state;
+       int spi = off / BPF_REG_SIZE - 1;
+       struct bpf_func_state *caller, *cur;
+       struct bpf_reg_state *arg;
+
+       if (state->no_stack_arg_load) {
+               verbose(env, "r11 load must be before any r11 store or call insn\n");
+               return -EINVAL;
+       }
+
+       if (spi + 1 > bpf_in_stack_arg_cnt(subprog)) {
+               verbose(env, "invalid read from stack arg off %d depth %d\n",
+                       off, bpf_in_stack_arg_cnt(subprog) * BPF_REG_SIZE);
+               return -EACCES;
+       }
+
+       caller = vstate->frame[vstate->curframe - 1];
+       arg = &caller->stack_arg_regs[spi];
+       cur = vstate->frame[vstate->curframe];
+       cur->regs[dst_regno] = *arg;
+       return 0;
+}
+
+static int check_outgoing_stack_args(struct bpf_verifier_env *env, struct bpf_func_state *caller,
+                                    int nargs)
+{
+       int i, spi;
+
+       for (i = MAX_BPF_FUNC_REG_ARGS; i < nargs; i++) {
+               spi = i - MAX_BPF_FUNC_REG_ARGS;
+               if (spi >= caller->out_stack_arg_cnt ||
+                   caller->stack_arg_regs[spi].type == NOT_INIT) {
+                       verbose(env, "callee expects %d args, stack arg%d is not initialized\n",
+                               nargs, spi + 1);
+                       return -EFAULT;
+               }
+       }
+
+       return 0;
+}
+
+static struct bpf_reg_state *get_func_arg_reg(struct bpf_func_state *caller,
+                                             struct bpf_reg_state *regs, int arg)
+{
+       if (arg < MAX_BPF_FUNC_REG_ARGS)
+               return &regs[arg + 1];
+
+       return &caller->stack_arg_regs[arg - MAX_BPF_FUNC_REG_ARGS];
+}
+
 static int check_map_access_type(struct bpf_verifier_env *env, struct bpf_reg_state *reg,
                                 int off, int size, enum bpf_access_type type)
 {
@@ -6217,10 +6344,20 @@ static int check_load_mem(struct bpf_verifier_env *env, struct bpf_insn *insn,
                          bool strict_alignment_once, bool is_ldsx,
                          bool allow_trust_mismatch, const char *ctx)
 {
+       struct bpf_verifier_state *vstate = env->cur_state;
+       struct bpf_func_state *state = vstate->frame[vstate->curframe];
        struct bpf_reg_state *regs = cur_regs(env);
        enum bpf_reg_type src_reg_type;
        int err;
 
+       /* Handle stack arg read */
+       if (is_stack_arg_ldx(insn)) {
+               err = check_reg_arg(env, insn->dst_reg, DST_OP_NO_MARK);
+               if (err)
+                       return err;
+               return check_stack_arg_read(env, state, insn->off, insn->dst_reg);
+       }
+
        /* check src operand */
        err = check_reg_arg(env, insn->src_reg, SRC_OP);
        if (err)
@@ -6249,10 +6386,20 @@ static int check_load_mem(struct bpf_verifier_env *env, struct bpf_insn *insn,
 static int check_store_reg(struct bpf_verifier_env *env, struct bpf_insn *insn,
                           bool strict_alignment_once)
 {
+       struct bpf_verifier_state *vstate = env->cur_state;
+       struct bpf_func_state *state = vstate->frame[vstate->curframe];
        struct bpf_reg_state *regs = cur_regs(env);
        enum bpf_reg_type dst_reg_type;
        int err;
 
+       /* Handle stack arg write */
+       if (is_stack_arg_stx(insn)) {
+               err = check_reg_arg(env, insn->src_reg, SRC_OP);
+               if (err)
+                       return err;
+               return check_stack_arg_write(env, state, insn->off, regs + insn->src_reg);
+       }
+
        /* check src1 operand */
        err = check_reg_arg(env, insn->src_reg, SRC_OP);
        if (err)
@@ -8860,6 +9007,15 @@ static void clear_caller_saved_regs(struct bpf_verifier_env *env,
        }
 }
 
+static void invalidate_outgoing_stack_args(const struct bpf_verifier_env *env,
+                                          struct bpf_func_state *state)
+{
+       int i, nslots = state->out_stack_arg_cnt;
+
+       for (i = 0; i < nslots; i++)
+               bpf_mark_reg_not_init(env, &state->stack_arg_regs[i]);
+}
+
 typedef int (*set_callee_state_fn)(struct bpf_verifier_env *env,
                                   struct bpf_func_state *caller,
                                   struct bpf_func_state *callee,
@@ -8922,6 +9078,7 @@ static int btf_check_func_arg_match(struct bpf_verifier_env *env, int subprog,
                                    struct bpf_reg_state *regs)
 {
        struct bpf_subprog_info *sub = subprog_info(env, subprog);
+       struct bpf_func_state *caller = cur_func(env);
        struct bpf_verifier_log *log = &env->log;
        u32 i;
        int ret;
@@ -8930,13 +9087,16 @@ static int btf_check_func_arg_match(struct bpf_verifier_env *env, int subprog,
        if (ret)
                return ret;
 
+       ret = check_outgoing_stack_args(env, caller, sub->arg_cnt);
+       if (ret)
+               return ret;
+
        /* check that BTF function arguments match actual types that the
         * verifier sees.
         */
        for (i = 0; i < sub->arg_cnt; i++) {
                argno_t argno = argno_from_arg(i + 1);
-               u32 regno = i + 1;
-               struct bpf_reg_state *reg = &regs[regno];
+               struct bpf_reg_state *reg = get_func_arg_reg(caller, regs, i);
                struct bpf_subprog_arg_info *arg = &sub->args[i];
 
                if (arg->arg_type == ARG_ANYTHING) {
@@ -9124,6 +9284,8 @@ static int check_func_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
                           int *insn_idx)
 {
        struct bpf_verifier_state *state = env->cur_state;
+       struct bpf_subprog_info *caller_info;
+       u16 callee_incoming, stack_arg_cnt;
        struct bpf_func_state *caller;
        int err, subprog, target_insn;
 
@@ -9166,6 +9328,7 @@ static int check_func_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
                /* mark global subprog for verifying after main prog */
                subprog_aux(env, subprog)->called = true;
                clear_caller_saved_regs(env, caller->regs);
+               invalidate_outgoing_stack_args(env, cur_func(env));
 
                /* All non-void global functions return a 64-bit SCALAR_VALUE. */
                if (!subprog_returns_void(env, subprog)) {
@@ -9177,6 +9340,16 @@ static int check_func_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
                return 0;
        }
 
+       /*
+        * Track caller's total stack arg count (incoming + max outgoing).
+        * This is needed so the JIT knows how much stack arg space to allocate.
+        */
+       caller_info = &env->subprog_info[caller->subprogno];
+       callee_incoming = bpf_in_stack_arg_cnt(&env->subprog_info[subprog]);
+       stack_arg_cnt = bpf_in_stack_arg_cnt(caller_info) + callee_incoming;
+       if (stack_arg_cnt > caller_info->stack_arg_cnt)
+               caller_info->stack_arg_cnt = stack_arg_cnt;
+
        /* for regular function entry setup new frame and continue
         * from that frame.
         */
@@ -9534,6 +9707,7 @@ static int prepare_func_exit(struct bpf_verifier_env *env, int *insn_idx)
         * bpf_throw, this will be done by copy_verifier_state for extra frames. */
        free_func_state(callee);
        state->frame[state->curframe--] = NULL;
+       invalidate_outgoing_stack_args(env, caller);
 
        /* for callbacks widen imprecise scalars to make programs like below verify:
         *
@@ -10160,6 +10334,7 @@ static int check_helper_call(struct bpf_verifier_env *env, struct bpf_insn *insn
                bpf_mark_reg_not_init(env, &regs[caller_saved[i]]);
                check_reg_arg(env, caller_saved[i], DST_OP_NO_MARK);
        }
+       invalidate_outgoing_stack_args(env, cur_func(env));
 
        /* helper call returns 64-bit value. */
        regs[BPF_REG_0].subreg_def = DEF_NOT_SUBREG;
@@ -12842,6 +13017,7 @@ static int check_kfunc_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
                bpf_mark_reg_not_init(env, &regs[regno]);
                regs[regno].subreg_def = DEF_NOT_SUBREG;
        }
+       invalidate_outgoing_stack_args(env, cur_func(env));
 
        /* Check return type */
        t = btf_type_skip_modifiers(desc_btf, meta.func_proto->type, NULL);
@@ -16961,6 +17137,14 @@ static int do_check_insn(struct bpf_verifier_env *env, bool *do_print_state)
                return check_store_reg(env, insn, false);
 
        case BPF_ST: {
+               /* Handle stack arg write (store immediate) */
+               if (is_stack_arg_st(insn)) {
+                       struct bpf_verifier_state *vstate = env->cur_state;
+                       struct bpf_func_state *state = vstate->frame[vstate->curframe];
+
+                       return check_stack_arg_write(env, state, insn->off, NULL);
+               }
+
                enum bpf_reg_type dst_reg_type;
 
                err = check_reg_arg(env, insn->dst_reg, SRC_OP);
@@ -16995,6 +17179,8 @@ static int do_check_insn(struct bpf_verifier_env *env, bool *do_print_state)
                                }
                        }
                        mark_reg_scratched(env, BPF_REG_0);
+                       if (bpf_in_stack_arg_cnt(&env->subprog_info[cur_func(env)->subprogno]))
+                               cur_func(env)->no_stack_arg_load = true;
                        if (insn->src_reg == BPF_PSEUDO_CALL)
                                return check_func_call(env, insn, &env->insn_idx);
                        if (insn->src_reg == BPF_PSEUDO_KFUNC_CALL)
@@ -18110,7 +18296,7 @@ static int do_check_common(struct bpf_verifier_env *env, int subprog)
                                goto out;
                        }
                }
-               for (i = BPF_REG_1; i <= sub->arg_cnt; i++) {
+               for (i = BPF_REG_1; i <= min_t(u32, sub->arg_cnt, MAX_BPF_FUNC_REG_ARGS); i++) {
                        arg = &sub->args[i - BPF_REG_1];
                        reg = &regs[i];
 
@@ -18153,6 +18339,12 @@ static int do_check_common(struct bpf_verifier_env *env, int subprog)
                                goto out;
                        }
                }
+               if (env->prog->type == BPF_PROG_TYPE_EXT && sub->arg_cnt > MAX_BPF_FUNC_REG_ARGS) {
+                       verbose(env, "freplace programs with >%d args not supported yet\n",
+                               MAX_BPF_FUNC_REG_ARGS);
+                       ret = -EINVAL;
+                       goto out;
+               }
        } else {
                /* if main BPF program has associated BTF info, validate that
                 * it's matching expected signature, and otherwise mark BTF
@@ -18160,8 +18352,11 @@ static int do_check_common(struct bpf_verifier_env *env, int subprog)
                 */
                if (env->prog->aux->func_info_aux) {
                        ret = btf_prepare_func_args(env, 0);
-                       if (ret || sub->arg_cnt != 1 || sub->args[0].arg_type != ARG_PTR_TO_CTX)
+                       if (ret || sub->arg_cnt != 1 || sub->args[0].arg_type != ARG_PTR_TO_CTX) {
                                env->prog->aux->func_info_aux[0].unreliable = true;
+                               sub->arg_cnt = 1;
+                               sub->stack_arg_cnt = 0;
+                       }
                }
 
                /* 1st arg to a function */