[PATCH bpf-next v2 10/11] bpf, x86: Pass a kfunc's 16-byte aligned stack argument 16-byte aligned

[Yonghong Song <yonghong.song@linux.dev>]

An argument aligned to 16 bytes, an __int128 or a struct holding one,
goes in a 16-byte aligned stack slot in the x86-64 calling convention.
The kernel keeps its stack only 8-byte aligned, being built with
-mpreferred-stack-boundary=3, yet a caller still realigns RSP before it
passes such an argument on the stack, and the callee can take the
alignment for granted. For

	u64 f(u64 a, u64 b, u64 c, u64 d, u64 e, u64 f, __int128 g)

both GCC and clang fold ((unsigned long)&g & 15) == 0 to true, and GCC
passes &g on as an aligned __int128 pointer without copying g.

The JIT places a kfunc's stack arguments at the offsets from RSP that
the convention gives them, holes included, but RSP itself is only 8-byte
aligned in a BPF program, so a kfunc can get such an argument at an
address that is not 16-byte aligned.

When a kfunc call passes an argument aligned to 16 bytes on the stack,
test RSP before the call, and if it is off by 8, copy the stack argument
slots to the first 16-byte aligned place below them. A push of each slot
does that in four bytes: it takes the address of its operand before it
moves RSP, so one push instruction, repeated, copies the slots from the
top down. After the call, put RSP back from RBP. The arm64 JIT keeps SP
16-byte aligned, so only x86-64 needs this.

bpf_jit_stack_slots_to_align(), next to bpf_jit_place_args(), returns
how many stack slots a call has to realign, from where its arguments
are placed. The x86-64 JIT now places a kfunc's arguments once, for it
and for bpf_jit_plan_slot_moves(), the half of bpf_jit_plan_arg_moves()
that plans the moves from a given placement. bpf_jit_plan_arg_moves()
keeps placing the arguments itself, for arm64.

Fixes: a3195057406f ("bpf, x86: Move kfunc arguments into the x86-64 calling convention")
Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
---
 arch/x86/net/bpf_jit_comp.c | 65 ++++++++++++++++++++++++++++++++++---
 include/linux/filter.h      |  5 +++
 kernel/bpf/core.c           | 50 +++++++++++++++++++++++-----
 3 files changed, 107 insertions(+), 13 deletions(-)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 9bf29dc39498..ad1e1187154f 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -1918,17 +1918,19 @@ static s32 x86_arg_stack_off(u8 pos)
 
 /*
  * Move the arguments the x86-64 ABI places somewhere other than the argument
- * slot the BPF calling convention gave them. @stack_base addresses the
+ * slot the BPF calling convention gave them, @pos_of_slot being where
+ * bpf_jit_place_args() put each of the @nslots slots. @stack_base addresses the
  * outgoing stack argument area from RBP. Return the number of emitted bytes.
  */
-static int emit_kfunc_arg_moves(const struct btf_func_model *fm, s32 stack_base, u8 **pprog)
+static int emit_kfunc_arg_moves(const u8 *pos_of_slot, u32 nslots, s32 stack_base,
+				u8 **pprog)
 {
 	struct bpf_jit_arg_move moves[BPF_JIT_MAX_ARG_MOVES];
 	const u8 nreg = x86_arg_abi.nr_arg_regs;
 	u8 *prog = *pprog, *start = prog;
 	u32 i, n;
 
-	n = bpf_jit_plan_arg_moves(&x86_arg_abi, fm, moves);
+	n = bpf_jit_plan_slot_moves(pos_of_slot, nslots, moves);
 
 	for (i = 0; i < n; i++) {
 		u8 dst = moves[i].dst, src = moves[i].src, reg;
@@ -1965,6 +1967,36 @@ static int emit_kfunc_arg_moves(const struct btf_func_model *fm, s32 stack_base,
 	return prog - start;
 }
 
+/*
+ * When RSP is not 16-byte aligned, copy the @n on-stack argument slots at RSP
+ * to the first 16-byte aligned place below them. A push takes the address of
+ * its memory operand before it moves RSP, so the same push copies each slot in
+ * turn, from the top down. The caller puts RSP back after the call. Return the
+ * number of emitted bytes.
+ */
+static int emit_kfunc_align_stack_args(u32 n, u8 **pprog)
+{
+	u8 *prog = *pprog, *start = prog, *skip;
+	/* a kfunc takes MAX_BPF_FUNC_ARGS slots at most: disp fits in a disp8 */
+	u8 disp = (n - 1) * 8;
+	u32 i;
+
+	EMIT4(0x40, 0xF6, 0xC4, 0x08);		/* test spl, 8 */
+	EMIT2(X86_JE, 0);			/* je over the copy */
+	skip = prog;
+	if (!(n & 1)) {
+		/* an even number of slots keeps RSP's alignment */
+		EMIT4(0x48, 0x83, 0xEC, 8);	/* sub rsp, 8 */
+		disp += 8;
+	}
+	for (i = 0; i < n; i++)
+		EMIT4(0xFF, 0x74, 0x24, disp);	/* push qword ptr [rsp + disp] */
+	skip[-1] = prog - skip;
+
+	*pprog = prog;
+	return prog - start;
+}
+
 /*
  * Rebase the __arena args of a kfunc call to arena kernel addresses,
  * rN = kern_vm_start + (u32)rN, with R12 holding kern_vm_start. A nullable
@@ -2953,6 +2985,8 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 			/* call */
 		case BPF_JMP | BPF_CALL: {
 			const struct btf_func_model *fm = NULL;
+			u8 pos_of_slot[MAX_BPF_FUNC_ARG_SLOTS];
+			u32 nslots, align_slots = 0;
 
 			func = (u8 *) __bpf_call_base + imm32;
 			if (src_reg == BPF_PSEUDO_CALL && tail_call_reachable) {
@@ -2969,8 +3003,20 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 				if (err < 0)
 					return err;
 				ip += err;
-				ip += emit_kfunc_arg_moves(fm, outgoing_arg_base -
-							   outgoing_rsp, &prog);
+				nslots = bpf_jit_place_args(&x86_arg_abi, fm, pos_of_slot);
+				ip += emit_kfunc_arg_moves(pos_of_slot, nslots,
+							   outgoing_arg_base - outgoing_rsp,
+							   &prog);
+				/*
+				 * The kernel keeps the stack only 8-byte aligned,
+				 * but a caller realigns RSP to pass an argument
+				 * aligned to 16 bytes on the stack, and the callee
+				 * may rely on its slot's alignment.
+				 */
+				align_slots = bpf_jit_stack_slots_to_align(&x86_arg_abi, fm,
+									   pos_of_slot);
+				if (align_slots)
+					ip += emit_kfunc_align_stack_args(align_slots, &prog);
 			}
 			if (func == (u8 *)arch_bpf_timed_may_goto) {
 				u32 fp = priv_frame_ptr ? X86_REG_R9 : BPF_REG_FP;
@@ -2993,6 +3039,15 @@ static int do_jit(struct bpf_verifier_env *env, struct bpf_prog *bpf_prog, int *
 				return -EINVAL;
 			if (priv_frame_ptr)
 				pop_r9(&prog);
+			if (align_slots) {
+				s32 off = outgoing_arg_base - outgoing_rsp;
+
+				/* lea rsp, [rbp + off], wherever the copy left it */
+				if (is_imm8(off))
+					EMIT4(0x48, 0x8D, 0x65, (u8)off);
+				else
+					EMIT3_off32(0x48, 0x8D, 0xA5, off);
+			}
 			/*
 			 * A kfunc returning more than 8 bytes hands the second
 			 * half back in RDX (the native ABI's second return reg),
diff --git a/include/linux/filter.h b/include/linux/filter.h
index 9339c6131f8f..55875f3a1510 100644
--- a/include/linux/filter.h
+++ b/include/linux/filter.h
@@ -1285,9 +1285,14 @@ struct bpf_jit_arg_move {
 	u8 src;
 };
 
+u32 bpf_jit_plan_slot_moves(const u8 *pos_of_slot, u32 nslots,
+			    struct bpf_jit_arg_move *moves);
 u32 bpf_jit_plan_arg_moves(const struct bpf_jit_arg_abi *abi,
 			   const struct btf_func_model *fm,
 			   struct bpf_jit_arg_move *moves);
+u32 bpf_jit_stack_slots_to_align(const struct bpf_jit_arg_abi *abi,
+				 const struct btf_func_model *fm,
+				 const u8 *pos_of_slot);
 u64 bpf_arch_uaddress_limit(void);
 void arch_bpf_stack_walk(bool (*consume_fn)(void *cookie, u64 ip, u64 sp, u64 bp), void *cookie);
 u64 arch_bpf_timed_may_goto(void);
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index 05c89396119a..3d6d504cdb5f 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -3379,15 +3379,10 @@ u32 bpf_jit_place_args(const struct bpf_jit_arg_abi *abi,
 	return slot;
 }
 
-u32 bpf_jit_plan_arg_moves(const struct bpf_jit_arg_abi *abi,
-			   const struct btf_func_model *fm,
-			   struct bpf_jit_arg_move *moves)
+u32 bpf_jit_plan_slot_moves(const u8 *pos_of_slot, u32 nslots,
+			    struct bpf_jit_arg_move *moves)
 {
-	u8 pos_of_slot[MAX_BPF_FUNC_ARG_SLOTS];
-	u32 nslots, n = 0, s, back;
-
-	nslots = bpf_jit_place_args(abi, fm, pos_of_slot);
-	back = nslots;
+	u32 n = 0, s, back = nslots;
 
 	/*
 	 * An argument is two eightbytes at most, so it frees one register at
@@ -3420,6 +3415,45 @@ u32 bpf_jit_plan_arg_moves(const struct bpf_jit_arg_abi *abi,
 	return n;
 }
 
+u32 bpf_jit_plan_arg_moves(const struct bpf_jit_arg_abi *abi,
+			   const struct btf_func_model *fm,
+			   struct bpf_jit_arg_move *moves)
+{
+	u8 pos_of_slot[MAX_BPF_FUNC_ARG_SLOTS];
+	u32 nslots;
+
+	nslots = bpf_jit_place_args(abi, fm, pos_of_slot);
+	return bpf_jit_plan_slot_moves(pos_of_slot, nslots, moves);
+}
+
+/*
+ * Return how many stack slots a call passes when one of its stack arguments is
+ * aligned to 16 bytes, or 0 otherwise. @pos_of_slot is where
+ * bpf_jit_place_args() put each argument slot. Such an argument goes in a
+ * 16-byte aligned stack slot, so a JIT whose stack is only 8-byte aligned has
+ * to realign that many slots before the call.
+ */
+u32 bpf_jit_stack_slots_to_align(const struct bpf_jit_arg_abi *abi,
+				 const struct btf_func_model *fm,
+				 const u8 *pos_of_slot)
+{
+	u32 i, k, slot = 0, top = 0;
+	bool align16 = false;
+
+	for (i = 0; i < fm->nr_args; i++) {
+		u32 nslots = btf_func_model_arg_slots(fm, i);
+
+		if ((fm->arg_flags[i] & BTF_FMODEL_ALIGN16_ARG) &&
+		    pos_of_slot[slot] >= abi->nr_arg_regs)
+			align16 = true;
+		for (k = 0; k < nslots; k++, slot++)
+			if (pos_of_slot[slot] >= abi->nr_arg_regs)
+				top = max_t(u32, top, pos_of_slot[slot] - abi->nr_arg_regs + 1);
+	}
+
+	return align16 ? top : 0;
+}
+
 bool __weak bpf_jit_supports_stack_args(void)
 {
 	return false;
-- 
2.53.0-Meta