From 5a34daaafd1c3de92b194031c00f75fce1417ebb Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Mon, 13 May 2024 13:18:18 -0700 Subject: [PATCH 1/7] Handle call local seperately from regular jump instructions Signed-off-by: Alan Jowett --- tests/factorial.data | 13 +++++++++++++ vm/ubpf_jit_arm64.c | 4 ++++ vm/ubpf_jit_support.c | 6 +++++- vm/ubpf_jit_support.h | 3 +++ vm/ubpf_jit_x86_64.c | 27 +++++++++++++++++++++++++-- vm/ubpf_jit_x86_64.h | 13 +++++++++++++ 6 files changed, 63 insertions(+), 3 deletions(-) create mode 100644 tests/factorial.data diff --git a/tests/factorial.data b/tests/factorial.data new file mode 100644 index 00000000..1d2c721e --- /dev/null +++ b/tests/factorial.data @@ -0,0 +1,13 @@ +-- asm +mov %r1, 1 +mov %r2, 20 +call local exponential +exit +exponential: +mul %r1, %r2 +sub %r2, 1 +jne %r2, 0, exponential +mov %r0, %r1 +exit +-- result +0x21C3677C82B40000 diff --git a/vm/ubpf_jit_arm64.c b/vm/ubpf_jit_arm64.c index c6710406..0896fd63 100644 --- a/vm/ubpf_jit_arm64.c +++ b/vm/ubpf_jit_arm64.c @@ -1189,6 +1189,10 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) *errmsg = ubpf_error("Too many LEA calculations."); break; } + case TooManyLocalCalls: { + *errmsg = ubpf_error("Too many local calls."); + break; + } case UnexpectedInstruction: { // errmsg set at time the error was detected because the message requires // information about the unexpected instruction. diff --git a/vm/ubpf_jit_support.c b/vm/ubpf_jit_support.c index ba96b9be..4f7b057f 100644 --- a/vm/ubpf_jit_support.c +++ b/vm/ubpf_jit_support.c @@ -34,9 +34,11 @@ initialize_jit_state_result(struct jit_state *state, struct ubpf_jit_result *com state->jumps = calloc(UBPF_MAX_INSTS, sizeof(state->jumps[0])); state->loads = calloc(UBPF_MAX_INSTS, sizeof(state->loads[0])); state->leas = calloc(UBPF_MAX_INSTS, sizeof(state->leas[0])); + state->local_calls = calloc(UBPF_MAX_INSTS, sizeof(state->local_calls[0])); state->num_jumps = 0; state->num_loads = 0; state->num_leas = 0; + state->num_local_calls = 0; state->jit_status = NoError; if (!state->pc_locs || !state->jumps || !state->loads || !state->leas) { @@ -59,6 +61,8 @@ release_jit_state_result(struct jit_state *state, struct ubpf_jit_result *compil state->loads = NULL; free(state->leas); state->leas = NULL; + free(state->local_calls); + state->local_calls = NULL; } void @@ -97,4 +101,4 @@ void emit_jump_target(struct jit_state* state, uint32_t jump_src) { fixup_jump_target(state->jumps, state->num_jumps, jump_src, state->offset); -} \ No newline at end of file +} diff --git a/vm/ubpf_jit_support.h b/vm/ubpf_jit_support.h index 19738975..58740ea5 100644 --- a/vm/ubpf_jit_support.h +++ b/vm/ubpf_jit_support.h @@ -33,6 +33,7 @@ enum JitProgress { TooManyJumps, TooManyLoads, TooManyLeas, + TooManyLocalCalls, NotEnoughSpace, UnexpectedInstruction, UnknownInstruction @@ -88,9 +89,11 @@ struct jit_state struct patchable_relative* jumps; struct patchable_relative* loads; struct patchable_relative* leas; + struct patchable_relative* local_calls; int num_jumps; int num_loads; int num_leas; + int num_local_calls; uint32_t stack_size; }; diff --git a/vm/ubpf_jit_x86_64.c b/vm/ubpf_jit_x86_64.c index 467d6daa..590dc048 100644 --- a/vm/ubpf_jit_x86_64.c +++ b/vm/ubpf_jit_x86_64.c @@ -121,7 +121,7 @@ emit_local_call(struct jit_state* state, uint32_t target_pc) emit_alu64_imm32(state, 0x81, 5, RSP, 4 * sizeof(uint64_t)); #endif emit1(state, 0xe8); // e8 is the opcode for a CALL - emit_jump_address_reloc(state, target_pc); + emit_local_call_address_reloc(state, target_pc); #if defined(_WIN32) /* Deallocate home register space - 4 registers */ emit_alu64_imm32(state, 0x81, 0, RSP, 4 * sizeof(uint64_t)); @@ -301,7 +301,6 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) } struct ebpf_inst inst = ubpf_fetch_instruction(vm, i); - state->pc_locs[i] = state->offset; int dst = map_register(inst.dst); int src = map_register(inst.src); @@ -314,6 +313,8 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) emit_alu64_imm32(state, 0x81, 5, RSP, 8); } + state->pc_locs[i] = state->offset; + switch (inst.opcode) { case EBPF_OP_ADD_IMM: emit_alu32_imm32(state, 0x81, 0, dst, inst.imm); @@ -739,6 +740,10 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) *errmsg = ubpf_error("Too many LEA calculations"); break; } + case TooManyLocalCalls: { + *errmsg = ubpf_error("Too many local calls"); + break; + } case UnexpectedInstruction: { // errmsg set at time the error was detected because the message requires // information about the unexpected instruction. @@ -938,6 +943,24 @@ resolve_patchable_relatives(struct jit_state* state) memcpy(offset_ptr, &rel, sizeof(uint32_t)); } + for (i = 0; i < state->num_local_calls; i++) { + struct patchable_relative local_call = state->local_calls[i]; + + int target_loc; + assert(local_call.target_offset == 0); + assert(local_call.target_pc != TARGET_PC_EXIT); + assert(local_call.target_pc != TARGET_PC_RETPOLINE); + + target_loc = state->pc_locs[local_call.target_pc]; + + /* Assumes call offset is at end of instruction */ + uint32_t rel = target_loc - (local_call.offset_loc + sizeof(uint32_t)); + rel -= 7; // For the "sub rsp, 8" instruction that is inserted before the call to align the stack. + + uint8_t* offset_ptr = &state->buf[local_call.offset_loc]; + memcpy(offset_ptr, &rel, sizeof(uint32_t)); + } + for (i = 0; i < state->num_loads; i++) { struct patchable_relative load = state->loads[i]; diff --git a/vm/ubpf_jit_x86_64.h b/vm/ubpf_jit_x86_64.h index 376e63c3..8cac1b8e 100644 --- a/vm/ubpf_jit_x86_64.h +++ b/vm/ubpf_jit_x86_64.h @@ -121,6 +121,19 @@ emit_jump_address_reloc(struct jit_state* state, int32_t target_pc) return target_address_offset; } +static uint32_t +emit_local_call_address_reloc(struct jit_state* state, int32_t target_pc) +{ + if (state->num_local_calls == UBPF_MAX_INSTS) { + state->jit_status = TooManyLocalCalls; + return 0; + } + uint32_t target_address_offset = state->offset; + emit_patchable_relative(state->offset, target_pc, 0, state->local_calls, state->num_local_calls++); + emit_4byte_offset_placeholder(state); + return target_address_offset; +} + static inline void emit_modrm(struct jit_state* state, int mod, int r, int m) { From e57df101c3b8af263f8d1f29b0b857afd98369e5 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Fri, 17 May 2024 15:45:13 -0700 Subject: [PATCH 2/7] Fix after local call prolog increased in size Signed-off-by: Alan Jowett --- vm/ubpf_jit_arm64.c | 38 +++++++++++++++++++++++++++++++++++--- vm/ubpf_jit_support.h | 1 + vm/ubpf_jit_x86_64.c | 7 ++++++- 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/vm/ubpf_jit_arm64.c b/vm/ubpf_jit_arm64.c index 8acf76a7..7f2396bc 100644 --- a/vm/ubpf_jit_arm64.c +++ b/vm/ubpf_jit_arm64.c @@ -338,7 +338,14 @@ static void emit_unconditionalbranch_immediate( struct jit_state* state, enum UnconditionalBranchImmediateOpcode op, int32_t target_pc) { - emit_patchable_relative(state->offset, target_pc, 0, state->jumps, state->num_jumps++); + struct patchable_relative* table = state->jumps; + int num_jumps = state->num_jumps; + if (op == UBR_BL && target_pc != TARGET_PC_ENTER) { + table = state->local_calls; + num_jumps = state->num_local_calls; + } + + emit_patchable_relative(state->offset, target_pc, 0, table, num_jumps); emit_instruction(state, op); } @@ -966,14 +973,20 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) // All checks for errors during the encoding of _this_ instruction // occur at the end of the loop. struct ebpf_inst inst = ubpf_fetch_instruction(vm, i); - state->pc_locs[i] = state->offset; if (i == 0 || vm->int_funcs[i]) { + size_t prolog_start = state->offset; emit_movewide_immediate(state, true, temp_register, ubpf_stack_usage_for_local_func(vm, i)); emit_addsub_immediate(state, true, AS_SUB, SP, SP, 16); emit_loadstorepair_immediate(state, LSP_STPX, temp_register, temp_register, SP, 0); + // Record the size of the prolog so that we can calculate offset when doing a local call. + if (state->bpf_function_prolog_size == 0) { + state->bpf_function_prolog_size = state->offset - prolog_start; + } } + state->pc_locs[i] = state->offset; + enum Registers dst = map_register(inst.dst); enum Registers src = map_register(inst.src); uint8_t opcode = inst.opcode; @@ -1371,6 +1384,25 @@ resolve_leas(struct jit_state* state) return true; } +static bool +resolve_local_calls(struct jit_state* state) +{ + for (unsigned i = 0; i < state->num_local_calls; ++i) { + struct patchable_relative local_call = state->local_calls[i]; + + int32_t target_loc; + assert(local_call.target_offset == 0); + assert(local_call.target_pc != TARGET_PC_EXIT); + assert(local_call.target_pc != TARGET_PC_RETPOLINE); + target_loc = state->pc_locs[local_call.target_pc]; + + int32_t rel = target_loc - local_call.offset_loc; + rel -= state->bpf_function_prolog_size; + resolve_branch_immediate(state, local_call.offset_loc, rel); + } + return true; +} + bool ubpf_jit_update_dispatcher_arm64( struct ubpf_vm* vm, external_function_dispatcher_t new_dispatcher, uint8_t* buffer, size_t size, uint32_t offset) @@ -1415,7 +1447,7 @@ ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, enum Jit goto out; } - if (!resolve_jumps(&state) || !resolve_loads(&state) || !resolve_leas(&state)) { + if (!resolve_jumps(&state) || !resolve_loads(&state) || !resolve_leas(&state) || !resolve_local_calls(&state)) { compile_result.errmsg = ubpf_error("Could not patch the relative addresses in the JIT'd code."); goto out; } diff --git a/vm/ubpf_jit_support.h b/vm/ubpf_jit_support.h index 3ea472ef..83ba66b8 100644 --- a/vm/ubpf_jit_support.h +++ b/vm/ubpf_jit_support.h @@ -97,6 +97,7 @@ struct jit_state int num_leas; int num_local_calls; uint32_t stack_size; + size_t bpf_function_prolog_size; // Count of bytes emitted at the start of the function. }; int diff --git a/vm/ubpf_jit_x86_64.c b/vm/ubpf_jit_x86_64.c index af24963b..3f2c6fff 100644 --- a/vm/ubpf_jit_x86_64.c +++ b/vm/ubpf_jit_x86_64.c @@ -342,6 +342,7 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) uint32_t target_pc = i + inst.offset + 1; if (i == 0 || vm->int_funcs[i]) { + size_t prolog_start = state->offset; uint16_t stack_usage = ubpf_stack_usage_for_local_func(vm, i); emit_alu64_imm32(state, 0x81, 5, RSP, 8); emit1(state, 0x48); @@ -349,6 +350,10 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) emit1(state, 0x04); // Mod: 00b Reg: 000b RM: 100b emit1(state, 0x24); // Scale: 00b Index: 100b Base: 100b emit4(state, stack_usage); + // Record the size of the prolog so that we can calculate offset when doing a local call. + if (state->bpf_function_prolog_size == 0) { + state->bpf_function_prolog_size = state->offset - prolog_start; + } } state->pc_locs[i] = state->offset; @@ -993,7 +998,7 @@ resolve_patchable_relatives(struct jit_state* state) /* Assumes call offset is at end of instruction */ uint32_t rel = target_loc - (local_call.offset_loc + sizeof(uint32_t)); - rel -= 7; // For the "sub rsp, 8" instruction that is inserted before the call to align the stack. + rel -= state->bpf_function_prolog_size; // For the prolog inserted at the start of every local call. uint8_t* offset_ptr = &state->buf[local_call.offset_loc]; memcpy(offset_ptr, &rel, sizeof(uint32_t)); From 0b580485cff49a47e6e11af1f98effe11c3c5843 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Fri, 17 May 2024 16:16:36 -0700 Subject: [PATCH 3/7] Zero initialize jit_state Signed-off-by: Alan Jowett --- vm/ubpf_jit_x86_64.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vm/ubpf_jit_x86_64.c b/vm/ubpf_jit_x86_64.c index 3f2c6fff..330dc3ea 100644 --- a/vm/ubpf_jit_x86_64.c +++ b/vm/ubpf_jit_x86_64.c @@ -1045,7 +1045,7 @@ resolve_patchable_relatives(struct jit_state* state) struct ubpf_jit_result ubpf_translate_x86_64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, enum JitMode jit_mode) { - struct jit_state state; + struct jit_state state = {0}; struct ubpf_jit_result compile_result; if (initialize_jit_state_result(&state, &compile_result, buffer, *size, jit_mode, &compile_result.errmsg) < 0) { From da4910b72a6082f1727a4f7581cec59e339b60f0 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Fri, 17 May 2024 16:39:33 -0700 Subject: [PATCH 4/7] Increment num_jumps after inserting into table Signed-off-by: Alan Jowett --- vm/ubpf_jit_arm64.c | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vm/ubpf_jit_arm64.c b/vm/ubpf_jit_arm64.c index 7f2396bc..b123ecc1 100644 --- a/vm/ubpf_jit_arm64.c +++ b/vm/ubpf_jit_arm64.c @@ -339,13 +339,13 @@ emit_unconditionalbranch_immediate( struct jit_state* state, enum UnconditionalBranchImmediateOpcode op, int32_t target_pc) { struct patchable_relative* table = state->jumps; - int num_jumps = state->num_jumps; + int* num_jumps = &state->num_jumps; if (op == UBR_BL && target_pc != TARGET_PC_ENTER) { table = state->local_calls; - num_jumps = state->num_local_calls; + num_jumps = &state->num_local_calls; } - emit_patchable_relative(state->offset, target_pc, 0, table, num_jumps); + emit_patchable_relative(state->offset, target_pc, 0, table, (*num_jumps)++); emit_instruction(state, op); } From e070f79fc1ba37a3e67c4f22961a53216d32f188 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Fri, 17 May 2024 17:18:34 -0700 Subject: [PATCH 5/7] Initialize bpf_function_prolog_size correctly Signed-off-by: Alan Jowett --- vm/ubpf_jit_support.c | 1 + vm/ubpf_jit_x86_64.c | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/vm/ubpf_jit_support.c b/vm/ubpf_jit_support.c index a685cbc7..8661be05 100644 --- a/vm/ubpf_jit_support.c +++ b/vm/ubpf_jit_support.c @@ -50,6 +50,7 @@ initialize_jit_state_result( state->num_local_calls = 0; state->jit_status = NoError; state->jit_mode = jit_mode; + state->bpf_function_prolog_size = 0; if (!state->pc_locs || !state->jumps || !state->loads || !state->leas) { *errmsg = ubpf_error("Could not allocate space needed to JIT compile eBPF program"); diff --git a/vm/ubpf_jit_x86_64.c b/vm/ubpf_jit_x86_64.c index 330dc3ea..3f2c6fff 100644 --- a/vm/ubpf_jit_x86_64.c +++ b/vm/ubpf_jit_x86_64.c @@ -1045,7 +1045,7 @@ resolve_patchable_relatives(struct jit_state* state) struct ubpf_jit_result ubpf_translate_x86_64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, enum JitMode jit_mode) { - struct jit_state state = {0}; + struct jit_state state; struct ubpf_jit_result compile_result; if (initialize_jit_state_result(&state, &compile_result, buffer, *size, jit_mode, &compile_result.errmsg) < 0) { From b24ea1ea36ee8017f3c4c958ffc1a77deeb7ec09 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Sat, 18 May 2024 10:15:25 -0700 Subject: [PATCH 6/7] Update ubpf_jit_arm64.c Assert that prolog size matches. --- vm/ubpf_jit_arm64.c | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vm/ubpf_jit_arm64.c b/vm/ubpf_jit_arm64.c index b123ecc1..0368c7e7 100644 --- a/vm/ubpf_jit_arm64.c +++ b/vm/ubpf_jit_arm64.c @@ -982,6 +982,8 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) // Record the size of the prolog so that we can calculate offset when doing a local call. if (state->bpf_function_prolog_size == 0) { state->bpf_function_prolog_size = state->offset - prolog_start; + } else { + assert(state->bpf_function_prolog_size == state->offset - prolog_start); } } From 711cfc065ef64654da01e51906e270e847a357e2 Mon Sep 17 00:00:00 2001 From: Alan Jowett Date: Sat, 18 May 2024 10:18:14 -0700 Subject: [PATCH 7/7] Update ubpf_jit_x86_64.c Assert that the prolog size matches. --- vm/ubpf_jit_x86_64.c | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vm/ubpf_jit_x86_64.c b/vm/ubpf_jit_x86_64.c index 3f2c6fff..4a50b526 100644 --- a/vm/ubpf_jit_x86_64.c +++ b/vm/ubpf_jit_x86_64.c @@ -353,6 +353,8 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) // Record the size of the prolog so that we can calculate offset when doing a local call. if (state->bpf_function_prolog_size == 0) { state->bpf_function_prolog_size = state->offset - prolog_start; + } else { + assert(state->bpf_function_prolog_size == state->offset - prolog_start); } }