wasmer_wasix/state/
func_env.rs

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