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
19const 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 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 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 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 if let Some(snapshot) = store_snapshot {
89 restore_store_snapshot(&mut store, &snapshot);
90 }
91
92 Ok((ctx, store))
93 }
94
95 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 pub fn data<'a>(&'a self, store: &'a impl AsStoreRef) -> &'a WasiEnv {
111 self.env.as_ref(store)
112 }
113
114 pub fn data_mut<'a>(&'a self, store: &'a mut impl AsStoreMut) -> &'a mut WasiEnv {
116 self.env.as_mut(store)
117 }
118
119 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 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 update_layout {
191 let new_layout = match stack_layout {
192 Some(layout) => layout,
193 None => {
194 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 if data_end >= stack_upper { 0 } else { data_end }
233 } else {
234 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 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 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 env.thread.set_memory_layout(layout.clone());
289
290 {
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 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 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 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 self.data(store).blocking_on_exit(process_exit_code);
349 }
350
351 #[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 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 rewind_state.is_none() {
427 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 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}