Fix compatibility with the Go runtime on Windows for exceptions

crates/wasmtime/src/runtime/vm/sys/windows/vectored_exceptions.rs

@@ -1,5 +1,92 @@
+//! Implementation of handling hardware traps generated by wasm (e.g. segfaults)
+//! on Windows.
+//!
+//! This module is implemented with Windows Vectored Exception Handling which
+//! is, I think, implemented on top of Structured Exception Handling (SEH). This
+//! is distinct from Unix signals where instead of a single global handler
+//! there's a list of vectored exception handlers which is managed by the
+//! Windows runtime. This list is sort of like a `VecDeque` where you can push
+//! on either end, and then you're able to remove any pushed entry later on.
+//!
+//! Windows's behavior here seems to first execute the ordered list of vectored
+//! exception handlers until one returns `EXCEPTION_CONTINUE_EXECUTION`. If this
+//! list is exhausted then it seems to go to default SEH routines which abort
+//! the process.
+//!
+//! Another interesting part, however, is that once an exception handler returns
+//! `EXCEPTION_CONTINUE_EXECUTION` Windows will then consult a similar deque of
+//! "continue handlers". These continue handlers have the same signature as the
+//! exception handlers and are managed with similar functions
+//! (`AddVectoredContinueHandler` instead of `AddVectoredExceptionHandler`). The
+//! difference here is that the first continue handler to return
+//! `EXCEPTION_CONTINUE_EXECUTION` will short-circuit the rest of the list. If
+//! none of them return `EXCEPTION_CONTINUE_EXECUTION` then the programs
+//! still resumes as normal.
+//!
+//! # Wasmtime's implementation
+//!
+//! Wasmtime installs both an exception handler and a continue handler. The
+//! purpose of the exception handler is to return `EXCEPTION_CONTINUE_EXECUTION`
+//! for any wasm exceptions that we want to catch (e.g. divide-by-zero, out of
+//! bounds memory accesses in wasm, `unreachable` via illegal instruction, etc).
+//! Note that this exception handler is installed at the front of the list to
+//! try to run it as soon as possible as, if we catch something, we want to
+//! bypass all other handlers.
+//!
+//! Wasmtime then also installs a continue handler, also at the front of the
+//! list, where the sole purpose of the continue handler is to also return
+//! `EXCEPTION_CONTINUE_EXECUTION` and bypass the rest of the continue handler
+//! list to get back to wasm ASAP. The reason for this is explained in the next
+//! section.
+//!
+//! To implement the continue handler in Wasmtime a thread-local variable
+//! `LAST_EXCEPTION_PC` is used here which is set during the exception handler
+//! and then tested during the continue handler. If it matches the current PC
+//! then it's assume that Wasmtime is the one that processed the exception and
+//! the `EXCEPTION_CONTINUE_EXECUTION` is returned.
+//!
+//! # Why both an exception and continue handler?
+//!
+//! All of Wasmtime's tests in this repository will pass if the continue handler
+//! is removed, so why have it? The primary reason at this time is integration
+//! with the Go runtime as discovered in the `wasmtime-go` embedding.
+//!
+//! Go's behavior for exceptions is:
+//!
+//! * An exception handler is installed at the front of the list of handlers
+//!   which looks for Go-originating exceptions. If one is found it returns
+//!   `EXCEPTION_CONTINUE_EXECUTION`, otherwise it forwards along with
+//!   `EXCEPTION_CONTINUE_SEARCH`. Wasmtime exceptions will properly go through
+//!   this handler and then hit Wasmtime's handler, so no issues yet.
+//!
+//! * Go then additionally installs *two* continue handlers. One at the front of
+//!   the list and one at the end. The continue handler at the front of the list
+//!   looks for Go-related exceptions dealing with things like
+//!   async/preemption/etc to resume execution back into Go. This means that the
+//!   handler will return `EXCEPTION_CONTINUE_EXECUTION` sometimes for
+//!   Go-specific reasons, and otherwise the handler returns
+//!   `EXCEPTION_CONTINUE_SEARCH`. As before this isn't a problem for Wasmtime
+//!   as nothing happens for non-Go-related exceptions.
+//!
+//! * The problem with Go is the second, final, continue handler. This will, by
+//!   default, abort the process for all exceptions whether or not they're Go
+//!   related. This seems to have some logic for whether or not Go was built as
+//!   a library or dylib but that seem to apply for Go-built executables (e.g.
+//!   `go test` in the wasmtime-go repository). This second handler is the
+//!   problematic one because in Wasmtime we "catch" the exception in the
+//!   exception handler function but then the process still aborts as all
+//!   continue handlers are run, including Go's abort-the-process handler.
+//!
+//! Thus the reason Wasmtime has a continue handler in addition to an exception
+//! handler. By installing a high-priority continue handler that pairs with the
+//! high-priority exception handler we can ensure that, for example, Go's
+//! fallback continue handler is never executed.
+//!
+//! This is all... a bit... roundabout. Sorry.
+
 use crate::prelude::*;
 use crate::runtime::vm::traphandlers::{TrapRegisters, TrapTest, tls};
+use std::cell::Cell;
 use std::ffi::c_void;
 use std::io;
 use windows_sys::Win32::Foundation::*;
@@ -9,25 +96,40 @@ use windows_sys::Win32::System::Diagnostics::Debug::*;
 pub type SignalHandler = Box<dyn Fn(*mut EXCEPTION_POINTERS) -> bool + Send + Sync>;
 
 pub struct TrapHandler {
-    handle: *mut c_void,
+    exception_handler: *mut c_void,
+    continue_handler: *mut c_void,
 }
 
 unsafe impl Send for TrapHandler {}
 unsafe impl Sync for TrapHandler {}
 
 impl TrapHandler {
     pub unsafe fn new(_macos_use_mach_ports: bool) -> TrapHandler {
-        // our trap handler needs to go first, so that we can recover from
+        // Our trap handler needs to go first, so that we can recover from
         // wasm faults and continue execution, so pass `1` as a true value
         // here.
-        let handle = unsafe { AddVectoredExceptionHandler(1, Some(exception_handler)) };
-        if handle.is_null() {
+        //
+        // Note that this is true for the "continue" handler as well since we
+        // want to short-circuit as many other continue handlers as we can on
+        // wasm exceptions.
+        let exception_handler = unsafe { AddVectoredExceptionHandler(1, Some(exception_handler)) };
+        if exception_handler.is_null() {
             panic!(
                 "failed to add exception handler: {}",
                 io::Error::last_os_error()
             );
         }
-        TrapHandler { handle }
+        let continue_handler = unsafe { AddVectoredContinueHandler(1, Some(continue_handler)) };
+        if continue_handler.is_null() {
+            panic!(
+                "failed to add continue handler: {}",
+                io::Error::last_os_error()
+            );
+        }
+        TrapHandler {
+            exception_handler,
+            continue_handler,
+        }
     }
 
     pub fn validate_config(&self, _macos_use_mach_ports: bool) {}
@@ -36,18 +138,36 @@ impl TrapHandler {
 impl Drop for TrapHandler {
     fn drop(&mut self) {
         unsafe {
-            let rc = RemoveVectoredExceptionHandler(self.handle);
+            let rc = RemoveVectoredExceptionHandler(self.exception_handler);
             if rc == 0 {
                 eprintln!(
                     "failed to remove exception handler: {}",
                     io::Error::last_os_error()
                 );
                 libc::abort();
             }
+            let rc = RemoveVectoredContinueHandler(self.continue_handler);
+            if rc == 0 {
+                eprintln!(
+                    "failed to remove continue handler: {}",
+                    io::Error::last_os_error()
+                );
+                libc::abort();
+            }
         }
     }
 }
 
+std::thread_local! {
+    static LAST_EXCEPTION_PC: Cell<usize> = const { Cell::new(0) };
+}
+
+/// Wasmtime's exception handler for Windows. See module docs for more.
+///
+/// # Safety
+///
+/// Invoked by Windows' vectored exception system; should not be called by
+/// anyone else.
 #[allow(
     clippy::cast_possible_truncation,
     reason = "too fiddly to handle and wouldn't help much anyway"
@@ -116,6 +236,7 @@ unsafe extern "system" fn exception_handler(exception_info: *mut EXCEPTION_POINT
             TrapTest::HandledByEmbedder => EXCEPTION_CONTINUE_EXECUTION,
             TrapTest::Trap(handler) => {
                 let context = unsafe { exception_info.ContextRecord.as_mut().unwrap() };
+                LAST_EXCEPTION_PC.with(|s| s.set(handler.pc));
                 cfg_if::cfg_if! {
                     if #[cfg(target_arch = "x86_64")] {
                         context.Rip = handler.pc as _;
@@ -139,3 +260,30 @@ unsafe extern "system" fn exception_handler(exception_info: *mut EXCEPTION_POINT
         }
     })
 }
+
+/// See module docs for more information on what this is doing.
+///
+/// # Safety
+///
+/// Invoked by Windows' vectored exception system; should not be called by
+/// anyone else.
+unsafe extern "system" fn continue_handler(exception_info: *mut EXCEPTION_POINTERS) -> i32 {
+    let context = unsafe { &(*(*exception_info).ContextRecord) };
+    let last_exception_pc = LAST_EXCEPTION_PC.with(|s| s.replace(0));
+
+    cfg_if::cfg_if! {
+        if #[cfg(target_arch = "x86_64")] {
+            let context_pc = context.Rip as usize;
+        } else if #[cfg(target_arch = "aarch64")] {
+            let context_pc = context.Pc as usize;
+        } else {
+            compile_error!("unsupported platform");
+        }
+    }
+
+    if last_exception_pc == context_pc {
+        EXCEPTION_CONTINUE_EXECUTION
+    } else {
+        EXCEPTION_CONTINUE_SEARCH
+    }
+}