Update JIT to handle overlapping instructions (#497)

The interpreter already supports functions that happen to be in more
than one function. This patch adds that support to the JIT.

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

tests/call0.data

@@ -0,0 +1,7 @@
+-- asm
+mov %r0, 0x12345678
+call local next
+next:
+exit
+-- result
+0x12345678

vm/ubpf_int.h

@@ -21,6 +21,7 @@
 #ifndef UBPF_INT_H
 #define UBPF_INT_H
 
+#include <stdbool.h>
 #include <stdint.h>
 #include <ubpf.h>
 #include "ebpf.h"
@@ -170,4 +171,20 @@ ubpf_stack_usage_for_local_func(const struct ubpf_vm* vm, uint16_t pc);
 bool
 ubpf_calculate_stack_usage_for_local_func(const struct ubpf_vm* vm, uint16_t pc, char** errmsg);
 
+/**
+ * @brief Determine whether an eBPF instruction has a fallthrough
+ *
+ * An eBPF instruction has a fallthrough unless the instruction performs
+ * unconditional change in control-flow. Currently, the only instruction
+ * that fits that description is the EXIT.
+ *
+ * @return True if the inst has a fallthrough; false, otherwise.
+ */
+static inline bool
+ubpf_instruction_has_fallthrough(const struct ebpf_inst inst)
+{
+    // The only instruction that does not have a fallthrough is the EXIT.
+    return inst.opcode != EBPF_OP_EXIT;
+}
+
 #endif

vm/ubpf_jit_arm64.c

@@ -22,6 +22,7 @@
  * [ArmARM-A H.a]: https://developer.arm.com/documentation/ddi0487/ha
  */
 
+#include <stdint.h>
 #define _GNU_SOURCE
 #include <stdio.h>
 #include <stdbool.h>
@@ -334,10 +335,11 @@ enum UnconditionalBranchImmediateOpcode
 };
 
 /* [ArmARM-A H.a]: C4.1.65: Unconditional branch (immediate).  */
-static void
+static uint32_t
 emit_unconditionalbranch_immediate(
     struct jit_state* state, enum UnconditionalBranchImmediateOpcode op, int32_t target_pc)
 {
+    uint32_t source_offset = state->offset;
     struct patchable_relative* table = state->jumps;
     int* num_jumps = &state->num_jumps;
     if (op == UBR_BL && target_pc != TARGET_PC_ENTER) {
@@ -347,6 +349,8 @@ emit_unconditionalbranch_immediate(
 
     emit_patchable_relative(state->offset, target_pc, 0, table, (*num_jumps)++);
     emit_instruction(state, op);
+
+    return source_offset;
 }
 
 enum Condition
@@ -974,6 +978,21 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
         // occur at the end of the loop.
         struct ebpf_inst inst = ubpf_fetch_instruction(vm, i);
 
+
+        // If a) the previous instruction could fallthrough to this instruction and
+        //    b) this instruction starts a local function, then
+        // we have to "jump around" the code that manipulates the stack!
+        uint32_t fallthrough_jump_source = 0;
+        bool fallthrough_jump_present = false;
+        if (i != 0 && vm->int_funcs[i]) {
+            struct ebpf_inst prev_inst = ubpf_fetch_instruction(vm, i - 1);
+            if (ubpf_instruction_has_fallthrough(prev_inst)) {
+                fallthrough_jump_source = emit_unconditionalbranch_immediate(state, UBR_B, 0);
+                fallthrough_jump_present = true;
+            }
+        }
+
+
         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));
@@ -987,6 +1006,10 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
             }
         }
 
+        if (fallthrough_jump_present) {
+            fixup_jump_target(state->jumps, state->num_jumps, fallthrough_jump_source, state->offset);
+        }
+
         state->pc_locs[i] = state->offset;
 
         enum Registers dst = map_register(inst.dst);

vm/ubpf_jit_support.c

@@ -77,13 +77,26 @@ release_jit_state_result(struct jit_state* state, struct ubpf_jit_result* compil
 }
 
 void
-emit_patchable_relative(
-    uint32_t offset, uint32_t target_pc, uint32_t manual_target_offset, struct patchable_relative* table, size_t index)
+emit_patchable_relative_ex(
+    uint32_t offset,
+    uint32_t target_pc,
+    uint32_t manual_target_offset,
+    struct patchable_relative* table,
+    size_t index,
+    bool near)
 {
     struct patchable_relative* jump = &table[index];
     jump->offset_loc = offset;
     jump->target_pc = target_pc;
     jump->target_offset = manual_target_offset;
+    jump->near = near;
+}
+
+void
+emit_patchable_relative(
+    uint32_t offset, uint32_t target_pc, uint32_t manual_target_offset, struct patchable_relative* table, size_t index)
+{
+    emit_patchable_relative_ex(offset, target_pc, manual_target_offset, table, index, false);
 }
 
 void

vm/ubpf_jit_support.h

@@ -49,6 +49,8 @@ struct patchable_relative
     uint32_t target_pc;
     /* ... the target_offset is set which overrides the automatic lookup. */
     uint32_t target_offset;
+    /* Whether or not this patchable relative is _near_. */
+    bool near;
 };
 
 /* Special values for target_pc in struct jump */
@@ -112,6 +114,39 @@ initialize_jit_state_result(
 void
 release_jit_state_result(struct jit_state* state, struct ubpf_jit_result* compile_result);
 
+/** @brief Add an entry to the given patchable relative table.
+ *
+ * Emitting an entry into the patchable relative table means that resolution of the target
+ * address can be postponed until all the instructions are emitted. Note: This function does
+ * not emit any instructions -- it simply updates metadata to guide resolution after code generation.
+ * _target_pc_ is in eBPF instruction units and _manual_target_offset_ is in JIT'd instruction
+ * units. In other words, setting _target_pc_ instead of _manual_target_offset_ will guide
+ * the resolution algorithm to find the JIT'd code that corresponds to the eBPF instruction
+ * (as the jump target); alternatively, setting _manual_target_offset_ will direct the
+ * resolution algorithm to find the JIT'd instruction at that offset (as the target).
+ *
+ * @param[in] offset The offset in the JIT'd code where the to-be-resolved target begins.
+ * @param[in] target_pc The offset of the eBPF instruction targeted by the jump.
+ * @param[in] manual_target_offset The offset of the JIT'd instruction targeted by the jump.
+ *                                 A non-zero value for this parameter overrides _target_pc_`.
+ * @param[in] table The relative patchable table to update.
+ * @param[in] index A spot in the _table_ to add/update according to the given parameters.
+ * @param[in] near Whether the target is relatively near the jump.
+ */
+void
+emit_patchable_relative_ex(
+    uint32_t offset,
+    uint32_t target_pc,
+    uint32_t manual_target_offset,
+    struct patchable_relative* table,
+    size_t index,
+    bool near);
+
+/** @brief Add an entry to the given patchable relative table.
+ *
+ * See emit_patchable_relative_ex. emit_patchable_relative's parameters have the same meaning
+ * but fixes the _near_ argument to false.
+ */
 void
 emit_patchable_relative(
     uint32_t offset, uint32_t target_pc, uint32_t manual_target_offset, struct patchable_relative* table, size_t index);

vm/ubpf_jit_x86_64.c

@@ -341,6 +341,19 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
         int src = map_register(inst.src);
         uint32_t target_pc = i + inst.offset + 1;
 
+        // If a) the previous instruction could fallthrough to this instruction and
+        //    b) this instruction starts a local function, then
+        // we have to "jump around" the code that manipulates the stack!
+        uint32_t fallthrough_jump_source = 0;
+        bool fallthrough_jump_present = false;
+        if (i != 0 && vm->int_funcs[i]) {
+            struct ebpf_inst prev_inst = ubpf_fetch_instruction(vm, i - 1);
+            if (ubpf_instruction_has_fallthrough(prev_inst)) {
+                fallthrough_jump_source = emit_near_jmp(state, 0);
+                fallthrough_jump_present = true;
+            }
+        }
+
         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);
@@ -350,14 +363,19 @@ 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;
             } else {
                 assert(state->bpf_function_prolog_size == state->offset - prolog_start);
             }
+
         }
 
+        if (fallthrough_jump_present) {
+            fixup_jump_target(state->jumps, state->num_jumps, fallthrough_jump_source, state->offset);
+        }
         state->pc_locs[i] = state->offset;
 
         switch (inst.opcode) {
@@ -778,8 +796,7 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
         }
 
         // If this is a ALU32 instruction, truncate the target register to 32 bits.
-        if (((inst.opcode & EBPF_CLS_MASK) == EBPF_CLS_ALU) &&
-            (inst.opcode & EBPF_ALU_OP_MASK) != 0xd0) {
+        if (((inst.opcode & EBPF_CLS_MASK) == EBPF_CLS_ALU) && (inst.opcode & EBPF_ALU_OP_MASK) != 0xd0) {
             emit_truncate_u32(state, dst);
         }
     }
@@ -994,11 +1011,29 @@ resolve_patchable_relatives(struct jit_state* state)
             target_loc = state->pc_locs[jump.target_pc];
         }
 
-        /* Assumes jump offset is at end of instruction */
-        uint32_t rel = target_loc - (jump.offset_loc + sizeof(uint32_t));
+        if (jump.near) {
+            /* When there is a near jump, we need to make sure that the target
+             * is within the proper limits. So, we start with a type that can
+             * hold values that are bigger than we'll ultimately need. If we
+             * went straight to the uint8_t, we couldn't tell if we overflowed.
+             */
+            int32_t rel = target_loc - (jump.offset_loc + sizeof(uint8_t));
+            if (!(-128 <= rel && rel < 128)) {
+                return false;
+            }
+            /* Now that we are sure the target is _near_ enough, we can move
+             * to the proper type.
+             */
+            int8_t rel8 = rel;
+            uint8_t* offset_ptr = &state->buf[jump.offset_loc];
+            *offset_ptr = rel8;
+        } else {
+            /* Assumes jump offset is at end of instruction */
+            uint32_t rel = target_loc - (jump.offset_loc + sizeof(uint32_t));
 
-        uint8_t* offset_ptr = &state->buf[jump.offset_loc];
-        memcpy(offset_ptr, &rel, sizeof(uint32_t));
+            uint8_t* offset_ptr = &state->buf[jump.offset_loc];
+            memcpy(offset_ptr, &rel, sizeof(uint32_t));
+        }
     }
 
     for (i = 0; i < state->num_local_calls; i++) {

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_near_jump_address_reloc(struct jit_state* state, int32_t target_pc)
+{
+    if (state->num_jumps == UBPF_MAX_INSTS) {
+        state->jit_status = TooManyJumps;
+        return 0;
+    }
+    uint32_t target_address_offset = state->offset;
+    emit_patchable_relative_ex(state->offset, target_pc, 0, state->jumps, state->num_jumps++, true /* near */);
+    emit1(state, 0x0);
+    return target_address_offset;
+}
+
 static uint32_t
 emit_local_call_address_reloc(struct jit_state* state, int32_t target_pc)
 {
@@ -402,11 +415,32 @@ emit_ret(struct jit_state* state)
     emit1(state, 0xc3);
 }
 
-static inline void
+/** @brief Emit a (32-bit) jump.
+ *
+ * @param[in] state The JIT state.
+ * @param[in] target_pc The PC to which to jump when this near
+ *                      jump is executed.
+ * @return The offset in the JIT'd code where the jump offset starts.
+ */
+static inline uint32_t
 emit_jmp(struct jit_state* state, uint32_t target_pc)
 {
     emit1(state, 0xe9);
-    emit_jump_address_reloc(state, target_pc);
+    return emit_jump_address_reloc(state, target_pc);
+}
+
+/** @brief Emit a near jump.
+ *
+ * @param[in] state The JIT state.
+ * @param[in] target_pc The PC to which to jump when this near
+ *                      jump is executed.
+ * @return The offset in the JIT'd code where the jump offset starts.
+ */
+static inline uint32_t
+emit_near_jmp(struct jit_state* state, uint32_t target_pc)
+{
+    emit1(state, 0xeb);
+    return emit_near_jump_address_reloc(state, target_pc);
 }
 
 static inline uint32_t