Move rewind buffer to after TLS area if it's at the bottom of the stack

lib/journal/src/concrete/archived.rs

@@ -1802,4 +1802,6 @@ 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,6 +616,8 @@ 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,
         }
     }
 }
@@ -627,6 +629,8 @@ 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,
         }
     }
 }
@@ -638,6 +642,8 @@ 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/wasi-types/src/wasix/mod.rs

@@ -26,4 +26,8 @@ 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/func_env.rs

@@ -25,6 +25,17 @@ 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>,
@@ -147,6 +158,18 @@ 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();
 
@@ -173,12 +196,34 @@ impl WasiFunctionEnv {
                 ));
             }
 
+            let tls_base = if let Some(tls_base) = tls_base {
+                match tls_base.get(store) {
+                    wasmer::Value::I32(a) => a as u64,
+                    wasmer::Value::I64(a) => a as u64,
+                    _ => DEFAULT_TLS_BASE,
+                }
+            } else {
+                DEFAULT_TLS_BASE
+            };
+
+            let tls_size = if let Some(tls_size) = tls_size {
+                match tls_size.get(store) {
+                    wasmer::Value::I32(a) => a as u64,
+                    wasmer::Value::I64(a) => a as u64,
+                    _ => DEFAULT_TLS_SIZE,
+                }
+            } else {
+                DEFAULT_TLS_SIZE
+            };
+
             // 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_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,6 +870,21 @@ 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> {
@@ -1137,14 +1152,12 @@ where
     let memory = unsafe { env.memory_view(&ctx) };
 
     // Write the addresses to the start of the stack space
-    let unwind_pointer = env.layout.stack_lower;
+    let unwind_pointer = get_unwind_buffer_start(env);
     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> {
         start: wasi_try_ok!(unwind_data_start.try_into().map_err(|_| Errno::Overflow)),
-        end: wasi_try_ok!(env
-            .layout
-            .stack_upper
+        end: wasi_try_ok!((env.layout.stack_upper - memory_stack.len() as u64)
             .try_into()
             .map_err(|_| Errno::Overflow)),
     };
@@ -1175,6 +1188,9 @@ 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(),
@@ -1198,6 +1214,7 @@ 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(),
@@ -1297,7 +1314,7 @@ pub fn rewind_ext<M: MemorySize>(
     };
 
     // Write the addresses to the start of the stack space
-    let rewind_pointer = env.layout.stack_lower;
+    let rewind_pointer = get_unwind_buffer_start(env);
     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,11 +65,16 @@ 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);