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

[Puranjay Mohan <puranjay@kernel.org>]

Xu Kuohai <xukuohai@huaweicloud.com> writes:

> On 8/11/2026 3:09 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 <puranjay@kernel.org>
>> ---
>>   Documentation/bpf/kfuncs.rst  |  6 +--
>>   arch/arm64/net/bpf_jit_comp.c | 85 ++++++++++++++++++++++++++++-------
>>   2 files changed, 73 insertions(+), 18 deletions(-)
>> 
>> diff --git a/Documentation/bpf/kfuncs.rst b/Documentation/bpf/kfuncs.rst
>> index 1004eb0bec617..d9cc2ab1cf018 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 7aad17a51f006..1cd327d213e34 100644
>> --- a/arch/arm64/net/bpf_jit_comp.c
>> +++ b/arch/arm64/net/bpf_jit_comp.c
>> @@ -2377,6 +2377,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))
>> @@ -2550,26 +2555,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);
>
> Maybe I'm missing something, do we need to validate whether the
> address in the src register is really inside the current bpf
> prog's arena?

The JIT can assume that it is a kernel address into the arena as it
comes from struct ops.

>
>> +}
>> +
>>   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;
>> +		}
>>   	}
>>   
>>   	/*
>> @@ -2585,6 +2622,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) {
>> @@ -2594,6 +2634,8 @@ 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);
>> +			if (arena_arg)
>> +				emit_arena_arg_conv(ctx, tmp, tmp, nullable, base_lo);
>
> I think it is still worth adding a brief comment to explain
> that this conversion never happens for original function calls,
> though we have a WARN_ON_ONCE below preventing this, and the commit
> message notes it.

Added the comment in v2.