External helpers (#403)

* Add Support For Online External Helper Function Dispatch Signed-off-by: Will Hawkins <hawkinsw@obs.cr> * fixup! Add Support For Online External Helper Function Dispatch Add arm64 JIT support. Signed-off-by: Will Hawkins <hawkinsw@obs.cr> * fixup! Add Support For Online External Helper Function Dispatch Add (hopefully) Windows JIT support. Signed-off-by: Will Hawkins <hawkinsw@obs.cr> * fixup! Add Support For Online External Helper Function Dispatch Add Doxygen documentation. Signed-off-by: Will Hawkins <hawkinsw@obs.cr> --------- Signed-off-by: Will Hawkins <hawkinsw@obs.cr> Co-authored-by: Alan Jowett <alanjo@microsoft.com>
iovisor · Mar 7, 2024 · 3ceee2f · 3ceee2f
1 parent 6c97bfd
commit 3ceee2f
Show file tree

Hide file tree

Showing 9 changed files with 210 additions and 39 deletions.
diff --git a/ubpf_plugin/test_helpers.h b/ubpf_plugin/test_helpers.h
@@ -2,6 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #pragma once
+#include "ubpf.h"
+#include <string.h>
 #include <cmath>
 #include <cstdint>
 #include <map>
@@ -73,12 +75,11 @@ unwind(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
     return a;
 }
 
-static std::map<uint32_t, uint64_t (*)(uint64_t r1, uint64_t r2, uint64_t r3, uint64_t r4, uint64_t r5)>
-    helper_functions = {
-        {0, gather_bytes},
-        {1, memfrob},
-        {2, no_op},
-        {3, sqrti},
-        {4, strcmp_ext},
-        {5, unwind},
+static std::map<uint32_t, external_function_t> helper_functions = {
+    {0, gather_bytes},
+    {1, memfrob},
+    {2, no_op},
+    {3, sqrti},
+    {4, strcmp_ext},
+    {5, unwind},
 };
diff --git a/ubpf_plugin/ubpf_plugin.cc b/ubpf_plugin/ubpf_plugin.cc
@@ -6,6 +6,7 @@
 // value of %r0 at the end of execution.
 // The program is intended to be used with the bpf conformance test suite.
 
+#include <cstdint>
 #include <cstring>
 #include <iostream>
 #include <memory>
@@ -21,6 +22,16 @@ extern "C"
 
 #include "test_helpers.h"
 
+uint64_t test_helpers_dispatcher(void *cookie, unsigned int idx, uint64_t p0, uint64_t p1,uint64_t p2,uint64_t p3, uint64_t p4) {
+    UNREFERENCED_PARAMETER(cookie);
+    return helper_functions[idx](p0, p1, p2, p3, p4);
+}
+
+bool test_helpers_validater(unsigned int idx, void *cookie) {
+    UNREFERENCED_PARAMETER(cookie);
+    return helper_functions.contains(idx);
+}
+
 /**
  * @brief Read in a string of hex bytes and return a vector of bytes.
  *
@@ -130,14 +141,18 @@ int main(int argc, char **argv)
         return 1;
     }
 
+/*
     for (auto &[key, value] : helper_functions)
     {
-        if (ubpf_register(vm.get(), key, "unnamed", reinterpret_cast<void *>(value)) != 0)
+        if (ubpf_register(vm.get(), key, "unnamed", value) != 0)
         {
             std::cerr << "Failed to register helper function" << std::endl;
             return 1;
         }
     }
+*/
+
+    ubpf_register_external_dispatcher(vm.get(), test_helpers_dispatcher, test_helpers_validater, NULL);
 
     if (ubpf_set_unwind_function_index(vm.get(), 5) != 0)
     {

diff --git a/vm/inc/ubpf.h b/vm/inc/ubpf.h
@@ -98,6 +98,23 @@ extern "C"
     void
     ubpf_set_error_print(struct ubpf_vm* vm, int (*error_printf)(FILE* stream, const char* format, ...));
 
+    /**
+     * @brief The type of an external function.
+     */
+    typedef uint64_t (*external_function_t)(uint64_t p0, uint64_t p1, uint64_t p2, uint64_t p3, uint64_t p4);
+
+    /**
+     * @brief Cast an external function to external_function_t
+     * Some external functions may not use all the parameters and, therefore,
+     * not match the external_function_t typedef. Use this for a conversion.
+     *
+     * @param[in] f The function to cast to match the signature of an
+     *              external function.
+     * @retval The external function, as external_function_t.
+     */
+    external_function_t
+    as_external_function_t(void* f);
+
     /**
      * @brief Register an external function.
      * The immediate field of a CALL instruction is an index into an array of
@@ -112,7 +129,43 @@ extern "C"
      * @retval -1 Failure.
      */
     int
-    ubpf_register(struct ubpf_vm* vm, unsigned int index, const char* name, void* fn);
+    ubpf_register(struct ubpf_vm* vm, unsigned int index, const char* name, external_function_t fn);
+
+    /**
+     * @brief The type of an external helper dispatcher function.
+     */
+    typedef uint64_t (*external_function_dispatcher_t)(
+        void* cookie, unsigned int index, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t);
+
+    /**
+     * @brief The type of an external helper validation function.
+     */
+    typedef bool (*external_function_validate_t)(unsigned int index, void* cookie);
+
+    /**
+     * @brief Register a function that dispatches to external helpers
+     * The immediate field of a CALL instruction is an index of a helper
+     * function to invoke. This API sets a callback that will choose the
+     * helper function to invoke (based on the index) and then invoke it.
+     * This API also sets a callback that the validator will use to determine
+     * if a given index is a valid external function.
+     *
+     * @param[in] vm The VM to register the function on.
+     * @param[in] dispatcher The callback that will dispatch to the external
+     *                       helper.
+     * @param[in] validater The callback that will validate that a given index
+     *                      is valid for an external helper.
+     * @param[in] cookie A pointer to some user-defined cookie that will be
+     *                   passed to the callbacks.
+     * @retval 0 Success.
+     * @retval -1 Failure.
+     */
+    int
+    ubpf_register_external_dispatcher(
+        struct ubpf_vm* vm,
+        external_function_dispatcher_t dispatcher,
+        external_function_validate_t validater,
+        void* cookie);
 
     /**
      * @brief Load code into a VM.

diff --git a/vm/test.c b/vm/test.c
@@ -516,14 +516,26 @@ bpf_map_delete_elem_impl(struct bpf_map* map, const void* key)
 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, 5, "unwind", unwind);
+    ubpf_register(vm, 0, "gather_bytes", as_external_function_t(gather_bytes));
+    ubpf_register(vm, 1, "memfrob", as_external_function_t(memfrob));
+    ubpf_register(vm, 2, "trash_registers", as_external_function_t(trash_registers));
+    ubpf_register(vm, 3, "sqrti", as_external_function_t(sqrti));
+    ubpf_register(vm, 4, "strcmp_ext", as_external_function_t(strcmp));
+    ubpf_register(vm, 5, "unwind", as_external_function_t(unwind));
     ubpf_set_unwind_function_index(vm, 5);
-    ubpf_register(vm, (unsigned int)(uintptr_t)bpf_map_lookup_elem, "bpf_map_lookup_elem", bpf_map_lookup_elem_impl);
-    ubpf_register(vm, (unsigned int)(uintptr_t)bpf_map_update_elem, "bpf_map_update_elem", bpf_map_update_elem_impl);
-    ubpf_register(vm, (unsigned int)(uintptr_t)bpf_map_delete_elem, "bpf_map_delete_elem", bpf_map_delete_elem_impl);
+    ubpf_register(
+        vm,
+        (unsigned int)(uintptr_t)bpf_map_lookup_elem,
+        "bpf_map_lookup_elem",
+        as_external_function_t(bpf_map_lookup_elem_impl));
+    ubpf_register(
+        vm,
+        (unsigned int)(uintptr_t)bpf_map_update_elem,
+        "bpf_map_update_elem",
+        as_external_function_t(bpf_map_update_elem_impl));
+    ubpf_register(
+        vm,
+        (unsigned int)(uintptr_t)bpf_map_delete_elem,
+        "bpf_map_delete_elem",
+        as_external_function_t(bpf_map_delete_elem_impl));
 }
diff --git a/vm/ubpf_int.h b/vm/ubpf_int.h
@@ -33,9 +33,15 @@ struct ubpf_vm
     uint16_t num_insts;
     ubpf_jit_fn jitted;
     size_t jitted_size;
+
     ext_func* ext_funcs;
     bool* int_funcs;
     const char** ext_func_names;
+
+    external_function_dispatcher_t dispatcher;
+    external_function_validate_t dispatcher_validate;
+    void* dispatcher_cookie;
+
     bool bounds_check_enabled;
     int (*error_printf)(FILE* stream, const char* format, ...);
     int (*translate)(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg);
@@ -68,6 +74,9 @@ char*
 ubpf_error(const char* fmt, ...);
 unsigned int
 ubpf_lookup_registered_function(struct ubpf_vm* vm, const char* name);
+uint64_t
+ubpf_dispatch_to_external_helper(
+    uint64_t p0, uint64_t p1, uint64_t p2, uint64_t p3, uint64_t p4, const struct ubpf_vm* vm, unsigned int idx);
 
 /**
  * @brief Fetch the instruction at the given index.

diff --git a/vm/ubpf_jit_arm64.c b/vm/ubpf_jit_arm64.c
@@ -157,7 +157,10 @@ 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_to(uint32_t amount, uint64_t boundary) { return (amount + (boundary - 1 )) & ~(boundary - 1); }
+static uint32_t inline align_to(uint32_t amount, uint64_t boundary)
+{
+    return (amount + (boundary - 1)) & ~(boundary - 1);
+}
 
 static void
 emit_bytes(struct jit_state* state, void* data, uint32_t len)
@@ -577,13 +580,21 @@ emit_jit_prologue(struct jit_state* state, size_t ubpf_stack_size)
 }
 
 static void
-emit_call(struct jit_state* state, uintptr_t func)
+emit_dispatched_external_helper_call(struct jit_state* state, struct ubpf_vm* vm, unsigned int idx)
 {
     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);
 
-    emit_movewide_immediate(state, true, temp_register, func);
+    // All parameters to the helper function are in the right spot
+    // for the dispatcher. All we need to do now is ...
+
+    // ... set up the final two parameters.
+    emit_movewide_immediate(state, true, R5, (uint64_t)vm);
+    emit_movewide_immediate(state, true, R6, idx);
+
+    // Call!
+    emit_movewide_immediate(state, true, temp_register, (uint64_t)ubpf_dispatch_to_external_helper);
     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.  */
@@ -1058,7 +1069,7 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
             break;
         case EBPF_OP_CALL:
             if (inst.src == 0) {
-                emit_call(state, (uintptr_t)vm->ext_funcs[inst.imm]);
+                emit_dispatched_external_helper_call(state, vm, 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);

diff --git a/vm/ubpf_jit_x86_64.c b/vm/ubpf_jit_x86_64.c
@@ -18,16 +18,15 @@
  * limitations under the License.
  */
 
+#define _GNU_SOURCE
+
 #include "ebpf.h"
 #include <stdint.h>
-#define _GNU_SOURCE
 #include <stdio.h>
 #include <stdlib.h>
 #include <stdbool.h>
 #include <unistd.h>
-#include <inttypes.h>
 #include <sys/mman.h>
-#include <errno.h>
 #include <assert.h>
 #include "ubpf_int.h"
 #include "ubpf_jit_x86_64.h"
@@ -607,7 +606,7 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
             /* We reserve RCX for shifts */
             if (inst.src == 0) {
                 emit_mov(state, RCX_ALT, RCX);
-                emit_call(state, vm->ext_funcs[inst.imm]);
+                emit_dispatched_external_helper_call(state, vm, 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);

diff --git a/vm/ubpf_jit_x86_64.h b/vm/ubpf_jit_x86_64.h
@@ -28,6 +28,9 @@
 #include <stdint.h>
 #include <string.h>
 
+#include "ubpf.h"
+#include "ubpf_int.h"
+
 #define RAX 0
 #define RCX 1
 #define RDX 2
@@ -373,7 +376,7 @@ emit_jmp(struct jit_state* state, uint32_t target_pc)
 }
 
 static inline void
-emit_call(struct jit_state* state, void* target)
+emit_dispatched_external_helper_call(struct jit_state* state, const struct ubpf_vm* vm, unsigned int idx)
 {
     /*
      * When we enter here, our stack is 16-byte aligned. Keep
@@ -386,16 +389,32 @@ emit_call(struct jit_state* state, void* target)
      */
     emit_alu64_imm32(state, 0x81, 5, RSP, sizeof(uint64_t));
 
+    emit_load_imm(state, RAX, idx);
+    emit_push(state, RAX);
+
+    emit_load_imm(state, RAX, (uint64_t)vm);
+    emit_push(state, RAX);
+
     /* Windows x64 ABI spills 5th parameter to stack */
     emit_push(state, map_register(5));
 
     /* Windows x64 ABI requires home register space.
      * Allocate home register space - 4 registers.
      */
     emit_alu64_imm32(state, 0x81, 5, RSP, 4 * sizeof(uint64_t));
+#else
+    // Save r9 -- I need it for a parameter!
+    emit_push(state, R9);
+
+    // Before it's a parameter, use it for a push.
+    emit_load_imm(state, R9, idx);
+    emit_push(state, R9);
+
+    emit_load_imm(state, R9, (uint64_t)vm);
 #endif
 
-    emit_load_imm(state, RAX, (uintptr_t)target);
+    emit_load_imm(state, RAX, (uintptr_t)ubpf_dispatch_to_external_helper);
+
 #ifndef UBPF_DISABLE_RETPOLINES
     emit1(state, 0xe8); // e8 is the opcode for a CALL
     emit_jump_target_address(state, TARGET_PC_RETPOLINE);
@@ -413,9 +432,15 @@ emit_call(struct jit_state* state, void* target)
     emit1(state, 0xd0);
 #endif
 
+    // The result is in RAX. Nothing to do there.
+    // Just rationalize the stack!
+
 #if defined(_WIN32)
-    /* Deallocate home register space + spilled register + alignment space - 5 registers */
-    emit_alu64_imm32(state, 0x81, 0, RSP, (4 + 1 + 1) * sizeof(uint64_t));
+    /* Deallocate home register space + 3 spilled parameters + alignment space */
+    emit_alu64_imm32(state, 0x81, 0, RSP, (4 + 3 + 1) * sizeof(uint64_t));
+#else
+    emit_pop(state, R9); // First one is a throw away (it's where our parameter was!)
+    emit_pop(state, R9); // This one is real!
 #endif
 }