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
22const 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 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 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 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 if let Some(snapshot) = store_snapshot {
92 restore_store_snapshot(&mut store, &snapshot);
93 }
94
95 Ok((ctx, store))
96 }
97
98 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 pub fn data<'a>(&'a self, store: &'a impl AsStoreRef) -> &'a WasiEnv {
114 self.env.as_ref(store)
115 }
116
117 pub fn data_mut<'a>(&'a self, store: &'a mut impl AsStoreMut) -> &'a mut WasiEnv {
119 self.env.as_mut(store)
120 }
121
122 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 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 update_layout {
190 let new_layout = match stack_layout {
191 Some(layout) => layout,
192 None => {
193 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 if data_end >= stack_upper { 0 } else { data_end }
232 } else {
233 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 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 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 env.thread.set_memory_layout(layout.clone());
288
289 {
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 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 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 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 self.data(store).blocking_on_exit(process_exit_code);
348 }
349
350 #[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 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 rewind_state.is_none() {
426 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 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}