From 2e13f7ce238892f583c4e9ac97d22254ea3aea4f Mon Sep 17 00:00:00 2001 From: Will Hawkins Date: Fri, 26 Apr 2024 23:27:20 -0400 Subject: [PATCH] fixup! [WIP] Make the JIT'd code completely portable. Add arm64 --- ...ubpf_test_jit_unexpected_instruction.input | 1 + .../ubpf_test_jit_unexpected_instruction.md | 0 custom_tests/srcs/test_helpers.h | 135 +++++++++++- custom_tests/srcs/ubpf_custom_test_support.cc | 1 + .../ubpf_test_jit_unexpected_instruction.cc | 48 +++++ custom_tests/srcs/ubpf_test_update_helpers.cc | 133 +++++++++--- vm/ubpf_jit_arm64.c | 192 +++++++++++++++--- vm/ubpf_vm.c | 4 +- 8 files changed, 441 insertions(+), 73 deletions(-) create mode 100644 custom_tests/data/ubpf_test_jit_unexpected_instruction.input create mode 100644 custom_tests/descrs/ubpf_test_jit_unexpected_instruction.md create mode 100644 custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc diff --git a/custom_tests/data/ubpf_test_jit_unexpected_instruction.input b/custom_tests/data/ubpf_test_jit_unexpected_instruction.input new file mode 100644 index 00000000..d5be5a56 --- /dev/null +++ b/custom_tests/data/ubpf_test_jit_unexpected_instruction.input @@ -0,0 +1 @@ +8f 00 00 00 01 00 00 00 \ No newline at end of file diff --git a/custom_tests/descrs/ubpf_test_jit_unexpected_instruction.md b/custom_tests/descrs/ubpf_test_jit_unexpected_instruction.md new file mode 100644 index 00000000..e69de29b diff --git a/custom_tests/srcs/test_helpers.h b/custom_tests/srcs/test_helpers.h index 594b6127..0b0bbd5f 100644 --- a/custom_tests/srcs/test_helpers.h +++ b/custom_tests/srcs/test_helpers.h @@ -109,20 +109,135 @@ updated_dispatcher_test_memfrob(uint64_t a, uint64_t b, uint64_t c, uint64_t d, return 43; } +static uint64_t +dispatcher_gather_bytes(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 44; + +} + +static uint64_t +updated_dispatcher_gather_bytes(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 45; +} + +static uint64_t +dispatcher_no_op(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 46; + +} + +static uint64_t +updated_dispatcher_no_op(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 47; +} + +static uint64_t +dispatcher_sqrti(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 48; + +} + +static uint64_t +updated_dispatcher_sqrti(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 49; +} + +static uint64_t +dispatcher_strcmp_ext(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 50; + +} + +static uint64_t +updated_dispatcher_strcmp_ext(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 51; +} + +static uint64_t +dispatcher_unwind(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 52; + +} + +static uint64_t +updated_dispatcher_unwind(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e) +{ + UNREFERENCED_PARAMETER(a); + UNREFERENCED_PARAMETER(b); + UNREFERENCED_PARAMETER(c); + UNREFERENCED_PARAMETER(d); + UNREFERENCED_PARAMETER(e); + return 53; +} + static std::map dispatch_test_helper_functions = { - {0, gather_bytes}, + {0, dispatcher_gather_bytes}, {1, dispatcher_test_memfrob }, - {2, no_op}, - {3, sqrti}, - {4, strcmp_ext}, - {5, unwind}, + {2, dispatcher_no_op}, + {3, dispatcher_sqrti}, + {4, dispatcher_strcmp_ext}, + {5, dispatcher_unwind}, }; static std::map updated_dispatch_test_helper_functions = { - {0, gather_bytes}, + {0, updated_dispatcher_gather_bytes}, {1, updated_dispatcher_test_memfrob }, - {2, no_op}, - {3, sqrti}, - {4, strcmp_ext}, - {5, unwind}, + {2, updated_dispatcher_no_op}, + {3, updated_dispatcher_sqrti}, + {4, updated_dispatcher_strcmp_ext}, + {5, updated_dispatcher_unwind}, }; \ No newline at end of file diff --git a/custom_tests/srcs/ubpf_custom_test_support.cc b/custom_tests/srcs/ubpf_custom_test_support.cc index 56380c82..6fc58e05 100644 --- a/custom_tests/srcs/ubpf_custom_test_support.cc +++ b/custom_tests/srcs/ubpf_custom_test_support.cc @@ -67,6 +67,7 @@ bool ubpf_setup_custom_test(ubpf_vm_up &vm, ubpf_jit_fn &jit_fn, std::string &error) { + jit_fn = nullptr; std::vector program = bytes_to_ebpf_inst(base16_decode(program_string)); char *error_s{nullptr}; diff --git a/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc b/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc new file mode 100644 index 00000000..daebf17c --- /dev/null +++ b/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft Corporation +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +extern "C" +{ +#include "ebpf.h" +#include "ubpf.h" +} + +#include "ubpf_custom_test_support.h" + +/** + * @brief This program reads BPF instructions from stdin and memory contents from + * the first agument. It then executes the BPF program and prints the + * value of %r0 at the end of execution. + */ +int main() +{ + std::string expected_error{"Failed to load program: unknown opcode 0x8f at PC 0" }; + ubpf_jit_fn jit_fn; + + std::string program_string; + std::getline(std::cin, program_string); + + std::vector program = bytes_to_ebpf_inst(base16_decode(program_string)); + + ubpf_vm_up vm(ubpf_create(), ubpf_destroy); + std::string error{}; + + if (!ubpf_setup_custom_test( + vm, + program_string, + custom_test_fixup_cb{[](ubpf_vm_up&, std::string& ) { + return true; + }}, + jit_fn, + error)) { + + if (jit_fn == nullptr && expected_error == error) + return 0; + } + + return 1; +} diff --git a/custom_tests/srcs/ubpf_test_update_helpers.cc b/custom_tests/srcs/ubpf_test_update_helpers.cc index 159426cd..85491a52 100644 --- a/custom_tests/srcs/ubpf_test_update_helpers.cc +++ b/custom_tests/srcs/ubpf_test_update_helpers.cc @@ -15,13 +15,21 @@ extern "C" { -#include "ebpf.h" #include "ubpf.h" } #include "ubpf_custom_test_support.h" #include "test_helpers.h" +struct HelperTestCase { + const char *testcase_name; + external_function_t helper_function1; + external_function_t helper_function2; + int index; + uint64_t result1; + uint64_t result2; +}; + /** * @brief This program reads BPF instructions from stdin and memory contents from * the first agument. It then executes the BPF program and prints the @@ -31,42 +39,105 @@ int main(int argc, char **argv) { std::vector args(argv, argv + argc); std::string program_string{}; - ubpf_jit_fn jit_fn; - uint64_t memory{0x123456789}; + bool success{true}; + + const char memfrob_testcase_name[] = "memfrob"; + const char gather_bytes_testcase_name[] = "gather bytes"; + const char sqrti_testcase_name[] = "sqrti"; + const char no_op_testcase_name[] = "no op"; + const char strcmp_testcase_name[] = "strcmp"; + const char unwind_testcase_name[] = "unwind"; + + std::vector test_cases{ + { + .testcase_name = memfrob_testcase_name, + .helper_function1 = dispatcher_test_memfrob, + .helper_function2 = updated_dispatcher_test_memfrob, + .index = 1, + .result1 = 42, + .result2 = 43 + }, + { + .testcase_name = gather_bytes_testcase_name, + .helper_function1 = dispatcher_gather_bytes, + .helper_function2 = updated_dispatcher_gather_bytes, + .index = 1, + .result1 = 44, + .result2 = 45 + }, + { + .testcase_name = no_op_testcase_name, + .helper_function1 = dispatcher_no_op, + .helper_function2 = updated_dispatcher_no_op, + .index = 1, + .result1 = 46, + .result2 = 47 + }, + { + .testcase_name = sqrti_testcase_name, + .helper_function1 = dispatcher_sqrti, + .helper_function2 = updated_dispatcher_sqrti, + .index = 1, + .result1 = 48, + .result2 = 49 + }, + { + .testcase_name = strcmp_testcase_name, + .helper_function1 = dispatcher_strcmp_ext, + .helper_function2 = updated_dispatcher_strcmp_ext, + .index = 1, + .result1 = 50, + .result2 = 51 + }, + { + .testcase_name = unwind_testcase_name, + .helper_function1 = dispatcher_unwind, + .helper_function2 = updated_dispatcher_unwind, + .index = 1, + .result1 = 52, + .result2 = 53 + } + }; std::getline(std::cin, program_string); - std::unique_ptr vm(ubpf_create(), ubpf_destroy); - std::string error{}; - if (!ubpf_setup_custom_test( - vm, - program_string, - [](ubpf_vm_up& vm, std::string& error) { - for (auto& [key, value] : dispatch_test_helper_functions) { - if (ubpf_register(vm.get(), key, "unnamed", value) != 0) { - error = "Failed to register helper function"; - return false; - } - } - return true; - }, - jit_fn, - error)) { - std::cerr << "Problem setting up custom test: " << error << std::endl; - return 1; - } + for (auto testcase : test_cases) { + ubpf_jit_fn jit_fn; + uint64_t memory{0x123456789}; + std::unique_ptr vm(ubpf_create(), ubpf_destroy); + std::string error{}; + if (!ubpf_setup_custom_test( + vm, + program_string, + [&testcase](ubpf_vm_up& vm, std::string& error) { + if (ubpf_register(vm.get(), testcase.index, "unnamed", testcase.helper_function1) != 0) { + error = "Failed to register helper function"; + return false; + } + return true; + }, + jit_fn, + error)) { + std::cerr << "Problem setting up custom test: " << error << std::endl; + return 1; + } - [[maybe_unused]] auto first_result = jit_fn(&memory, sizeof(uint64_t)); - std::cout << "first_result: " << first_result << "\n"; + [[maybe_unused]] auto first_result = jit_fn(&memory, sizeof(uint64_t)); - for (auto& [key, value] : updated_dispatch_test_helper_functions) { - if (ubpf_register(vm.get(), key, "unnamed", value) != 0) { - error = "Failed to register helper function"; + if (ubpf_register(vm.get(), testcase.index, "unnamed", testcase.helper_function2) != 0) { + std::cout << "Failed to register helper function\n"; return 1; } - } - [[maybe_unused]] auto second_result = jit_fn(&memory, sizeof(uint64_t)); - std::cout << "second_result: " << second_result << "\n"; - return !(first_result == 42); + [[maybe_unused]] auto second_result = jit_fn(&memory, sizeof(uint64_t)); + + auto current_success{(first_result == testcase.result1 && second_result == testcase.result2)}; + if (!current_success) { + std::cout << "There was a failure with test " << testcase.testcase_name << ": " << + testcase.result1 << " != " << first_result << " or " << + testcase.result2 << " != " << second_result << "!\n"; + } + success &= current_success; + } + return success ? 0 : 1; } diff --git a/vm/ubpf_jit_arm64.c b/vm/ubpf_jit_arm64.c index 4949b003..1fa3355e 100644 --- a/vm/ubpf_jit_arm64.c +++ b/vm/ubpf_jit_arm64.c @@ -28,9 +28,7 @@ #include #include #include -#include #include -#include #include #include "ubpf_int.h" @@ -42,6 +40,7 @@ #define TARGET_PC_EXIT ~UINT32_C(0) #define TARGET_PC_ENTER (~UINT32_C(0) & 0x0101) #define TARGET_PC_EXTERNAL_DISPATCHER (~UINT32_C(0) & 0x1010) +#define TARGET_LOAD_HELPER_TABLE (~UINT32_C(0) & 0x101010) // This is guaranteed to be an illegal A64 instruction. #define BAD_OPCODE ~UINT32_C(0) @@ -50,6 +49,15 @@ struct patchable_relative { uint32_t offset_loc; uint32_t target_pc; + uint32_t target_offset; + uint32_t adr; +}; + +enum JitError { + NoError, + NotEnoughSpace, + UnexpectedInstruction, + UnknownInstruction }; struct jit_state @@ -61,7 +69,9 @@ struct jit_state uint32_t exit_loc; uint32_t entry_loc; uint32_t dispatcher_loc; + uint32_t helper_table_loc; uint32_t unwind_loc; + enum JitError jit_error; struct patchable_relative* jumps; struct patchable_relative* loads; int num_jumps; @@ -117,6 +127,8 @@ static enum Registers temp_register = R24; static enum Registers temp_div_register = R25; // Temp register for load/store offsets static enum Registers offset_register = R26; +// Special register for external dispatcher context. +static enum Registers VOLATILE_CTXT = R26; // Number of eBPF registers #define REGISTER_MAP_SIZE 11 @@ -169,8 +181,11 @@ static uint32_t inline align_to(uint32_t amount, uint64_t boundary) static void emit_bytes(struct jit_state* state, void* data, uint32_t len) { - assert(len <= state->size); - assert(state->offset <= state->size - len); + if (!(len <= state->size || state->offset <= state->size - len)) { + state->jit_error = NotEnoughSpace; + return; + } + if ((state->offset + len) > state->size) { state->offset = state->size; return; @@ -277,6 +292,16 @@ emit_loadstore_literal( emit_instruction(state, op | reg_op_base | rt); } +static void +emit_adr(struct jit_state *state, uint32_t offset, enum Registers rd) +{ + uint32_t immlo = (offset & 0x00000003) << 30; + uint32_t immhi = (offset & 0x00ffffe0) << 5; + uint32_t instr = 0x10000000 | immhi | immlo | rd; + emit_instruction(state, instr); + +} + enum LoadStorePairOpcode { // op V L @@ -351,15 +376,25 @@ enum UnconditionalBranchImmediateOpcode UBR_BL = 0x94000000U, // 1001_0100_0000_0000_0000_0000_0000_0000 }; -static void +static int note_jump(struct jit_state* state, uint32_t target_pc) { + int index = state->num_jumps++; + if (state->num_jumps == UBPF_MAX_INSTS) { - return; + return -1; } - struct patchable_relative* jump = &state->jumps[state->num_jumps++]; + struct patchable_relative* jump = &state->jumps[index]; jump->offset_loc = state->offset; jump->target_pc = target_pc; + return index; +} + +static void +note_jump_offset(struct jit_state *state, int index, uint32_t target_offset) +{ + struct patchable_relative* to_update = &state->jumps[index]; + to_update->target_offset = target_offset; } static void @@ -373,6 +408,18 @@ note_load(struct jit_state* state, uint32_t target_pc) load->target_pc = target_pc; } +static void +note_adr(struct jit_state* state, uint32_t offset) +{ + if (state->num_loads == UBPF_MAX_INSTS) { + return; + } + struct patchable_relative* load = &state->loads[state->num_loads++]; + load->offset_loc = state->offset; + load->target_pc = offset; + load->adr = offset; +} + /* [ArmARM-A H.a]: C4.1.65: Unconditional branch (immediate). */ static void @@ -411,11 +458,12 @@ enum ConditionalBranchImmediateOpcode }; /* [ArmARM-A H.a]: C4.1.65: Conditional branch (immediate). */ -static void +static int emit_conditionalbranch_immediate(struct jit_state* state, enum Condition cond, uint32_t target_pc) { - note_jump(state, target_pc); + int index = note_jump(state, target_pc); emit_instruction(state, BR_Bcond | (0 << 5) | cond); + return index; } enum CompareBranchOpcode @@ -577,6 +625,16 @@ update_load_literal(struct jit_state* state, uint32_t instr_offset, int32_t targ memcpy(state->buf + instr_offset, &instr, sizeof(uint32_t)); } +static void +update_adr(struct jit_state* state, uint32_t instr_offset, int32_t immediate) +{ + uint32_t instr; + uint32_t immhi = (immediate & 0x00ffffff) << 5; + memcpy(&instr, state->buf + instr_offset, sizeof(uint32_t)); + instr |= immhi; + memcpy(state->buf + instr_offset, &instr, sizeof(uint32_t)); +} + /* Generate the function prologue. * * We set the stack to look like: @@ -609,6 +667,9 @@ emit_jit_prologue(struct jit_state* state, size_t ubpf_stack_size) /* Setup UBPF frame pointer. */ emit_addsub_immediate(state, true, AS_ADD, map_register(10), SP, state->stack_size); + /* Copy R0 to the volatile context for safe keeping. */ + emit_logical_register(state, true, LOG_ORR, VOLATILE_CTXT, RZ, R0); + emit_unconditionalbranch_immediate(state, UBR_BL, TARGET_PC_ENTER); emit_unconditionalbranch_immediate(state, UBR_B, TARGET_PC_EXIT); state->entry_loc = state->offset; @@ -619,6 +680,10 @@ emit_dispatched_external_helper_call(struct jit_state* state, struct ubpf_vm* vm { (void)vm; + // There are two paths through the function: + // 1. There is an external dispatcher registered. If so, we prioritize that. + // 2. We fall back to the regular registered helper. + uint32_t stack_movement = align_to(8, 16); emit_addsub_immediate(state, true, AS_SUB, SP, SP, stack_movement); emit_loadstore_immediate(state, LS_STRX, R30, SP, 0); @@ -626,14 +691,39 @@ emit_dispatched_external_helper_call(struct jit_state* state, struct ubpf_vm* vm // All parameters to the helper function are in the right spot // for the dispatcher. All we need to do now is ... + // Call! + note_load(state, TARGET_PC_EXTERNAL_DISPATCHER); + emit_loadstore_literal(state, LS_LDRL, temp_register); + + // Check whether temp_register is empty. + emit_addsub_immediate(state, true, AS_SUBS, temp_register, temp_register, 0); + + // Jump if we are ... + int jump_index = emit_conditionalbranch_immediate(state, COND_NE, 0); + + // Otherwise, load the target into the register. + + emit_movewide_immediate(state, true, R5, idx); + emit_movewide_immediate(state, true, R6, 3); + emit_dataprocessing_twosource(state, true, DP2_LSLV, R5, R5, R6); + + emit_movewide_immediate(state, true, temp_register, 0); + note_adr(state, TARGET_LOAD_HELPER_TABLE); + emit_adr(state, 0, temp_register); + emit_addsub_register(state, true, AS_ADD, temp_register, temp_register, R5); + emit_loadstore_immediate(state, LS_LDRX, temp_register, temp_register, 0); + + // And now we, too, are ... + + // ... ready to roll. + + note_jump_offset(state, jump_index, state->offset); + // ... set up the final two parameters. emit_movewide_immediate(state, true, R5, idx); // TODO - emit_movewide_immediate(state, true, R6, 0); + emit_logical_register(state, true, LOG_ORR, R6, RZ, VOLATILE_CTXT); - // Call! - note_load(state, TARGET_PC_EXTERNAL_DISPATCHER); - emit_loadstore_literal(state, LS_LDRL, temp_register); emit_unconditionalbranch_register(state, BR_BLR, temp_register); /* On exit need to move result from r0 to whichever register we've mapped EBPF r0 to. */ @@ -705,6 +795,16 @@ emit_dispatched_external_helper_address(struct jit_state *state, uint64_t dispat return helper_address; } +static uint32_t +emit_helper_table(struct jit_state* state, struct ubpf_vm* vm) { + + uint32_t helper_table_address_target = state->offset; + for (int i = 0; iext_funcs[i], sizeof(uint64_t)); + } + return helper_table_address_target; +} + static bool is_imm_op(struct ebpf_inst const* inst) { @@ -979,6 +1079,9 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) emit_jit_prologue(state, UBPF_STACK_SIZE); for (i = 0; i < vm->num_insts; i++) { + + // 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; @@ -989,6 +1092,9 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) int sixty_four = is_alu64_op(&inst); + // If this is an operation with an immediate operand (and that immediate + // operand is _not_ simple), then we convert the operation to the equivalent + // register version after moving the immediate into a temporary register. if (is_imm_op(&inst) && !is_simple_imm(&inst)) { emit_movewide_immediate(state, sixty_four, temp_register, (int64_t)inst.imm); src = temp_register; @@ -1196,16 +1302,30 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg) case EBPF_OP_RSH64_IMM: case EBPF_OP_ARSH64_IMM: *errmsg = ubpf_error("Unexpected instruction at PC %d: opcode %02x, immediate %08x", i, opcode, inst.imm); - return -1; + state->jit_error = UnexpectedInstruction; default: *errmsg = ubpf_error("Unknown instruction at PC %d: opcode %02x", i, opcode); - return -1; + state->jit_error = UnknownInstruction; + } + + switch (state->jit_error) { + case NotEnoughSpace: + *errmsg = ubpf_error("Target buffer too small"); + return -1; + case UnknownInstruction: + case UnexpectedInstruction: + // Error messages for this case are set when discovered. + return -1; + case NoError: + // just keep going! + break; } } emit_jit_epilogue(state); state->dispatcher_loc = emit_dispatched_external_helper_address(state, (uint64_t)vm->dispatcher); + state->helper_table_loc = emit_helper_table(state, vm); return 0; } @@ -1233,7 +1353,9 @@ resolve_jumps(struct jit_state* state) struct patchable_relative jump = state->jumps[i]; int32_t target_loc; - if (jump.target_pc == TARGET_PC_EXIT) { + if (jump.target_offset != 0) { + target_loc = jump.target_offset; + } else if (jump.target_pc == TARGET_PC_EXIT) { target_loc = state->exit_loc; } else if (jump.target_pc == TARGET_PC_ENTER) { target_loc = state->entry_loc; @@ -1256,6 +1378,8 @@ resolve_loads(struct jit_state* state) int32_t target_loc; if (jump.target_pc == TARGET_PC_EXTERNAL_DISPATCHER) { target_loc = state->dispatcher_loc; + } else if (jump.target_pc == TARGET_LOAD_HELPER_TABLE) { + target_loc = state->helper_table_loc; } else { return false; } @@ -1263,7 +1387,11 @@ resolve_loads(struct jit_state* state) int32_t rel = target_loc - jump.offset_loc; assert(rel % 4 == 0); rel >>= 2; - update_load_literal(state, jump.offset_loc, rel); + if (jump.adr != 0) { + update_adr(state, jump.offset_loc, rel); + } else { + update_load_literal(state, jump.offset_loc, rel); + } } return true; } @@ -1271,21 +1399,26 @@ resolve_loads(struct jit_state* state) 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) { UNUSED_PARAMETER(vm); - UNUSED_PARAMETER(new_dispatcher); - UNUSED_PARAMETER(buffer); - UNUSED_PARAMETER(size); - UNUSED_PARAMETER(offset); + uint64_t jit_upper_bound = (uint64_t)buffer + size; + void *dispatcher_address = (void*)((uint64_t)buffer + offset); + if ((uint64_t)dispatcher_address + sizeof(void*) < jit_upper_bound) { + memcpy(dispatcher_address, &new_dispatcher, sizeof(void*)); + return true; + } + return false; } bool ubpf_jit_update_helper_arm64(struct ubpf_vm* vm, ext_func new_helper, unsigned int idx, uint8_t* buffer, size_t size, uint32_t offset) { UNUSED_PARAMETER(vm); - UNUSED_PARAMETER(new_helper); - UNUSED_PARAMETER(idx); - UNUSED_PARAMETER(buffer); - UNUSED_PARAMETER(size); - UNUSED_PARAMETER(offset); + uint64_t jit_upper_bound = (uint64_t)buffer + size; + + void* dispatcher_address = (void*)((uint64_t)buffer + offset + (8 * idx)); + if ((uint64_t)dispatcher_address + sizeof(void*) < jit_upper_bound) { + memcpy(dispatcher_address, &new_helper, sizeof(void*)); + return true; + } return false; } @@ -1322,11 +1455,6 @@ ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size) goto out; } - if (state.offset == state.size) { - compile_result.errmsg = ubpf_error("Target buffer too small"); - goto out; - } - if (!resolve_jumps(&state) || !resolve_loads(&state)) { compile_result.errmsg = ubpf_error("Could not patch the relative addresses in the JIT'd code."); goto out; @@ -1334,6 +1462,8 @@ ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size) compile_result.compile_result = UBPF_JIT_COMPILE_SUCCESS; *size = state.offset; + compile_result.external_dispatcher_offset = state.dispatcher_loc; + compile_result.external_helper_offset = state.helper_table_loc; out: free(state.pc_locs); diff --git a/vm/ubpf_vm.c b/vm/ubpf_vm.c index f8d38fcb..2246db3b 100644 --- a/vm/ubpf_vm.c +++ b/vm/ubpf_vm.c @@ -98,7 +98,9 @@ ubpf_create(void) vm->jit_update_dispatcher = ubpf_jit_update_dispatcher_x86_64; vm->jit_update_helper = ubpf_jit_update_helper_x86_64; #elif defined(__aarch64__) || defined(_M_ARM64) - vm->translate = ubpf_translate_arm64; + vm->jit_translate = ubpf_translate_arm64; + vm->jit_update_dispatcher = ubpf_jit_update_dispatcher_arm64; + vm->jit_update_helper = ubpf_jit_update_helper_arm64; #else vm->translate = ubpf_translate_null; #endif