From 50de1c47a41d4031f6002969e71dc6954dedb6b2 Mon Sep 17 00:00:00 2001
From: Puranjay Mohan <puranjay@kernel.org>
Date: Thu, 13 Aug 2026 12:03:48 -0700
Subject: [PATCH 1/7] bpf, arm64: Fix stack-passed arguments for indirect
 trampolines

save_args() reads stack-passed arguments relative to FP assuming the
trampoline is entered through the fentry call from a traced function, in
which case both the parent frame (FP/x9) and the traced function frame
(FP/LR) are saved before FP is set, so the arguments start at FP + 32.

An indirect trampoline for a struct_ops callback is entered through a
function pointer (blr), so only the FP/LR frame is pushed and the
arguments start at FP + 16, not FP + 32. Every stack-passed argument of
a struct_ops callback with more than eight argument slots is read two
slots off.

This went unnoticed because no struct_ops member passed arguments on the
stack until bpf_testmod_ops3::test_arena_stack, added by
commit 2d4de9a493a0 ("selftests/bpf: Test stack-passed struct_ops arena arguments").
That member covers this on arm64 once the JIT gains arena argument
support later in this series. Pass is_struct_ops into save_args() and
pick the offset accordingly, mirroring the x86 fix.

Fixes: 9014cf56f13d ("bpf, arm64: Support up to 12 function arguments")
Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
Reviewed-by: Xu Kuohai <xukuohai@huawei.com>
Link: https://lore.kernel.org/bpf/20260813190356.335181-2-puranjay@kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/arm64/net/bpf_jit_comp.c | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index 74b4083791da..7938b3422d3c 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -2525,9 +2525,8 @@ static void clear_garbage(struct jit_ctx *ctx, int reg, int effective_bytes)
 }
 
 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)
+		      const struct btf_func_model *m, const struct arg_aux *a,
+		      bool for_call_origin, bool is_struct_ops)
 {
 	int i;
 	int reg;
@@ -2547,7 +2546,15 @@ static void save_args(struct jit_ctx *ctx, int bargs_off, int oargs_off,
 		bargs_off += 8;
 	}
 
-	soff = 32; /* on stack arguments start from FP + 32 */
+	/*
+	 * On-stack arguments start above the frame(s) pushed by the trampoline
+	 * prologue. Entered through the fentry call from a traced function, the
+	 * prologue saves both the parent (FP/x9) and the traced function
+	 * (FP/LR) frames, so the arguments start at FP + 32. A struct_ops
+	 * callback is called indirectly and only the FP/LR frame is saved, so
+	 * they start at FP + 16.
+	 */
+	soff = is_struct_ops ? 16 : 32;
 	doff = (for_call_origin ? oargs_off : bargs_off);
 
 	/* save on stack arguments */
@@ -2737,7 +2744,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);
+	save_args(ctx, bargs_off, oargs_off, m, a, false, is_struct_ops);
 
 	/* save callee saved registers */
 	emit(A64_STR64I(A64_R(19), A64_SP, regs_off), ctx);
@@ -2786,7 +2793,7 @@ 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);
+		save_args(ctx, bargs_off, oargs_off, m, a, true, is_struct_ops);
 		/* call original func */
 		emit(A64_LDR64I(A64_R(10), A64_SP, retaddr_off), ctx);
 		emit(A64_ADR(A64_LR, AARCH64_INSN_SIZE * 2), ctx);

From f4adb983ef3d7a23d7f67834bc12cec528312899 Mon Sep 17 00:00:00 2001
From: Tejun Heo <tj@kernel.org>
Date: Thu, 13 Aug 2026 12:03:49 -0700
Subject: [PATCH 2/7] arm64: insn: Add encoder for ADD/SUB (extended register)

The insn library encodes the immediate and shifted-register forms of
ADD/SUB but not the extended-register form. The BPF JIT wants it to
rebase a 32-bit arena offset onto the arena kernel base in a single
instruction, add xN, xBASE, wN, uxtw, instead of a separate zero-extend
followed by a plain add.

Add aarch64_insn_gen_add_sub_extended_reg(), modeled on the
shifted-register generator. The option and imm3 fields occupy the same
bits as the shifted form's shift amount, so they are encoded through the
existing IMM_6 field type. The opt field in bits 23:22 is part of the
opcode here rather than a shift type, and any value other than 00 is
unallocated, so the decode masks cover it.

Note that register 31 does not mean the same thing in the two forms: in
the extended-register encoding it is SP for Rn, and for Rd unless the
instruction sets the flags, while it stays XZR for Rm. Callers porting a
shifted-register site that passes A64_ZR need to be aware of that, so
say so above the function.

Signed-off-by: Tejun Heo <tj@kernel.org>
Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
Reviewed-by: Xu Kuohai <xukuohai@huawei.com>
Link: https://lore.kernel.org/bpf/20260813190356.335181-3-puranjay@kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/arm64/include/asm/insn.h | 23 ++++++++++++++
 arch/arm64/lib/insn.c         | 60 +++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/arch/arm64/include/asm/insn.h b/arch/arm64/include/asm/insn.h
index cc0702fa64a7..1ce75a442638 100644
--- a/arch/arm64/include/asm/insn.h
+++ b/arch/arm64/include/asm/insn.h
@@ -205,6 +205,18 @@ enum aarch64_insn_adsb_type {
 	AARCH64_INSN_ADSB_SUB_SETFLAGS
 };
 
+/* option field of add/sub (extended register) */
+enum aarch64_insn_extend_type {
+	AARCH64_INSN_EXTEND_UXTB,
+	AARCH64_INSN_EXTEND_UXTH,
+	AARCH64_INSN_EXTEND_UXTW,
+	AARCH64_INSN_EXTEND_UXTX,
+	AARCH64_INSN_EXTEND_SXTB,
+	AARCH64_INSN_EXTEND_SXTH,
+	AARCH64_INSN_EXTEND_SXTW,
+	AARCH64_INSN_EXTEND_SXTX,
+};
+
 enum aarch64_insn_movewide_type {
 	AARCH64_INSN_MOVEWIDE_ZERO,
 	AARCH64_INSN_MOVEWIDE_KEEP,
@@ -378,6 +390,10 @@ __AARCH64_INSN_FUNCS(add,	0x7F200000, 0x0B000000)
 __AARCH64_INSN_FUNCS(adds,	0x7F200000, 0x2B000000)
 __AARCH64_INSN_FUNCS(sub,	0x7F200000, 0x4B000000)
 __AARCH64_INSN_FUNCS(subs,	0x7F200000, 0x6B000000)
+__AARCH64_INSN_FUNCS(add_ext,	0x7FE00000, 0x0B200000)
+__AARCH64_INSN_FUNCS(adds_ext,	0x7FE00000, 0x2B200000)
+__AARCH64_INSN_FUNCS(sub_ext,	0x7FE00000, 0x4B200000)
+__AARCH64_INSN_FUNCS(subs_ext,	0x7FE00000, 0x6B200000)
 __AARCH64_INSN_FUNCS(madd,	0x7FE08000, 0x1B000000)
 __AARCH64_INSN_FUNCS(msub,	0x7FE08000, 0x1B008000)
 __AARCH64_INSN_FUNCS(udiv,	0x7FE0FC00, 0x1AC00800)
@@ -637,6 +653,13 @@ u32 aarch64_insn_gen_add_sub_shifted_reg(enum aarch64_insn_register dst,
 					 int shift,
 					 enum aarch64_insn_variant variant,
 					 enum aarch64_insn_adsb_type type);
+u32 aarch64_insn_gen_add_sub_extended_reg(enum aarch64_insn_register dst,
+					  enum aarch64_insn_register src,
+					  enum aarch64_insn_register reg,
+					  enum aarch64_insn_extend_type extend,
+					  int shift,
+					  enum aarch64_insn_variant variant,
+					  enum aarch64_insn_adsb_type type);
 u32 aarch64_insn_gen_data1(enum aarch64_insn_register dst,
 			   enum aarch64_insn_register src,
 			   enum aarch64_insn_variant variant,
diff --git a/arch/arm64/lib/insn.c b/arch/arm64/lib/insn.c
index 37ce75f7f1f0..e70ac0238515 100644
--- a/arch/arm64/lib/insn.c
+++ b/arch/arm64/lib/insn.c
@@ -986,6 +986,66 @@ u32 aarch64_insn_gen_add_sub_shifted_reg(enum aarch64_insn_register dst,
 	return aarch64_insn_encode_immediate(AARCH64_INSN_IMM_6, insn, shift);
 }
 
+/*
+ * Unlike the shifted-register form, register 31 is not XZR everywhere here:
+ * it encodes SP for @src, and for @dst too unless @type sets the flags. Only
+ * @reg keeps the XZR meaning.
+ */
+u32 aarch64_insn_gen_add_sub_extended_reg(enum aarch64_insn_register dst,
+					  enum aarch64_insn_register src,
+					  enum aarch64_insn_register reg,
+					  enum aarch64_insn_extend_type extend,
+					  int shift,
+					  enum aarch64_insn_variant variant,
+					  enum aarch64_insn_adsb_type type)
+{
+	u32 insn;
+
+	switch (type) {
+	case AARCH64_INSN_ADSB_ADD:
+		insn = aarch64_insn_get_add_ext_value();
+		break;
+	case AARCH64_INSN_ADSB_SUB:
+		insn = aarch64_insn_get_sub_ext_value();
+		break;
+	case AARCH64_INSN_ADSB_ADD_SETFLAGS:
+		insn = aarch64_insn_get_adds_ext_value();
+		break;
+	case AARCH64_INSN_ADSB_SUB_SETFLAGS:
+		insn = aarch64_insn_get_subs_ext_value();
+		break;
+	default:
+		pr_err("%s: unknown add/sub encoding %d\n", __func__, type);
+		return AARCH64_BREAK_FAULT;
+	}
+
+	switch (variant) {
+	case AARCH64_INSN_VARIANT_32BIT:
+		break;
+	case AARCH64_INSN_VARIANT_64BIT:
+		insn |= AARCH64_INSN_SF_BIT;
+		break;
+	default:
+		pr_err("%s: unknown variant encoding %d\n", __func__, variant);
+		return AARCH64_BREAK_FAULT;
+	}
+
+	if (shift < 0 || shift > 4) {
+		pr_err("%s: invalid shift encoding %d\n", __func__, shift);
+		return AARCH64_BREAK_FAULT;
+	}
+
+	insn = aarch64_insn_encode_register(AARCH64_INSN_REGTYPE_RD, insn, dst);
+
+	insn = aarch64_insn_encode_register(AARCH64_INSN_REGTYPE_RN, insn, src);
+
+	insn = aarch64_insn_encode_register(AARCH64_INSN_REGTYPE_RM, insn, reg);
+
+	/* option in bits [15:13] and imm3 in [12:10] together fill IMM_6 */
+	return aarch64_insn_encode_immediate(AARCH64_INSN_IMM_6, insn,
+					     (extend << 3) | shift);
+}
+
 u32 aarch64_insn_gen_data1(enum aarch64_insn_register dst,
 			   enum aarch64_insn_register src,
 			   enum aarch64_insn_variant variant,

From 760cb40cfd6a4e5781e4a32b46b1ef49950fd271 Mon Sep 17 00:00:00 2001
From: Puranjay Mohan <puranjay@kernel.org>
Date: Thu, 13 Aug 2026 12:03:50 -0700
Subject: [PATCH 3/7] bpf, arm64: JIT __arena kfunc argument rebasing

Implement arena argument rebasing for kfunc calls on arm64. x28 already
holds kern_vm_start whenever the prog has an arena, and the newly added
extended-register add zero-extends the 32-bit arena offset in place, so
an unconditional argument costs a single instruction emitted right
before the call:

  add xN, x28, wN, uxtw

A nullable argument first truncates into wN so that a zero offset leaves
xN holding a real NULL, then tests it and jumps over the add:

  mov wN, wN
  cbz wN, 1f
  add xN, x28, wN, uxtw
1:

The rebase is native code generated after constant blinding has run on
the BPF instruction stream, so blinding never sees it and needs no
special handling. The emitted count depends only on the kfunc model, so
it is identical across JIT passes.

bpf_jit_supports_arena_args() is not flipped yet; that happens when the
struct_ops trampoline side is in place as well.

Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
Reviewed-by: Xu Kuohai <xukuohai@huawei.com>
Link: https://lore.kernel.org/bpf/20260813190356.335181-4-puranjay@kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 arch/arm64/net/bpf_jit.h      | 11 +++++++++
 arch/arm64/net/bpf_jit_comp.c | 42 +++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/arch/arm64/net/bpf_jit.h b/arch/arm64/net/bpf_jit.h
index d13de4222cfb..b2fe6e6dcf44 100644
--- a/arch/arm64/net/bpf_jit.h
+++ b/arch/arm64/net/bpf_jit.h
@@ -243,6 +243,17 @@
 /* Rn - Rm; set condition flags */
 #define A64_CMP(sf, Rn, Rm) A64_SUBS(sf, A64_ZR, Rn, Rm)
 
+/* Add/subtract (extended register) */
+#define A64_ADDSUB_EREG(sf, Rd, Rn, Rm, ext, shift, type) \
+	aarch64_insn_gen_add_sub_extended_reg(Rd, Rn, Rm, \
+		AARCH64_INSN_EXTEND_##ext, shift, A64_VARIANT(sf), \
+		AARCH64_INSN_ADSB_##type)
+/* Rd = Rn + (EXT(Rm) << shift) */
+#define A64_ADD_EXT(sf, Rd, Rn, Rm, ext, shift) \
+	A64_ADDSUB_EREG(sf, Rd, Rn, Rm, ext, shift, ADD)
+/* Rd = Rn + (u32)Rm */
+#define A64_ADD_UXTW(Rd, Rn, Rm) A64_ADD_EXT(1, Rd, Rn, Rm, UXTW, 0)
+
 /* Data-processing (1 source) */
 #define A64_DATA1(sf, Rd, Rn, type) aarch64_insn_gen_data1(Rd, Rn, \
 	A64_VARIANT(sf), AARCH64_INSN_DATA1_##type)
diff --git a/arch/arm64/net/bpf_jit_comp.c b/arch/arm64/net/bpf_jit_comp.c
index 7938b3422d3c..e31490c0e331 100644
--- a/arch/arm64/net/bpf_jit_comp.c
+++ b/arch/arm64/net/bpf_jit_comp.c
@@ -1256,6 +1256,43 @@ static void emit_stack_arg_store_imm(s32 imm, s16 bpf_off, const u8 tmp, struct
 	}
 }
 
+/*
+ * Rebase the __arena args of a kfunc call to arena kernel addresses,
+ * xN = kern_vm_start + (u32)xN, with the arena base register holding
+ * kern_vm_start. A nullable arg preserves NULL by skipping the add, tested
+ * on the truncated value as arena NULL is offset 0.
+ */
+static int emit_kfunc_arena_args(struct jit_ctx *ctx, const struct bpf_insn *insn)
+{
+	const u8 arena_vm_base = bpf2a64[ARENA_VM_START];
+	const struct btf_func_model *fm;
+	int i;
+
+	fm = bpf_jit_find_kfunc_model(ctx->prog, insn);
+	if (!fm)
+		return -EINVAL;
+
+	for (i = 0; i < min_t(int, fm->nr_args, MAX_BPF_FUNC_REG_ARGS); i++) {
+		const u8 reg = bpf2a64[BPF_REG_1 + i];
+		u8 flags = fm->arg_flags[i];
+
+		if (!(flags & BTF_FMODEL_ARENA_ARG))
+			continue;
+		if (WARN_ON_ONCE(!ctx->arena_vm_start))
+			return -EINVAL;
+
+		if (flags & BTF_FMODEL_NULLABLE_ARG) {
+			/* 32-bit mov clears the upper 32 bits */
+			emit(A64_MOV(0, reg, reg), ctx);
+			/* skip the add so that NULL stays NULL */
+			emit(A64_CBZ(0, reg, 2), ctx);
+		}
+		emit(A64_ADD_UXTW(reg, arena_vm_base, reg), ctx);
+	}
+
+	return 0;
+}
+
 /* JITs an eBPF instruction.
  * Returns:
  * 0  - successfully JITed an 8-byte eBPF instruction.
@@ -1678,6 +1715,11 @@ static int build_insn(const struct bpf_verifier_env *env, const struct bpf_insn
 					    &func_addr, &func_addr_fixed);
 		if (ret < 0)
 			return ret;
+		if (insn->src_reg == BPF_PSEUDO_KFUNC_CALL) {
+			ret = emit_kfunc_arena_args(ctx, insn);
+			if (ret < 0)
+				return ret;
+		}
 		emit_call(func_addr, ctx);
 		/*
 		 * Call to arch_bpf_timed_may_goto() is emitted by the

From bb5bad6a78b6334ab9f8ab99cdb92bae11d9acea Mon Sep 17 00:00:00 2001
From: Puranjay Mohan <puranjay@kernel.org>
Date: Thu, 13 Aug 2026 12:03:51 -0700
Subject: [PATCH 4/7] bpf, arm64: Convert struct_ops arena arguments in the
 trampoline

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>
Reviewed-by: Xu Kuohai <xukuohai@huawei.com>
Link: https://lore.kernel.org/bpf/20260813190356.335181-5-puranjay@kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 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 10e725cbe64c..85f73e0bbd0f 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 e31490c0e331..c18e005a41db 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);

From 1c5bc60f957d853ec7183d79419360979dddd421 Mon Sep 17 00:00:00 2001
From: Puranjay Mohan <puranjay@kernel.org>
Date: Thu, 13 Aug 2026 12:03:52 -0700
Subject: [PATCH 5/7] selftests/bpf: Add arm64 JIT-sequence tests for __arena
 kfunc arguments

Pin the arm64 counterparts of the x86-64 rebase sequences: the single
extended-register add for an unconditional argument, the nullable
truncate-test-and-skip variant, and all five argument registers in one
call. The nullable cases use a local label so the branch is pinned to
the instruction right after the add, and the label line does not spell
out the call because arm64 emits either a direct bl or a materialize-
and-blr pair depending on the distance to the kfunc.

Note that on arm64 an unconditional argument is one instruction with
nothing to anchor it against, so arena_arg_jit_rebase alone cannot tell
the two forms apart; it only requires that nothing is emitted between
the rebase and the call. The args5 test is what pins the distinction,
since its four consecutive adds leave no room for a nullable
truncate-and-branch pair between them.

Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
Acked-by: Xu Kuohai <xukuohai@huawei.com>
Link: https://lore.kernel.org/bpf/20260813190356.335181-6-puranjay@kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../selftests/bpf/progs/arena_kfunc_jit.c     | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tools/testing/selftests/bpf/progs/arena_kfunc_jit.c b/tools/testing/selftests/bpf/progs/arena_kfunc_jit.c
index c9b918662616..b5a01cbc33a7 100644
--- a/tools/testing/selftests/bpf/progs/arena_kfunc_jit.c
+++ b/tools/testing/selftests/bpf/progs/arena_kfunc_jit.c
@@ -32,6 +32,10 @@ __jited("	movl	%edi, %edi")
 __jited("	addq	%r12, %rdi")
 __jited("...")
 __jited("	callq	{{.*}}")
+__arch_arm64
+__jited("...")
+__jited("	add	x0, x28, w0, uxtw")
+__jited("	{{(bl|mov)	.*}}")
 __success
 int arena_arg_jit_rebase(void *ctx)
 {
@@ -48,6 +52,12 @@ __jited("	testl	%edi, %edi")
 __jited("	je	L0")
 __jited("	addq	%r12, %rdi")
 __jited("L0:	callq	{{.*}}")
+__arch_arm64
+__jited("...")
+__jited("	mov	w0, w0")
+__jited("	cbz	w0, L0")
+__jited("	add	x0, x28, w0, uxtw")
+__jited("L0:	{{.*}}")
 __success
 int arena_arg_jit_nullable(void *ctx)
 {
@@ -72,6 +82,16 @@ __jited("	testl	%r8d, %r8d")
 __jited("	je	L0")
 __jited("	addq	%r12, %r8")
 __jited("L0:	callq	{{.*}}")
+__arch_arm64
+__jited("...")
+__jited("	add	x0, x28, w0, uxtw")
+__jited("	add	x1, x28, w1, uxtw")
+__jited("	add	x2, x28, w2, uxtw")
+__jited("	add	x3, x28, w3, uxtw")
+__jited("	mov	w4, w4")
+__jited("	cbz	w4, L0")
+__jited("	add	x4, x28, w4, uxtw")
+__jited("L0:	{{.*}}")
 __success
 int arena_arg_jit_args5(void *ctx)
 {

From 05a3575f225c81712b925c28b1caa1c24ac6ed3b Mon Sep 17 00:00:00 2001
From: Puranjay Mohan <puranjay@kernel.org>
Date: Thu, 13 Aug 2026 12:03:53 -0700
Subject: [PATCH 6/7] selftests/bpf: Enable __arena argument tests on arm64

The arena kfunc and struct_ops argument tests were restricted to x86-64
because it was the only JIT that implemented the conversions. arm64 does
now, so let them run there too: tag every program in arena_kfunc.c with
__arch_arm64 in addition to __arch_x86_64, and widen the __x86_64__
guards in the struct_ops arena test.

Without this the tests report SKIP on arm64 rather than exercising the
newly added JIT support.

Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
Acked-by: Xu Kuohai <xukuohai@huawei.com>
Link: https://lore.kernel.org/bpf/20260813190356.335181-7-puranjay@kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../selftests/bpf/prog_tests/test_struct_ops_arena.c   | 10 +++++-----
 tools/testing/selftests/bpf/progs/arena_kfunc.c        |  9 +++++++++
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/test_struct_ops_arena.c b/tools/testing/selftests/bpf/prog_tests/test_struct_ops_arena.c
index 940ec2cda0d5..7f9f54ba3fbe 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_struct_ops_arena.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_struct_ops_arena.c
@@ -6,7 +6,7 @@
 #include "struct_ops_arena_attach.skel.h"
 #include "struct_ops_arena_fail.skel.h"
 
-#if defined(__x86_64__)
+#if defined(__x86_64__) || defined(__aarch64__)
 /*
  * Attach callbacks with __arena and __arena__nullable arguments and drive
  * them through the bpf_testmod_ops3_call_test_arena*() kfuncs.
@@ -111,11 +111,11 @@ static void arena_arg_attach(void)
 void serial_test_struct_ops_arena(void)
 {
 	/*
-	 * Arena struct_ops arguments need JIT support, currently x86-64 only.
-	 * Elsewhere verification fails with "JIT does not support arena
-	 * arguments", so the programs cannot even load.
+	 * Arena struct_ops arguments need JIT support, currently x86-64 and
+	 * arm64 only. Elsewhere verification fails with "JIT does not support
+	 * arena arguments", so the programs cannot even load.
 	 */
-#if defined(__x86_64__)
+#if defined(__x86_64__) || defined(__aarch64__)
 	if (test__start_subtest("arena_arg"))
 		arena_arg();
 	if (test__start_subtest("arena_arg_fail"))
diff --git a/tools/testing/selftests/bpf/progs/arena_kfunc.c b/tools/testing/selftests/bpf/progs/arena_kfunc.c
index cdcea889da58..bf0d304e0e59 100644
--- a/tools/testing/selftests/bpf/progs/arena_kfunc.c
+++ b/tools/testing/selftests/bpf/progs/arena_kfunc.c
@@ -27,6 +27,7 @@ volatile u64 stash;
 
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __success __retval(0)
 int arena_arg_forms(void *ctx)
 {
@@ -70,6 +71,7 @@ int arena_arg_forms(void *ctx)
  */
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __success __retval(0)
 int arena_arg_rebase(void *ctx)
 {
@@ -111,6 +113,7 @@ int arena_arg_rebase(void *ctx)
 
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __success __retval(0)
 int arena_args5(void *ctx)
 {
@@ -142,6 +145,7 @@ int arena_args5(void *ctx)
 
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __success __retval(0)
 int arena_arg_mixed(void *ctx)
 {
@@ -169,6 +173,7 @@ int arena_arg_mixed(void *ctx)
 /* kernel-side faults on unpopulated pages recover via the scratch page */
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __success __retval(0)
 int arena_arg_unpopulated(void *ctx)
 {
@@ -189,6 +194,7 @@ int arena_arg_unpopulated(void *ctx)
 
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __failure __msg("arena pointer requires a program with an associated arena")
 int arena_arg_no_arena(void *ctx)
 {
@@ -198,6 +204,7 @@ int arena_arg_no_arena(void *ctx)
 
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __failure __msg("is not a pointer to arena or scalar")
 int arena_arg_bad_reg(void *ctx)
 {
@@ -213,6 +220,7 @@ int arena_arg_bad_reg(void *ctx)
 	defined(__BPF_FEATURE_STACK_ARGUMENT)
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __failure __msg("arena pointer cannot be a stack argument")
 int arena_arg_stack(void *ctx)
 {
@@ -223,6 +231,7 @@ int arena_arg_stack(void *ctx)
 #else
 SEC("syscall")
 __arch_x86_64
+__arch_arm64
 __description("arena_arg_stack: not supported, dummy test")
 __success
 int arena_arg_stack(void *ctx)

From 197d34b169435a447830ddd915d596431b9f311a Mon Sep 17 00:00:00 2001
From: Puranjay Mohan <puranjay@kernel.org>
Date: Thu, 13 Aug 2026 12:03:54 -0700
Subject: [PATCH 7/7] selftests/bpf: Test a multi-slot argument before a
 struct_ops arena argument

The trampoline reads the __arena flag from the btf_func_model per
argument but stores the ctx one register slot at a time, so the two only
line up if every preceding argument occupies exactly one slot. Every
arena-bearing member of bpf_testmod_ops3 takes single-slot arguments, so
nothing exercises the mapping and a mis-indexed arg_flags lookup would
go unnoticed on any architecture.

Add test_arena_multislot(), whose first argument is a 16-byte struct
passed by value. It fills ctx[0] and ctx[1], putting the arena pointer
at argument index one but slot two. The callback checks both halves of
the struct before dereferencing ctx[2], so a JIT that walks registers
instead of arguments converts the wrong slot and fails the test.

Signed-off-by: Puranjay Mohan <puranjay@kernel.org>
Acked-by: Xu Kuohai <xukuohai@huawei.com>
Link: https://lore.kernel.org/bpf/20260813190356.335181-8-puranjay@kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../selftests/bpf/progs/struct_ops_arena.c    | 24 +++++++++++++++++++
 .../selftests/bpf/test_kmods/bpf_testmod.c    | 15 ++++++++++++
 .../selftests/bpf/test_kmods/bpf_testmod.h    |  8 +++++++
 .../bpf/test_kmods/bpf_testmod_kfunc.h        |  1 +
 4 files changed, 48 insertions(+)

diff --git a/tools/testing/selftests/bpf/progs/struct_ops_arena.c b/tools/testing/selftests/bpf/progs/struct_ops_arena.c
index ba04c73d8d96..8aa8639df91f 100644
--- a/tools/testing/selftests/bpf/progs/struct_ops_arena.c
+++ b/tools/testing/selftests/bpf/progs/struct_ops_arena.c
@@ -59,11 +59,28 @@ int test_arena_stack_cb(unsigned long long *ctx)
 	return 0;
 }
 
+SEC("struct_ops/test_arena_multislot")
+int test_arena_multislot_cb(unsigned long long *ctx)
+{
+	u64 __arena *ptr = (u64 __arena *)ctx[2];
+
+	arena_touch++;
+	/*
+	 * The 16-byte struct occupies ctx[0] and ctx[1], so @ptr is argument
+	 * one but slot two. Getting that wrong hands the callback a scalar.
+	 */
+	if (ctx[0] != 11 || ctx[1] != 22)
+		return 0xbad;
+	*ptr += 1;
+	return 0;
+}
+
 SEC(".struct_ops.link")
 struct bpf_testmod_ops3 testmod_arena = {
 	.test_arena = (void *)test_arena_cb,
 	.test_arena_nullable = (void *)test_arena_nullable_cb,
 	.test_arena_stack = (void *)test_arena_stack_cb,
+	.test_arena_multislot = (void *)test_arena_multislot_cb,
 };
 
 SEC("syscall")
@@ -109,6 +126,13 @@ int trigger(void *ctx)
 	if (*val != 44)
 		return 9;
 
+	/* a multi-slot arg precedes the arena pointer here */
+	ret = bpf_testmod_ops3_call_test_arena_multislot((u64 *)val);
+	if (ret)
+		return 10;
+	if (*val != 45)
+		return 11;
+
 	bpf_arena_free_pages(&arena, (void __arena *)val, 1);
 #endif
 	return 0;
diff --git a/tools/testing/selftests/bpf/test_kmods/bpf_testmod.c b/tools/testing/selftests/bpf/test_kmods/bpf_testmod.c
index a6133f7521f3..9366a3c578f1 100644
--- a/tools/testing/selftests/bpf/test_kmods/bpf_testmod.c
+++ b/tools/testing/selftests/bpf/test_kmods/bpf_testmod.c
@@ -402,12 +402,19 @@ static int bpf_testmod_ops3__test_arena_stack(u64 a, u64 b, u64 c, u64 d,
 	return 0;
 }
 
+static int bpf_testmod_ops3__test_arena_multislot(struct bpf_testmod_arena_pair p,
+						  u64 *ptr__arena)
+{
+	return 0;
+}
+
 static struct bpf_testmod_ops3 __bpf_testmod_ops3 = {
 	.test_1 = bpf_testmod_test_3,
 	.test_2 = bpf_testmod_test_4,
 	.test_arena = bpf_testmod_ops3__test_arena,
 	.test_arena_nullable = bpf_testmod_ops3__test_arena_nullable,
 	.test_arena_stack = bpf_testmod_ops3__test_arena_stack,
+	.test_arena_multislot = bpf_testmod_ops3__test_arena_multislot,
 };
 
 static void bpf_testmod_test_struct_ops3(void)
@@ -441,6 +448,13 @@ __bpf_kfunc int bpf_testmod_ops3_call_test_arena_stack(u64 *ptr__arena)
 	return st_ops3->test_arena_stack(1, 2, 3, 4, 5, 6, 7, 8, ptr__arena);
 }
 
+__bpf_kfunc int bpf_testmod_ops3_call_test_arena_multislot(u64 *ptr__arena)
+{
+	struct bpf_testmod_arena_pair p = { .a = 11, .b = 22 };
+
+	return st_ops3->test_arena_multislot(p, ptr__arena);
+}
+
 struct bpf_testmod_btf_type_tag_1 {
 	int a;
 };
@@ -852,6 +866,7 @@ BTF_ID_FLAGS(func, bpf_testmod_ops3_call_test_2)
 BTF_ID_FLAGS(func, bpf_testmod_ops3_call_test_arena)
 BTF_ID_FLAGS(func, bpf_testmod_ops3_call_test_arena_nullable)
 BTF_ID_FLAGS(func, bpf_testmod_ops3_call_test_arena_stack)
+BTF_ID_FLAGS(func, bpf_testmod_ops3_call_test_arena_multislot)
 BTF_ID_FLAGS(func, bpf_kfunc_get_default_trusted_ptr_test);
 BTF_ID_FLAGS(func, bpf_kfunc_put_default_trusted_ptr_test);
 BTF_KFUNCS_END(bpf_testmod_common_kfunc_ids)
diff --git a/tools/testing/selftests/bpf/test_kmods/bpf_testmod.h b/tools/testing/selftests/bpf/test_kmods/bpf_testmod.h
index 33f2af5b7085..210b919290cc 100644
--- a/tools/testing/selftests/bpf/test_kmods/bpf_testmod.h
+++ b/tools/testing/selftests/bpf/test_kmods/bpf_testmod.h
@@ -103,6 +103,12 @@ struct bpf_testmod_ops2 {
 	int (*test_1)(void);
 };
 
+/* 16 bytes, so it takes two argument slots when passed by value */
+struct bpf_testmod_arena_pair {
+	u64 a;
+	u64 b;
+};
+
 struct bpf_testmod_ops3 {
 	int (*test_1)(void);
 	int (*test_2)(void);
@@ -112,6 +118,8 @@ struct bpf_testmod_ops3 {
 	/* enough leading args to force @ptr onto the stack on x86 and arm64 */
 	int (*test_arena_stack)(u64 a, u64 b, u64 c, u64 d, u64 e, u64 f,
 				u64 g, u64 h, u64 *ptr);
+	/* a multi-slot leading arg, so @ptr is not at the slot its arg index suggests */
+	int (*test_arena_multislot)(struct bpf_testmod_arena_pair p, u64 *ptr);
 };
 
 struct st_ops_args {
diff --git a/tools/testing/selftests/bpf/test_kmods/bpf_testmod_kfunc.h b/tools/testing/selftests/bpf/test_kmods/bpf_testmod_kfunc.h
index c4383acb53c1..7d81070eefe7 100644
--- a/tools/testing/selftests/bpf/test_kmods/bpf_testmod_kfunc.h
+++ b/tools/testing/selftests/bpf/test_kmods/bpf_testmod_kfunc.h
@@ -123,6 +123,7 @@ void bpf_testmod_test_mod_kfunc(int i) __ksym;
 int bpf_testmod_ops3_call_test_arena(__u64 *ptr__arena) __ksym;
 int bpf_testmod_ops3_call_test_arena_nullable(__u64 *ptr__arena__nullable) __ksym;
 int bpf_testmod_ops3_call_test_arena_stack(__u64 *ptr__arena) __ksym;
+int bpf_testmod_ops3_call_test_arena_multislot(__u64 *ptr__arena) __ksym;
 
 __u64 bpf_kfunc_call_test1(struct sock *sk, __u32 a, __u64 b,
 				__u32 c, __u64 d) __ksym;