Merge pull request #88 from Alan-Jowett/issue87

Implement unwind on success semantics to support bpf_tail_call.

bpf_tail_call has special semantics. If the call returns success, then ubpf needs to store the return value and unwind the stack so to the caller.

Resolves: #87 

Signed-off-by: Alan Jowett Alan.Jowett@microsoft.com

tests/call_unwind.data

@@ -0,0 +1,9 @@
+-- asm
+mov r1, 0
+call 5
+mov r0, 2
+exit
+-- result
+0x0
+-- no register offset
+call instruction

tests/call_unwind_fail.data

@@ -0,0 +1,9 @@
+-- asm
+mov r1, -1
+call 5
+mov r0, 2
+exit
+-- result
+0x2
+-- no register offset
+call instruction

vm/inc/ubpf.h

@@ -112,4 +112,14 @@ ubpf_jit_fn ubpf_compile(struct ubpf_vm *vm, char **errmsg);
  * message will be stored in 'errmsg' and should be freed by the caller.
  */
 int ubpf_translate(struct ubpf_vm *vm, uint8_t *buffer, size_t *size, char **errmsg);
+
+/*
+ * Instruct the uBPF runtime to apply unwind-on-success semantics to a helper
+ * function. If the function returns 0, the uBPF runtime will end execution of
+ * the eBPF program and immediately return control to the caller. This is used
+ * for implementing function like the "bpf_tail_call" helper. 
+ * Returns 0 on success, -1 on if there is already an unwind helper set.
+ */
+int ubpf_set_unwind_function_index(struct ubpf_vm *vm, unsigned int idx);
+
 #endif

vm/test.c

@@ -226,6 +226,12 @@ sqrti(uint32_t x)
     return sqrt(x);
 }
 
+static uint64_t
+unwind(uint64_t i)
+{
+    return i;
+}
+
 static void
 register_functions(struct ubpf_vm *vm)
 {
@@ -234,4 +240,6 @@ register_functions(struct ubpf_vm *vm)
     ubpf_register(vm, 2, "trash_registers", trash_registers);
     ubpf_register(vm, 3, "sqrti", sqrti);
     ubpf_register(vm, 4, "strcmp_ext", strcmp);
+    ubpf_register(vm, 5, "unwind", unwind);
+    ubpf_set_unwind_function_index(vm, 5);
 }

vm/ubpf_int.h

@@ -32,6 +32,7 @@ struct ubpf_vm {
     const char **ext_func_names;
     bool bounds_check_enabled;
     int (*error_printf)(FILE* stream, const char* format, ...);
+    int unwind_stack_extension_index;
 };
 
 char *ubpf_error(const char *fmt, ...);

vm/ubpf_jit_x86_64.c

@@ -408,6 +408,10 @@ translate(struct ubpf_vm *vm, struct jit_state *state, char **errmsg)
             /* 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);
+            }
             break;
         case EBPF_OP_EXIT:
             if (i != vm->num_insts - 1) {

vm/ubpf_jit_x86_64.h

@@ -61,6 +61,7 @@ struct jit_state {
     uint32_t *pc_locs;
     uint32_t exit_loc;
     uint32_t div_by_zero_loc;
+    uint32_t unwind_loc;
     struct jump *jumps;
     int num_jumps;
 };

vm/ubpf_vm.c

@@ -66,6 +66,8 @@ ubpf_create(void)
 
     vm->bounds_check_enabled = true;
     vm->error_printf = fprintf;
+
+    vm->unwind_stack_extension_index = -1;
     return vm;
 }
 
@@ -90,6 +92,17 @@ ubpf_register(struct ubpf_vm *vm, unsigned int idx, const char *name, void *fn)
 
     vm->ext_funcs[idx] = (ext_func)fn;
     vm->ext_func_names[idx] = name;
+
+    return 0;
+}
+
+int ubpf_set_unwind_function_index(struct ubpf_vm *vm, unsigned int idx)
+{
+    if (vm->unwind_stack_extension_index != -1) {
+        return -1;
+    }
+
+    vm->unwind_stack_extension_index = idx;
     return 0;
 }
 
@@ -568,6 +581,11 @@ ubpf_exec(const struct ubpf_vm *vm, void *mem, size_t mem_len, uint64_t* bpf_ret
             return 0;
         case EBPF_OP_CALL:
             reg[0] = vm->ext_funcs[inst.imm](reg[1], reg[2], reg[3], reg[4], reg[5]);
+            // Unwind the stack if unwind extension returns success.
+            if (inst.imm == vm->unwind_stack_extension_index && reg[0] == 0) {
+                *bpf_return_value = reg[0];
+                return 0;
+            }
             break;
         }
     }