[PATCH bpf-next v2 3/6] bpf: x86: Support in-register struct arguments

[Yonghong Song <yhs@fb.com>]

In C, struct value can be passed as a function argument.
For small structs, struct value may be passed in
one or more registers. For trampoline based bpf programs,
This would cause complication since one-to-one mapping between
function argument and arch argument register is not valid
any more.

To support struct value argument and make bpf programs
easy to write, the bpf program function parameter is
changed from struct type to a pointer to struct type.
The following is a simplified example.

In one of later selftests, we have a bpf_testmod function:
    struct bpf_testmod_struct_arg_2 {
        long a;
        long b;
    };
    noinline int
    bpf_testmod_test_struct_arg_2(int a, struct bpf_testmod_struct_arg_2 b, int c) {
        bpf_testmod_test_struct_arg_result = a + b.a + b.b + c;
        return bpf_testmod_test_struct_arg_result;
    }

When a bpf program is attached to the bpf_testmod function
bpf_testmod_test_struct_arg_2(), the bpf program may look like
    SEC("fentry/bpf_testmod_test_struct_arg_2")
    int BPF_PROG(test_struct_arg_3, int a, struct bpf_testmod_struct_arg_2 *b, int c)
    {
        t2_a = a;
        t2_b_a = b->a;
        t2_b_b = b->b;
        t2_c = c;
        return 0;
    }

Basically struct value becomes a pointer to the struct.
The trampoline stack will be increased to store the stack values and
the pointer to these values will be saved in the stack slot corresponding
to that argument. For x86_64, the struct size is limited up to 16 bytes
so the struct can fit in one or two registers. The struct size of more
than 16 bytes is not supported now as our current use case is
for sockptr_t in the argument. We could handle such large struct's
in the future if we have concrete use cases.

The main changes are in save_regs() and restore_regs(). The following
illustrated the trampoline asm codes for save_regs() and restore_regs().
save_regs():
    /* first argument */
    mov    DWORD PTR [rbp-0x18],edi
    /* second argument: struct, save actual values and put the pointer to the slot */
    lea    rax,[rbp-0x40]
    mov    QWORD PTR [rbp-0x10],rax
    mov    QWORD PTR [rbp-0x40],rsi
    mov    QWORD PTR [rbp-0x38],rdx
    /* third argument */
    mov    DWORD PTR [rbp-0x8],esi
restore_regs():
    mov    edi,DWORD PTR [rbp-0x18]
    mov    rsi,QWORD PTR [rbp-0x40]
    mov    rdx,QWORD PTR [rbp-0x38]
    mov    esi,DWORD PTR [rbp-0x8]

Signed-off-by: Yonghong Song <yhs@fb.com>
---
 arch/x86/net/bpf_jit_comp.c | 137 +++++++++++++++++++++++++++++-------
 1 file changed, 110 insertions(+), 27 deletions(-)

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 2657b58001cf..2fa7d694c559 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -37,6 +37,11 @@ static u8 *emit_code(u8 *ptr, u32 bytes, unsigned int len)
 #define EMIT3(b1, b2, b3)	EMIT((b1) + ((b2) << 8) + ((b3) << 16), 3)
 #define EMIT4(b1, b2, b3, b4)   EMIT((b1) + ((b2) << 8) + ((b3) << 16) + ((b4) << 24), 4)
 
+#define EMIT2_UPDATE_PROG(bytes, len) \
+	do { *prog = emit_code(*prog, bytes, len); } while (0)
+#define EMIT4_UPDATE_PROG(b1, b2, b3, b4) \
+	EMIT2_UPDATE_PROG((b1) + ((b2) << 8) + ((b3) << 16) + ((b4) << 24), 4)
+
 #define EMIT1_off32(b1, off) \
 	do { EMIT1(b1); EMIT(off, 4); } while (0)
 #define EMIT2_off32(b1, b2, off) \
@@ -1749,36 +1754,92 @@ st:			if (is_imm8(insn->off))
 }
 
 static void save_regs(const struct btf_func_model *m, u8 **prog, int nr_args,
-		      int regs_off)
+		      int regs_off, int struct_val_off)
 {
-	int i;
-	/* Store function arguments to stack.
-	 * For a function that accepts two pointers the sequence will be:
-	 * mov QWORD PTR [rbp-0x10],rdi
-	 * mov QWORD PTR [rbp-0x8],rsi
-	 */
-	for (i = 0; i < min(nr_args, 6); i++)
-		emit_stx(prog, bytes_to_bpf_size(m->arg_size[i]),
-			 BPF_REG_FP,
-			 i == 5 ? X86_REG_R9 : BPF_REG_1 + i,
-			 -(regs_off - i * 8));
+	int i, j, curr_reg_idx, curr_s_stack_off, s_arg_nregs;
+
+	curr_reg_idx = curr_s_stack_off = 0;
+	for (i = 0; i < nr_args; i++) {
+		if (!(m->arg_flags[i] & BTF_FMODEL_STRUCT_ARG)) {
+			/* Store function arguments to stack.
+			 * For a function that accepts two pointers the sequence will be:
+			 * mov QWORD PTR [rbp-0x10],rdi
+			 * mov QWORD PTR [rbp-0x8],rsi
+			 */
+			emit_stx(prog, bytes_to_bpf_size(m->arg_size[i]),
+				 BPF_REG_FP,
+				 curr_reg_idx == 5 ? X86_REG_R9 : BPF_REG_1 + curr_reg_idx,
+				 -(regs_off - i * 8));
+			curr_reg_idx++;
+			continue;
+		}
+
+		/* lea rax, [rbp - struct_val_off] */
+		EMIT4_UPDATE_PROG(0x48, 0x8D, 0x45, -(struct_val_off - curr_s_stack_off * 8));
+		/* The struct value is converted to a pointer argument.
+		 * mov QWORD PTR [rbp - s_arg_idx * 8],rax
+		 */
+		emit_stx(prog, bytes_to_bpf_size(8),
+			 BPF_REG_FP, BPF_REG_0, -(regs_off - i * 8));
+
+		/* Save struct registers to stack.
+		 * For example, argument 1 (second argument) size is 16 which occupies two
+		 * registers, these two register values will be saved in stack.
+		 * mov QWORD PTR [rbp-0x40],rsi
+		 * mov QWORD PTR [rbp-0x38],rdx
+		 */
+		s_arg_nregs = (m->arg_size[i] + 7) / 8;
+		for (j = 0; j < s_arg_nregs; j++) {
+			emit_stx(prog, bytes_to_bpf_size(8),
+				 BPF_REG_FP,
+				 curr_reg_idx == 5 ? X86_REG_R9 : BPF_REG_1 + curr_reg_idx,
+				 -(struct_val_off - curr_s_stack_off * 8));
+			curr_reg_idx++;
+			curr_s_stack_off++;
+		}
+	}
 }
 
 static void restore_regs(const struct btf_func_model *m, u8 **prog, int nr_args,
-			 int regs_off)
+			 int regs_off, int struct_val_off)
 {
-	int i;
+	int i, j, curr_reg_idx, curr_s_stack_off, s_arg_nregs;
+
+	curr_reg_idx = curr_s_stack_off = 0;
+	for (i = 0; i < nr_args; i++) {
+		if (!(m->arg_flags[i] & BTF_FMODEL_STRUCT_ARG)) {
+			/* Restore struct values from stack.
+			 * For example, argument 1 (second argument) size is 16 which occupies two
+			 * registers, these two register values will be restored to the original
+			 * registers.
+			 * mov rsi,QWORD PTR [rbp-0x40]
+			 * mov rdx,QWORD PTR [rbp-0x38]
+			 */
+			emit_ldx(prog, bytes_to_bpf_size(m->arg_size[i]),
+				 curr_reg_idx  == 5 ? X86_REG_R9 : BPF_REG_1 + curr_reg_idx,
+				 BPF_REG_FP,
+				 -(regs_off - i * 8));
+			curr_reg_idx++;
+			continue;
+		}
 
-	/* Restore function arguments from stack.
-	 * For a function that accepts two pointers the sequence will be:
-	 * EMIT4(0x48, 0x8B, 0x7D, 0xF0); mov rdi,QWORD PTR [rbp-0x10]
-	 * EMIT4(0x48, 0x8B, 0x75, 0xF8); mov rsi,QWORD PTR [rbp-0x8]
-	 */
-	for (i = 0; i < min(nr_args, 6); i++)
-		emit_ldx(prog, bytes_to_bpf_size(m->arg_size[i]),
-			 i == 5 ? X86_REG_R9 : BPF_REG_1 + i,
-			 BPF_REG_FP,
-			 -(regs_off - i * 8));
+		/* Restore struct values from stack.
+		 * For example, argument 1 (second argument) size is 16 which occupies two
+		 * registers, these two register values will be restored to the original
+		 * registers.
+		 * mov rsi,QWORD PTR [rbp-0x40]
+		 * mov rdx,QWORD PTR [rbp-0x38]
+		 */
+		s_arg_nregs = (m->arg_size[i] + 7) / 8;
+		for (j = 0; j < s_arg_nregs; j++) {
+			emit_ldx(prog, bytes_to_bpf_size(8),
+				 curr_reg_idx == 5 ? X86_REG_R9 : BPF_REG_1 + curr_reg_idx,
+				 BPF_REG_FP,
+				 -(struct_val_off - curr_s_stack_off * 8));
+			curr_reg_idx++;
+			curr_s_stack_off++;
+		}
+	}
 }
 
 static int invoke_bpf_prog(const struct btf_func_model *m, u8 **pprog,
@@ -2020,6 +2081,7 @@ int arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *image, void *i
 	struct bpf_tramp_links *fentry = &tlinks[BPF_TRAMP_FENTRY];
 	struct bpf_tramp_links *fexit = &tlinks[BPF_TRAMP_FEXIT];
 	struct bpf_tramp_links *fmod_ret = &tlinks[BPF_TRAMP_MODIFY_RETURN];
+	int struct_val_off, extra_nregs = 0;
 	u8 **branches = NULL;
 	u8 *prog;
 	bool save_ret;
@@ -2028,6 +2090,20 @@ int arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *image, void *i
 	if (nr_args > 6)
 		return -ENOTSUPP;
 
+	for (i = 0; i < MAX_BPF_FUNC_ARGS; i++) {
+		if (m->arg_flags[i] & BTF_FMODEL_STRUCT_ARG) {
+			/* Only support up to 16 bytes struct which should keep
+			 * values in registers.
+			 */
+			if (m->arg_size[i] > 16)
+				return -ENOTSUPP;
+
+			extra_nregs += (m->arg_size[i] + 7) / 8 - 1;
+		}
+	}
+	if (nr_args + extra_nregs > 6)
+		return -ENOTSUPP;
+
 	/* Generated trampoline stack layout:
 	 *
 	 * RBP + 8         [ return address  ]
@@ -2066,6 +2142,13 @@ int arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *image, void *i
 	stack_size += (sizeof(struct bpf_tramp_run_ctx) + 7) & ~0x7;
 	run_ctx_off = stack_size;
 
+	/* For structure argument */
+	for (i = 0; i < MAX_BPF_FUNC_ARGS; i++) {
+		if (m->arg_flags[i] & BTF_FMODEL_STRUCT_ARG)
+			stack_size += (m->arg_size[i] + 7) & ~0x7;
+	}
+	struct_val_off = stack_size;
+
 	if (flags & BPF_TRAMP_F_SKIP_FRAME) {
 		/* skip patched call instruction and point orig_call to actual
 		 * body of the kernel function.
@@ -2101,7 +2184,7 @@ int arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *image, void *i
 		emit_stx(&prog, BPF_DW, BPF_REG_FP, BPF_REG_0, -ip_off);
 	}
 
-	save_regs(m, &prog, nr_args, regs_off);
+	save_regs(m, &prog, nr_args, regs_off, struct_val_off);
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
 		/* arg1: mov rdi, im */
@@ -2131,7 +2214,7 @@ int arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *image, void *i
 	}
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
-		restore_regs(m, &prog, nr_args, regs_off);
+		restore_regs(m, &prog, nr_args, regs_off, struct_val_off);
 
 		if (flags & BPF_TRAMP_F_ORIG_STACK) {
 			emit_ldx(&prog, BPF_DW, BPF_REG_0, BPF_REG_FP, 8);
@@ -2172,7 +2255,7 @@ int arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *image, void *i
 		}
 
 	if (flags & BPF_TRAMP_F_RESTORE_REGS)
-		restore_regs(m, &prog, nr_args, regs_off);
+		restore_regs(m, &prog, nr_args, regs_off, struct_val_off);
 
 	/* This needs to be done regardless. If there were fmod_ret programs,
 	 * the return value is only updated on the stack and still needs to be
-- 
2.30.2