wasmer_wasix/state/
func_env.rs

1use tracing::trace;
2use wasmer::{
3    AsStoreMut, AsStoreRef, ExportError, FunctionEnv, FunctionEnvMut, Imports, Instance, Memory,
4    Module, Store,
5};
6use wasmer_wasix_types::{wasi::ExitCode, wasix::WasiMemoryLayout};
7
8#[allow(unused_imports)]
9use crate::os::task::thread::RewindResultType;
10#[cfg(feature = "journal")]
11use crate::syscalls::restore_snapshot;
12use crate::{
13    RewindStateOption, StoreSnapshot, WasiEnv, WasiError, WasiModuleInstanceHandles,
14    WasiRuntimeError, WasiThreadError,
15    runtime::task_manager::SpawnMemoryTypeOrStore,
16    state::WasiModuleTreeHandles,
17    utils::{get_wasi_version, get_wasi_versions, store::restore_store_snapshot},
18};
19
20use super::Linker;
21
22/// The default stack size for WASIX - the number itself is the default that compilers
23/// have used in the past when compiling WASM apps.
24///
25/// (this is only used for programs that have no stack pointer)
26const DEFAULT_STACK_SIZE: u64 = 1_048_576u64;
27const DEFAULT_STACK_BASE: u64 = DEFAULT_STACK_SIZE;
28
29#[derive(Clone, Debug)]
30pub struct WasiFunctionEnv {
31    pub env: FunctionEnv<WasiEnv>,
32}
33
34impl WasiFunctionEnv {
35    pub fn new(store: &mut impl AsStoreMut, env: WasiEnv) -> Self {
36        Self {
37            env: FunctionEnv::new(store, env),
38        }
39    }
40
41    // Creates a new environment context on a new store
42    pub fn new_with_store(
43        module: Module,
44        env: WasiEnv,
45        store_snapshot: Option<StoreSnapshot>,
46        spawn_type: SpawnMemoryTypeOrStore,
47        update_layout: bool,
48        call_initialize: bool,
49        parent_linker_and_ctx: Option<(Linker, &mut FunctionEnvMut<WasiEnv>)>,
50    ) -> Result<(Self, Store), WasiThreadError> {
51        // Create a new store and put the memory object in it
52        // (but only if it has imported memory)
53        let (memory, store): (Option<wasmer::Memory>, Option<wasmer::Store>) = match spawn_type {
54            SpawnMemoryTypeOrStore::New => (None, None),
55            SpawnMemoryTypeOrStore::Type(mut ty) => {
56                ty.shared = true;
57
58                let mut store = env.runtime.new_store();
59
60                // Note: If memory is shared, maximum needs to be set in the
61                // browser otherwise creation will fail.
62                let _ = ty.maximum.get_or_insert(wasmer_types::Pages::max_value());
63
64                let mem = Memory::new(&mut store, ty).map_err(|err| {
65                    tracing::error!(
66                        error = &err as &dyn std::error::Error,
67                        memory_type=?ty,
68                        "could not create memory",
69                    );
70                    WasiThreadError::MemoryCreateFailed(err)
71                })?;
72                (Some(mem), Some(store))
73            }
74            SpawnMemoryTypeOrStore::StoreAndMemory(s, m) => (m, Some(s)),
75        };
76
77        let mut store = store.unwrap_or_else(|| env.runtime().new_store());
78
79        let (_, ctx) = env.instantiate(
80            module,
81            &mut store,
82            memory,
83            update_layout,
84            call_initialize,
85            parent_linker_and_ctx,
86        )?;
87
88        // FIXME: shouldn't this happen _before_ instantiating, so the startup code in the instance
89        // has access to the globals?
90        // Set all the globals
91        if let Some(snapshot) = store_snapshot {
92            restore_store_snapshot(&mut store, &snapshot);
93        }
94
95        Ok((ctx, store))
96    }
97
98    /// Get an `Imports` for a specific version of WASI detected in the module.
99    pub fn import_object(
100        &self,
101        store: &mut impl AsStoreMut,
102        module: &Module,
103    ) -> Result<Imports, WasiError> {
104        let wasi_version = get_wasi_version(module, false).ok_or(WasiError::UnknownWasiVersion)?;
105        Ok(crate::generate_import_object_from_env(
106            store,
107            &self.env,
108            wasi_version,
109        ))
110    }
111
112    /// Gets a reference to the WasiEnvironment
113    pub fn data<'a>(&'a self, store: &'a impl AsStoreRef) -> &'a WasiEnv {
114        self.env.as_ref(store)
115    }
116
117    /// Gets a mutable- reference to the host state in this context.
118    pub fn data_mut<'a>(&'a self, store: &'a mut impl AsStoreMut) -> &'a mut WasiEnv {
119        self.env.as_mut(store)
120    }
121
122    /// Initializes the WasiEnv using the instance exports
123    /// (this must be executed before attempting to use it)
124    /// (as the stores can not by themselves be passed between threads we can store the module
125    ///  in a thread-local variables and use it later - for multithreading)
126    // FIXME: this probably doesn't work with WASIX modules, since they import their memories?
127    pub fn initialize(
128        &mut self,
129        store: &mut impl AsStoreMut,
130        instance: Instance,
131    ) -> Result<(), ExportError> {
132        let exported_memory = instance
133            .exports
134            .iter()
135            .filter_map(|(_, export)| {
136                if let wasmer::Extern::Memory(memory) = export {
137                    Some(memory.clone())
138                } else {
139                    None
140                }
141            })
142            .next()
143            .ok_or_else(|| ExportError::Missing("No exported memory found".to_string()))?;
144
145        self.initialize_handles_and_layout(
146            store,
147            instance.clone(),
148            WasiModuleTreeHandles::Static(WasiModuleInstanceHandles::new(
149                exported_memory,
150                store,
151                instance,
152                None,
153            )),
154            None,
155            true,
156        )
157    }
158
159    /// Initializes the WasiEnv using the instance exports and a provided optional memory
160    /// (this must be executed before attempting to use it)
161    /// (as the stores can not by themselves be passed between threads we can store the module
162    ///  in a thread-local variables and use it later - for multithreading)
163    // FIXME: Move this code to somewhere that makes sense (in WasiEnv?)
164    pub fn initialize_handles_and_layout(
165        &mut self,
166        store: &mut impl AsStoreMut,
167        instance: Instance,
168        handles: WasiModuleTreeHandles,
169        stack_layout: Option<WasiMemoryLayout>,
170        update_layout: bool,
171    ) -> Result<(), ExportError> {
172        let is_wasix_module = crate::utils::is_wasix_module(instance.module());
173
174        let new_inner = handles;
175
176        let main_module_handles = new_inner.main_module_instance_handles();
177        let stack_pointer = main_module_handles.stack_pointer.clone();
178        let data_end = main_module_handles.data_end.clone();
179        let stack_low = main_module_handles.stack_low.clone();
180        let stack_high = main_module_handles.stack_high.clone();
181        let tls_base = main_module_handles.tls_base.clone();
182
183        let env = self.data_mut(store);
184        env.set_inner(new_inner);
185
186        env.state.fs.set_is_wasix(is_wasix_module);
187
188        // If the stack offset and size is not set then do so
189        if update_layout {
190            let new_layout = match stack_layout {
191                Some(layout) => layout,
192                None => {
193                    // Set the base stack
194                    let stack_upper = if let Some(stack_high) = stack_high {
195                        match stack_high.get(store) {
196                            wasmer::Value::I32(a) => a as u64,
197                            wasmer::Value::I64(a) => a as u64,
198                            _ => DEFAULT_STACK_BASE,
199                        }
200                    } else if let Some(stack_pointer) = stack_pointer {
201                        match stack_pointer.get(store) {
202                            wasmer::Value::I32(a) => a as u64,
203                            wasmer::Value::I64(a) => a as u64,
204                            _ => DEFAULT_STACK_BASE,
205                        }
206                    } else {
207                        DEFAULT_STACK_BASE
208                    };
209
210                    if stack_upper == 0 {
211                        return Err(ExportError::Missing(
212                            "stack_high or stack_pointer is not set to the upper stack range"
213                                .to_string(),
214                        ));
215                    }
216
217                    let mut stack_lower = if let Some(stack_low) = stack_low {
218                        match stack_low.get(store) {
219                            wasmer::Value::I32(a) => a as u64,
220                            wasmer::Value::I64(a) => a as u64,
221                            _ => 0,
222                        }
223                    } else if let Some(data_end) = data_end {
224                        let data_end = match data_end.get(store) {
225                            wasmer::Value::I32(a) => a as u64,
226                            wasmer::Value::I64(a) => a as u64,
227                            _ => 0,
228                        };
229                        // It's possible for the data section to be above the stack, we check for that here and
230                        // if it is, we'll assume the stack starts at address 0
231                        if data_end >= stack_upper { 0 } else { data_end }
232                    } else {
233                        // clang-16 and higher generate the `__stack_low` global, and it can be exported with
234                        // `-Wl,--export=__stack_low`. clang-15 generates `__data_end`, which should be identical
235                        // and can be exported if `__stack_low` is not available.
236                        if self.data(store).will_use_asyncify() {
237                            tracing::warn!(
238                                "Missing both __stack_low and __data_end exports, unwinding may cause memory corruption"
239                            );
240                        }
241                        0
242                    };
243
244                    if stack_lower >= stack_upper {
245                        if self.data(store).will_use_asyncify() {
246                            tracing::warn!(
247                                "Detected lower end of stack to be above higher end, ignoring stack_lower; \
248                                unwinding may cause memory corruption"
249                            );
250                        }
251                        stack_lower = 0;
252                    }
253
254                    // Note: the TLS base global may not be initialized at the point when this
255                    // code runs, so we take a zero value to mean it wasn't initialized and we
256                    // don't know it. It's never actually zero for statically-linked, non-PIE
257                    // modules.
258                    let tls_base = if let Some(tls_base) = tls_base {
259                        match tls_base.get(store) {
260                            wasmer::Value::I32(a) => a as u64,
261                            wasmer::Value::I64(a) => a as u64,
262                            _ => 0,
263                        }
264                    } else {
265                        0
266                    };
267
268                    WasiMemoryLayout {
269                        stack_lower,
270                        stack_upper,
271                        stack_size: stack_upper - stack_lower,
272                        guard_size: 0,
273                        tls_base: if tls_base == 0 { None } else { Some(tls_base) },
274                    }
275                }
276            };
277
278            // Update the stack layout which is need for asyncify
279            let env = self.data_mut(store);
280            let tid = env.tid();
281            let layout = &mut env.layout;
282            layout.stack_upper = new_layout.stack_upper;
283            layout.stack_lower = new_layout.stack_lower;
284            layout.stack_size = layout.stack_upper - layout.stack_lower;
285
286            // Replace the thread object itself
287            env.thread.set_memory_layout(layout.clone());
288
289            // Replace the thread object with this new layout
290            {
291                let mut guard = env.process.lock();
292                guard
293                    .threads
294                    .values_mut()
295                    .filter(|t| t.tid() == tid)
296                    .for_each(|t| t.set_memory_layout(layout.clone()))
297            }
298        }
299        tracing::trace!("initializing with layout {:?}", self.data(store).layout);
300
301        Ok(())
302    }
303
304    /// Like `import_object` but containing all the WASI versions detected in
305    /// the module.
306    pub fn import_object_for_all_wasi_versions(
307        &self,
308        store: &mut impl AsStoreMut,
309        module: &Module,
310    ) -> Result<Imports, WasiError> {
311        let wasi_versions =
312            get_wasi_versions(module, false).ok_or(WasiError::UnknownWasiVersion)?;
313
314        let mut resolver = Imports::new();
315        for version in wasi_versions.iter() {
316            let new_import_object =
317                crate::generate_import_object_from_env(store, &self.env, *version);
318            for ((n, m), e) in new_import_object.into_iter() {
319                resolver.define(&n, &m, e);
320            }
321        }
322
323        Ok(resolver)
324    }
325
326    /// # Safety
327    ///
328    /// This function should only be called from within a syscall
329    /// as it can potentially execute local thread variable cleanup
330    /// code
331    pub fn on_exit(&self, store: &mut impl AsStoreMut, process_exit_code: Option<ExitCode>) {
332        trace!(
333            "wasi[{}:{}]::on_exit",
334            self.data(store).pid(),
335            self.data(store).tid()
336        );
337
338        if let Some(linker) = self.data(store).inner().linker().cloned() {
339            // Note: this call will also process pending dl operations, hence unblocking
340            // other threads that may be waiting for this one to pick the operation up
341            if let Err(e) = linker.shutdown_instance_group(&mut self.env.clone().into_mut(store)) {
342                tracing::warn!("Failed to shutdown linker instance group: {e:?}");
343            }
344        }
345
346        // Cleans up all the open files (if this is the main thread)
347        self.data(store).blocking_on_exit(process_exit_code);
348    }
349
350    /// Bootstraps this main thread and context with any journals that
351    /// may be present
352    ///
353    /// # Safety
354    ///
355    /// This function manipulates the memory of the process and thus must be executed
356    /// by the WASM process thread itself.
357    ///
358    #[allow(clippy::result_large_err)]
359    #[allow(unused_variables, unused_mut)]
360    #[tracing::instrument(skip_all)]
361    pub unsafe fn bootstrap(
362        &self,
363        mut store: &'_ mut impl AsStoreMut,
364    ) -> Result<RewindStateOption, WasiRuntimeError> {
365        tracing::debug!("bootstrap start");
366
367        #[allow(unused_mut)]
368        let mut rewind_state = None;
369
370        #[cfg(feature = "journal")]
371        {
372            // If there are journals we need to restore then do so (this will
373            // prevent the initialization function from running
374            let restore_ro_journals = self
375                .data(&store)
376                .runtime
377                .read_only_journals()
378                .collect::<Vec<_>>();
379            let restore_w_journals = self
380                .data(&store)
381                .runtime
382                .writable_journals()
383                .collect::<Vec<_>>();
384            if !restore_ro_journals.is_empty() || !restore_w_journals.is_empty() {
385                tracing::trace!("replaying journal=true");
386                self.data_mut(&mut store).replaying_journal = true;
387
388                for journal in restore_ro_journals {
389                    let ctx = self.env.clone().into_mut(&mut store);
390                    let rewind = match unsafe { restore_snapshot(ctx, journal.as_ref(), true) } {
391                        Ok(r) => r,
392                        Err(err) => {
393                            tracing::trace!("replaying journal=false (err={:?})", err);
394                            self.data_mut(&mut store).replaying_journal = false;
395                            return Err(err);
396                        }
397                    };
398                    rewind_state = rewind.map(|rewind| (rewind, RewindResultType::RewindRestart));
399                }
400
401                for journal in restore_w_journals {
402                    let ctx = self.env.clone().into_mut(&mut store);
403                    let rewind = match unsafe {
404                        restore_snapshot(ctx, journal.as_ref().as_dyn_readable_journal(), true)
405                    } {
406                        Ok(r) => r,
407                        Err(err) => {
408                            tracing::trace!("replaying journal=false (err={:?})", err);
409                            self.data_mut(&mut store).replaying_journal = false;
410                            return Err(err);
411                        }
412                    };
413                    rewind_state = rewind.map(|rewind| (rewind, RewindResultType::RewindRestart));
414                }
415
416                tracing::trace!("replaying journal=false");
417                self.data_mut(&mut store).replaying_journal = false;
418            }
419
420            // If there is no rewind state then the journal is being replayed
421            // and hence we do not need to write an init module event
422            //
423            // But otherwise we need to notify the journal of the module hash
424            // so that recompiled modules will restart
425            if rewind_state.is_none() {
426                // The first event we save is an event that records the module hash.
427                // Note: This is used to detect if an incorrect journal is used on the wrong
428                // process or if a process has been recompiled
429                let wasm_hash = Box::from(self.data(&store).process.module_hash.as_bytes());
430                let mut ctx = self.env.clone().into_mut(&mut store);
431                crate::journal::JournalEffector::save_event(
432                    &mut ctx,
433                    crate::journal::JournalEntry::InitModuleV1 { wasm_hash },
434                )
435                .map_err(|err| {
436                    WasiRuntimeError::Runtime(wasmer::RuntimeError::new(format!(
437                        "journal failed to save the module initialization event - {err}"
438                    )))
439                })?;
440            } else {
441                // Otherwise we should emit a clear ethereal event
442                let mut ctx = self.env.clone().into_mut(&mut store);
443                crate::journal::JournalEffector::save_event(
444                    &mut ctx,
445                    crate::journal::JournalEntry::ClearEtherealV1,
446                )
447                .map_err(|err| {
448                    WasiRuntimeError::Runtime(wasmer::RuntimeError::new(format!(
449                        "journal failed to save clear ethereal event - {err}",
450                    )))
451                })?;
452            }
453        }
454
455        tracing::debug!("bootstrap complete");
456
457        Ok(rewind_state)
458    }
459}