wasmtime platform: make init_traps catch #UD

This makes wasmtime_init_traps actually register an exception handler,
and adds handling for #UD to recognize it as a trap that should be
forwraded to wasmtime.  More traps will need to be added in the future
in order to ensure correctness.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/wasm_runtime/src/component.rs

@@ -33,6 +33,8 @@ use spin::Mutex;
 use wasmtime::component::{Component, Instance, Linker};
 use wasmtime::{Config, Engine, Store};
 
+use crate::platform;
+
 static CUR_ENGINE: Mutex<Option<Engine>> = Mutex::new(None);
 static CUR_LINKER: Mutex<Option<Linker<()>>> = Mutex::new(None);
 static CUR_STORE: Mutex<Option<Store<()>>> = Mutex::new(None);
@@ -74,6 +76,8 @@ fn load_wasm_module(function_call: &FunctionCall) -> Result<Vec<u8>> {
 
 #[no_mangle]
 pub extern "C" fn hyperlight_main() {
+    platform::register_page_fault_handler();
+
     let mut config = Config::new();
     config.memory_reservation(0);
     config.memory_guard_size(0);

src/wasm_runtime/src/module.rs

@@ -32,7 +32,7 @@ use hyperlight_guest_bin::host_comm::print_output_with_host_print;
 use spin::Mutex;
 use wasmtime::{Config, Engine, Linker, Module, Store, Val};
 
-use crate::{hostfuncs, marshal, wasip1};
+use crate::{hostfuncs, marshal, platform, wasip1};
 
 static CUR_ENGINE: Mutex<Option<Engine>> = Mutex::new(None);
 static CUR_LINKER: Mutex<Option<Linker<()>>> = Mutex::new(None);
@@ -141,6 +141,8 @@ fn load_wasm_module(function_call: &FunctionCall) -> Result<Vec<u8>> {
 #[no_mangle]
 #[allow(clippy::fn_to_numeric_cast)] // GuestFunctionDefinition expects a function pointer as i64
 pub extern "C" fn hyperlight_main() {
+    platform::register_page_fault_handler();
+
     register_function(GuestFunctionDefinition::new(
         "PrintOutput".to_string(),
         vec![ParameterType::String],

src/wasm_runtime/src/platform.rs

@@ -16,7 +16,9 @@ limitations under the License.
 
 use alloc::alloc::{alloc, dealloc, Layout};
 use core::ffi::c_void;
-use core::sync::atomic::{AtomicPtr, Ordering};
+use core::sync::atomic::{AtomicPtr, AtomicU64, Ordering};
+
+use hyperlight_guest_bin::exceptions::handler;
 
 // Wasmtime Embedding Interface
 
@@ -64,10 +66,52 @@ pub extern "C" fn wasmtime_page_size() -> usize {
 #[allow(non_camel_case_types)] // we didn't choose the name!
 type wasmtime_trap_handler_t =
     extern "C" fn(ip: usize, fp: usize, has_faulting_addr: bool, faulting_addr: usize);
+static WASMTIME_REQUESTED_TRAP_HANDLER: AtomicU64 = AtomicU64::new(0);
+fn wasmtime_trap_handler(
+    exception_number: u64,
+    info: *mut handler::ExceptionInfo,
+    ctx: *mut handler::Context,
+    _page_fault_address: u64,
+) -> bool {
+    let requested_handler = WASMTIME_REQUESTED_TRAP_HANDLER.load(Ordering::Relaxed);
+    if requested_handler != 0 {
+        #[allow(clippy::collapsible_if)] // We will add more cases
+        if exception_number == 6 {
+            // #UD
+            // we assume that handle_trap always longjmp's away, so don't bother
+            // setting up a terribly proper stack frame
+            unsafe {
+                let orig_rip = (&raw mut (*info).rip).read_volatile();
+                (&raw mut (*info).rip).write_volatile(requested_handler);
+                // TODO: This only works on amd64 sysv
+                (&raw mut (*ctx).gprs[9]).write_volatile(orig_rip);
+                let orig_rbp = (&raw mut (*ctx).gprs[8]).read_volatile();
+                (&raw mut (*ctx).gprs[10]).write_volatile(orig_rbp);
+                (&raw mut (*ctx).gprs[11]).write_volatile(0);
+                (&raw mut (*ctx).gprs[12]).write_volatile(0);
+            }
+            return true;
+        }
+        // TODO: Add handlers for any other traps that wasmtime needs
+    }
+    false
+}
 
-// TODO: Correctly handle traps.
 #[no_mangle]
-pub extern "C" fn wasmtime_init_traps(_handler: wasmtime_trap_handler_t) -> i32 {
+pub extern "C" fn wasmtime_init_traps(handler: wasmtime_trap_handler_t) -> i32 {
+    WASMTIME_REQUESTED_TRAP_HANDLER.store(handler as usize as u64, Ordering::Relaxed);
+    // On amd64, vector 6 is #UD
+    // See AMD64 Architecture Programmer's Manual, Volume 2
+    //    §8.2 Vectors, p. 245
+    //      Table 8-1: Interrupt Vector Source and Cause
+    handler::handlers[6].store(wasmtime_trap_handler as usize as u64, Ordering::Release);
+    // TODO: Add handlers for any other traps that wasmtime needs,
+    //       probably including at least some floating-point
+    //       exceptions
+    // TODO: Ensure that invalid accesses to mprotect()'d regions also
+    //       need to trap, although those will need to go through the
+    //       page fault handler instead of using this handler that
+    //       takes over the exception.
     0
 }