Re: [PATCH bpf-next v2 4/7] bpf, arm64: Convert struct_ops arena arguments in the trampoline

On 8/14/2026 3:03 AM, Puranjay Mohan wrote:
> Implement the struct_ops arena argument conversion on arm64. save_args()
> receives the arena base from bpf_tramp_arena_base() and consults the
> btf_func_model argument flags as it copies each native argument into the
> BPF ctx, routing a marked argument through x10 with the low half of the
> base materialized once into x11:
> 
>    sub w10, wsrc, w11    /* truncate and clear the upper 32 bits */
>    str x10, [sp, #slot]
> 
> A nullable argument tests the full 64-bit kernel pointer first:
> 
>    mov x10, xsrc
>    cbz x10, 1f
>    sub w10, w10, w11
> 1:
>    str x10, [sp, #slot]
> 
> The 32-bit subtraction is sufficient since (u32)(kaddr - base) ==
> (u32)kaddr - (u32)base, and it clears the upper half as the JITs require
> of arena pointer registers. Stack-passed arguments already reload
> through x10, so only the subtraction (and the NULL test) is inserted
> there.
> 
> The register loop now walks arguments rather than registers so that the
> per-argument flags line up with the slots a multi-slot argument occupies;
> the sequence of stores is otherwise unchanged. bpf_tramp_arena_base()
> returns a base only for a single-program struct_ops indirect trampoline,
> so a tracing trampoline emits exactly what it did before and never
> touches x11. The size probe reruns the same emission with the same model
> and nodes, so the image size matches by construction.
> 
> Conversion must never reach the original function, which takes kernel
> addresses. That holds because BPF_TRAMP_F_INDIRECT is incompatible with
> BPF_TRAMP_F_CALL_ORIG, so pass 0 rather than the base to the call-origin
> save_args() and assert the flag combination the same way x86 does,
> rather than leaving the invariant to a comment.
> 
> With both the kfunc and struct_ops directions implemented, flip
> bpf_jit_supports_arena_args() on for arm64 and drop the x86-64-only
> qualifier from the kfunc documentation.
> 
> Signed-off-by: Puranjay Mohan <[email protected]>
> ---
>   Documentation/bpf/kfuncs.rst  |  6 +--
>   arch/arm64/net/bpf_jit_comp.c | 92 +++++++++++++++++++++++++++++------
>   2 files changed, 80 insertions(+), 18 deletions(-)
> 
> diff --git a/Documentation/bpf/kfuncs.rst b/Documentation/bpf/kfuncs.rst
> index 10e725cbe64cd..85f73e0bbd0ff 100644
> --- a/Documentation/bpf/kfuncs.rst
> +++ b/Documentation/bpf/kfuncs.rst
> @@ -301,9 +301,9 @@ An example is given below::
>           }
>   
>   Calling such a kfunc requires the program to use an arena map and a JIT with
> -arena argument support (currently x86-64); verification fails otherwise. The
> -program can pass any value without compromising the kernel. A value that does
> -not point into the arena is a program bug.
> +arena argument support (currently x86-64 and arm64); verification fails
> +otherwise. The program can pass any value without compromising the kernel. A
> +value that does not point into the arena is a program bug.
>   
>   The suffixes have the same meaning on the arguments of struct_ops stub
>   functions, with the conversion running in the opposite direction. The
> diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
> index e31490c0e331c..c18e005a41dbe 100644
> --- a/arch/arm64/net/bpf_jit_comp.c
> +++ b/arch/arm64/net/bpf_jit_comp.c
> @@ -2393,6 +2393,11 @@ bool bpf_jit_supports_stack_args(void)
>   	return true;
>   }
>   
> +bool bpf_jit_supports_arena_args(void)
> +{
> +	return true;
> +}
> +
>   void *bpf_arch_text_copy(void *dst, void *src, size_t len)
>   {
>   	if (!aarch64_insn_copy(dst, src, len))
> @@ -2566,26 +2571,58 @@ static void clear_garbage(struct jit_ctx *ctx, int reg, int effective_bytes)
>   	}
>   }
>   
> +/*
> + * Convert an arena kernel address into the arena pointer form on its way into
> + * the BPF ctx, dst = (u32)(src - kern_vm_start), with @base_lo holding the low
> + * 32 bits of kern_vm_start. A nullable arg preserves NULL, tested on the full
> + * 64-bit kernel pointer. The 32-bit subtraction both truncates and clears the
> + * upper half, so the stored value satisfies the JIT invariant for arena
> + * pointer registers.
> + */
> +static void emit_arena_arg_conv(struct jit_ctx *ctx, u8 dst, u8 src, bool nullable, u8 base_lo)
> +{
> +	if (nullable) {
> +		if (dst != src)
> +			emit(A64_MOV(1, dst, src), ctx);
> +		/* skip the subtraction so that NULL stays NULL */
> +		emit(A64_CBZ(1, dst, 2), ctx);
> +		src = dst;
> +	}
> +	emit(A64_SUB(0, dst, src, base_lo), ctx);
> +}
> +
>   static void save_args(struct jit_ctx *ctx, int bargs_off, int oargs_off,
>   		      const struct btf_func_model *m, const struct arg_aux *a,
> -		      bool for_call_origin, bool is_struct_ops)
> +		      bool for_call_origin, bool is_struct_ops, u64 arena_base)
>   {
> -	int i;
> -	int reg;
> -	int doff;
> -	int soff;
> -	int slots;
>   	u8 tmp = bpf2a64[TMP_REG_1];
> +	u8 base_lo = bpf2a64[TMP_REG_2];
> +	int i, reg, doff, soff, slots;
> +
> +	/* only the low 32 bits of the base take part in the subtraction */
> +	if (arena_base)
> +		emit_a64_mov_i(0, base_lo, (s32)(u32)arena_base, ctx);
>   
>   	/* store arguments to the stack for the bpf program, or restore
>   	 * arguments from stack for the original function
>   	 */
> -	for (reg = 0; reg < a->regs_for_args; reg++) {
> -		emit(for_call_origin ?
> -		     A64_LDR64I(reg, A64_SP, bargs_off) :
> -		     A64_STR64I(reg, A64_SP, bargs_off),
> -		     ctx);
> -		bargs_off += 8;
> +	for (i = 0, reg = 0; i < a->args_in_regs; i++) {
> +		bool arena_arg = arena_base && (m->arg_flags[i] & BTF_FMODEL_ARENA_ARG);
> +		bool nullable = m->arg_flags[i] & BTF_FMODEL_NULLABLE_ARG;
> +
> +		slots = (m->arg_size[i] + 7) / 8;
> +		while (slots-- > 0) {
> +			if (for_call_origin) {
> +				emit(A64_LDR64I(reg, A64_SP, bargs_off), ctx);
> +			} else if (arena_arg) {
> +				emit_arena_arg_conv(ctx, tmp, reg, nullable, base_lo);
> +				emit(A64_STR64I(tmp, A64_SP, bargs_off), ctx);
> +			} else {
> +				emit(A64_STR64I(reg, A64_SP, bargs_off), ctx);
> +			}
> +			reg++;
> +			bargs_off += 8;
> +		}
>   	}
>   
>   	/*
> @@ -2601,6 +2638,9 @@ static void save_args(struct jit_ctx *ctx, int bargs_off, int oargs_off,
>   
>   	/* save on stack arguments */
>   	for (i = a->args_in_regs; i < m->nr_args; i++) {
> +		bool arena_arg = arena_base && (m->arg_flags[i] & BTF_FMODEL_ARENA_ARG);
> +		bool nullable = m->arg_flags[i] & BTF_FMODEL_NULLABLE_ARG;
> +
>   		slots = (m->arg_size[i] + 7) / 8;
>   		/* verifier ensures arg_size <= 16, so slots equals 1 or 2 */
>   		while (slots-- > 0) {
> @@ -2610,6 +2650,15 @@ static void save_args(struct jit_ctx *ctx, int bargs_off, int oargs_off,
>   			 */
>   			if (slots == 0 && !for_call_origin)
>   				clear_garbage(ctx, tmp, m->arg_size[i] % 8);
> +			/*
> +			 * No guard on for_call_origin here: only the indirect
> +			 * trampoline is given a base, and it never calls the
> +			 * original function, so arguments are never converted
> +			 * on their way back out to it. See the WARN_ON_ONCE()
> +			 * in prepare_trampoline().
> +			 */
> +			if (arena_arg)
> +				emit_arena_arg_conv(ctx, tmp, tmp, nullable, base_lo);
>   			emit(A64_STR64I(tmp, A64_SP, doff), ctx);
>   			soff += 8;
>   			doff += 8;
> @@ -2669,8 +2718,21 @@ static int prepare_trampoline(struct jit_ctx *ctx, struct bpf_tramp_image *im,
>   	bool is_struct_ops = is_struct_ops_tramp(fentry);
>   	int cookie_off, cookie_cnt, cookie_bargs_off;
>   	int fsession_cnt = bpf_fsession_cnt(tnodes);
> +	u64 arena_base;
>   	u64 func_meta;
>   
> +	/*
> +	 * F_INDIRECT is only compatible with F_RET_FENTRY_RET, it is explicitly
> +	 * incompatible with F_CALL_ORIG | F_SKIP_FRAME | F_IP_ARG because
> +	 * @func_addr. Arena conversion relies on this: bpf_tramp_arena_base()
> +	 * only returns a base for the indirect trampoline, which therefore
> +	 * never calls the original function with converted arguments.
> +	 */
> +	WARN_ON_ONCE((flags & BPF_TRAMP_F_INDIRECT) &&
> +		     (flags & ~(BPF_TRAMP_F_INDIRECT | BPF_TRAMP_F_RET_FENTRY_RET)));
> +
> +	arena_base = bpf_tramp_arena_base(m, tnodes, flags);
> +
>   	/* trampoline stack layout:
>   	 *                    [ parent ip         ]
>   	 *                    [ FP                ]
> @@ -2786,7 +2848,7 @@ static int prepare_trampoline(struct jit_ctx *ctx, struct bpf_tramp_image *im,
>   	store_func_meta(ctx, func_meta, func_meta_off);
>   
>   	/* save args for bpf */
> -	save_args(ctx, bargs_off, oargs_off, m, a, false, is_struct_ops);
> +	save_args(ctx, bargs_off, oargs_off, m, a, false, is_struct_ops, arena_base);
>   
>   	/* save callee saved registers */
>   	emit(A64_STR64I(A64_R(19), A64_SP, regs_off), ctx);
> @@ -2834,8 +2896,8 @@ static int prepare_trampoline(struct jit_ctx *ctx, struct bpf_tramp_image *im,
>   	}
>   
>   	if (flags & BPF_TRAMP_F_CALL_ORIG) {
> -		/* save args for original func */
> -		save_args(ctx, bargs_off, oargs_off, m, a, true, is_struct_ops);
> +		/* the original func takes kernel addresses, never converted ones */
> +		save_args(ctx, bargs_off, oargs_off, m, a, true, is_struct_ops, 0);
>   		/* call original func */
>   		emit(A64_LDR64I(A64_R(10), A64_SP, retaddr_off), ctx);
>   		emit(A64_ADR(A64_LR, AARCH64_INSN_SIZE * 2), ctx);

Reviewed-by: Xu Kuohai <[email protected]>