From 3d264586e982645d7cae37c97f2d5a119ef2e47b Mon Sep 17 00:00:00 2001 From: Arshia Ghafoori Date: Mon, 6 May 2024 20:26:23 +0330 Subject: [PATCH] Move rewind buffer to after TLS area if it's at the bottom of the stack --- lib/journal/src/concrete/archived.rs | 2 + lib/journal/src/concrete/archived_from.rs | 6 +++ lib/wasi-types/src/wasix/mod.rs | 4 ++ lib/wasix/src/state/func_env.rs | 45 ++++++++++++++++++++ lib/wasix/src/syscalls/mod.rs | 27 +++++++++--- lib/wasix/src/syscalls/wasix/thread_spawn.rs | 5 +++ 6 files changed, 84 insertions(+), 5 deletions(-) diff --git a/lib/journal/src/concrete/archived.rs b/lib/journal/src/concrete/archived.rs index d03b8583dd0..85f66df6c2d 100644 --- a/lib/journal/src/concrete/archived.rs +++ b/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, } diff --git a/lib/journal/src/concrete/archived_from.rs b/lib/journal/src/concrete/archived_from.rs index 853973061bf..bec9070c5cb 100644 --- a/lib/journal/src/concrete/archived_from.rs +++ b/lib/journal/src/concrete/archived_from.rs @@ -616,6 +616,8 @@ impl From 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 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, } } } diff --git a/lib/wasi-types/src/wasix/mod.rs b/lib/wasi-types/src/wasix/mod.rs index 8f601ed578b..b494d23936a 100644 --- a/lib/wasi-types/src/wasix/mod.rs +++ b/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, } diff --git a/lib/wasix/src/state/func_env.rs b/lib/wasix/src/state/func_env.rs index 6eff4e5166b..c12a71b2c9d 100644 --- a/lib/wasix/src/state/func_env.rs +++ b/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, @@ -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()); diff --git a/lib/wasix/src/syscalls/mod.rs b/lib/wasix/src/syscalls/mod.rs index 2cd424ea9a0..d658006f1bb 100644 --- a/lib/wasix/src/syscalls/mod.rs +++ b/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 { @@ -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>() as u64); let unwind_data = __wasi_asyncify_t:: { 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( }; // 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>() as u64); let rewind_data_end = rewind_data_start + (rewind_stack.len() as u64); diff --git a/lib/wasix/src/syscalls/wasix/thread_spawn.rs b/lib/wasix/src/syscalls/wasix/thread_spawn.rs index f30d2d772fa..7d792f19a3e 100644 --- a/lib/wasix/src/syscalls/wasix/thread_spawn.rs +++ b/lib/wasix/src/syscalls/wasix/thread_spawn.rs @@ -65,11 +65,16 @@ pub fn thread_spawn_internal_from_wasi( 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);