Add support for disjoint helper id address space

Signed-off-by: Alan Jowett <alan.jowett@microsoft.com>

vm/inc/ubpf.h

@@ -33,6 +33,8 @@
 struct ubpf_vm;
 typedef uint64_t (*ubpf_jit_fn)(void *mem, size_t mem_len);
 
+typedef uint64_t (*ubpf_helper_fn)(uint64_t arg0, uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4);
+
 struct ubpf_vm *ubpf_create(void);
 void ubpf_destroy(struct ubpf_vm *vm);
 
@@ -64,7 +66,21 @@ void ubpf_set_error_print(struct ubpf_vm *vm, int (*error_printf)(FILE* stream,
  *
  * Returns 0 on success, -1 on error.
  */
-int ubpf_register(struct ubpf_vm *vm, unsigned int idx, const char *name, void *fn);
+int ubpf_register(struct ubpf_vm *vm, unsigned int idx, const char *name, ubpf_helper_fn fn);
+
+/*
+ * Register a callback function used to resolve external functions.
+ * 
+ * The immediate field of a CALL instruction is an identifier that denotes
+ * which helper function should be invoked. 
+ * 
+ * 'context' A opaque pointer used when invoking the resolver function.
+ * 'helper_resolver' A pointer to a function that resolves helper identifiers
+ * to function addresses.
+ * 
+ * Returns 0 on success, -1 on error.
+ */
+int ubpf_register_helper_resolver(struct ubpf_vm *vm, void* context, ubpf_helper_fn (*helper_resolver)(void* context, int32_t helper_identifier));
 
 /*
  * Load code into a VM

vm/test.c

@@ -31,6 +31,13 @@
 void ubpf_set_register_offset(int x);
 static void *readfile(const char *path, size_t maxlen, size_t *len);
 static void register_functions(struct ubpf_vm *vm);
+static ubpf_helper_fn helper_resolver(void* helper_resolver_context, int32_t helper_id);
+typedef struct {
+    uint32_t helper_function_count;
+    ubpf_helper_fn helper_functions[];
+}helper_function_dispatch_table_t;
+
+static helper_function_dispatch_table_t helper_function_dispatch_table;
 
 static void usage(const char *name)
 {
@@ -105,6 +112,8 @@ int main(int argc, char **argv)
 
     register_functions(vm);
 
+    ubpf_register_helper_resolver(vm, &helper_function_dispatch_table, helper_resolver);
+
     /* 
      * The ELF magic corresponds to an RSH instruction with an offset,
      * which is invalid.
@@ -225,12 +234,34 @@ sqrti(uint32_t x)
     return sqrt(x);
 }
 
+static helper_function_dispatch_table_t helper_function_dispatch_table = {
+    5,
+    {
+        (ubpf_helper_fn)gather_bytes,
+        (ubpf_helper_fn)memfrob,
+        (ubpf_helper_fn)trash_registers,
+        (ubpf_helper_fn)sqrti,
+        (ubpf_helper_fn)strcmp,
+    },
+};
+
+
 static void
 register_functions(struct ubpf_vm *vm)
 {
-    ubpf_register(vm, 0, "gather_bytes", gather_bytes);
-    ubpf_register(vm, 1, "memfrob", memfrob);
-    ubpf_register(vm, 2, "trash_registers", trash_registers);
-    ubpf_register(vm, 3, "sqrti", sqrti);
-    ubpf_register(vm, 4, "strcmp_ext", strcmp);
+    ubpf_register(vm, 0, "gather_bytes", (ubpf_helper_fn)gather_bytes);
+    ubpf_register(vm, 1, "memfrob", (ubpf_helper_fn)memfrob);
+    ubpf_register(vm, 2, "trash_registers", (ubpf_helper_fn)trash_registers);
+    ubpf_register(vm, 3, "sqrti", (ubpf_helper_fn)sqrti);
+    ubpf_register(vm, 4, "strcmp_ext", (ubpf_helper_fn)strcmp);
 }
+
+static ubpf_helper_fn helper_resolver(void* helper_resolver_context, int32_t helper_id)
+{
+    helper_function_dispatch_table_t* helper_functions = (helper_function_dispatch_table_t*)helper_resolver_context;
+    if (helper_id >= 0 && helper_id <= helper_functions->helper_function_count) {
+        return helper_functions->helper_functions[helper_id];
+    } else {
+        return NULL;
+    }
+}

vm/ubpf_int.h

@@ -21,20 +21,22 @@
 #include "ebpf.h"
 
 struct ebpf_inst;
-typedef uint64_t (*ext_func)(uint64_t arg0, uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4);
 
 struct ubpf_vm {
     struct ebpf_inst *insts;
     uint16_t num_insts;
     ubpf_jit_fn jitted;
     size_t jitted_size;
-    ext_func *ext_funcs;
+    ubpf_helper_fn *ext_funcs;
     const char **ext_func_names;
     bool bounds_check_enabled;
     int (*error_printf)(FILE* stream, const char* format, ...);
+    void* helper_resolver_context;
+    ubpf_helper_fn (*helper_resolver)(void* context, int32_t helper_identifier);
 };
 
 char *ubpf_error(const char *fmt, ...);
 unsigned int ubpf_lookup_registered_function(struct ubpf_vm *vm, const char *name);
+ubpf_helper_fn ubpf_resolve_helper_function(const struct ubpf_vm *vm, int32_t helper_id);
 
 #endif

vm/ubpf_jit_x86_64.c

@@ -398,11 +398,17 @@ translate(struct ubpf_vm *vm, struct jit_state *state, char **errmsg)
             emit_cmp(state, src, dst);
             emit_jcc(state, 0x8e, target_pc);
             break;
-        case EBPF_OP_CALL:
+        case EBPF_OP_CALL:{
+            ubpf_helper_fn target = ubpf_resolve_helper_function(vm, inst.imm);
+            if (!target) {
+                *errmsg = ubpf_error("Unknown helper function identifier at PC %d: helper_id %02x", i, inst.imm);
+                return -1;
+            }
             /* We reserve RCX for shifts */
             emit_mov(state, RCX_ALT, RCX);
-            emit_call(state, vm->ext_funcs[inst.imm]);
+            emit_call(state, target);
             break;
+        }
         case EBPF_OP_EXIT:
             if (i != vm->num_insts - 1) {
                 emit_jmp(state, TARGET_PC_EXIT);

vm/ubpf_vm.c

@@ -82,13 +82,13 @@ ubpf_destroy(struct ubpf_vm *vm)
 }
 
 int
-ubpf_register(struct ubpf_vm *vm, unsigned int idx, const char *name, void *fn)
+ubpf_register(struct ubpf_vm *vm, unsigned int idx, const char *name, ubpf_helper_fn fn)
 {
     if (idx >= MAX_EXT_FUNCS) {
         return -1;
     }
 
-    vm->ext_funcs[idx] = (ext_func)fn;
+    vm->ext_funcs[idx] = fn;
     vm->ext_func_names[idx] = name;
     return 0;
 }
@@ -106,6 +106,30 @@ ubpf_lookup_registered_function(struct ubpf_vm *vm, const char *name)
     return -1;
 }
 
+int ubpf_register_helper_resolver(struct ubpf_vm *vm, void* context, ubpf_helper_fn (*helper_resolver)(void* context, int32_t helper_identifier))
+{
+    if (vm->helper_resolver) {
+        return -1;
+    }
+
+    vm->helper_resolver = helper_resolver;
+    vm->helper_resolver_context = context;
+    return 0;
+}
+
+ubpf_helper_fn
+ubpf_resolve_helper_function(const struct ubpf_vm *vm, int32_t helper_id)
+{
+    if (vm->helper_resolver) {
+        return vm->helper_resolver(vm->helper_resolver_context, helper_id);
+    }
+    else if (helper_id >= 0 && helper_id <= MAX_EXT_FUNCS) {
+        return vm->ext_funcs[helper_id];
+    } else {
+        return NULL;
+    }
+}
+
 int
 ubpf_load(struct ubpf_vm *vm, const void *code, uint32_t code_len, char **errmsg)
 {
@@ -565,7 +589,7 @@ ubpf_exec(const struct ubpf_vm *vm, void *mem, size_t mem_len)
         case EBPF_OP_EXIT:
             return reg[0];
         case EBPF_OP_CALL:
-            reg[0] = vm->ext_funcs[inst.imm](reg[1], reg[2], reg[3], reg[4], reg[5]);
+            reg[0] = ubpf_resolve_helper_function(vm, inst.imm)(reg[1], reg[2], reg[3], reg[4], reg[5]);
             break;
         }
     }
@@ -705,17 +729,21 @@ validate(const struct ubpf_vm *vm, const struct ebpf_inst *insts, uint32_t num_i
             }
             break;
 
-        case EBPF_OP_CALL:
-            if (inst.imm < 0 || inst.imm >= MAX_EXT_FUNCS) {
-                *errmsg = ubpf_error("invalid call immediate at PC %d", i);
-                return false;
-            }
-            if (!vm->ext_funcs[inst.imm]) {
-                *errmsg = ubpf_error("call to nonexistent function %u at PC %d", inst.imm, i);
+        case EBPF_OP_CALL:{
+            ubpf_helper_fn target = ubpf_resolve_helper_function(vm, inst.imm);
+            if (!target){
+                if (inst.imm < 0 || inst.imm >= MAX_EXT_FUNCS) {
+                    *errmsg = ubpf_error("invalid call immediate at PC %d", i);
+                }
+                else {
+                    *errmsg = ubpf_error("call to nonexistent function %u at PC %d", inst.imm, i);
+                    return false;
+                }
                 return false;
             }
             break;
-
+        }
+
         case EBPF_OP_EXIT:
             break;