Implement signal-handling for MMU epochs and an epoch-ending method on Store.

epoch_mmu_trap_via_signal_handler() test demonstrates the signal handler's detection of epoch interrupts and shows a do-nothing version of the asm trampoline returns control successfully to wasm.

Author: erikrose

crates/wasmtime/src/runtime/code_memory.rs

@@ -311,6 +311,18 @@ impl CodeMemory {
         &self.mmap[self.trap_data.clone()]
     }
 
+    /// Returns the address at which to resume after the given epoch check.
+    pub fn return_address_for_epoch_check(&self, check_offset: u32) -> Option<*const ()> {
+        let section = unsafe {
+            std::slice::from_raw_parts(
+                self.mmap[self.epoch_check_data.clone()].as_ptr() as *const u8,
+                self.epoch_check_data.len(),
+            )
+        };
+        return_offset_for_epoch_check(section, check_offset)
+            .map(|offset| (self.text().as_ptr().addr() + offset as usize) as *const ())
+    }
+
     /// Publishes the internal ELF image to be ready for execution.
     ///
     /// This method can only be when the image is not published (its

crates/wasmtime/src/runtime/store.rs

@@ -1178,6 +1178,12 @@ impl<T> Store<T> {
         self.inner.epoch_deadline_callback(Box::new(callback));
     }
 
+    /// Iff epoch_interruption_via_mmu() is on, end the current epoch, causing
+    /// Wasm code belonging to this store to yield to a different task.
+    pub fn end_mmu_epoch(&self) {
+        self.inner.vm_store_context().protect_interrupt_page();
+    }
+
     /// Set an exception as the currently pending exception, and
     /// return an error that propagates the throw.
     ///

crates/wasmtime/src/runtime/vm/sys/unix/signals.rs

@@ -1,7 +1,9 @@
 //! Trap handling on Unix based on POSIX signals.
 
 use crate::prelude::*;
+use crate::runtime::module::lookup_code;
 use crate::runtime::vm::traphandlers::{TrapRegisters, TrapTest, tls};
+use core::arch::naked_asm;
 use std::cell::RefCell;
 use std::io;
 use std::mem;
@@ -18,6 +20,10 @@ static mut PREV_SIGBUS: libc::sigaction = UNINIT_SIGACTION;
 static mut PREV_SIGILL: libc::sigaction = UNINIT_SIGACTION;
 static mut PREV_SIGFPE: libc::sigaction = UNINIT_SIGACTION;
 
+// From signal.h. Not yet exposed in libc. Value is valid for Linux and Mac; not
+// sure about elsewhere.
+#[cfg(all(target_os = "linux"))]
+const SEGV_ACCERR: libc::c_int = 2;
 pub struct TrapHandler;
 
 impl TrapHandler {
@@ -128,6 +134,26 @@ now.
     }
 }
 
+/// Switches tasks in response to a signal thrown under MMU-based epoch
+/// interruption.
+///
+/// Saves register state, makes a host call to switch tasks, restores state, and
+/// returns.
+#[unsafe(naked)]
+unsafe extern "C" fn task_switch_trampoline(_vmctx: usize) {
+    naked_asm!(
+        "
+        // Save regs.
+        // Call hostcall to do task switch, passing in vmctx.
+        // Restore regs (including r10).
+
+        // Resume right after the load instruction that triggered the signal
+        // handler.
+        jmp r10
+        "
+    );
+}
+
 unsafe extern "C" fn trap_handler(
     signum: libc::c_int,
     siginfo: *mut libc::siginfo_t,
@@ -148,6 +174,37 @@ unsafe extern "C" fn trap_handler(
             None => return false,
         };
 
+        // Check for segfaults meant as cues to end an epoch.
+        #[cfg(all(target_os = "linux", target_arch = "x86_64"))]
+        if signum == libc::SIGSEGV && unsafe { (*siginfo).si_code } == SEGV_ACCERR {
+            // Compare it with the offsets of epoch-check instructions as stored
+            // in the object file.
+            let ucontext = unsafe { &mut *(context as *mut libc::ucontext_t) };
+            let pc = ucontext.uc_mcontext.gregs[libc::REG_RIP as usize] as usize;
+            // Now things get expensive: we call lookup_code(), which takes a global lock.
+            if let Some((code_memory, offset_within_code)) = lookup_code(pc) {
+                // We're within a 'target_arch = "x86_64"', so we can just treat
+                // the stored little-endians as native u32s.
+                if let Some(return_address) = code_memory.return_address_for_epoch_check(
+                    offset_within_code
+                        .try_into()
+                        .expect("epoch-check location should fit in 32 bits"),
+                ) {
+                    // It is an epoch check. Arrange to resume at asm trampoline
+                    // after signal handler exits:
+                    ucontext.uc_mcontext.gregs[libc::REG_RIP as usize] =
+                        task_switch_trampoline as *const () as i64;
+                    // Put original resumption address in R10:
+                    ucontext.uc_mcontext.gregs[libc::REG_R10 as usize] = return_address as i64;
+                    // Trampoline can expect the vmctx in RDI.
+
+                    return true;
+                }
+            }
+            // Else it is an ordinary trap or the epochchecks section is somehow
+            // missing from the binary; continue on.
+        }
+
         // If we hit an exception while handling a previous trap, that's
         // quite bad, so bail out and let the system handle this
         // recursive segfault.

tests/all/epoch_mmu.rs

@@ -1,7 +1,7 @@
 #![cfg(not(miri))]
 
 use object::{LittleEndian, Object, ObjectSection, U32Bytes};
-use wasmtime::{Config, Engine};
+use wasmtime::{Config, Engine, Instance, Module, Store};
 use wasmtime_environ::obj::ELF_WASMTIME_EPOCH_CHECKS;
 
 /// Asserts that each epoch-check offset encoded into the binary points to the
@@ -60,3 +60,40 @@ fn epoch_check_offsets() {
         "Neither check's load instruction uses R12 of RSP as its source, so all length bits should be 0."
     );
 }
+
+/// Runs a wasm function that loops forever with MMU-based epoch interruption
+/// enabled. Incrementing the epoch past the deadline triggers the signal
+/// handler (`trap_handler`) which converts the SIGSEGV into a trap.
+#[test]
+fn epoch_mmu_trap_via_signal_handler() {
+    let mut config = Config::new();
+    config.epoch_interruption_via_mmu(true);
+    let engine = Engine::new(&config).unwrap();
+    let module = Module::new(
+        &engine,
+        r#"(module
+             (memory 0)
+             (func (export "answer") (result i32)
+                i32.const 42
+             )
+           )"#,
+    )
+    .unwrap();
+
+    // Trap as soon as the first epoch check is encountered, in the function
+    // prologue. Recall that MMU-based epochs don't operate based on a numeric
+    // deadline but on an external entity protecting the memory page, typically
+    // on a timer.
+    let mut store = Store::new(&engine, ());
+    store.epoch_deadline_trap(); // Allegedly the default.
+    // Protect that page:
+    store.end_mmu_epoch();
+
+    let instance = Instance::new(&mut store, &module, &[]).unwrap();
+    let func = instance
+        .get_typed_func::<(), i32>(&mut store, "answer")
+        .unwrap();
+
+    let result = func.call(&mut store, ()).unwrap();
+    assert_eq!(result, 42);
+}