[stack-switching] Fix ASan-host interop

crates/cranelift/src/alias_region.rs

@@ -85,6 +85,11 @@ enum VmType {
     VMGcHeader,
     VMDrcHeader,
     VMCopyingHeader,
+    #[allow(
+        dead_code,
+        reason = "generated uniformly for all VM types via `for_each_vm_type!`"
+    )]
+    VMAsanStackSupport,
 }
 
 /// A key that uniquely identifies an alias region across an entire compilation.

crates/cranelift/src/func_environ/stack_switching/instructions.rs

@@ -30,10 +30,10 @@ fn emit_stack_switch<'a>(
     store_context_ptr: ir::Value,
     load_context_ptr: ir::Value,
     payload: ir::Value,
-    asan_target_csi: impl FnOnce(
+    asan_stack_information: impl FnOnce(
         &mut crate::func_environ::FuncEnvironment<'a>,
         &mut FunctionBuilder,
-    ) -> ir::Value,
+    ) -> (ir::Value, ir::Value),
 ) -> ir::Value {
     if !env.compiler.tunables().asan_stack_switching {
         return builder
@@ -44,7 +44,7 @@ fn emit_stack_switch<'a>(
     // ASan-aware stack switching.
     // The `asan_target_csi` is provided as a function to lazily load
     // the necessary ASan bookkeeping.
-    let target_csi = asan_target_csi(env, builder);
+    let (source_csi, target_csi) = asan_stack_information(env, builder);
     let pointer_type = env.pointer_type();
     let slot = env.get_or_create_asan_fake_stack_slot(builder);
     let fake_stack_save = builder.ins().stack_addr(pointer_type, slot, 0);
@@ -60,7 +60,7 @@ fn emit_stack_switch<'a>(
     );
     builder.ins().call(
         asan_start_switch_fiber,
-        &[vmctx, fake_stack_save, target_csi],
+        &[vmctx, fake_stack_save, source_csi, target_csi],
     );
 
     let result = builder
@@ -74,7 +74,7 @@ fn emit_stack_switch<'a>(
     );
     builder
         .ins()
-        .call(asan_finish_switch_fiber, &[vmctx, fake_stack]);
+        .call(asan_finish_switch_fiber, &[vmctx, fake_stack, source_csi]);
 
     result
 }
@@ -1914,7 +1914,12 @@ fn translate_resume_impl<'a>(
             control_context_ptr,
             control_context_ptr,
             resume_payload,
-            |env, builder| last_ancestor.common_stack_information(env, builder).address,
+            |env, builder| {
+                (
+                    parent_csi.address,
+                    last_ancestor.common_stack_information(env, builder).address,
+                )
+            },
         );
 
         // At this point we know nothing about the continuation that just
@@ -2234,9 +2239,14 @@ pub(crate) fn translate_suspend<'a>(
         control_context_ptr,
         suspend_payload,
         |env, builder| {
-            handler_stack_chain
-                .get_common_stack_information(env, builder)
-                .address
+            (
+                end_of_chain_contref
+                    .common_stack_information(env, builder)
+                    .address,
+                handler_stack_chain
+                    .get_common_stack_information(env, builder)
+                    .address,
+            )
         },
     );
 
@@ -2511,9 +2521,14 @@ pub(crate) fn translate_switch<'a>(
             tmp_control_context,
             switch_payload,
             |env, builder| {
-                switchee_contref_last_ancestor
-                    .common_stack_information(env, builder)
-                    .address
+                (
+                    switcher_contref_last_ancestor
+                        .common_stack_information(env, builder)
+                        .address,
+                    switchee_contref_last_ancestor
+                        .common_stack_information(env, builder)
+                        .address,
+                )
             },
         )
     };

crates/environ/src/builtin.rs

@@ -191,9 +191,14 @@ macro_rules! foreach_builtin_function {
             // Helpers used around stack switch instructions in
             // ASan-instrumented builds.
             #[cfg(feature = "stack-switching")]
-            asan_start_switch_fiber(vmctx: vmctx, fake_stack_save: pointer, target_csi: pointer);
+            asan_start_switch_fiber(
+                vmctx: vmctx,
+                fake_stack_save: pointer,
+                source_csi: pointer,
+                target_csi: pointer
+            );
             #[cfg(feature = "stack-switching")]
-            asan_finish_switch_fiber(vmctx: vmctx, fake_stack: pointer);
+            asan_finish_switch_fiber(vmctx: vmctx, fake_stack: pointer, current_csi: pointer);
         }
     };
 }

crates/environ/src/vmoffsets.rs

@@ -71,6 +71,9 @@ macro_rules! define_vm_type_offsets {
     (@size ($p:expr) VMStackLimits) => { u32::from(($p).vm_stack_limits().size()) };
     (@size ($p:expr) VMHostArray) => { u32::from(($p).vm_host_array().size()) };
     (@size ($p:expr) VMPayloads) => { u32::from(($p).vm_payloads().size()) };
+    (@size ($p:expr) VMAsanStackSupport) => {
+        u32::from(($p).vm_asan_stack_support().size())
+    };
     (@size ($p:expr) VMCommonStackInformation) => {
         u32::from(($p).vm_common_stack_information().size())
     };
@@ -117,6 +120,9 @@ macro_rules! define_vm_type_offsets {
     (@align ($p:expr) VMStackLimits) => { u32::from(($p).vm_stack_limits().align()) };
     (@align ($p:expr) VMHostArray) => { u32::from(($p).vm_host_array().align()) };
     (@align ($p:expr) VMPayloads) => { u32::from(($p).vm_payloads().align()) };
+    (@align ($p:expr) VMAsanStackSupport) => {
+        u32::from(($p).vm_asan_stack_support().align())
+    };
     (@align ($p:expr) VMCommonStackInformation) => {
         u32::from(($p).vm_common_stack_information().align())
     };

crates/environ/src/vmtypes.rs

@@ -617,6 +617,28 @@ macro_rules! for_each_vm_type {
                 pub gc_ref_data: Option<VmPtr<u8>>,
             }
 
+            /// AddressSanitizer bookkeeping associated with a stack.
+            #[derive(Debug, Clone, Default)]
+            #[repr(C)]
+            #[snake_name = vm_asan_stack_support]
+            pub struct VMAsanStackSupport {
+                /// Bottom of the usable stack range reported to ASan,
+                /// or `None` until ASan discovers an initial stack's
+                /// bounds. Technically only needed by the runtime,
+                /// but we are using `VmPtr` here to avoid introducing
+                /// macro rules for `*mut u8` just for the sake of ASan
+                /// builds.
+                pub stack_bottom: Option<VmPtr<u8>>,
+
+                /// Size of the usable stack range reported to ASan.
+                pub stack_size: usize,
+
+                /// The stack whose bounds ASan will report when the next
+                /// stack-switch handshake is completed on this stack.
+                /// `None` when no handshake is pending.
+                pub pending_source_csi: Option<VmPtr<VMCommonStackInformation>>,
+            }
+
             /// The information saved for every stack, whether it is a
             /// continuation's or the initial stack's.
             #[derive(Debug, Clone)]
@@ -638,16 +660,9 @@ macro_rules! for_each_vm_type {
                 /// The index within `handlers` of the first `switch` handler.
                 pub first_switch_handler_index: u32,
 
-                /// Bottom of the usable stack range reported to ASan,
-                /// or `None` until ASan discovers an initial stack's
-                /// bounds.  Technically only needed by the runtime,
-                /// but we are using `VmPtr` here to avoid introducing
-                /// macro rules for `*mut u8` just for the sake of ASan
-                /// builds.
-                pub asan_stack_bottom: Option<VmPtr<u8>>,
-
-                /// Size of the usable stack range reported to ASan.
-                pub asan_stack_size: usize,
+                /// AddressSanitizer bookkeeping for this stack.
+                #[aggregate]
+                pub asan: VMAsanStackSupport,
             }
 
             /// A continuation.
@@ -663,13 +678,13 @@ macro_rules! for_each_vm_type {
                 #[aggregate]
                 pub parent_chain: VMStackChain,
 
+                /// Revision counter.
+                pub revision: usize,
+
                 /// The end of this continuation's parent chain while it is
                 /// `Suspended` or `Fresh`, and `None` while it is running.
                 pub last_ancestor: Option<VmPtr<VMContRef>>,
 
-                /// Revision counter.
-                pub revision: usize,
-
                 /// The stack this continuation runs on.
                 #[aggregate]
                 pub stack: VMContinuationStack,

crates/wasmtime/src/runtime/store.rs

@@ -2166,11 +2166,11 @@ at https://bytecodealliance.org/security.
             let asan_range = stack
                 .asan_range()
                 .expect("supported continuation stacks have a usable range");
-            continuation.common_stack_information.asan_stack_bottom = Some(vm::VmPtr::from(
+            continuation.common_stack_information.asan.stack_bottom = Some(vm::VmPtr::from(
                 NonNull::new(asan_range.start as *mut u8)
                     .expect("a continuation stack's ASan range must have a non-null bottom"),
             ));
-            continuation.common_stack_information.asan_stack_size = asan_range.len();
+            continuation.common_stack_information.asan.stack_size = asan_range.len();
         }
         continuation.stack = stack;
         let ptr = continuation.deref_mut() as *mut VMContRef;

crates/wasmtime/src/runtime/vm/libcalls.rs

@@ -1195,18 +1195,26 @@ unsafe fn asan_start_switch_fiber(
     _store: &mut dyn VMStore,
     _instance: InstanceId,
     fake_stack_save: *mut u8,
+    source_csi: *mut u8,
     target_csi: *mut u8,
 ) {
-    unsafe { crate::vm::stack_switching::asan::start_switch_fiber(fake_stack_save, target_csi) }
+    unsafe {
+        crate::vm::stack_switching::asan::start_switch_fiber(
+            fake_stack_save,
+            source_csi,
+            target_csi,
+        )
+    }
 }
 
 #[cfg(feature = "stack-switching")]
 unsafe fn asan_finish_switch_fiber(
     _store: &mut dyn VMStore,
     _instance: InstanceId,
     fake_stack: *mut u8,
+    current_csi: *mut u8,
 ) {
-    unsafe { crate::vm::stack_switching::asan::finish_switch_fiber(fake_stack) }
+    unsafe { crate::vm::stack_switching::asan::finish_switch_fiber(fake_stack, current_csi) }
 }
 
 #[cfg(feature = "gc")]

crates/wasmtime/src/runtime/vm/stack_switching.rs

@@ -90,8 +90,7 @@ impl VMCommonStackInformation {
             state: VMStackState::Running,
             handlers: VMHostArray::empty(),
             first_switch_handler_index: 0,
-            asan_stack_bottom: None,
-            asan_stack_size: 0,
+            asan: Default::default(),
         }
     }
 }
@@ -158,8 +157,7 @@ impl VMContRef {
             state,
             handlers,
             first_switch_handler_index: 0,
-            asan_stack_bottom: None,
-            asan_stack_size: 0,
+            asan: Default::default(),
         };
         let parent_chain = VMStackChain::Absent;
         let last_ancestor = None;