diff --git a/CMakeLists.txt b/CMakeLists.txt
index d110ef58..a537ef34 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -29,6 +29,7 @@ add_subdirectory("vm")
         STEP_TARGETS build)
 
 if(UBPF_ENABLE_TESTS)
+  add_subdirectory("custom_tests")
   add_subdirectory("ubpf_plugin")
   if (NOT UBPF_SKIP_EXTERNAL)
   endif()
diff --git a/cmake/settings.cmake b/cmake/settings.cmake
index 3a06504f..0f2f0afb 100644
--- a/cmake/settings.cmake
+++ b/cmake/settings.cmake
@@ -17,7 +17,8 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
       -Wall
       -Werror
       -Iinc
-      -O2
+      -O0
+      -g
       -Wunused-parameter
       -fPIC
     )
diff --git a/custom_tests/CMakeLists.txt b/custom_tests/CMakeLists.txt
new file mode 100644
index 00000000..a5cc6c2f
--- /dev/null
+++ b/custom_tests/CMakeLists.txt
@@ -0,0 +1,120 @@
+# Copyright (c) Microsoft Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+set(CMAKE_CXX_STANDARD 20)
+
+file(GLOB test_descr_files ${CMAKE_SOURCE_DIR}/custom_tests/descrs/*.md)
+
+add_library(ubpf_custom_test_support srcs/ubpf_custom_test_support.cc)
+
+target_link_libraries(
+    ubpf_custom_test_support
+    ubpf
+    ubpf_settings
+)
+
+target_include_directories(ubpf_custom_test_support PUBLIC ".srcs/")
+target_include_directories(ubpf_custom_test_support PRIVATE
+    "${CMAKE_SOURCE_DIR}/vm"
+    "${CMAKE_BINARY_DIR}/vm"
+    "${CMAKE_SOURCE_DIR}/vm/inc"
+    "${CMAKE_BINARY_DIR}/vm/inc"
+)
+
+set(QEMU_RUNNER "")
+if(CMAKE_SYSTEM_PROCESSOR STREQUAL aarch64 AND (NOT CMAKE_HOST_SYSTEM_PROCESSOR STREQUAL aarch64))
+	set(QEMU_RUNNER qemu-aarch64 -L /usr/aarch64-linux-gnu)
+endif()
+message(WARNING "qemu_runner: ${QEMU_RUNNER}")
+
+foreach(test_file ${test_descr_files})
+    get_filename_component(test_name ${test_file} NAME_WE)
+    set(test_source_path "${CMAKE_SOURCE_DIR}/custom_tests/srcs/${test_name}.cc")
+    message(WARNING "test_name: ${test_name}")
+    message(WARNING "test_source_path: ${test_source_path}")
+ 
+    add_executable(
+        ${test_name}
+        ${test_source_path}
+    )
+    target_include_directories(${test_name} PRIVATE
+        "${CMAKE_SOURCE_DIR}/vm"
+        "${CMAKE_BINARY_DIR}/vm"
+        "${CMAKE_SOURCE_DIR}/vm/inc"
+        "${CMAKE_BINARY_DIR}/vm/inc"
+    )
+    target_link_libraries(
+        ${test_name}
+        ubpf
+        ubpf_custom_test_support
+        ubpf_settings
+    )
+    set(potential_input_file ${CMAKE_SOURCE_DIR}/custom_tests/data/${test_name}.input)
+    if (EXISTS ${potential_input_file})
+        list(JOIN QEMU_RUNNER " " QEMU_RUNNER_STR)
+        add_test(
+            NAME ${test_name}-Custom
+	    COMMAND sh -c "cat ${potential_input_file} | ${QEMU_RUNNER_STR} $<TARGET_FILE:${test_name}>"
+        )
+    else()
+        add_test(
+            NAME ${test_name}-Custom
+	    COMMAND ${QEMU_RUNNER} $<TARGET_FILE:${test_name}>
+        )
+    endif()
+    message(WARNING "test_source: ${test_source}")
+endforeach()
+
+#if(NOT BPF_CONFORMANCE_RUNNER)
+#    set(BPF_CONFORMANCE_RUNNER ${CMAKE_BINARY_DIR}/external/bpf_conformance/bin/bpf_conformance_runner)
+#else()
+#    message(STATUS "Using custom bpf_conformance_runner: ${BPF_CONFORMANCE_RUNNER}")
+#endif()
+
+#if(CMAKE_SYSTEM_PROCESSOR STREQUAL aarch64 AND (NOT CMAKE_HOST_SYSTEM_PROCESSOR STREQUAL aarch64))
+#    set(PLUGIN_JIT --plugin_path ${CMAKE_BINARY_DIR}/bin/run-jit.sh)
+#    set(PLUGIN_INTERPRET --plugin_path ${CMAKE_BINARY_DIR}/bin/run-interpret.sh)
+#else()
+#    if(PLATFORM_WINDOWS)
+#        set(PLATFORM_EXECUTABLE_EXTENSION ".exe")
+#    else()
+#        set(PLATFORM_EXECUTABLE_EXTENSION "")
+#    endif()
+#    set(PLUGIN_JIT --plugin_path ${CMAKE_BINARY_DIR}/bin/ubpf_plugin${PLATFORM_EXECUTABLE_EXTENSION} --plugin_options --jit)
+#    set(PLUGIN_INTERPRET --plugin_path ${CMAKE_BINARY_DIR}/bin/ubpf_plugin${PLATFORM_EXECUTABLE_EXTENSION} --plugin_options --interpret)
+#endif()
+
+## Add all names of tests that are expected to fail to the TESTS_EXPECTED_TO_FAIL list
+#list(APPEND TESTS_EXPECTED_TO_FAIL "duplicate_label")
+## TODO: remove this once we have a proper implementation of interlocked operations
+## and support for calling local functions.
+#list(APPEND TESTS_EXPECTED_TO_FAIL "lock")
+
+#foreach(file ${files})
+#    unset(EXPECT_FAILURE)
+#    foreach(to_fail ${TESTS_EXPECTED_TO_FAIL})
+#        if(NOT EXPECT_FAILURE)
+#            string(REGEX MATCH "${to_fail}" EXPECT_FAILURE "${file}")
+#            if(EXPECT_FAILURE)
+#                message(STATUS "Expecting ${file} test to fail.")
+#            endif()
+#        endif()
+#    endforeach()
+#    add_test(
+#        NAME ${file}-JIT
+#        COMMAND ${BPF_CONFORMANCE_RUNNER} --test_file_path ${file} ${PLUGIN_JIT}
+#    )
+
+#    if(EXPECT_FAILURE)
+#        set_tests_properties(${file}-JIT PROPERTIES WILL_FAIL TRUE)
+#    endif()
+
+#    add_test(
+#        NAME ${file}-Interpreter
+#        COMMAND ${BPF_CONFORMANCE_RUNNER} --test_file_path ${file} ${PLUGIN_INTERPRET}
+#    )
+
+#    if(EXPECT_FAILURE)
+#        set_tests_properties(${file}-Interpreter PROPERTIES WILL_FAIL TRUE)
+#    endif()
+#endforeach()
diff --git a/custom_tests/README.md b/custom_tests/README.md
new file mode 100644
index 00000000..5ce054c4
--- /dev/null
+++ b/custom_tests/README.md
@@ -0,0 +1,35 @@
+## Writing a uBPF Custom Tests
+
+Custom tests are enabled by creating two (2) or three (3) different files in the `custom_tests` directory.
+
+### Files Of a uBPF Custom Test
+
+#### Description Files
+
+The first file to create is the Description File. The Description File is a file with a `.md` extension that resides in the `descrs` directory. The purpose of this file is to identify the name of the test (everything before the `.md` extension) and provide a place to document the purpose of the test.
+
+#### Source Files
+
+The second file to create is the Source File. The Source file should reside in the `srcs` directory and have a name that matches its Description File (with the `.cc` extension rather than the `.md` extension).
+
+#### Input Files
+
+The final file is optional. The Input File resides in the `data` directory and should have the same name as the other two (2) files but with an `.input` extension rather than  `.cc` or `.md` for the Source and Description File respectively. If present, the contents of this file will be given to the executed custom test over standard input.
+
+### Building
+
+The Source Files for a custom test are compiled using C++20 and are saved as an executable named according to the name of the test in the CMake build directory.
+
+### Return Values
+
+All successful tests should return `0`. All failing tests should return something other than `0`.
+
+### Supporting Libraries
+
+To reduce the boilerplate needed to write custom tests, there is a custom test library with several helpful functions. These functions are documented in the library's header file (`custom_tests/srcs/ubpf_custom_test_support.h`).
+
+### Putting It Together
+
+After describing the test's purpose in a Markdown syntax in a file named, say, `test_example.md` and stored in the `descrs` directory, you can write the test's Source Code (in C++20) and give it the name `test_example.cc` in the `srcs` directory. If the test needs input, you can save that input in the tests Input File (`test_input.input`) in the `data` directory.
+
+Because all the files are present, this test will be run when the CTest target is invoked. Because there the optional `test_input.input` file is present, the contents of that file will be given to the executable via standard input.
\ No newline at end of file
diff --git a/custom_tests/data/ubpf_test_external_dispatcher_context_overwrite.input b/custom_tests/data/ubpf_test_external_dispatcher_context_overwrite.input
new file mode 100644
index 00000000..45c5a332
--- /dev/null
+++ b/custom_tests/data/ubpf_test_external_dispatcher_context_overwrite.input
@@ -0,0 +1 @@
+b7  01  00  00  01  02  03  04  85  00  00  00  01  00  00  00  95  00  00  00  00  00  00  00
diff --git a/custom_tests/data/ubpf_test_external_dispatcher_simple_context.input b/custom_tests/data/ubpf_test_external_dispatcher_simple_context.input
new file mode 100644
index 00000000..12873c00
--- /dev/null
+++ b/custom_tests/data/ubpf_test_external_dispatcher_simple_context.input
@@ -0,0 +1 @@
+85  00  00  00  01  00  00  00 95  00  00  00  00  00  00  00
diff --git a/custom_tests/data/ubpf_test_jit_unexpected_instruction.input b/custom_tests/data/ubpf_test_jit_unexpected_instruction.input
new file mode 100644
index 00000000..d5be5a56
--- /dev/null
+++ b/custom_tests/data/ubpf_test_jit_unexpected_instruction.input
@@ -0,0 +1 @@
+8f  00  00  00  01  00  00  00
\ No newline at end of file
diff --git a/custom_tests/data/ubpf_test_update_dispatcher.input b/custom_tests/data/ubpf_test_update_dispatcher.input
new file mode 100644
index 00000000..12873c00
--- /dev/null
+++ b/custom_tests/data/ubpf_test_update_dispatcher.input
@@ -0,0 +1 @@
+85  00  00  00  01  00  00  00 95  00  00  00  00  00  00  00
diff --git a/custom_tests/data/ubpf_test_update_helpers.input b/custom_tests/data/ubpf_test_update_helpers.input
new file mode 100644
index 00000000..12873c00
--- /dev/null
+++ b/custom_tests/data/ubpf_test_update_helpers.input
@@ -0,0 +1 @@
+85  00  00  00  01  00  00  00 95  00  00  00  00  00  00  00
diff --git a/custom_tests/descrs/ubpf_test_external_dispatcher_context_overwrite.md b/custom_tests/descrs/ubpf_test_external_dispatcher_context_overwrite.md
new file mode 100644
index 00000000..e69de29b
diff --git a/custom_tests/descrs/ubpf_test_external_dispatcher_simple_context.md b/custom_tests/descrs/ubpf_test_external_dispatcher_simple_context.md
new file mode 100644
index 00000000..e69de29b
diff --git a/custom_tests/descrs/ubpf_test_jit_buffer_too_small.md b/custom_tests/descrs/ubpf_test_jit_buffer_too_small.md
new file mode 100644
index 00000000..e69de29b
diff --git a/custom_tests/descrs/ubpf_test_jit_unexpected_instruction.md b/custom_tests/descrs/ubpf_test_jit_unexpected_instruction.md
new file mode 100644
index 00000000..e69de29b
diff --git a/custom_tests/descrs/ubpf_test_update_dispatcher.md b/custom_tests/descrs/ubpf_test_update_dispatcher.md
new file mode 100644
index 00000000..e69de29b
diff --git a/custom_tests/descrs/ubpf_test_update_helpers.md b/custom_tests/descrs/ubpf_test_update_helpers.md
new file mode 100644
index 00000000..e69de29b
diff --git a/custom_tests/srcs/test_helpers.h b/custom_tests/srcs/test_helpers.h
new file mode 100644
index 00000000..1d190865
--- /dev/null
+++ b/custom_tests/srcs/test_helpers.h
@@ -0,0 +1,225 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+#include "ubpf.h"
+#include <string.h>
+#include <cmath>
+#include <cstdint>
+#include <map>
+
+#if !defined(UNREFERENCED_PARAMETER)
+#define UNREFERENCED_PARAMETER(P) (void)(P)
+#endif
+
+static uint64_t
+gather_bytes(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    return ((uint64_t)(a & 0xff) << 32) | ((uint64_t)(b & 0xff) << 24) | ((uint64_t)(c & 0xff) << 16) |
+           ((uint64_t)(d & 0xff) << 8) | (e & 0xff);
+};
+
+static uint64_t
+memfrob(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+
+    uint8_t* p = reinterpret_cast<uint8_t*>(a);
+    for (uint64_t i = 0; i < b; i++) {
+        p[i] ^= 42;
+    }
+    return 0;
+};
+
+;
+
+static uint64_t
+no_op(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+
+    return 0;
+}
+
+static uint64_t
+sqrti(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+
+    return static_cast<uint64_t>(std::sqrt(a));
+}
+
+static uint64_t
+strcmp_ext(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return strcmp(reinterpret_cast<char*>(a), reinterpret_cast<char*>(b));
+}
+
+static uint64_t
+unwind(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return a;
+}
+
+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},
+};
+
+static uint64_t
+dispatcher_test_memfrob(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 42;
+
+}
+
+static uint64_t
+updated_dispatcher_test_memfrob(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 43;
+}
+
+static uint64_t
+dispatcher_gather_bytes(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 44;
+
+}
+
+static uint64_t
+updated_dispatcher_gather_bytes(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 45;
+}
+
+static uint64_t
+dispatcher_no_op(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 46;
+
+}
+
+static uint64_t
+updated_dispatcher_no_op(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 47;
+}
+
+static uint64_t
+dispatcher_sqrti(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 48;
+
+}
+
+static uint64_t
+updated_dispatcher_sqrti(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 49;
+}
+
+static uint64_t
+dispatcher_strcmp_ext(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 50;
+
+}
+
+static uint64_t
+updated_dispatcher_strcmp_ext(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 51;
+}
+
+static uint64_t
+dispatcher_unwind(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 52;
+
+}
+
+static uint64_t
+updated_dispatcher_unwind(uint64_t a, uint64_t b, uint64_t c, uint64_t d, uint64_t e)
+{
+    UNREFERENCED_PARAMETER(a);
+    UNREFERENCED_PARAMETER(b);
+    UNREFERENCED_PARAMETER(c);
+    UNREFERENCED_PARAMETER(d);
+    UNREFERENCED_PARAMETER(e);
+    return 53;
+}
diff --git a/custom_tests/srcs/ubpf_custom_test_support.cc b/custom_tests/srcs/ubpf_custom_test_support.cc
new file mode 100644
index 00000000..6fc58e05
--- /dev/null
+++ b/custom_tests/srcs/ubpf_custom_test_support.cc
@@ -0,0 +1,110 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#include "ubpf_int.h"
+#include <cstdint>
+#include <cstring>
+#include <vector>
+#include <string>
+#include <sstream>
+#include <memory>
+#include <optional>
+
+extern "C"
+{
+#include "ebpf.h"
+#include "ubpf.h"
+}
+
+
+#include "ubpf_custom_test_support.h"
+
+/**
+ * @brief Read in a string of hex bytes and return a vector of bytes.
+ *
+ * @param[in] input String containing hex bytes.
+ * @return Vector of bytes.
+ */
+std::vector<uint8_t>
+base16_decode(const std::string &input)
+{
+    std::vector<uint8_t> output;
+    std::stringstream ss(input);
+    std::string value;
+    output.reserve(input.size() / 3);
+    while (std::getline(ss, value, ' '))
+    {
+        try
+        {
+            output.push_back(static_cast<uint8_t>(std::stoi(value, nullptr, 16)));
+        }
+        catch (...)
+        {
+            // Ignore invalid values.
+        }
+    }
+    return output;
+}
+
+/**
+ * @brief Convert a vector of bytes to a vector of ebpf_inst.
+ *
+ * @param[in] bytes Vector of bytes.
+ * @return Vector of ebpf_inst.
+ */
+std::vector<ebpf_inst>
+bytes_to_ebpf_inst(std::vector<uint8_t> bytes)
+{
+    std::vector<ebpf_inst> instructions(bytes.size() / sizeof(ebpf_inst));
+    memcpy(instructions.data(), bytes.data(), bytes.size());
+    return instructions;
+}
+
+
+bool ubpf_setup_custom_test(ubpf_vm_up &vm,
+                       const std::string program_string,
+                       std::optional<custom_test_fixup_cb> fixup_f,
+                       ubpf_jit_fn &jit_fn,
+                       std::string &error)
+{
+    jit_fn = nullptr;
+    std::vector<ebpf_inst> program = bytes_to_ebpf_inst(base16_decode(program_string));
+    char *error_s{nullptr};
+
+    if (vm == nullptr)
+    {
+        error = "Failed to create VM";
+        return false;
+    }
+
+    if (ubpf_set_unwind_function_index(vm.get(), 5) != 0)
+    {
+        error = "Failed to set unwind function index";
+        return false;
+    }
+
+    if (fixup_f.has_value())
+    {
+        if (!(fixup_f.value())(vm, error)) {
+            return false;
+        }
+    }
+
+    if (ubpf_load(vm.get(), program.data(), static_cast<uint32_t>(program.size() * sizeof(ebpf_inst)), &error_s) != 0)
+    {
+        error = "Failed to load program: " + std::string{error_s};
+        free(error_s);
+        return false;
+    }
+
+    jit_fn = ubpf_compile(vm.get(), &error_s);
+    if (jit_fn == nullptr)
+    {
+        error = "Failed to compile: " + std::string{error_s};
+        free(error_s);
+        return false;
+    }
+
+    free(error_s);
+    return true;
+}
diff --git a/custom_tests/srcs/ubpf_custom_test_support.h b/custom_tests/srcs/ubpf_custom_test_support.h
new file mode 100644
index 00000000..8957c5c6
--- /dev/null
+++ b/custom_tests/srcs/ubpf_custom_test_support.h
@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+#include <string>
+#include <optional>
+#include <functional>
+
+extern "C"
+{
+#include "ebpf.h"
+#include "ubpf.h"
+}
+
+#define UNREFERENCED_PARAMETER (void)
+
+/**
+ * @brief Read in a string of hex bytes and return a vector of bytes.
+ *
+ * @param[in] input String containing hex bytes.
+ * @return Vector of bytes.
+ */
+std::vector<uint8_t>
+base16_decode(const std::string &input);
+
+
+/**
+ * @brief Convert a vector of bytes to a vector of ebpf_inst.
+ *
+ * @param[in] bytes Vector of bytes.
+ * @return Vector of ebpf_inst.
+ */
+std::vector<ebpf_inst>
+bytes_to_ebpf_inst(std::vector<uint8_t> bytes);
+
+
+using ubpf_vm_up = std::unique_ptr<ubpf_vm, decltype(&ubpf_destroy)>;
+using custom_test_fixup_cb = std::function<bool(ubpf_vm_up &, std::string &error)>;
+
+
+/**
+ * @brief Do the common necessary work to setup a custom test.
+ *
+ * @param[in] vm The VM for which to prepare the test.
+ * @param[in] program_string A string of raw bytes that make up the eBPF program to execute under this test.
+ * @param[in] fixup_f A function that will be invoked after the program is loaded and before it is compiled.
+ * @param[out] jit_fn A function that can be invoked to run the jit'd program.
+ * @param[out] error A string containing the error message (if any) generated during custom test configuration.
+ * @return True or false depending on whether setting up the custom test succeeded.
+ */
+bool ubpf_setup_custom_test(ubpf_vm_up &vm,
+                       const std::string program_string,
+                       std::optional<custom_test_fixup_cb> fixup_f,
+                       ubpf_jit_fn &jit_fn,
+                       std::string &error);
+
diff --git a/custom_tests/srcs/ubpf_test_external_dispatcher_context_overwrite.cc b/custom_tests/srcs/ubpf_test_external_dispatcher_context_overwrite.cc
new file mode 100644
index 00000000..20e65a64
--- /dev/null
+++ b/custom_tests/srcs/ubpf_test_external_dispatcher_context_overwrite.cc
@@ -0,0 +1,76 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+// This program reads BPF instructions from stdin and memory contents from
+// the first agument. It then executes the BPF program and prints the
+// value of %r0 at the end of execution.
+// The program is intended to be used with the bpf conformance test suite.
+
+#include <cstdint>
+#include <iostream>
+#include <memory>
+#include <vector>
+#include <string>
+
+extern "C"
+{
+#include "ubpf.h"
+}
+
+#include "ubpf_custom_test_support.h"
+
+uint64_t *cookie_pointer_value{nullptr};
+uint64_t
+external_dispatcher(uint64_t p0, uint64_t p1, uint64_t p2, uint64_t p3, uint64_t p4, unsigned int idx, void* cookie)
+{
+    UNREFERENCED_PARAMETER(p0);
+    UNREFERENCED_PARAMETER(p1);
+    UNREFERENCED_PARAMETER(p2);
+    UNREFERENCED_PARAMETER(p3);
+    UNREFERENCED_PARAMETER(p4);
+    UNREFERENCED_PARAMETER(idx);
+    uint64_t* ccookie = (uint64_t*)cookie;
+    cookie_pointer_value = ccookie;
+    return 1;
+}
+
+bool
+external_dispatcher_validater(unsigned int idx, const struct ubpf_vm* cookie)
+{
+    UNREFERENCED_PARAMETER(idx);
+    UNREFERENCED_PARAMETER(cookie);
+    return true;
+}
+
+int main(int argc, char **argv)
+{
+    std::vector<std::string> args(argv, argv + argc);
+    std::string program_string{};
+    ubpf_jit_fn jit_fn;
+    uint64_t memory{0x123456789};
+
+    std::getline(std::cin, program_string);
+
+    std::unique_ptr<ubpf_vm, decltype(&ubpf_destroy)> vm(ubpf_create(), ubpf_destroy);
+    std::string error{};
+    if (!ubpf_setup_custom_test(
+            vm,
+            program_string,
+            [](ubpf_vm_up& vm, std::string &error) {
+                if (ubpf_register_external_dispatcher(vm.get(), external_dispatcher, external_dispatcher_validater) < 0) {
+                    error = "Failed to register external dispatcher.";
+                    return false;
+                }
+                return true;
+            },
+            jit_fn,
+            error)) {
+        std::cerr << "Problem setting up custom test: " << error << std::endl;
+        return 1;
+    }
+
+    [[maybe_unused]] auto result = jit_fn(&memory, sizeof(uint64_t));
+
+    // Ultimately we want the cookie pointer that we got in the dispatcher to match what we sent in!
+    return !(cookie_pointer_value == &memory);
+}
diff --git a/custom_tests/srcs/ubpf_test_external_dispatcher_simple_context.cc b/custom_tests/srcs/ubpf_test_external_dispatcher_simple_context.cc
new file mode 100644
index 00000000..20e65a64
--- /dev/null
+++ b/custom_tests/srcs/ubpf_test_external_dispatcher_simple_context.cc
@@ -0,0 +1,76 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+// This program reads BPF instructions from stdin and memory contents from
+// the first agument. It then executes the BPF program and prints the
+// value of %r0 at the end of execution.
+// The program is intended to be used with the bpf conformance test suite.
+
+#include <cstdint>
+#include <iostream>
+#include <memory>
+#include <vector>
+#include <string>
+
+extern "C"
+{
+#include "ubpf.h"
+}
+
+#include "ubpf_custom_test_support.h"
+
+uint64_t *cookie_pointer_value{nullptr};
+uint64_t
+external_dispatcher(uint64_t p0, uint64_t p1, uint64_t p2, uint64_t p3, uint64_t p4, unsigned int idx, void* cookie)
+{
+    UNREFERENCED_PARAMETER(p0);
+    UNREFERENCED_PARAMETER(p1);
+    UNREFERENCED_PARAMETER(p2);
+    UNREFERENCED_PARAMETER(p3);
+    UNREFERENCED_PARAMETER(p4);
+    UNREFERENCED_PARAMETER(idx);
+    uint64_t* ccookie = (uint64_t*)cookie;
+    cookie_pointer_value = ccookie;
+    return 1;
+}
+
+bool
+external_dispatcher_validater(unsigned int idx, const struct ubpf_vm* cookie)
+{
+    UNREFERENCED_PARAMETER(idx);
+    UNREFERENCED_PARAMETER(cookie);
+    return true;
+}
+
+int main(int argc, char **argv)
+{
+    std::vector<std::string> args(argv, argv + argc);
+    std::string program_string{};
+    ubpf_jit_fn jit_fn;
+    uint64_t memory{0x123456789};
+
+    std::getline(std::cin, program_string);
+
+    std::unique_ptr<ubpf_vm, decltype(&ubpf_destroy)> vm(ubpf_create(), ubpf_destroy);
+    std::string error{};
+    if (!ubpf_setup_custom_test(
+            vm,
+            program_string,
+            [](ubpf_vm_up& vm, std::string &error) {
+                if (ubpf_register_external_dispatcher(vm.get(), external_dispatcher, external_dispatcher_validater) < 0) {
+                    error = "Failed to register external dispatcher.";
+                    return false;
+                }
+                return true;
+            },
+            jit_fn,
+            error)) {
+        std::cerr << "Problem setting up custom test: " << error << std::endl;
+        return 1;
+    }
+
+    [[maybe_unused]] auto result = jit_fn(&memory, sizeof(uint64_t));
+
+    // Ultimately we want the cookie pointer that we got in the dispatcher to match what we sent in!
+    return !(cookie_pointer_value == &memory);
+}
diff --git a/custom_tests/srcs/ubpf_test_jit_buffer_too_small.cc b/custom_tests/srcs/ubpf_test_jit_buffer_too_small.cc
new file mode 100644
index 00000000..9a34760a
--- /dev/null
+++ b/custom_tests/srcs/ubpf_test_jit_buffer_too_small.cc
@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#include <vector>
+#include <string>
+
+extern "C"
+{
+#include "ebpf.h"
+#include "ubpf.h"
+}
+
+#include "ubpf_custom_test_support.h"
+
+/**
+ * @brief This program reads BPF instructions from stdin and memory contents from
+ * the first agument. It then executes the BPF program and prints the
+ * value of %r0 at the end of execution.
+ */
+int main()
+{
+    std::string expected_error{"Failed to compile: Target buffer too small"};
+    std::string program_string{"95 00 00 00 00 00 00 00"};
+    ubpf_jit_fn jit_fn;
+
+    std::vector<ebpf_inst> program = bytes_to_ebpf_inst(base16_decode(program_string));
+
+    ubpf_vm_up vm(ubpf_create(), ubpf_destroy);
+    std::string error{};
+    char *error_s{nullptr};
+
+    if (!ubpf_setup_custom_test(
+            vm,
+            program_string,
+            custom_test_fixup_cb{[](ubpf_vm_up& vm, std::string& error) {
+                if (ubpf_set_jit_code_size(vm.get(), 1) < 0) {
+                    error = "Could not set the jit code size.";
+                    return false;
+                }
+                return true;
+            }},
+            jit_fn,
+            error)) {
+        free(error_s);
+        if (jit_fn == nullptr && expected_error == error)
+            return 0;
+    }
+
+    free(error_s);
+    return 1;
+}
diff --git a/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc b/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc
new file mode 100644
index 00000000..daebf17c
--- /dev/null
+++ b/custom_tests/srcs/ubpf_test_jit_unexpected_instruction.cc
@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#include <vector>
+#include <string>
+#include <iostream>
+
+extern "C"
+{
+#include "ebpf.h"
+#include "ubpf.h"
+}
+
+#include "ubpf_custom_test_support.h"
+
+/**
+ * @brief This program reads BPF instructions from stdin and memory contents from
+ * the first agument. It then executes the BPF program and prints the
+ * value of %r0 at the end of execution.
+ */
+int main()
+{
+    std::string expected_error{"Failed to load program: unknown opcode 0x8f at PC 0" };
+    ubpf_jit_fn jit_fn;
+
+    std::string program_string;
+    std::getline(std::cin, program_string);
+
+    std::vector<ebpf_inst> program = bytes_to_ebpf_inst(base16_decode(program_string));
+
+    ubpf_vm_up vm(ubpf_create(), ubpf_destroy);
+    std::string error{};
+
+    if (!ubpf_setup_custom_test(
+            vm,
+            program_string,
+            custom_test_fixup_cb{[](ubpf_vm_up&, std::string& ) {
+                return true;
+            }},
+            jit_fn,
+            error)) {
+
+        if (jit_fn == nullptr && expected_error == error)
+            return 0;
+    }
+
+    return 1;
+}
diff --git a/custom_tests/srcs/ubpf_test_update_dispatcher.cc b/custom_tests/srcs/ubpf_test_update_dispatcher.cc
new file mode 100644
index 00000000..849f844c
--- /dev/null
+++ b/custom_tests/srcs/ubpf_test_update_dispatcher.cc
@@ -0,0 +1,110 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+// This program reads BPF instructions from stdin and memory contents from
+// the first agument. It then executes the BPF program and prints the
+// value of %r0 at the end of execution.
+// The program is intended to be used with the bpf conformance test suite.
+
+#include "ubpf_int.h"
+#include <cstdint>
+#include <iostream>
+#include <memory>
+#include <vector>
+#include <string>
+
+extern "C"
+{
+#include "ubpf.h"
+}
+
+#include "ubpf_custom_test_support.h"
+
+const uint64_t dispatcher_test_dispatcher_failure{40};
+const uint64_t dispatcher_test_dispatcher_success{42};
+const uint64_t updated_dispatcher_test_dispatcher_failure{41};
+const uint64_t updated_dispatcher_test_dispatcher_success{43};
+
+uint64_t
+dispatcher_test_dispatcher(
+    uint64_t p0, uint64_t p1, uint64_t p2, uint64_t p3, uint64_t p4, unsigned int idx, void* cookie)
+{
+    UNREFERENCED_PARAMETER(p0);
+    UNREFERENCED_PARAMETER(p1);
+    UNREFERENCED_PARAMETER(p2);
+    UNREFERENCED_PARAMETER(p3);
+    UNREFERENCED_PARAMETER(p4);
+    UNREFERENCED_PARAMETER(cookie);
+    if (idx != 1) {
+        return dispatcher_test_dispatcher_failure;
+    }
+    return dispatcher_test_dispatcher_success;
+}
+
+uint64_t
+updated_dispatcher_test_dispatcher(
+    uint64_t p0, uint64_t p1, uint64_t p2, uint64_t p3, uint64_t p4, unsigned int idx, void* cookie)
+{
+    UNREFERENCED_PARAMETER(p0);
+    UNREFERENCED_PARAMETER(p1);
+    UNREFERENCED_PARAMETER(p2);
+    UNREFERENCED_PARAMETER(p3);
+    UNREFERENCED_PARAMETER(p4);
+    UNREFERENCED_PARAMETER(cookie);
+    if (idx != 1) {
+        return updated_dispatcher_test_dispatcher_failure;
+    }
+    return updated_dispatcher_test_dispatcher_success;
+}
+
+bool
+test_helpers_validater(unsigned int idx, const struct ubpf_vm* vm)
+{
+    UNREFERENCED_PARAMETER(idx);
+    UNREFERENCED_PARAMETER(vm);
+    return true;
+}
+
+int
+main(int argc, char** argv)
+{
+    std::vector<std::string> args(argv, argv + argc);
+    std::string program_string;
+    std::string memory_string;
+
+    std::getline(std::cin, program_string);
+
+    ubpf_jit_fn jit_fn;
+    uint64_t memory{0x123456789};
+    std::unique_ptr<ubpf_vm, decltype(&ubpf_destroy)> vm(ubpf_create(), ubpf_destroy);
+    std::string error{};
+    if (!ubpf_setup_custom_test(
+            vm,
+            program_string,
+            [](ubpf_vm_up& vm, std::string& error) {
+                if (ubpf_register_external_dispatcher(vm.get(), dispatcher_test_dispatcher, test_helpers_validater)) {
+                    error = "Failed to register the external dispatcher function";
+                    return false;
+                }
+                return true;
+            },
+            jit_fn,
+            error)) {
+        std::cerr << "Problem setting up custom test: " << error << std::endl;
+        return 1;
+    }
+
+    auto first_result = jit_fn(&memory, sizeof(uint64_t));
+
+    if (ubpf_register_external_dispatcher(vm.get(), updated_dispatcher_test_dispatcher, test_helpers_validater)) {
+        std::cout << "Failed to register updated dispatcher function\n";
+        return 1;
+    }
+
+    auto second_result = jit_fn(&memory, sizeof(uint64_t));
+
+    auto current_success{
+        (first_result == dispatcher_test_dispatcher_success &&
+         second_result == updated_dispatcher_test_dispatcher_success)};
+    return current_success ? 0 : 1;
+}
diff --git a/custom_tests/srcs/ubpf_test_update_helpers.cc b/custom_tests/srcs/ubpf_test_update_helpers.cc
new file mode 100644
index 00000000..85491a52
--- /dev/null
+++ b/custom_tests/srcs/ubpf_test_update_helpers.cc
@@ -0,0 +1,143 @@
+// Copyright (c) Microsoft Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+// This program reads BPF instructions from stdin and memory contents from
+// the first agument. It then executes the BPF program and prints the
+// value of %r0 at the end of execution.
+// The program is intended to be used with the bpf conformance test suite.
+
+#include "ubpf_int.h"
+#include <cstdint>
+#include <iostream>
+#include <memory>
+#include <vector>
+#include <string>
+
+extern "C"
+{
+#include "ubpf.h"
+}
+
+#include "ubpf_custom_test_support.h"
+#include "test_helpers.h"
+
+struct HelperTestCase {
+    const char *testcase_name;
+    external_function_t helper_function1;
+    external_function_t helper_function2;
+    int index;
+    uint64_t result1;
+    uint64_t result2;
+};
+
+/**
+ * @brief This program reads BPF instructions from stdin and memory contents from
+ * the first agument. It then executes the BPF program and prints the
+ * value of %r0 at the end of execution.
+ */
+int main(int argc, char **argv)
+{
+    std::vector<std::string> args(argv, argv + argc);
+    std::string program_string{};
+    bool success{true};
+
+    const char memfrob_testcase_name[] = "memfrob";
+    const char gather_bytes_testcase_name[] = "gather bytes";
+    const char sqrti_testcase_name[] = "sqrti";
+    const char no_op_testcase_name[] = "no op";
+    const char strcmp_testcase_name[] = "strcmp";
+    const char unwind_testcase_name[] = "unwind";
+
+    std::vector<HelperTestCase> test_cases{
+        {
+            .testcase_name = memfrob_testcase_name,
+            .helper_function1 = dispatcher_test_memfrob,
+            .helper_function2 = updated_dispatcher_test_memfrob,
+            .index = 1,
+            .result1 = 42,
+            .result2 = 43
+        },
+        {
+            .testcase_name = gather_bytes_testcase_name,
+            .helper_function1 = dispatcher_gather_bytes,
+            .helper_function2 = updated_dispatcher_gather_bytes,
+            .index = 1,
+            .result1 = 44,
+            .result2 = 45
+        },
+        {
+            .testcase_name = no_op_testcase_name,
+            .helper_function1 = dispatcher_no_op,
+            .helper_function2 = updated_dispatcher_no_op,
+            .index = 1,
+            .result1 = 46,
+            .result2 = 47
+        },
+        {
+            .testcase_name = sqrti_testcase_name,
+            .helper_function1 = dispatcher_sqrti,
+            .helper_function2 = updated_dispatcher_sqrti,
+            .index = 1,
+            .result1 = 48,
+            .result2 = 49
+        },
+        {
+            .testcase_name = strcmp_testcase_name,
+            .helper_function1 = dispatcher_strcmp_ext,
+            .helper_function2 = updated_dispatcher_strcmp_ext,
+            .index = 1,
+            .result1 = 50,
+            .result2 = 51
+        },
+        {
+            .testcase_name = unwind_testcase_name,
+            .helper_function1 = dispatcher_unwind,
+            .helper_function2 = updated_dispatcher_unwind,
+            .index = 1,
+            .result1 = 52,
+            .result2 = 53
+        }
+    };
+
+    std::getline(std::cin, program_string);
+
+    for (auto testcase : test_cases) {
+        ubpf_jit_fn jit_fn;
+        uint64_t memory{0x123456789};
+        std::unique_ptr<ubpf_vm, decltype(&ubpf_destroy)> vm(ubpf_create(), ubpf_destroy);
+        std::string error{};
+        if (!ubpf_setup_custom_test(
+                vm,
+                program_string,
+                [&testcase](ubpf_vm_up& vm, std::string& error) {
+                        if (ubpf_register(vm.get(), testcase.index, "unnamed", testcase.helper_function1) != 0) {
+                            error = "Failed to register helper function";
+                            return false;
+                        }
+                    return true;
+                },
+                jit_fn,
+                error)) {
+            std::cerr << "Problem setting up custom test: " << error << std::endl;
+            return 1;
+        }
+
+        [[maybe_unused]] auto first_result = jit_fn(&memory, sizeof(uint64_t));
+
+        if (ubpf_register(vm.get(), testcase.index, "unnamed", testcase.helper_function2) != 0) {
+            std::cout << "Failed to register helper function\n";
+            return 1;
+        }
+
+        [[maybe_unused]] auto second_result = jit_fn(&memory, sizeof(uint64_t));
+
+        auto current_success{(first_result == testcase.result1 && second_result == testcase.result2)};
+        if (!current_success) {
+            std::cout << "There was a failure with test " << testcase.testcase_name << ": " << 
+            testcase.result1 << " != " << first_result << " or "  <<
+            testcase.result2 << " != " << second_result << "!\n";
+        }
+        success &= current_success;
+    }
+    return success ? 0 : 1;
+}
diff --git a/external/bpf_conformance b/external/bpf_conformance
index 41ec3192..6356e3ee 160000
--- a/external/bpf_conformance
+++ b/external/bpf_conformance
@@ -1 +1 @@
-Subproject commit 41ec319242a4e598850650158dbaa91397df45a4
+Subproject commit 6356e3ee03b2c1a70fd6839a85713974bab999b3
diff --git a/ubpf_plugin/ubpf_plugin.cc b/ubpf_plugin/ubpf_plugin.cc
index 3accc64e..80859c27 100644
--- 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 "ubpf_int.h"
 #include <cstdint>
 #include <cstring>
 #include <iostream>
@@ -27,8 +28,8 @@ uint64_t test_helpers_dispatcher(uint64_t p0, uint64_t p1,uint64_t p2,uint64_t p
     return helper_functions[idx](p0, p1, p2, p3, p4);
 }
 
-bool test_helpers_validater(unsigned int idx, void *cookie) {
-    UNREFERENCED_PARAMETER(cookie);
+bool test_helpers_validater(unsigned int idx, const struct ubpf_vm *vm) {
+    UNREFERENCED_PARAMETER(vm);
     return helper_functions.contains(idx);
 }
 
@@ -141,7 +142,7 @@ int main(int argc, char **argv)
         return 1;
     }
 
-/*
+    /*
     for (auto &[key, value] : helper_functions)
     {
         if (ubpf_register(vm.get(), key, "unnamed", value) != 0)
@@ -150,9 +151,9 @@ int main(int argc, char **argv)
             return 1;
         }
     }
-*/
+    */
 
-    ubpf_register_external_dispatcher(vm.get(), test_helpers_dispatcher, test_helpers_validater, NULL);
+    ubpf_register_external_dispatcher(vm.get(), test_helpers_dispatcher, test_helpers_validater);
 
     if (ubpf_set_unwind_function_index(vm.get(), 5) != 0)
     {
diff --git a/vm/inc/ubpf.h b/vm/inc/ubpf.h
index 81c92a1e..6e969f0e 100644
--- a/vm/inc/ubpf.h
+++ b/vm/inc/ubpf.h
@@ -140,7 +140,7 @@ extern "C"
     /**
      * @brief The type of an external helper validation function.
      */
-    typedef bool (*external_function_validate_t)(unsigned int index, void* cookie);
+    typedef bool (*external_function_validate_t)(unsigned int index, const struct ubpf_vm* vm);
 
     /**
      * @brief Register a function that dispatches to external helpers
@@ -155,8 +155,6 @@ extern "C"
      *                       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.
      */
@@ -164,8 +162,7 @@ extern "C"
     ubpf_register_external_dispatcher(
         struct ubpf_vm* vm,
         external_function_dispatcher_t dispatcher,
-        external_function_validate_t validater,
-        void* cookie);
+        external_function_validate_t validater);
 
     /**
      * @brief Load code into a VM.
@@ -397,6 +394,18 @@ extern "C"
     int
     ubpf_register_data_bounds_check(struct ubpf_vm* vm, void* user_context, ubpf_bounds_check bounds_check);
 
+    /**
+     * @brief Set a bounds check function for the VM.
+     *
+     * @param[in] vm The VM to set the bounds check function for.
+     * @param[in] user_context The user context to pass to the bounds check function.
+     * @param[in] bounds_check The bounds check function.
+     * @retval 0 Success.
+     * @retval -1 Failure.
+     */
+    int
+    ubpf_set_jit_code_size(struct ubpf_vm* vm, size_t code_size);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/vm/ubpf_int.h b/vm/ubpf_int.h
index 861e6a80..9c5311ba 100644
--- a/vm/ubpf_int.h
+++ b/vm/ubpf_int.h
@@ -24,15 +24,33 @@
 #include <ubpf.h>
 #include "ebpf.h"
 
+#define UNUSED_PARAMETER(x) ((void)x)
+
 struct ebpf_inst;
 typedef uint64_t (*ext_func)(uint64_t arg0, uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4);
 
+typedef enum {
+    UBPF_JIT_COMPILE_SUCCESS,
+    UBPF_JIT_COMPILE_FAILURE,
+} upbf_jit_result_t;
+
+struct ubpf_jit_result {
+    uint32_t external_dispatcher_offset;
+    uint32_t external_helper_offset;
+    upbf_jit_result_t compile_result;
+    char *errmsg;
+};
+
+#define MAX_EXT_FUNCS 64
+
 struct ubpf_vm
 {
     struct ebpf_inst* insts;
     uint16_t num_insts;
     ubpf_jit_fn jitted;
     size_t jitted_size;
+    size_t jitter_buffer_size;
+    struct ubpf_jit_result jitted_result;
 
     ext_func* ext_funcs;
     bool* int_funcs;
@@ -40,11 +58,12 @@ struct ubpf_vm
 
     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);
+    struct ubpf_jit_result (*jit_translate)(struct ubpf_vm* vm, uint8_t* buffer, size_t* size);
+    bool (*jit_update_dispatcher)(struct ubpf_vm* vm, external_function_dispatcher_t new_dispatcher, uint8_t* buffer, size_t size, uint32_t offset);
+    bool (*jit_update_helper)(struct ubpf_vm* vm, ext_func new_helper, unsigned int idx, uint8_t* buffer, size_t size, uint32_t offset);
     int unwind_stack_extension_index;
     uint64_t pointer_secret;
     ubpf_data_relocation data_relocation_function;
@@ -63,12 +82,24 @@ struct ubpf_stack_frame
 };
 
 /* The various JIT targets.  */
-int
-ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg);
-int
-ubpf_translate_x86_64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg);
-int
-ubpf_translate_null(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg);
+
+// arm64
+struct ubpf_jit_result
+ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size);
+bool ubpf_jit_update_dispatcher_arm64(struct ubpf_vm* vm, external_function_dispatcher_t new_dispatcher, uint8_t* buffer, size_t size, uint32_t offset);
+bool ubpf_jit_update_helper_arm64(struct ubpf_vm* vm, ext_func new_helper, unsigned int idx, uint8_t* buffer, size_t size, uint32_t offset);
+
+//x86_64
+struct ubpf_jit_result
+ubpf_translate_x86_64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size);
+bool ubpf_jit_update_dispatcher_x86_64(struct ubpf_vm* vm, external_function_dispatcher_t new_dispatcher, uint8_t* buffer, size_t size, uint32_t offset);
+bool ubpf_jit_update_helper_x86_64(struct ubpf_vm* vm, ext_func new_helper, unsigned int idx, uint8_t* buffer, size_t size, uint32_t offset);
+
+//uhm, hello?
+struct ubpf_jit_result
+ubpf_translate_null(struct ubpf_vm* vm, uint8_t* buffer, size_t* size);
+bool ubpf_jit_update_dispatcher_null(struct ubpf_vm* vm, external_function_dispatcher_t new_dispatcher, uint8_t* buffer, size_t size, uint32_t offset);
+bool ubpf_jit_update_helper_null(struct ubpf_vm* vm, ext_func new_helper, unsigned int idx, uint8_t* buffer, size_t size, uint32_t offset);
 
 char*
 ubpf_error(const char* fmt, ...);
diff --git a/vm/ubpf_jit.c b/vm/ubpf_jit.c
index 0e7b10fe..58a0e9cb 100644
--- a/vm/ubpf_jit.c
+++ b/vm/ubpf_jit.c
@@ -23,31 +23,65 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <stdbool.h>
+#include <string.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"
 
-#define UNUSED(x) ((void)x)
 
 int
 ubpf_translate(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg)
 {
-    return vm->translate(vm, buffer, size, errmsg);
+    struct ubpf_jit_result jit_result = vm->jit_translate(vm, buffer, size);
+    vm->jitted_result = jit_result;
+    if (jit_result.errmsg) {
+        *errmsg = jit_result.errmsg;
+    }
+    return jit_result.compile_result == UBPF_JIT_COMPILE_SUCCESS ? 0 : -1;
 }
 
-int
-ubpf_translate_null(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg)
+struct ubpf_jit_result
+ubpf_translate_null(struct ubpf_vm* vm, uint8_t* buffer, size_t* size)
 {
+    struct ubpf_jit_result compile_result;
+    compile_result.compile_result = UBPF_JIT_COMPILE_FAILURE;
+    compile_result.external_dispatcher_offset = 0;
+
     /* NULL JIT target - just returns an error. */
-    UNUSED(vm);
-    UNUSED(buffer);
-    UNUSED(size);
-    *errmsg = ubpf_error("Code can not be JITed on this target.");
-    return -1;
+    UNUSED_PARAMETER(vm);
+    UNUSED_PARAMETER(buffer);
+    UNUSED_PARAMETER(size);
+    compile_result.errmsg = ubpf_error("Code can not be JITed on this target.");
+    return compile_result;
+}
+
+bool ubpf_jit_update_dispatcher_null(struct ubpf_vm* vm, external_function_dispatcher_t new_dispatcher, uint8_t* buffer, size_t size, uint32_t offset)
+{
+    UNUSED_PARAMETER(vm);
+    UNUSED_PARAMETER(new_dispatcher);
+    UNUSED_PARAMETER(buffer);
+    UNUSED_PARAMETER(size);
+    UNUSED_PARAMETER(offset);
+    return false;
+}
+
+bool ubpf_jit_update_helper_null(struct ubpf_vm* vm, ext_func new_helper, unsigned int idx, uint8_t* buffer, size_t size, uint32_t offset)
+{
+    UNUSED_PARAMETER(vm);
+    UNUSED_PARAMETER(new_helper);
+    UNUSED_PARAMETER(idx);
+    UNUSED_PARAMETER(buffer);
+    UNUSED_PARAMETER(size);
+    UNUSED_PARAMETER(offset);
+    return false;
+}
+
+int
+ubpf_set_jit_code_size(struct ubpf_vm* vm, size_t code_size) 
+{
+    vm->jitter_buffer_size = code_size;
+    return 0;
 }
 
 ubpf_jit_fn
@@ -68,7 +102,7 @@ ubpf_compile(struct ubpf_vm* vm, char** errmsg)
         return NULL;
     }
 
-    jitted_size = 65536;
+    jitted_size = vm->jitter_buffer_size;
     buffer = calloc(jitted_size, 1);
     if (buffer == NULL) {
         *errmsg = ubpf_error("internal uBPF error: calloc failed: %s\n", strerror(errno));
diff --git a/vm/ubpf_jit_arm64.c b/vm/ubpf_jit_arm64.c
index e3749364..3ce6622a 100644
--- a/vm/ubpf_jit_arm64.c
+++ b/vm/ubpf_jit_arm64.c
@@ -28,9 +28,7 @@
 #include <stdbool.h>
 #include <string.h>
 #include <unistd.h>
-#include <inttypes.h>
 #include <sys/mman.h>
-#include <errno.h>
 #include <assert.h>
 #include "ubpf_int.h"
 
@@ -42,6 +40,7 @@
 #define TARGET_PC_EXIT ~UINT32_C(0)
 #define TARGET_PC_ENTER (~UINT32_C(0) & 0x0101)
 #define TARGET_PC_EXTERNAL_DISPATCHER (~UINT32_C(0) & 0x1010)
+#define TARGET_LOAD_HELPER_TABLE (~UINT32_C(0) & 0x101010)
 
 // This is guaranteed to be an illegal A64 instruction.
 #define BAD_OPCODE ~UINT32_C(0)
@@ -50,6 +49,15 @@ struct patchable_relative
 {
     uint32_t offset_loc;
     uint32_t target_pc;
+    uint32_t target_offset;
+    uint32_t adr;
+};
+
+enum JitError {
+    NoError,
+    NotEnoughSpace,
+    UnexpectedInstruction,
+    UnknownInstruction
 };
 
 struct jit_state
@@ -61,7 +69,9 @@ struct jit_state
     uint32_t exit_loc;
     uint32_t entry_loc;
     uint32_t dispatcher_loc;
+    uint32_t helper_table_loc;
     uint32_t unwind_loc;
+    enum JitError jit_error;
     struct patchable_relative* jumps;
     struct patchable_relative* loads;
     int num_jumps;
@@ -117,6 +127,8 @@ static enum Registers temp_register = R24;
 static enum Registers temp_div_register = R25;
 // Temp register for load/store offsets
 static enum Registers offset_register = R26;
+// Special register for external dispatcher context.
+static enum Registers VOLATILE_CTXT = R26;
 
 // Number of eBPF registers
 #define REGISTER_MAP_SIZE 11
@@ -169,8 +181,11 @@ static uint32_t inline align_to(uint32_t amount, uint64_t boundary)
 static void
 emit_bytes(struct jit_state* state, void* data, uint32_t len)
 {
-    assert(len <= state->size);
-    assert(state->offset <= state->size - len);
+    if (!(len <= state->size && state->offset <= state->size - len)) {
+        state->jit_error = NotEnoughSpace;
+        return;
+    }
+
     if ((state->offset + len) > state->size) {
         state->offset = state->size;
         return;
@@ -277,6 +292,16 @@ emit_loadstore_literal(
     emit_instruction(state, op | reg_op_base | rt);
 }
 
+static void
+emit_adr(struct jit_state *state, uint32_t offset, enum Registers rd)
+{
+    uint32_t immlo = (offset & 0x00000003) << 30;
+    uint32_t immhi = (offset & 0x00ffffe0) << 5;
+    uint32_t instr = 0x10000000 | immhi | immlo | rd;
+    emit_instruction(state, instr);
+
+}
+
 enum LoadStorePairOpcode
 {
     // op    V    L
@@ -351,15 +376,25 @@ enum UnconditionalBranchImmediateOpcode
     UBR_BL = 0x94000000U, // 1001_0100_0000_0000_0000_0000_0000_0000
 };
 
-static void
+static int
 note_jump(struct jit_state* state, uint32_t target_pc)
 {
+    int index = state->num_jumps++;
+
     if (state->num_jumps == UBPF_MAX_INSTS) {
-        return;
+        return -1;
     }
-    struct patchable_relative* jump = &state->jumps[state->num_jumps++];
+    struct patchable_relative* jump = &state->jumps[index];
     jump->offset_loc = state->offset;
     jump->target_pc = target_pc;
+    return index;
+}
+
+static void
+note_jump_offset(struct jit_state *state, int index, uint32_t target_offset)
+{
+    struct patchable_relative* to_update = &state->jumps[index];
+    to_update->target_offset = target_offset;
 }
 
 static void
@@ -373,6 +408,18 @@ note_load(struct jit_state* state, uint32_t target_pc)
     load->target_pc = target_pc;
 }
 
+static void
+note_adr(struct jit_state* state, uint32_t offset)
+{
+    if (state->num_loads == UBPF_MAX_INSTS) {
+        return;
+    }
+    struct patchable_relative* load = &state->loads[state->num_loads++];
+    load->offset_loc = state->offset;
+    load->target_pc = offset;
+    load->adr = offset;
+}
+
 
 /* [ArmARM-A H.a]: C4.1.65: Unconditional branch (immediate).  */
 static void
@@ -411,11 +458,12 @@ enum ConditionalBranchImmediateOpcode
 };
 
 /* [ArmARM-A H.a]: C4.1.65: Conditional branch (immediate).  */
-static void
+static int
 emit_conditionalbranch_immediate(struct jit_state* state, enum Condition cond, uint32_t target_pc)
 {
-    note_jump(state, target_pc);
+    int index = note_jump(state, target_pc);
     emit_instruction(state, BR_Bcond | (0 << 5) | cond);
+    return index;
 }
 
 enum CompareBranchOpcode
@@ -577,6 +625,16 @@ update_load_literal(struct jit_state* state, uint32_t instr_offset, int32_t targ
     memcpy(state->buf + instr_offset, &instr, sizeof(uint32_t));
 }
 
+static void
+update_adr(struct jit_state* state, uint32_t instr_offset, int32_t immediate)
+{
+    uint32_t instr;
+    uint32_t immhi = (immediate & 0x00ffffff) << 5;
+    memcpy(&instr, state->buf + instr_offset, sizeof(uint32_t));
+    instr |= immhi;
+    memcpy(state->buf + instr_offset, &instr, sizeof(uint32_t));
+}
+
 /* Generate the function prologue.
  *
  * We set the stack to look like:
@@ -609,6 +667,9 @@ emit_jit_prologue(struct jit_state* state, size_t ubpf_stack_size)
     /* Setup UBPF frame pointer. */
     emit_addsub_immediate(state, true, AS_ADD, map_register(10), SP, state->stack_size);
 
+    /* Copy R0 to the volatile context for safe keeping. */
+    emit_logical_register(state, true, LOG_ORR, VOLATILE_CTXT, RZ, R0);
+
     emit_unconditionalbranch_immediate(state, UBR_BL, TARGET_PC_ENTER);
     emit_unconditionalbranch_immediate(state, UBR_B, TARGET_PC_EXIT);
     state->entry_loc = state->offset;
@@ -617,6 +678,12 @@ emit_jit_prologue(struct jit_state* state, size_t ubpf_stack_size)
 static void
 emit_dispatched_external_helper_call(struct jit_state* state, struct ubpf_vm* vm, unsigned int idx)
 {
+    (void)vm;
+
+    // There are two paths through the function:
+    // 1. There is an external dispatcher registered. If so, we prioritize that.
+    // 2. We fall back to the regular registered helper.
+
     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);
@@ -624,13 +691,39 @@ emit_dispatched_external_helper_call(struct jit_state* state, struct ubpf_vm* vm
     // All parameters to the helper function are in the right spot
     // for the dispatcher. All we need to do now is ...
 
+       // Call!
+    note_load(state, TARGET_PC_EXTERNAL_DISPATCHER);
+    emit_loadstore_literal(state, LS_LDRL, temp_register);
+
+    // Check whether temp_register is empty.
+    emit_addsub_immediate(state, true, AS_SUBS, temp_register, temp_register, 0);
+
+    // Jump if we are ...
+    int jump_index = emit_conditionalbranch_immediate(state, COND_NE, 0);
+
+    // Otherwise, load the target into the register.
+
+    emit_movewide_immediate(state, true, R5, idx);
+    emit_movewide_immediate(state, true, R6, 3);
+    emit_dataprocessing_twosource(state, true, DP2_LSLV, R5, R5, R6);
+
+    emit_movewide_immediate(state, true, temp_register, 0);
+    note_adr(state, TARGET_LOAD_HELPER_TABLE);
+    emit_adr(state, 0, temp_register);
+    emit_addsub_register(state, true, AS_ADD, temp_register, temp_register, R5);
+    emit_loadstore_immediate(state, LS_LDRX, temp_register, temp_register, 0);
+
+    // And now we, too, are ...
+
+    // ... ready to roll.
+
+    note_jump_offset(state, jump_index, state->offset);
+
     // ... set up the final two parameters.
     emit_movewide_immediate(state, true, R5, idx);
-    emit_movewide_immediate(state, true, R6, (uint64_t)vm->dispatcher_cookie);
+    // TODO
+    emit_logical_register(state, true, LOG_ORR, R6, RZ, VOLATILE_CTXT);
 
-    // Call!
-    note_load(state, TARGET_PC_EXTERNAL_DISPATCHER);
-    emit_loadstore_literal(state, LS_LDRL, temp_register);
     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.  */
@@ -702,6 +795,16 @@ emit_dispatched_external_helper_address(struct jit_state *state, uint64_t dispat
     return helper_address;
 }
 
+static uint32_t
+emit_helper_table(struct jit_state* state, struct ubpf_vm* vm) {
+
+    uint32_t helper_table_address_target = state->offset;
+    for (int i = 0; i<MAX_EXT_FUNCS; i++) {
+        emit_bytes(state, &vm->ext_funcs[i], sizeof(uint64_t));
+    }
+    return helper_table_address_target;
+}
+
 static bool
 is_imm_op(struct ebpf_inst const* inst)
 {
@@ -976,6 +1079,9 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
     emit_jit_prologue(state, UBPF_STACK_SIZE);
 
     for (i = 0; i < vm->num_insts; i++) {
+
+        // All checks for errors during the encoding of _this_ instruction
+        // occur at the end of the loop.
         struct ebpf_inst inst = ubpf_fetch_instruction(vm, i);
         state->pc_locs[i] = state->offset;
 
@@ -986,6 +1092,9 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
 
         int sixty_four = is_alu64_op(&inst);
 
+        // If this is an operation with an immediate operand (and that immediate
+        // operand is _not_ simple), then we convert the operation to the equivalent
+        // register version after moving the immediate into a temporary register.
         if (is_imm_op(&inst) && !is_simple_imm(&inst)) {
             emit_movewide_immediate(state, sixty_four, temp_register, (int64_t)inst.imm);
             src = temp_register;
@@ -1193,16 +1302,30 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
         case EBPF_OP_RSH64_IMM:
         case EBPF_OP_ARSH64_IMM:
             *errmsg = ubpf_error("Unexpected instruction at PC %d: opcode %02x, immediate %08x", i, opcode, inst.imm);
-            return -1;
+            state->jit_error = UnexpectedInstruction;
         default:
             *errmsg = ubpf_error("Unknown instruction at PC %d: opcode %02x", i, opcode);
-            return -1;
+            state->jit_error = UnknownInstruction;
+        }
+
+        switch (state->jit_error) {
+            case NotEnoughSpace:
+                *errmsg = ubpf_error("Target buffer too small");
+                return -1;
+            case UnknownInstruction:
+            case UnexpectedInstruction:
+                // Error messages for this case are set when discovered.
+                return -1;
+            case NoError:
+                // just keep going!
+                break;
         }
     }
 
     emit_jit_epilogue(state);
 
     state->dispatcher_loc =  emit_dispatched_external_helper_address(state, (uint64_t)vm->dispatcher);
+    state->helper_table_loc = emit_helper_table(state, vm);
 
     return 0;
 }
@@ -1230,7 +1353,9 @@ resolve_jumps(struct jit_state* state)
         struct patchable_relative jump = state->jumps[i];
 
         int32_t target_loc;
-        if (jump.target_pc == TARGET_PC_EXIT) {
+        if (jump.target_offset != 0) {
+            target_loc = jump.target_offset;
+        } else if (jump.target_pc == TARGET_PC_EXIT) {
             target_loc = state->exit_loc;
         } else if (jump.target_pc == TARGET_PC_ENTER) {
             target_loc = state->entry_loc;
@@ -1253,6 +1378,8 @@ resolve_loads(struct jit_state* state)
         int32_t target_loc;
         if (jump.target_pc == TARGET_PC_EXTERNAL_DISPATCHER) {
             target_loc = state->dispatcher_loc;
+        } else if (jump.target_pc == TARGET_LOAD_HELPER_TABLE) {
+            target_loc = state->helper_table_loc;
         } else {
             return false;
         }
@@ -1260,17 +1387,50 @@ resolve_loads(struct jit_state* state)
         int32_t rel = target_loc - jump.offset_loc;
         assert(rel % 4 == 0);
         rel >>= 2;
-        update_load_literal(state, jump.offset_loc, rel);
+        if (jump.adr != 0) {
+            update_adr(state, jump.offset_loc, rel);
+        } else {
+            update_load_literal(state, jump.offset_loc, rel);
+        }
     }
     return true;
 }
 
+bool ubpf_jit_update_dispatcher_arm64(struct ubpf_vm* vm, external_function_dispatcher_t new_dispatcher, uint8_t* buffer, size_t size, uint32_t offset)
+{
+    UNUSED_PARAMETER(vm);
+    uint64_t jit_upper_bound = (uint64_t)buffer + size;
+    void *dispatcher_address = (void*)((uint64_t)buffer + offset);
+    if ((uint64_t)dispatcher_address + sizeof(void*) < jit_upper_bound) {
+        memcpy(dispatcher_address, &new_dispatcher, sizeof(void*));
+        return true;
+    }
+
+    return false;
+}
+
+bool ubpf_jit_update_helper_arm64(struct ubpf_vm* vm, ext_func new_helper, unsigned int idx, uint8_t* buffer, size_t size, uint32_t offset)
+{
+    UNUSED_PARAMETER(vm);
+    uint64_t jit_upper_bound = (uint64_t)buffer + size;
+
+    void* dispatcher_address = (void*)((uint64_t)buffer + offset + (8 * idx));
+    if ((uint64_t)dispatcher_address + sizeof(void*) < jit_upper_bound) {
+        memcpy(dispatcher_address, &new_helper, sizeof(void*));
+        return true;
+    }
+    return false;
+}
 
-int
-ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg)
+struct ubpf_jit_result
+ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size)
 {
     struct jit_state state;
-    int result = -1;
+    struct ubpf_jit_result compile_result;
+
+    compile_result.compile_result = UBPF_JIT_COMPILE_FAILURE;
+    compile_result.errmsg = NULL;
+    compile_result.external_dispatcher_offset = 0;
 
     state.offset = 0;
     state.size = *size;
@@ -1280,37 +1440,36 @@ ubpf_translate_arm64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** e
     state.loads = calloc(UBPF_MAX_INSTS, sizeof(state.loads[0]));
     state.num_jumps = 0;
     state.num_loads = 0;
+    state.jit_error = NoError;
+
 
     if (!state.pc_locs || !state.jumps) {
-        *errmsg = ubpf_error("Out of memory");
+        compile_result.errmsg = ubpf_error("Out of memory");
         goto out;
     }
 
-    if (translate(vm, &state, errmsg) < 0) {
+    if (translate(vm, &state, &compile_result.errmsg) < 0) {
         goto out;
     }
 
     if (state.num_jumps == UBPF_MAX_INSTS) {
-        *errmsg = ubpf_error("Excessive number of jump targets");
-        goto out;
-    }
-
-    if (state.offset == state.size) {
-        *errmsg = ubpf_error("Target buffer too small");
+        compile_result.errmsg = ubpf_error("Excessive number of jump targets");
         goto out;
     }
 
     if (!resolve_jumps(&state) || !resolve_loads(&state)) {
-        *errmsg = ubpf_error("Could not patch the relative addresses in the JIT'd code.");
+        compile_result.errmsg = ubpf_error("Could not patch the relative addresses in the JIT'd code.");
         goto out;
     }
 
-    result = 0;
+    compile_result.compile_result = UBPF_JIT_COMPILE_SUCCESS;
     *size = state.offset;
+    compile_result.external_dispatcher_offset = state.dispatcher_loc;
+    compile_result.external_helper_offset = state.helper_table_loc;
 
 out:
     free(state.pc_locs);
     free(state.jumps);
     free(state.loads);
-    return result;
+    return compile_result;
 }
diff --git a/vm/ubpf_jit_x86_64.c b/vm/ubpf_jit_x86_64.c
index 45e7e67e..e9bfa791 100644
--- a/vm/ubpf_jit_x86_64.c
+++ b/vm/ubpf_jit_x86_64.c
@@ -134,6 +134,16 @@ emit_dispatched_external_helper_address(struct jit_state* state, struct ubpf_vm*
     return external_helper_address_target;
 }
 
+static uint32_t
+emit_helper_table(struct jit_state* state, struct ubpf_vm* vm) {
+
+    uint32_t helper_table_address_target = state->offset;
+    for (int i = 0; i<MAX_EXT_FUNCS; i++) {
+        emit8(state, (uint64_t)vm->ext_funcs[i]);
+    }
+    return helper_table_address_target;
+}
+
 static uint32_t
 emit_retpoline(struct jit_state* state)
 {
@@ -200,6 +210,23 @@ ubpf_set_register_offset(int x)
     }
 }
 
+/*
+ * In order to make it so that the generated code is completely standalone, all the necessary
+ * function pointers for external helpers are embedded in the jitted code. The layout looks like:
+ *
+ *                 state->buffer: CODE
+ *                                CODE
+ *                                CODE
+ *                                ...
+ *                                CODE
+ *                                External Helper External Dispatcher Function Pointer (8 bytes, maybe NULL)
+ *                                External Helper Function Pointer Idx 0 (8 bytes, maybe NULL)
+ *                                External Helper Function Pointer Idx 1 (8 bytes, maybe NULL)
+ *                                ...
+ *                                External Helper Function Pointer Idx MAX_EXT_FUNCS-1 (8 bytes, maybe NULL)
+ * state->buffer + state->offset: 
+ */
+
 static int
 translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
 {
@@ -215,6 +242,11 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
         emit_mov(state, platform_parameter_registers[0], map_register(BPF_REG_1));
     }
 
+    /* Move the platform parameter register to the (volatile) register
+     * that holds the pointer to the context.
+     */
+    emit_mov(state, platform_parameter_registers[0], VOLATILE_CTXT);
+
     /*
      * Assuming that the stack is 16-byte aligned right before
      * the call insn that brought us to this code, when
@@ -614,7 +646,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_dispatched_external_helper_call(state, vm, inst.imm);
+                emit_dispatched_external_helper_call(state, 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);
@@ -708,6 +740,7 @@ translate(struct ubpf_vm* vm, struct jit_state* state, char** errmsg)
 
     state->retpoline_loc = emit_retpoline(state);
     state->dispatcher_loc = emit_dispatched_external_helper_address(state, vm);
+    state->helper_table_loc = emit_helper_table(state, vm);
 
     return 0;
 }
@@ -866,6 +899,8 @@ resolve_patchable_relatives(struct jit_state* state)
         int target_loc;
         if (load.target_pc == TARGET_PC_EXTERNAL_DISPATCHER) {
             target_loc = state->dispatcher_loc;
+        } else if (load.target_pc == TARGET_LOAD_HELPER_TABLE) {
+            target_loc = state->helper_table_loc;
         } else {
             target_loc = -1;
             return false;
@@ -880,11 +915,15 @@ resolve_patchable_relatives(struct jit_state* state)
     return true;
 }
 
-int
-ubpf_translate_x86_64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char** errmsg)
+struct ubpf_jit_result
+ubpf_translate_x86_64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size)
 {
     struct jit_state state;
-    int result = -1;
+    struct ubpf_jit_result compile_result;
+
+    compile_result.compile_result = UBPF_JIT_COMPILE_FAILURE;
+    compile_result.errmsg = NULL;
+    compile_result.external_dispatcher_offset = 0;
 
     state.offset = 0;
     state.size = *size;
@@ -896,35 +935,63 @@ ubpf_translate_x86_64(struct ubpf_vm* vm, uint8_t* buffer, size_t* size, char**
     state.num_loads = 0;
 
     if (!state.pc_locs || !state.jumps) {
-        *errmsg = ubpf_error("Out of memory");
+        compile_result.errmsg = ubpf_error("Out of memory");
         goto out;
     }
 
-    if (translate(vm, &state, errmsg) < 0) {
+    if (translate(vm, &state, &compile_result.errmsg) < 0) {
         goto out;
     }
 
     if (state.num_jumps == UBPF_MAX_INSTS) {
-        *errmsg = ubpf_error("Excessive number of jump targets");
+        compile_result.errmsg = ubpf_error("Excessive number of jump targets");
         goto out;
     }
 
     if (state.offset == state.size) {
-        *errmsg = ubpf_error("Target buffer too small");
+        compile_result.errmsg = ubpf_error("Target buffer too small");
         goto out;
     }
 
     if (!resolve_patchable_relatives(&state)) {
-        *errmsg = ubpf_error("Could not patch the relative addresses in the JIT'd code.");
+        compile_result.errmsg = ubpf_error("Could not patch the relative addresses in the JIT'd code.");
         goto out;
     }
 
-    result = 0;
+    compile_result.compile_result = UBPF_JIT_COMPILE_SUCCESS;
+    compile_result.external_dispatcher_offset = state.dispatcher_loc;
+    compile_result.external_helper_offset = state.helper_table_loc;
     *size = state.offset;
 
 out:
     free(state.pc_locs);
     free(state.jumps);
     free(state.loads);
-    return result;
+    return compile_result;
+}
+
+bool ubpf_jit_update_dispatcher_x86_64(struct ubpf_vm* vm, external_function_dispatcher_t new_dispatcher, uint8_t* buffer, size_t size, uint32_t offset)
+{
+    UNUSED_PARAMETER(vm);
+    uint64_t jit_upper_bound = (uint64_t)buffer + size;
+    void *dispatcher_address = (void*)((uint64_t)buffer + offset);
+    if ((uint64_t)dispatcher_address + sizeof(void*) < jit_upper_bound) {
+        memcpy(dispatcher_address, &new_dispatcher, sizeof(void*));
+        return true;
+    }
+
+    return false;
+}
+
+bool ubpf_jit_update_helper_x86_64(struct ubpf_vm* vm, ext_func new_helper, unsigned int idx, uint8_t* buffer, size_t size, uint32_t offset)
+{
+    UNUSED_PARAMETER(vm);
+    uint64_t jit_upper_bound = (uint64_t)buffer + size;
+
+    void* dispatcher_address = (void*)((uint64_t)buffer + offset + (8 * idx));
+    if ((uint64_t)dispatcher_address + sizeof(void*) < jit_upper_bound) {
+        memcpy(dispatcher_address, &new_helper, sizeof(void*));
+        return true;
+    }
+    return false;
 }
diff --git a/vm/ubpf_jit_x86_64.h b/vm/ubpf_jit_x86_64.h
index 7c093d70..9fe48458 100644
--- a/vm/ubpf_jit_x86_64.h
+++ b/vm/ubpf_jit_x86_64.h
@@ -29,7 +29,6 @@
 #include <string.h>
 
 #include "ubpf.h"
-#include "ubpf_int.h"
 
 #define RAX 0
 #define RCX 1
@@ -49,6 +48,8 @@
 #define R14 14
 #define R15 15
 
+#define VOLATILE_CTXT 11
+
 enum operand_size
 {
     S8,
@@ -59,8 +60,12 @@ enum operand_size
 
 struct patchable_relative
 {
+    /* Where in the instruction stream should this relative address be patched. */
     uint32_t offset_loc;
+    /* Which PC should this target. The ultimate offset will be determined
+     * automatically unless ... */
     uint32_t target_pc;
+    /* ... the target_offset is set which overrides the automatic lookup. */
     uint32_t target_offset;
 };
 
@@ -68,6 +73,7 @@ struct patchable_relative
 #define TARGET_PC_EXIT -1
 #define TARGET_PC_RETPOLINE -3
 #define TARGET_PC_EXTERNAL_DISPATCHER -4
+#define TARGET_LOAD_HELPER_TABLE -5
 
 struct jit_state
 {
@@ -79,6 +85,7 @@ struct jit_state
     uint32_t unwind_loc;
     uint32_t retpoline_loc;
     uint32_t dispatcher_loc;
+    uint32_t helper_table_loc;
     struct patchable_relative* jumps;
     struct patchable_relative* loads;
     int num_jumps;
@@ -88,7 +95,6 @@ struct jit_state
 static inline void
 emit_bytes(struct jit_state* state, void* data, uint32_t len)
 {
-    assert(state->offset <= state->size - len);
     if ((state->offset + len) > state->size) {
         state->offset = state->size;
         return;
@@ -121,16 +127,32 @@ emit8(struct jit_state* state, uint64_t x)
     emit_bytes(state, &x, sizeof(x));
 }
 
+static void
+emit_4byte_offset_placeholder(struct jit_state *state) {
+    emit4(state, 0);
+}
+
+
+// TODO: Document here.
 static inline void
+emit_patchable_relative(uint32_t offset, uint32_t target_pc, uint32_t manual_target_offset, struct patchable_relative *table, size_t index)
+{
+    struct patchable_relative* jump = &table[index];
+    jump->offset_loc = offset;
+    jump->target_pc = target_pc;
+    jump->target_offset = manual_target_offset;
+}
+
+static uint32_t
 emit_jump_target_address(struct jit_state* state, int32_t target_pc)
 {
     if (state->num_jumps == UBPF_MAX_INSTS) {
-        return;
+        return 0;
     }
-    struct patchable_relative* jump = &state->jumps[state->num_jumps++];
-    jump->offset_loc = state->offset;
-    jump->target_pc = target_pc;
-    emit4(state, 0);
+    uint32_t target_address_offset = state->offset;
+    emit_patchable_relative(state->offset, target_pc, 0, state->jumps, state->num_jumps++);
+    emit_4byte_offset_placeholder(state);
+    return target_address_offset;
 }
 
 static inline void
@@ -139,14 +161,13 @@ emit_jump_target_offset(struct jit_state* state, uint32_t jump_loc, uint32_t jum
     if (state->num_jumps == UBPF_MAX_INSTS) {
         return;
     }
-    struct patchable_relative* jump = &state->jumps[state->num_jumps++];
-    jump->offset_loc = jump_loc;
-    jump->target_offset = jump_state_offset;
+    emit_patchable_relative(jump_loc, 0, jump_state_offset, state->jumps, state->num_jumps++);
 }
 
 static inline void
 emit_modrm(struct jit_state* state, int mod, int r, int m)
 {
+    // Only the top 2 bits of the mod should be used.
     assert(!(mod & ~0xc0));
     emit1(state, (mod & 0xc0) | ((r & 7) << 3) | (m & 7));
 }
@@ -292,12 +313,12 @@ emit_cmp32(struct jit_state* state, int src, int dst)
     emit_alu32(state, 0x39, src, dst);
 }
 
-static inline void
+static inline uint32_t
 emit_jcc(struct jit_state* state, int code, int32_t target_pc)
 {
     emit1(state, 0x0f);
     emit1(state, code);
-    emit_jump_target_address(state, target_pc);
+    return emit_jump_target_address(state, target_pc);
 }
 
 /* Load [src + offset] into dst */
@@ -332,20 +353,20 @@ emit_load_imm(struct jit_state* state, int dst, int64_t imm)
     }
 }
 
-/* Load sign-extended immediate into register */
-static inline void
-emit_load_relative(struct jit_state* state, int target_pc)
+static uint32_t
+emit_load_relative(struct jit_state* state, int dst, int relative_load_tgt)
 {
     if (state->num_loads == UBPF_MAX_INSTS) {
-        return;
+        return 0;
     }
-    emit1(state, 0x48);
+
+    emit_rex(state, 1, 0, 0, 0);
     emit1(state, 0x8b);
-    emit1(state, 0x05);
-    struct patchable_relative* load = &state->loads[state->num_loads++];
-    load->offset_loc = state->offset;
-    load->target_pc = target_pc;
-    emit4(state, 0);
+    emit_modrm(state, 0, dst, 0x05);
+    uint32_t load_target_offset = state->offset;
+    emit_patchable_relative(state->offset, relative_load_tgt, 0, state->loads, state->num_loads++);
+    emit_4byte_offset_placeholder(state);
+    return load_target_offset;
 }
 
 /* Store register src to [dst + offset] */
@@ -396,45 +417,160 @@ emit_jmp(struct jit_state* state, uint32_t target_pc)
 }
 
 static inline void
-emit_dispatched_external_helper_call(struct jit_state* state, const struct ubpf_vm* vm, unsigned int idx)
+emit_dispatched_external_helper_call(struct jit_state* state, unsigned int idx)
 {
     /*
+     * Note: We do *not* have to preserve any x86-64 registers here ...
+     * ... according to the SystemV ABI: rbx (eBPF6),
+     *                                   r13 (eBPF7),
+     *                                   r14 (eBPF8),
+     *                                   r15 (eBPF9), and
+     *                                   rbp (eBPF10) are all preserved.
+     * ... according to the Windows ABI: r15 (eBPF6)
+     *                                   rdi (eBPF7),
+     *                                   rsi (eBPF8),
+     *                                   rbx (eBPF9), and
+     *                                   rbp (eBPF10) are all preserved.
+     *
      * When we enter here, our stack is 16-byte aligned. Keep
      * it that way!
      */
 
+    /*
+     * There are two things that could happen:
+     * 1. The user has registered an external dispatcher and we need to
+     *    send control there to invoke an external helper.
+     * 2. The user is relying on the default dispatcher to pass control
+     *    to the registered external helper.
+     * To determine which action to take, we will first consider the 8
+     * bytes at TARGET_PC_EXTERNAL_DISPATCHER. If those 8 bytes have an
+     * address, that represents the address of the user-registered external
+     * dispatcher and we pass control there. That function signature looks like
+     * uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, unsigned int index, void* cookie
+     * so we make sure that the arguments are done properly depending on the abi.
+     *
+     * If there is no external dispatcher registered, the user is expected
+     * to have registered a handler with us for the helper with index idx.
+     * There is a table of MAX_ function pointers starting at TARGET_LOAD_HELPER_TABLE.
+     * Each of those functions has a signature that looks like 
+     * uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, void* cookie
+     * We load the appropriate function pointer by using idx to index it and then
+     * make sure that the arguments are set properly depending on the abi.
+     */
+
+    // Save r9 -- I need it for a parameter!
+    emit_push(state, R9);
+    // Save register where volatile context is stored.
+    emit_push(state, VOLATILE_CTXT);
+    // ^^ Stack is aligned here.
+
 #if defined(_WIN32)
-    /* We have to create a little space to keep our 16-byte
-     * alignment happy
+    /* Because we may need 24 bytes on the stack but at least 16, we have to take 32
+     * to keep alignment happy. We may ultimately need it all, but we certainly
+     * need 16! Later, though, there is a push that always happens (MARKER2), so
+     * we only allocate 24 here.
      */
-    emit_alu64_imm32(state, 0x81, 5, RSP, sizeof(uint64_t));
+    emit_alu64_imm32(state, 0x81, 5, RSP, 3*sizeof(uint64_t));
+#endif
 
-    emit_load_imm(state, RAX, (uint64_t)vm->dispatcher_cookie);
-    emit_push(state, RAX);
+    emit_load_relative(state, RAX, TARGET_PC_EXTERNAL_DISPATCHER);
+    // cmp rax, 0
+    emit_cmp_imm32(state, RAX, 0);
+    // jne skip_default_dispatcher_label
+    uint32_t skip_default_dispatcher_source = emit_jcc(state, 0x85, 0);
+
+    // Default dispatcher:
+
+    // Load the address of the helper function from the table.
+    // mov rax, idx
+    emit_alu32(state, 0xc7, 0, RAX);
+    emit4(state, idx);
+    // shl rax, 3 (i.e., multiply the index by 8 because addresses are that size on x86-64)
+    emit_alu64_imm8(state, 0xc1, 4, RAX, 3);
+
+    // lea r9, [rip + HELPER TABLE ADDRESS]
+    emit_rex(state, 1, 1, 0, 0);
+    emit1(state, 0x8d);
+    emit_modrm(state, 0, R9, 0x05);
+    emit_patchable_relative(state->offset, TARGET_LOAD_HELPER_TABLE, 0, state->loads, state->num_loads++);
+    emit_4byte_offset_placeholder(state);
+
+    // add rax, r9
+    emit_alu64(state, 0x01, R9, RAX);
+    // load rax, [rax]
+    emit_load(state, S64, RAX, RAX, 0);
+
+    // There is no index for the registered helper function. They just get
+    // 5 arguments and a context, which becomes the 6th argument to the function ...
+#if defined(_WIN32)
+    // and spills to the stack on Windows.
+    // mov qword [rsp], VOLATILE_CTXT
+    emit1(state, 0x4c);
+    emit1(state, 0x89);
+    emit1(state, 0x5c);
+    emit1(state, 0x24);
+    emit1(state, 0x00);
+#else
+    // and goes in R9 on SystemV.
+    emit_mov(state, VOLATILE_CTXT, R9);
+#endif
+
+    // jmp call_label
+    emit1(state, 0xe9);
+    uint32_t skip_external_dispatcher_source = state->offset;
+    emit_4byte_offset_placeholder(state);
+
+    // External dispatcher:
+
+    // skip_default_dispatcher_label:
+    emit_jump_target_offset(state, skip_default_dispatcher_source, state->offset);
+
+    // Using an external dispatcher. They get a total of 7 arguments. The
+    // 6th argument is the index of the function to call which ...
+
+
+#if defined(_WIN32)
+    // and spills to the stack on Windows.
+
+    // mov qword [rsp + 8], VOLATILE_CTXT
+    emit1(state, 0x4c);
+    emit1(state, 0x89);
+    emit1(state, 0x5c);
+    emit1(state, 0x24);
+    emit1(state, 0x08);
+
+    emit_load_imm(state, VOLATILE_CTXT, (uint64_t)idx);
+
+    // mov qword [rsp + 8], VOLATILE_CTXT
+    emit1(state, 0x4c);
+    emit1(state, 0x89);
+    emit1(state, 0x5c);
+    emit1(state, 0x24);
+    emit1(state, 0x00);
+#else
+    // and goes in R9 on SystemV.
+    emit_load_imm(state, R9, (uint64_t)idx);
+    // And the 7th is already spilled to the stack in the right spot because
+    // we wanted to save it -- cool (see MARKER1, above).
+
+    // Intentional no-op for 7th argument.
+#endif
+
+    // Control flow converges for call:
 
-    emit_load_imm(state, RAX, idx);
-    emit_push(state, RAX);
+    // call_label:
+    emit_jump_target_offset(state, skip_external_dispatcher_source, state->offset);
 
-    /* Windows x64 ABI spills 5th parameter to stack */
+#if defined(_WIN32)
+    /* Windows x64 ABI spills 5th parameter to stack (MARKER2) */
     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, (uint64_t)vm->dispatcher_cookie);
-    emit_push(state, R9);
-
-    emit_load_imm(state, R9, (uint64_t)idx);
 #endif
 
-    emit_load_relative(state, TARGET_PC_EXTERNAL_DISPATCHER);
-
 #ifndef UBPF_DISABLE_RETPOLINES
     emit1(state, 0xe8); // e8 is the opcode for a CALL
     emit_jump_target_address(state, TARGET_PC_RETPOLINE);
@@ -456,12 +592,12 @@ emit_dispatched_external_helper_call(struct jit_state* state, const struct ubpf_
     // Just rationalize the stack!
 
 #if defined(_WIN32)
-    /* Deallocate home register space + 3 spilled parameters + alignment space */
+    /* Deallocate home register space + (up to ) 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
+
+    emit_pop(state, VOLATILE_CTXT); // Restore register where volatile context is stored.
+    emit_pop(state, R9); // Restore r9.
 }
 
 #endif
diff --git a/vm/ubpf_vm.c b/vm/ubpf_vm.c
index 8e9341d0..a7db73a4 100644
--- a/vm/ubpf_vm.c
+++ b/vm/ubpf_vm.c
@@ -30,9 +30,9 @@
 #include "ubpf_int.h"
 #include <unistd.h>
 
-#define MAX_EXT_FUNCS 64
 #define SHIFT_MASK_32_BIT(X) ((X) & 0x1f)
 #define SHIFT_MASK_64_BIT(X) ((X) & 0x3f)
+#define DEFAULT_JITTER_BUFFER_SIZE 65536
 
 static bool
 validate(const struct ubpf_vm* vm, const struct ebpf_inst* insts, uint32_t num_insts, char** errmsg);
@@ -65,20 +65,9 @@ ubpf_set_error_print(struct ubpf_vm* vm, int (*error_printf)(FILE* stream, const
 }
 
 static uint64_t
-ubpf_default_external_dispatcher(uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4, uint64_t arg5, unsigned int index, void* cookie)
+ubpf_default_external_dispatcher(uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4, uint64_t arg5, unsigned int index, external_function_t *external_fns)
 {
-    struct ubpf_vm *vm = (struct ubpf_vm*)cookie;
-    return vm->ext_funcs[index](arg1, arg2, arg3, arg4, arg5);
-}
-
-static bool
-ubpf_default_external_validator(unsigned int index, void* cookie)
-{
-    struct ubpf_vm *vm = (struct ubpf_vm*)cookie;
-    if (index < MAX_EXT_FUNCS) {
-        return vm->ext_funcs[index] != NULL;
-    }
-    return false;
+    return external_fns[index](arg1, arg2, arg3, arg4, arg5);
 }
 
 struct ubpf_vm*
@@ -105,19 +94,20 @@ ubpf_create(void)
     vm->error_printf = fprintf;
 
 #if defined(__x86_64__) || defined(_M_X64)
-    vm->translate = ubpf_translate_x86_64;
+    vm->jit_translate = ubpf_translate_x86_64;
+    vm->jit_update_dispatcher = ubpf_jit_update_dispatcher_x86_64;
+    vm->jit_update_helper = ubpf_jit_update_helper_x86_64;
 #elif defined(__aarch64__) || defined(_M_ARM64)
-    vm->translate = ubpf_translate_arm64;
+    vm->jit_translate = ubpf_translate_arm64;
+    vm->jit_update_dispatcher = ubpf_jit_update_dispatcher_arm64;
+    vm->jit_update_helper = ubpf_jit_update_helper_arm64;
 #else
     vm->translate = ubpf_translate_null;
 #endif
     vm->unwind_stack_extension_index = -1;
 
-    // By default, we will set an internal function to be the dispatcher.
-    // If the user wants to override it, that's great (see ubpf_register_external_dispatcher).
-    vm->dispatcher = ubpf_default_external_dispatcher;
-    vm->dispatcher_validate = ubpf_default_external_validator;
-    vm->dispatcher_cookie = vm;
+    vm->jitted_result.compile_result = UBPF_JIT_COMPILE_FAILURE;
+    vm->jitter_buffer_size = DEFAULT_JITTER_BUFFER_SIZE;
     return vm;
 }
 
@@ -148,19 +138,51 @@ ubpf_register(struct ubpf_vm* vm, unsigned int idx, const char* name, external_f
     vm->ext_funcs[idx] = (ext_func)fn;
     vm->ext_func_names[idx] = name;
 
-    return 0;
+    int success = 0;
+
+    if (vm->jitted_result.compile_result == UBPF_JIT_COMPILE_SUCCESS) {
+        if (mprotect(vm->jitted, vm->jitted_size, PROT_READ | PROT_WRITE) < 0) {
+            return -1;
+        }
+
+        // Now, update!
+        if (!vm->jit_update_helper(vm, fn, idx, (uint8_t*)vm->jitted, vm->jitted_size, vm->jitted_result.external_helper_offset)) {
+            // Can't immediately stop here because we have unprotected memory!
+            success = -1;
+        }
+
+        if (mprotect(vm->jitted, vm->jitted_size, PROT_READ | PROT_EXEC) < 0) {
+            return -1;
+        }
+    }
+    return success;
 }
 
 int
 ubpf_register_external_dispatcher(
-    struct ubpf_vm* vm, external_function_dispatcher_t dispatcher, external_function_validate_t validater, void* cookie)
+    struct ubpf_vm* vm, external_function_dispatcher_t dispatcher, external_function_validate_t validater)
 {
     vm->dispatcher = dispatcher;
     vm->dispatcher_validate = validater;
-    vm->dispatcher_cookie = cookie;
 
-    /* TODO: If the code is already JIT'd, update the dispatcher's address. */
-    return 0;
+    int success = 0;
+
+    if (vm->jitted_result.compile_result == UBPF_JIT_COMPILE_SUCCESS) {
+        if (mprotect(vm->jitted, vm->jitted_size, PROT_READ | PROT_WRITE) < 0) {
+            return -1;
+        }
+
+        // Now, update!
+        if (!vm->jit_update_dispatcher(vm, dispatcher, (uint8_t*)vm->jitted, vm->jitted_size, vm->jitted_result.external_dispatcher_offset)) {
+            // Can't immediately stop here because we have unprotected memory!
+            success = -1;
+        }
+
+        if (mprotect(vm->jitted, vm->jitted_size, PROT_READ | PROT_EXEC) < 0) {
+            return -1;
+        }
+    }
+    return success;
 }
 
 int
@@ -187,16 +209,6 @@ ubpf_lookup_registered_function(struct ubpf_vm* vm, const char* name)
     return -1;
 }
 
-bool
-ubpf_validate_external_helper(const struct ubpf_vm* vm, unsigned int idx)
-{
-    if (vm->dispatcher_validate) {
-        return vm->dispatcher_validate(idx, vm->dispatcher_cookie);
-    }
-
-    return vm->ext_funcs[idx] != NULL;
-}
-
 int
 ubpf_load(struct ubpf_vm* vm, const void* code, uint32_t code_len, char** errmsg)
 {
@@ -342,6 +354,7 @@ ubpf_exec(const struct ubpf_vm* vm, void* mem, size_t mem_len, uint64_t* bpf_ret
     uint64_t _reg[16];
     uint64_t ras_index = 0;
     int return_value = -1;
+    void *external_dispatcher_cookie = mem;
 
 // Windows Kernel mode limits stack usage to 12K, so we need to allocate it dynamically.
 #if defined(NTDDI_VERSION) && defined(WINNT)
@@ -899,9 +912,11 @@ ubpf_exec(const struct ubpf_vm* vm, void* mem, size_t mem_len, uint64_t* bpf_ret
             // program was assembled with the same endianess as the host machine.
             if (inst.src == 0) {
                 // Handle call by address to external function.
-                // reg[0] = vm->ext_funcs[inst.imm](reg[1], reg[2], reg[3], reg[4], reg[5]);
-                reg[0] = vm->dispatcher(reg[1], reg[2], reg[3], reg[4], reg[5], inst.imm, (void*)vm->dispatcher_cookie);
-                // Unwind the stack if unwind extension returns success.
+                if (vm->dispatcher != NULL) {
+                    reg[0] = vm->dispatcher(reg[1], reg[2], reg[3], reg[4], reg[5], inst.imm, external_dispatcher_cookie);
+                } else {
+                    reg[0] = ubpf_default_external_dispatcher(reg[1], reg[2], reg[3], reg[4], reg[5], inst.imm, vm->ext_funcs);
+                }
                 if (inst.imm == vm->unwind_stack_extension_index && reg[0] == 0) {
                     *bpf_return_value = reg[0];
                     return_value = 0;
@@ -1111,7 +1126,7 @@ validate(const struct ubpf_vm* vm, const struct ebpf_inst* insts, uint32_t num_i
                     *errmsg = ubpf_error("invalid call immediate at PC %d", i);
                     return false;
                 }
-                if ((vm->dispatcher != NULL && !vm->dispatcher_validate(inst.imm, vm->dispatcher_cookie)) ||
+                if ((vm->dispatcher != NULL && !vm->dispatcher_validate(inst.imm, vm)) ||
                     (vm->dispatcher == NULL && !vm->ext_funcs[inst.imm])) {
                     *errmsg = ubpf_error("call to nonexistent function %u at PC %d", inst.imm, i);
                     return false;