Do the proper move from the BPF calling convention to the x86-64 calling
convention to satisfy the native requirement.

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 | 73 +++++++++++++++++++++++++++++++++++--
 1 file changed, 70 insertions(+), 3 deletions(-)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index bba351944202..0496607a7003 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -1839,6 +1839,60 @@ 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. */
+		if (src == BPF_JIT_ARG_TMP) {
+			reg = AUX_REG;
+		} else if (src_mem) {
+			reg = dst_mem || dst == BPF_JIT_ARG_TMP ? 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 +1904,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 +2897,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 +4413,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.52.0