Implement local calls in JIT/interpreter. (#271)

Signed-off-by: Will Hawkins <hawkinsw@obs.cr>

ubpf_plugin/CMakeLists.txt

@@ -45,7 +45,7 @@ endif()
 foreach(file ${files})
     # TODO: remove this once we have a proper implementation of interlocked operations
     # and support for calling local functions.
-    string(REGEX MATCH "(lock|call_local)" EXPECT_FAILURE "${file}")
+    string(REGEX MATCH "(lock)" EXPECT_FAILURE "${file}")
     add_test(
         NAME ${file}-JIT
         COMMAND ${BPF_CONFORMANCE_RUNNER} --test_file_path ${file} ${PLUGIN_JIT}

vm/ebpf.h

@@ -33,6 +33,22 @@ struct ebpf_inst
     int32_t imm;
 };
 
+enum bpf_register
+{
+    BPF_REG_0 = 0,
+    BPF_REG_1,
+    BPF_REG_2,
+    BPF_REG_3,
+    BPF_REG_4,
+    BPF_REG_5,
+    BPF_REG_6,
+    BPF_REG_7,
+    BPF_REG_8,
+    BPF_REG_9,
+    BPF_REG_10,
+    _BPF_REG_MAX,
+};
+
 #define EBPF_CLS_MASK 0x07
 #define EBPF_ALU_OP_MASK 0xf0
 #define EBPF_JMP_OP_MASK 0xf0

vm/inc/ubpf.h

@@ -40,10 +40,17 @@ extern "C"
 #endif
 
 /**
- * @brief Default stack size for the VM.
+ * @brief Default stack size for the VM. Must be divisible by 16.
  */
 #if !defined(UBPF_STACK_SIZE)
 #define UBPF_STACK_SIZE 512
+#endif
+
+/**
+ * @brief Default maximum number of nested calls in the VM.
+ */
+#if !defined(UBPF_MAX_CALL_DEPTH)
+#define UBPF_MAX_CALL_DEPTH 10
 #endif
 
     /**

vm/ubpf_int.h

@@ -49,6 +49,12 @@ struct ubpf_vm
 #endif
 };
 
+struct ubpf_stack_frame
+{
+    uint16_t return_address;
+    uint64_t saved_registers[4];
+};
+
 /* The various JIT targets.  */
 int
 ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg);

vm/ubpf_jit_arm64.c

@@ -40,6 +40,7 @@
 
 /* Special values for target_pc in struct jump */
 #define TARGET_PC_EXIT ~UINT32_C(0)
+#define TARGET_PC_ENTER (~UINT32_C(0) & 0x0101)
 
 // This is guaranteed to be an illegal A64 instruction.
 #define BAD_OPCODE ~UINT32_C(0)
@@ -57,6 +58,7 @@ struct jit_state
     uint32_t size;
     uint32_t* pc_locs;
     uint32_t exit_loc;
+    uint32_t entry_loc;
     uint32_t unwind_loc;
     struct jump* jumps;
     int num_jumps;
@@ -155,6 +157,8 @@ emit_movewide_immediate(struct jit_state* state, bool sixty_four, enum Registers
 static void
 divmod(struct jit_state* state, uint8_t opcode, int rd, int rn, int rm);
 
+static uint32_t inline align_stack_amount(uint32_t amount) { return (amount + 15) & ~15U; }
+
 static void
 emit_bytes(struct jit_state* state, void* data, uint32_t len)
 {
@@ -547,11 +551,12 @@ static void
 emit_function_prologue(struct jit_state* state, size_t ubpf_stack_size)
 {
     uint32_t register_space = _countof(callee_saved_registers) * 8 + 2 * 8;
-    state->stack_size = (ubpf_stack_size + register_space + 15) & ~15U;
+    state->stack_size = align_stack_amount(ubpf_stack_size + register_space);
     emit_addsub_immediate(state, true, AS_SUB, SP, SP, state->stack_size);
 
     /* Set up frame */
     emit_loadstorepair_immediate(state, LSP_STPX, R29, R30, SP, 0);
+    /* In ARM64 calling convention, R29 is the frame pointer. */
     emit_addsub_immediate(state, true, AS_ADD, R29, SP, 0);
 
     /* Save callee saved registers */
@@ -563,11 +568,19 @@ emit_function_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);
+
+    emit_unconditionalbranch_immediate(state, UBR_BL, TARGET_PC_ENTER);
+    emit_unconditionalbranch_immediate(state, UBR_B, TARGET_PC_EXIT);
+    state->entry_loc = state->offset;
 }
 
 static void
 emit_call(struct jit_state* state, uintptr_t func)
 {
+    uint32_t stack_movement = align_stack_amount(8);
+    emit_addsub_immediate(state, true, AS_SUB, SP, SP, stack_movement);
+    emit_loadstore_immediate(state, LS_STRX, R30, SP, 0);
+
     emit_movewide_immediate(state, true, temp_register, func);
     emit_unconditonalbranch_register(state, BR_BLR, temp_register);
 
@@ -576,6 +589,25 @@ emit_call(struct jit_state* state, uintptr_t func)
     if (dest != R0) {
         emit_logical_register(state, true, LOG_ORR, dest, RZ, R0);
     }
+
+    emit_loadstore_immediate(state, LS_LDRX, R30, SP, 0);
+    emit_addsub_immediate(state, true, AS_ADD, SP, SP, stack_movement);
+}
+
+static void
+emit_local_call(struct jit_state* state, uint32_t target_pc)
+{
+    uint32_t stack_movement = align_stack_amount(40);
+    emit_addsub_immediate(state, true, AS_SUB, SP, SP, stack_movement);
+    emit_loadstore_immediate(state, LS_STRX, R30, SP, 0);
+    emit_loadstorepair_immediate(state, LSP_STPX, map_register(6), map_register(7), SP, 8);
+    emit_loadstorepair_immediate(state, LSP_STPX, map_register(8), map_register(9), SP, 24);
+    note_jump(state, target_pc);
+    emit_unconditionalbranch_immediate(state, UBR_BL, target_pc);
+    emit_loadstore_immediate(state, LS_LDRX, R30, SP, 0);
+    emit_loadstorepair_immediate(state, LSP_LDPX, map_register(6), map_register(7), SP, 8);
+    emit_loadstorepair_immediate(state, LSP_LDPX, map_register(8), map_register(9), SP, 24);
+    emit_addsub_immediate(state, true, AS_ADD, SP, SP, stack_movement);
 }
 
 static void
@@ -588,6 +620,9 @@ emit_function_epilogue(struct jit_state* state)
         emit_logical_register(state, true, LOG_ORR, R0, RZ, map_register(0));
     }
 
+    /* We could be anywhere in the stack if we excepted. Get our head right. */
+    emit_addsub_immediate(state, true, AS_ADD, SP, R29, 0);
+
     /* Restore callee-saved registers).  */
     size_t i;
     for (i = 0; i < _countof(callee_saved_registers); i += 2) {
@@ -1020,16 +1055,21 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
             emit_conditionalbranch_immediate(state, to_condition(opcode), target_pc);
             break;
         case EBPF_OP_CALL:
-            emit_call(state, (uintptr_t)vm->ext_funcs[inst.imm]);
-            if (inst.imm == vm->unwind_stack_extension_index) {
-                emit_addsub_immediate(state, true, AS_SUBS, RZ, map_register(0), 0);
-                emit_conditionalbranch_immediate(state, COND_EQ, TARGET_PC_EXIT);
+            if (inst.src == 0) {
+                emit_call(state, (uintptr_t)vm->ext_funcs[inst.imm]);
+                if (inst.imm == vm->unwind_stack_extension_index) {
+                    emit_addsub_immediate(state, true, AS_SUBS, RZ, map_register(0), 0);
+                    emit_conditionalbranch_immediate(state, COND_EQ, TARGET_PC_EXIT);
+                }
+            } else if (inst.src == 1) {
+                uint32_t call_target = i + inst.imm + 1;
+                emit_local_call(state, call_target);
+            } else {
+                emit_unconditionalbranch_immediate(state, UBR_B, TARGET_PC_EXIT);
             }
             break;
         case EBPF_OP_EXIT:
-            if (i != vm->num_insts - 1) {
-                emit_unconditionalbranch_immediate(state, UBR_B, TARGET_PC_EXIT);
-            }
+            emit_unconditonalbranch_register(state, BR_RET, R30);
             break;
 
         case EBPF_OP_STXW:
@@ -1122,6 +1162,8 @@ resolve_jumps(struct jit_state* state)
         int32_t target_loc;
         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;
         } else {
             target_loc = state->pc_locs[jump.target_pc];
         }

vm/ubpf_jit_x86_64.c

@@ -18,6 +18,8 @@
  * limitations under the License.
  */
 
+#include "ebpf.h"
+#include <stdint.h>
 #define _GNU_SOURCE
 #include <stdio.h>
 #include <stdlib.h>
@@ -94,8 +96,27 @@ static int register_map[REGISTER_MAP_SIZE] = {
 static int
 map_register(int r)
 {
-    assert(r < REGISTER_MAP_SIZE);
-    return register_map[r % REGISTER_MAP_SIZE];
+    assert(r < _BPF_REG_MAX);
+    return register_map[r % _BPF_REG_MAX];
+}
+
+static inline void
+emit_local_call(struct jit_state* state, uint32_t target_pc)
+{
+    /*
+     * Pushing 4 * 8 = 32 bytes will maintain the invariant
+     * that the stack is 16-byte aligned.
+     */
+    emit_push(state, map_register(BPF_REG_6));
+    emit_push(state, map_register(BPF_REG_7));
+    emit_push(state, map_register(BPF_REG_8));
+    emit_push(state, map_register(BPF_REG_9));
+    emit1(state, 0xe8); // e8 is the opcode for a CALL
+    emit_jump_offset(state, target_pc);
+    emit_pop(state, map_register(BPF_REG_9));
+    emit_pop(state, map_register(BPF_REG_8));
+    emit_pop(state, map_register(BPF_REG_7));
+    emit_pop(state, map_register(BPF_REG_6));
 }
 
 /* For testing, this changes the mapping between x86 and eBPF registers */
@@ -133,14 +154,44 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
 
     /* Move first platform parameter register into register 1 */
     if (map_register(1) != platform_parameter_registers[0]) {
-        emit_mov(state, platform_parameter_registers[0], map_register(1));
+        emit_mov(state, platform_parameter_registers[0], map_register(BPF_REG_1));
     }
 
-    /* Copy stack pointer to R10 */
-    emit_mov(state, RSP, map_register(10));
+    /*
+     * Copy stack pointer to R10. We will also use RBP to restore
+     * the stack pointer after execution.
+     */
+    emit_mov(state, RSP, map_register(BPF_REG_10));
+
+    uint64_t ubpf_stack_configuration_delta = UBPF_STACK_SIZE;
+    /*
+     * Assuming that the stack is 16-byte aligned when
+     * we start executing the jit'd code, we need to maintain
+     * that invariant. The UBPF_STACK_SIZE is guaranteed to be
+     * divisible by 16. However, if we pushed an odd number of
+     * registers on the stack when we are saving state (see above),
+     * then we have violated the invariant. So, we'll have to
+     * make an additional tweak to rectify the situation.
+     */
+    if (sizeof(platform_nonvolatile_registers) % 16 != 0) {
+        ubpf_stack_configuration_delta += 0x8;
+    }
 
     /* Allocate stack space */
-    emit_alu64_imm32(state, 0x81, 5, RSP, UBPF_STACK_SIZE);
+    emit_alu64_imm32(state, 0x81, 5, RSP, ubpf_stack_configuration_delta);
+
+    /*
+     * Use a call to set up a place where we can land after eBPF program's
+     * final EXIT call. This will make JIT of BPF EXIT call easier in the
+     * presence of calls to local functions.
+     */
+    emit1(state, 0xe8);
+    emit4(state, 5);
+    /*
+     * We jump over this instruction in the first place; return here
+     * after the eBPF program is finished executing.
+     */
+    emit_jmp(state, TARGET_PC_EXIT);
 
     for (i = 0; i < vm->num_insts; i++) {
         struct ebpf_inst inst = ubpf_fetch_instruction(vm, i);
@@ -489,17 +540,20 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
             break;
         case EBPF_OP_CALL:
             /* We reserve RCX for shifts */
-            emit_mov(state, RCX_ALT, RCX);
-            emit_call(state, vm->ext_funcs[inst.imm]);
-            if (inst.imm == vm->unwind_stack_extension_index) {
-                emit_cmp_imm32(state, map_register(0), 0);
-                emit_jcc(state, 0x84, TARGET_PC_EXIT);
+            if (inst.src == 0) {
+                emit_mov(state, RCX_ALT, RCX);
+                emit_call(state, vm->ext_funcs[inst.imm]);
+                if (inst.imm == vm->unwind_stack_extension_index) {
+                    emit_cmp_imm32(state, map_register(BPF_REG_0), 0);
+                    emit_jcc(state, 0x84, TARGET_PC_EXIT);
+                }
+            } else if (inst.src == 1) {
+                uint32_t target_pc = i + inst.imm + 1;
+                emit_local_call(state, target_pc);
             }
             break;
         case EBPF_OP_EXIT:
-            if (i != vm->num_insts - 1) {
-                emit_jmp(state, TARGET_PC_EXIT);
-            }
+            emit_ret(state);
             break;
 
         case EBPF_OP_LDXW:
@@ -558,12 +612,12 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
     state->exit_loc = state->offset;
 
     /* Move register 0 into rax */
-    if (map_register(0) != RAX) {
-        emit_mov(state, map_register(0), RAX);
+    if (map_register(BPF_REG_0) != RAX) {
+        emit_mov(state, map_register(BPF_REG_0), RAX);
     }
 
-    /* Deallocate stack space */
-    emit_alu64_imm32(state, 0x81, 0, RSP, UBPF_STACK_SIZE);
+    /* Deallocate stack space by restoring RSP from RBP. */
+    emit_mov(state, map_register(BPF_REG_10), RSP);
 
     /* Restore platform non-volatile registers */
     for (i = 0; i < _countof(platform_nonvolatile_registers); i++) {

vm/ubpf_jit_x86_64.h

@@ -343,6 +343,12 @@ emit_store_imm32(struct jit_state* state, enum operand_size size, int dst, int32
     }
 }
 
+static inline void
+emit_ret(struct jit_state* state)
+{
+    emit1(state, 0xc3);
+}
+
 static inline void
 emit_call(struct jit_state* state, void* target)
 {
@@ -355,10 +361,22 @@ emit_call(struct jit_state* state, void* target)
     emit_alu64_imm32(state, 0x81, 5, RSP, 4 * sizeof(uint64_t));
 #endif
 
+    /*
+     * Because we do not push/pop here, the invariant is maintained
+     * that the stack pointer is 16-byte aligned.
+     */
+
     /* TODO use direct call when possible */
     emit_load_imm(state, RAX, (uintptr_t)target);
     /* callq *%rax */
     emit1(state, 0xff);
+    // ModR/M byte: b11010000b = xd
+    //               ^
+    //               register-direct addressing.
+    //                 ^
+    //                 opcode extension (2)
+    //                    ^
+    //                    rax is register 0
     emit1(state, 0xd0);
 
 #if defined(_WIN32)