Add option to invoke callout prior to each instruction (#563)

Signed-off-by: Alan Jowett <[email protected]>
Co-authored-by: Alan Jowett <[email protected]>

custom_tests/data/ubpf_test_debug_function.input

@@ -0,0 +1 @@
+85  00  00  00  01  00  00  00 95  00  00  00  00  00  00  00

custom_tests/descrs/ubpf_test_debug_function.md

@@ -0,0 +1,3 @@
+## Test Description
+
+This test verifies that a debug call-out can be registered to inspect the state of the vm prior to each instruction.

custom_tests/srcs/ubpf_test_debug_function.cc

@@ -0,0 +1,107 @@
+// Copyright (c) Will Hawkins
+// SPDX-License-Identifier: Apache-2.0
+
+#include <cstdint>
+#include <iostream>
+#include <memory>
+#include <stdint.h>
+#include <vector>
+#include <string>
+
+extern "C"
+{
+#include "ubpf.h"
+}
+
+#include "ubpf_custom_test_support.h"
+
+typedef struct _vm_state {
+    int pc;
+    std::vector<uint64_t> registers;
+    std::vector<uint8_t> stack;
+} vm_state_t;
+
+void
+debug_callout(void* context, int program_counter, const uint64_t registers[16], const uint8_t* stack_start, size_t stack_length)
+{
+    std::vector<vm_state_t>* vm_states = static_cast<std::vector<vm_state_t>*>(context);
+    vm_state_t vm_state{};
+
+    vm_state.pc = program_counter;
+    for (int i = 0; i < 16; i++) {
+        vm_state.registers.push_back(registers[i]);
+    }
+    for (size_t i = 0; i < stack_length; i++) {
+        vm_state.stack.push_back(stack_start[i]);
+    }
+
+    vm_states->push_back(vm_state);
+}
+
+uint64_t test_function_1(uint64_t r1, uint64_t r2, uint64_t r3, uint64_t r4, uint64_t r5)
+{
+    return r1 + r2 + r3 + r4 + r5;
+}
+
+int
+main(int argc, char** argv)
+{
+    std::string program_string{};
+    std::string error{};
+    ubpf_jit_fn jit_fn;
+
+    std::vector<vm_state_t> vm_states;
+
+    if (!get_program_string(argc, argv, program_string, error)) {
+        std::cerr << error << std::endl;
+        return 1;
+    }
+
+    uint64_t memory{0x123456789};
+
+    std::unique_ptr<ubpf_vm, decltype(&ubpf_destroy)> vm(ubpf_create(), ubpf_destroy);
+    if (!ubpf_setup_custom_test(
+            vm,
+            program_string,
+            [](ubpf_vm_up& vm, std::string& error) {
+                int retval = ubpf_register(vm.get(), 1, "test_function_1", test_function_1);
+                if (retval < 0) {
+                    error = "Problem registering test function retval=" + std::to_string(retval);
+                    return false;
+                }
+                return true;
+            },
+            jit_fn,
+            error)) {
+        std::cerr << "Problem setting up custom test: " << error << std::endl;
+        return 1;
+    }
+
+    if (ubpf_register_debug_fn(vm.get(), &vm_states, debug_callout) < 0) {
+        std::cerr << "Problem registering debug function" << std::endl;
+        return 1;
+    }
+
+    uint64_t bpf_return_value;
+    if (ubpf_exec(vm.get(), &memory, sizeof(memory), &bpf_return_value)) {
+        std::cerr << "Problem executing program" << std::endl;
+        return 1;
+    }
+
+    if (vm_states.empty()) {
+        std::cerr << "No debug callouts were made" << std::endl;
+        return 1;
+    }
+
+    for (auto& vm_state : vm_states) {
+        std::cout << "Program Counter: " << vm_state.pc << std::endl;
+        for (int i = 0; i < 16; i++) {
+            std::cout << "Register " << i << ": " << vm_state.registers[i] << std::endl;
+        }
+        std::cout << "Stack: ";
+        for (auto& stack_byte : vm_state.stack) {
+            std::cout << std::hex << static_cast<int>(stack_byte) << " ";
+        }
+        std::cout << std::endl;
+    }
+}

vm/inc/ubpf.h

@@ -562,6 +562,34 @@ extern "C"
      */
     bool
     ubpf_toggle_undefined_behavior_check(struct ubpf_vm* vm, bool enable);
+
+    /**
+     * @brief A function to invoke before each instruction.
+     *
+     * @param[in, out] context Context passed in to ubpf_register_debug_fn.
+     * @param[in] program_counter Current instruction pointer.
+     * @param[in] registers Array of 11 registers representing the VM state.
+     * @param[in] stack_start Pointer to the beginning of the stack.
+     * @param[in] stack_length Size of the stack in bytes.
+     */
+    typedef void (*ubpf_debug_fn)(
+        void* context,
+        int program_counter,
+        const uint64_t registers[16],
+        const uint8_t* stack_start,
+        size_t stack_length);
+
+    /**
+     * @brief Add option to invoke a debug function before each instruction.
+     * Note: This only applies to the interpreter and not the JIT.
+     *
+     * @param[in] vm VM to add the option to.
+     * @param[in] debug_fn Function to invoke before each instruction. Pass NULL to remove the function.
+     * @return 0 on success.
+     * @return -1 on failure.
+     */
+    int
+    ubpf_register_debug_fn(struct ubpf_vm* vm, void* context, ubpf_debug_fn debug_function);
 #ifdef __cplusplus
 }
 #endif

vm/ubpf_int.h

@@ -99,6 +99,8 @@ struct ubpf_vm
     ubpf_bounds_check bounds_check_function;
     void* bounds_check_user_data;
     int instruction_limit;
+    void* debug_function_context; ///< Context pointer that is passed to the debug function.
+    ubpf_debug_fn debug_function; ///< Debug function that is called before each instruction.
 #ifdef DEBUG
     uint64_t* regs;
 #endif

vm/ubpf_vm.c

@@ -677,6 +677,11 @@ ubpf_exec_ex(
 
         struct ebpf_inst inst = ubpf_fetch_instruction(vm, pc++);
 
+        // Invoke the debug function to allow the user to inspect the state of the VM if it is enabled.
+        if (vm->debug_function) {
+            vm->debug_function(vm->debug_function_context, cur_pc, reg, stack_start, stack_length);
+        }
+
         if (!ubpf_validate_shadow_register(vm, &shadow_registers, inst)) {
             return_value = -1;
             goto cleanup;
@@ -1893,3 +1898,15 @@ ubpf_register_stack_usage_calculator(struct ubpf_vm* vm, stack_usage_calculator_
     vm->stack_usage_calculator = calculator;
     return 0;
 }
+int
+ubpf_register_debug_fn(struct ubpf_vm* vm, void* context, ubpf_debug_fn debug_function)
+{
+    if ((vm->debug_function != NULL && debug_function != NULL) ||
+        (vm->debug_function == NULL && debug_function == NULL)) {
+        return -1;
+    }
+
+    vm->debug_function = debug_function;
+    vm->debug_function_context = context;
+    return 0;
+}