[PATCH bpf-next v4 10/15] bpf, x86: Move kfunc arguments into the x86-64 calling convention

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

The x86-64 allows 6 register arguments and then stack arguments, but
backfilling the 6th register is possible, as in the following example:

	/* s is a 16-byte struct */
	void kfunc(u64 a, u64 b, u64 c, u64 d, u64 e, struct big s, u64 f)

BPF puts s in slots 5 and 6, R9 and the first stack slot, and f in slot
7, while x86-64 puts s in the two stack slots and f in R9. Emit the
moves bpf_jit_plan_arg_moves() plans to bridge the two.

Only one argument ever moves down, as an argument frees one register at
most, so carrying that single value in the scratch register past its own
destination is enough.

In addition, the arena argument walk counts eightbytes rather than
parameters, as an argument may take two registers.

Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
---
 arch/x86/net/bpf_jit_comp.c | 78 +++++++++++++++++++++++++++++++++++--
 1 file changed, 75 insertions(+), 3 deletions(-)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index bba351944202..2671e4118d00 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -1839,6 +1839,65 @@ static int emit_spectre_bhb_barrier(u8 **pprog, u8 *ip,
 	return 0;
 }
 
+static const struct bpf_jit_arg_abi x86_arg_abi = {
+	.nr_arg_regs		= 6,
+	.backfill_after_stack	= true,
+	.even_stack_align	= true,
+};
+
+static const u8 x86_arg_reg[] = {
+	BPF_REG_1, BPF_REG_2, BPF_REG_3, BPF_REG_4, BPF_REG_5, X86_REG_R9,
+};
+
+/*
+ * Move the arguments the x86-64 ABI places somewhere other than the argument
+ * slot the BPF calling convention gave them. @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)
+{
+	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);
+
+	for (i = 0; i < n; i++) {
+		u8 dst = moves[i].dst, src = moves[i].src, reg;
+		bool dst_mem = dst != BPF_JIT_ARG_TMP && dst >= nreg;
+		bool src_mem = src != BPF_JIT_ARG_TMP && src >= nreg;
+
+		/*
+		 * Take the value into a register: the one it belongs in, the
+		 * scratch when it is carried past its own destination, and
+		 * BPF_REG_AX only to pass one stack slot to another.
+		 */
+		if (src == BPF_JIT_ARG_TMP) {
+			reg = AUX_REG;
+		} else if (src_mem) {
+			reg = dst == BPF_JIT_ARG_TMP ? AUX_REG :
+			      dst_mem ? BPF_REG_AX : x86_arg_reg[dst];
+			emit_ldx(&prog, BPF_DW, reg, BPF_REG_FP,
+				 stack_base + (src - nreg) * 8);
+		} else {
+			reg = x86_arg_reg[src];
+		}
+
+		/* And leave it where the argument belongs. */
+		if (dst == BPF_JIT_ARG_TMP)
+			emit_mov_reg(&prog, true, AUX_REG, reg);
+		else if (dst_mem)
+			emit_stx(&prog, BPF_DW, BPF_REG_FP, reg,
+				 stack_base + (dst - nreg) * 8);
+		else if (reg != x86_arg_reg[dst])
+			emit_mov_reg(&prog, true, x86_arg_reg[dst], reg);
+	}
+
+	*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
@@ -1850,11 +1909,17 @@ static int emit_kfunc_arena_args(struct bpf_prog *bpf_prog,
 {
 	u8 *prog = *pprog;
 	u8 *start = prog;
-	int i;
+	int i, slot;
 
-	for (i = 0; i < min_t(int, fm->nr_args, MAX_BPF_FUNC_REG_ARGS); i++) {
+	for (i = 0, slot = 0; i < fm->nr_args; i++) {
+		u32 arg_regs = (fm->arg_size[i] + 7) / 8;
 		u8 flags = fm->arg_flags[i];
-		u32 reg = BPF_REG_1 + i;
+		u32 reg;
+
+		if (slot + arg_regs > MAX_BPF_FUNC_REG_ARGS)
+			break;
+		reg = BPF_REG_1 + slot;
+		slot += arg_regs;
 
 		if (!(flags & BTF_FMODEL_ARENA_ARG))
 			continue;
@@ -2837,6 +2902,8 @@ 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);
 			}
 			if (priv_frame_ptr) {
 				push_r9(&prog);
@@ -4351,6 +4418,11 @@ bool bpf_jit_supports_kfunc_ret_reg_pair(void)
 	return true;
 }
 
+const struct bpf_jit_arg_abi *bpf_jit_arg_abi(void)
+{
+	return &x86_arg_abi;
+}
+
 bool bpf_jit_supports_stack_args(void)
 {
 	return true;
-- 
2.53.0-Meta