Merge pull request #4658 from wasmerio/fix/rewind-buffer-take-3-stack…

…-low

Third solution to the rewind buffer issue: Read lower end of stack from __stack_low or __data_end

lib/journal/src/concrete/archived.rs

@@ -1802,6 +1802,4 @@ pub struct JournalWasiMemoryLayout {
     pub stack_lower: u64,
     pub guard_size: u64,
     pub stack_size: u64,
-    pub tls_base: u64,
-    pub tls_size: u64,
 }

lib/journal/src/concrete/archived_from.rs

@@ -616,8 +616,6 @@ impl From<JournalWasiMemoryLayout> for WasiMemoryLayout {
             stack_lower: value.stack_lower,
             guard_size: value.guard_size,
             stack_size: value.stack_size,
-            tls_base: value.tls_base,
-            tls_size: value.tls_size,
         }
     }
 }
@@ -629,8 +627,6 @@ impl From<&'_ ArchivedJournalWasiMemoryLayout> for WasiMemoryLayout {
             stack_lower: value.stack_lower,
             guard_size: value.guard_size,
             stack_size: value.stack_size,
-            tls_base: value.tls_base,
-            tls_size: value.tls_size,
         }
     }
 }
@@ -642,8 +638,6 @@ impl From<WasiMemoryLayout> for JournalWasiMemoryLayout {
             stack_lower: value.stack_lower,
             guard_size: value.guard_size,
             stack_size: value.stack_size,
-            tls_base: value.tls_base,
-            tls_size: value.tls_size,
         }
     }
 }

lib/journal/src/concrete/log_file.rs

@@ -345,8 +345,6 @@ mod tests {
                     stack_lower: 1024,
                     guard_size: 16,
                     stack_size: 1024,
-                    tls_base: 16,
-                    tls_size: 32,
                 },
                 start: wasmer_wasix_types::wasix::ThreadStartType::MainThread,
             })
@@ -376,8 +374,6 @@ mod tests {
                     stack_lower: 1024,
                     guard_size: 16,
                     stack_size: 1024,
-                    tls_base: 16,
-                    tls_size: 32,
                 },
                 start: wasmer_wasix_types::wasix::ThreadStartType::MainThread,
             })
@@ -428,8 +424,6 @@ mod tests {
                     stack_lower: 1024,
                     guard_size: 16,
                     stack_size: 1024,
-                    tls_base: 16,
-                    tls_size: 32,
                 },
                 start: wasmer_wasix_types::wasix::ThreadStartType::MainThread,
             })
@@ -477,8 +471,6 @@ mod tests {
                     stack_lower: 1024,
                     guard_size: 16,
                     stack_size: 1024,
-                    tls_base: 16,
-                    tls_size: 32,
                 },
                 start: wasmer_wasix_types::wasix::ThreadStartType::MainThread,
             })

lib/journal/src/concrete/tests.rs

@@ -75,8 +75,6 @@ pub fn test_record_set_thread() {
             stack_lower: 1024,
             guard_size: 16,
             stack_size: 1024,
-            tls_base: 16,
-            tls_size: 32,
         },
         start: wasmer_wasix_types::wasix::ThreadStartType::MainThread,
     });

lib/wasi-types/src/wasix/mod.rs

@@ -26,8 +26,4 @@ pub struct WasiMemoryLayout {
     pub guard_size: u64,
     /// Total size of the stack
     pub stack_size: u64,
-    /// Base address of the TLS area
-    pub tls_base: u64,
-    /// Total size of the TLS area
-    pub tls_size: u64,
 }

lib/wasix/src/state/env.rs

@@ -59,6 +59,15 @@ pub struct WasiInstanceHandles {
     /// Points to the current location of the memory stack pointer
     pub(crate) stack_pointer: Option<Global>,
 
+    /// Points to the end of the data section
+    pub(crate) data_end: Option<Global>,
+
+    /// Points to the lower end of the stack
+    pub(crate) stack_low: Option<Global>,
+
+    /// Points to the higher end of the stack
+    pub(crate) stack_high: Option<Global>,
+
     /// Main function that will be invoked (name = "_start")
     #[derivative(Debug = "ignore")]
     pub(crate) start: Option<TypedFunction<(), ()>>,
@@ -140,6 +149,21 @@ impl WasiInstanceHandles {
                 .get_global("__stack_pointer")
                 .map(|a| a.clone())
                 .ok(),
+            data_end: instance
+                .exports
+                .get_global("__data_end")
+                .map(|a| a.clone())
+                .ok(),
+            stack_low: instance
+                .exports
+                .get_global("__stack_low")
+                .map(|a| a.clone())
+                .ok(),
+            stack_high: instance
+                .exports
+                .get_global("__stack_high")
+                .map(|a| a.clone())
+                .ok(),
             start: instance.exports.get_typed_function(store, "_start").ok(),
             initialize: instance
                 .exports

lib/wasix/src/state/func_env.rs

@@ -25,17 +25,6 @@ use crate::{
 const DEFAULT_STACK_SIZE: u64 = 1_048_576u64;
 const DEFAULT_STACK_BASE: u64 = DEFAULT_STACK_SIZE;
 
-/// The default size and base address of the TLS area for modules that
-/// don't export __tls_base and __tls_size. This is the default value
-/// used by clang-15 at the time of writing this code.
-const DEFAULT_TLS_BASE: u64 = 1024u64;
-
-/// This is merely an attempt at a guess. We're avoiding the __pthread
-/// struct that each thread has (hoping it comes first), but there's
-/// no telling how many thread statics a module has and how much space
-/// they take.
-const DEFAULT_TLS_SIZE: u64 = 1024u64;
-
 #[derive(Clone, Debug)]
 pub struct WasiFunctionEnv {
     pub env: FunctionEnv<WasiEnv>,
@@ -158,20 +147,11 @@ impl WasiFunctionEnv {
             |v| Ok(v.clone()),
         )?;
 
-        let tls_base = instance
-            .exports
-            .get_global("__tls_base")
-            .map(|a| a.clone())
-            .ok();
-
-        let tls_size = instance
-            .exports
-            .get_global("__tls_size")
-            .map(|a| a.clone())
-            .ok();
-
         let new_inner = WasiInstanceHandles::new(memory, store, instance);
         let stack_pointer = new_inner.stack_pointer.clone();
+        let data_end = new_inner.data_end.clone();
+        let stack_low = new_inner.stack_low.clone();
+        let stack_high = new_inner.stack_high.clone();
 
         let env = self.data_mut(store);
         env.set_inner(new_inner);
@@ -181,7 +161,13 @@ impl WasiFunctionEnv {
         // If the stack offset and size is not set then do so
         if update_layout {
             // Set the base stack
-            let stack_base = if let Some(stack_pointer) = stack_pointer {
+            let stack_base = if let Some(stack_high) = stack_high {
+                match stack_high.get(store) {
+                    wasmer::Value::I32(a) => a as u64,
+                    wasmer::Value::I64(a) => a as u64,
+                    _ => DEFAULT_STACK_BASE,
+                }
+            } else if let Some(stack_pointer) = stack_pointer {
                 match stack_pointer.get(store) {
                     wasmer::Value::I32(a) => a as u64,
                     wasmer::Value::I64(a) => a as u64,
@@ -190,40 +176,48 @@ impl WasiFunctionEnv {
             } else {
                 DEFAULT_STACK_BASE
             };
+
             if stack_base == 0 {
                 return Err(ExportError::Missing(
-                    "stack_pointer is not set to the upper stack range".to_string(),
+                    "stack_high or stack_pointer is not set to the upper stack range".to_string(),
                 ));
             }
 
-            let tls_base = if let Some(tls_base) = tls_base {
-                match tls_base.get(store) {
+            let mut stack_lower = if let Some(stack_low) = stack_low {
+                match stack_low.get(store) {
                     wasmer::Value::I32(a) => a as u64,
                     wasmer::Value::I64(a) => a as u64,
-                    _ => DEFAULT_TLS_BASE,
+                    _ => 0,
                 }
-            } else {
-                DEFAULT_TLS_BASE
-            };
-
-            let tls_size = if let Some(tls_size) = tls_size {
-                match tls_size.get(store) {
+            } else if let Some(data_end) = data_end {
+                match data_end.get(store) {
                     wasmer::Value::I32(a) => a as u64,
                     wasmer::Value::I64(a) => a as u64,
-                    _ => DEFAULT_TLS_SIZE,
+                    _ => 0,
                 }
             } else {
-                DEFAULT_TLS_SIZE
+                // clang-16 and higher generate the `__stack_low` global, and it can be exported with
+                // `-Wl,--export=__stack_low`. clang-15 generates `__data_end`, which should be identical
+                // and can be exported if `__stack_low` is not available.
+                tracing::warn!("Missing both __stack_low and __data_end exports, unwinding may cause memory corruption");
+                0
             };
 
+            if stack_lower >= stack_base {
+                tracing::warn!(
+                    "Detected lower end of stack to be above higher end, ignoring stack_lower; \
+                    unwinding may cause memory corruption"
+                );
+                stack_lower = 0;
+            }
+
             // Update the stack layout which is need for asyncify
             let env = self.data_mut(store);
             let tid = env.tid();
             let layout = &mut env.layout;
             layout.stack_upper = stack_base;
+            layout.stack_lower = stack_lower;
             layout.stack_size = layout.stack_upper - layout.stack_lower;
-            layout.tls_base = tls_base;
-            layout.tls_size = tls_size;
 
             // Replace the thread object itself
             env.thread.set_memory_layout(layout.clone());

lib/wasix/src/syscalls/mod.rs

@@ -870,21 +870,6 @@ pub(crate) fn get_stack_upper(env: &WasiEnv) -> u64 {
     env.layout.stack_upper
 }
 
-fn get_unwind_buffer_start(env: &WasiEnv) -> u64 {
-    let stack_lower = env.layout.stack_lower;
-    let stack_upper = env.layout.stack_upper;
-    let tls_base = env.layout.tls_base;
-
-    if stack_upper.saturating_sub(tls_base) < tls_base.saturating_sub(stack_lower) {
-        // tls_base is closer to stack_upper, so we assume TLS lives at the top of the stack.
-        // This is the case for threads spawned by wasix-libc.
-        stack_lower
-    } else {
-        // Otherwise, we assume TLS lives on the bottom, in which case we have to "skip" it.
-        tls_base + env.layout.tls_size
-    }
-}
-
 pub(crate) unsafe fn get_memory_stack_pointer(
     ctx: &mut FunctionEnvMut<'_, WasiEnv>,
 ) -> Result<u64, String> {
@@ -1152,7 +1137,7 @@ where
     let memory = unsafe { env.memory_view(&ctx) };
 
     // Write the addresses to the start of the stack space
-    let unwind_pointer = get_unwind_buffer_start(env);
+    let unwind_pointer = env.layout.stack_lower;
     let unwind_data_start =
         unwind_pointer + (std::mem::size_of::<__wasi_asyncify_t<M::Offset>>() as u64);
     let unwind_data = __wasi_asyncify_t::<M::Offset> {
@@ -1188,9 +1173,6 @@ where
     trace!(
         stack_upper = env.layout.stack_upper,
         stack_lower = env.layout.stack_lower,
-        tls_base = env.layout.tls_base,
-        tls_size = env.layout.tls_size,
-        %unwind_pointer,
         "wasi[{}:{}]::unwinding (used_stack_space={} total_stack_space={})",
         ctx.data().pid(),
         ctx.data().tid(),
@@ -1214,7 +1196,6 @@ where
         let unwind_stack_finish: u64 = unwind_data_result.start.into();
         let unwind_size = unwind_stack_finish - unwind_stack_begin;
         trace!(
-            %unwind_pointer,
             "wasi[{}:{}]::unwound (memory_stack_size={} unwind_size={})",
             ctx.data().pid(),
             ctx.data().tid(),
@@ -1314,7 +1295,7 @@ pub fn rewind_ext<M: MemorySize>(
     };
 
     // Write the addresses to the start of the stack space
-    let rewind_pointer = get_unwind_buffer_start(env);
+    let rewind_pointer = env.layout.stack_lower;
     let rewind_data_start =
         rewind_pointer + (std::mem::size_of::<__wasi_asyncify_t<M::Offset>>() as u64);
     let rewind_data_end = rewind_data_start + (rewind_stack.len() as u64);

lib/wasix/src/syscalls/wasix/thread_spawn.rs

@@ -65,16 +65,11 @@ pub fn thread_spawn_internal_from_wasi<M: MemorySize>(
         let tls_base: u64 = start.tls_base.try_into().map_err(|_| Errno::Overflow)?;
         let stack_lower = stack_upper - stack_size;
 
-        // TLS size is constant
-        let tls_size = env.layout.tls_size;
-
         WasiMemoryLayout {
             stack_upper,
             stack_lower,
             guard_size,
             stack_size,
-            tls_base,
-            tls_size,
         }
     };
     tracing::trace!("spawn with layout {:?}", layout);