jonaskruckenberg.de
Next Generation WASM Microkernel Operating System
1// Copyright 2025 Jonas Kruckenberg
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8mod host;
9mod typed;
10
11use crate::arch;
12use crate::mem::VirtualAddress;
13use crate::util::zip_eq::IteratorExt;
14use crate::wasm::indices::VMSharedTypeIndex;
15use crate::wasm::store::{StoreOpaque, Stored};
16use crate::wasm::trap_handler::{Trap, TrapReason};
17use crate::wasm::types::FuncType;
18use crate::wasm::values::Val;
19use crate::wasm::vm::{
20 ExportedFunction, VMArrayCallHostFuncContext, VMFuncRef, VMFunctionImport, VMOpaqueContext,
21 VMVal, VmPtr,
22};
23use crate::wasm::{MAX_WASM_STACK, Module, Store};
24use alloc::boxed::Box;
25use alloc::sync::Arc;
26use anyhow::ensure;
27use core::ffi::c_void;
28use core::mem;
29use core::ptr::NonNull;
30pub use host::{HostFunc, IntoFunc};
31pub use typed::{TypedFunc, WasmParams, WasmResults, WasmTy};
32
33#[derive(Clone, Copy, Debug)]
34pub struct Func(pub(super) Stored<FuncData>);
35#[derive(Debug)]
36pub struct FuncData {
37 kind: FuncKind,
38}
39
40#[derive(Debug)]
41
42enum FuncKind {
43 StoreOwned { export: ExportedFunction },
44 SharedHost(Arc<HostFunc>),
45 Host(Box<HostFunc>),
46}
47
48impl Func {
49 pub fn wrap<T, Params, Results>(
50 store: &mut Store<T>,
51 func: impl IntoFunc<T, Params, Results>,
52 ) -> TypedFunc<Params, Results>
53 where
54 Params: WasmParams,
55 Results: WasmResults,
56 {
57 let (func, ty) = HostFunc::wrap(store.engine(), func);
58
59 let stored = store.add_function(FuncData {
60 kind: FuncKind::Host(Box::new(func)),
61 });
62
63 // Safety: the Rust generics ensure this is safe
64 unsafe { TypedFunc::new_unchecked(Self(stored), ty) }
65 }
66
67 pub fn typed<Params, Results>(
68 self,
69 store: &StoreOpaque,
70 ) -> crate::Result<TypedFunc<Params, Results>>
71 where
72 Params: WasmParams,
73 Results: WasmResults,
74 {
75 let ty = self.ty(store);
76 Params::typecheck(store.engine(), ty.params())?;
77 Results::typecheck(store.engine(), ty.results())?;
78
79 // Safety: the Rust generics ensure this is safe
80 Ok(unsafe { TypedFunc::new_unchecked(self, ty) })
81 }
82
83 pub fn call(
84 self,
85 store: &mut StoreOpaque,
86 params: &[Val],
87 results: &mut [Val],
88 ) -> crate::Result<()> {
89 // Do the typechecking. Notice how `TypedFunc::call` is essentially the same function
90 // minus this typechecking? Yeah. That's the benefit of the typed function.
91 let ty = self.ty(store);
92
93 let params_ = ty.params().zip_eq(params);
94 for (expected, param) in params_ {
95 let found = param.ty(store)?;
96 ensure!(
97 expected.matches(&found),
98 "Type mismatch. Expected `{expected:?}`, but found `{found:?}`"
99 );
100 }
101
102 ensure!(
103 results.len() >= ty.results().len(),
104 "Results slice too small. Need space for at least {}, but got only {}",
105 ty.results().len(),
106 results.len()
107 );
108
109 // Safety: we have checked the types above, we're safe to proceed
110 unsafe { self.call_unchecked(store, params, results) }
111 }
112
113 /// Calls the given function with the provided arguments and places the results in the provided
114 /// results slice.
115 ///
116 /// # Errors
117 ///
118 /// TODO
119 ///
120 /// # Safety
121 ///
122 /// It is up to the caller to ensure the provided arguments are of the correct types and that
123 /// the `results` slice has enough space to hold the results of the function.
124 pub unsafe fn call_unchecked(
125 self,
126 store: &mut StoreOpaque,
127 params: &[Val],
128 results: &mut [Val],
129 ) -> crate::Result<()> {
130 // This function mainly performs the lowering and lifting of VMVal values from and to Rust.
131 // Because - unlike TypedFunc - we don't have compile-time knowledge about the function type,
132 // we use a heap allocated vec (obtained through `store.take_wasm_vmval_storage()`) to store
133 // our parameters into and read results from.
134 //
135 // This is obviously a little less efficient, but it's not that big of a deal.
136
137 // take out the argument storage from the store
138 let mut values_vec = store.take_wasm_vmval_storage();
139 debug_assert!(values_vec.is_empty());
140
141 // resize the vec so we can be sure that params and results will fit.
142 let values_vec_size = params.len().max(results.len());
143 values_vec.resize_with(values_vec_size, || VMVal::v128(0));
144
145 // copy the arguments into the storage vec
146 for (arg, slot) in params.iter().zip(&mut values_vec) {
147 // Safety: the store stays alive for the duration of the call
148 unsafe {
149 *slot = arg.to_vmval(store)?;
150 }
151 }
152
153 // Safety: func refs obtained from our store are always usable by us.
154 let func_ref = unsafe { self.vm_func_ref(store).as_ref() };
155
156 // do the actual call
157 // Safety: at this point we have typechecked, we have allocated enough memory for the params
158 // and results, and obtained a valid func ref to call.
159 unsafe {
160 do_call(store, func_ref, &mut values_vec)?;
161 }
162
163 // copy the results out of the storage
164 let func_ty = self.ty(store);
165 for ((i, slot), vmval) in results.iter_mut().enumerate().zip(&values_vec) {
166 let ty = func_ty.result(i).unwrap();
167 // Safety: caller has to ensure safety
168 *slot = unsafe { Val::from_vmval(store, *vmval, ty) };
169 }
170
171 // clean up and return the argument storage
172 values_vec.truncate(0);
173 store.return_wasm_vmval_storage(values_vec);
174
175 Ok(())
176 }
177
178 pub fn ty(self, store: &StoreOpaque) -> FuncType {
179 FuncType::from_shared_type_index(store.engine(), self.type_index(store))
180 }
181
182 pub fn matches_ty(self, store: &StoreOpaque, ty: FuncType) -> bool {
183 let actual_ty = self.ty(store);
184 actual_ty.matches(&ty)
185 }
186
187 pub(super) fn type_index(self, store: &StoreOpaque) -> VMSharedTypeIndex {
188 // Safety: TODO
189 unsafe { self.vm_func_ref(store).as_ref().type_index }
190 }
191
192 pub(super) unsafe fn from_exported_function(
193 store: &mut StoreOpaque,
194 export: ExportedFunction,
195 ) -> Self {
196 let stored = store.add_function(FuncData {
197 kind: FuncKind::StoreOwned { export },
198 });
199 Self(stored)
200 }
201
202 pub(super) fn as_vmfunction_import(
203 self,
204 store: &mut StoreOpaque,
205 module: &Module,
206 ) -> VMFunctionImport {
207 let f = self.vm_func_ref(store);
208
209 // Safety: TODO
210 unsafe {
211 VMFunctionImport {
212 wasm_call: f.as_ref().wasm_call.unwrap_or_else(|| {
213 // Assert that this is a array-call function, since those
214 // are the only ones that could be missing a `wasm_call`
215 // trampoline.
216 let _ = VMArrayCallHostFuncContext::from_opaque(f.as_ref().vmctx.as_non_null());
217
218 let sig = self.type_index(store);
219
220 let ptr = module.wasm_to_array_trampoline(sig).expect(
221 "if the wasm is importing a function of a given type, it must have the \
222 type's trampoline",
223 );
224
225 VmPtr::from(ptr)
226 }),
227 array_call: f.as_ref().array_call,
228 vmctx: f.as_ref().vmctx,
229 }
230 }
231 }
232
233 pub(super) fn comes_from_same_store(self, store: &StoreOpaque) -> bool {
234 store.has_function(self.0)
235 }
236
237 pub(super) unsafe fn from_vm_func_ref(
238 store: &mut StoreOpaque,
239 func_ref: NonNull<VMFuncRef>,
240 ) -> Self {
241 // Safety: ensured by caller
242 unsafe {
243 debug_assert!(func_ref.as_ref().type_index != VMSharedTypeIndex::default());
244 Func::from_exported_function(store, ExportedFunction { func_ref })
245 }
246 }
247
248 pub(super) fn vm_func_ref(self, store: &StoreOpaque) -> NonNull<VMFuncRef> {
249 match &store[self.0].kind {
250 FuncKind::StoreOwned { export } => export.func_ref,
251 FuncKind::SharedHost(func) => func.func_ref(),
252 FuncKind::Host(func) => func.func_ref(),
253 }
254 }
255
256 pub(super) unsafe fn from_vmval(store: &mut StoreOpaque, raw: *mut c_void) -> Option<Self> {
257 // Safety: ensured by caller
258 unsafe { Some(Func::from_vm_func_ref(store, NonNull::new(raw.cast())?)) }
259 }
260
261 /// Extracts the raw value of this `Func`, which is owned by `store`.
262 ///
263 /// This function returns a value that's suitable for writing into the
264 /// `funcref` field of the [`VMVal`] structure.
265 ///
266 /// # Safety
267 ///
268 /// The returned value is only valid for as long as the store is alive and
269 /// this function is properly rooted within it. Additionally this function
270 /// should not be liberally used since it's a very low-level knob.
271 pub(super) unsafe fn to_vmval(self, store: &mut StoreOpaque) -> *mut c_void {
272 self.vm_func_ref(store).as_ptr().cast()
273 }
274}
275
276pub(super) unsafe fn do_call(
277 store: &mut StoreOpaque,
278 func_ref: &VMFuncRef,
279 params_and_results: &mut [VMVal],
280) -> crate::Result<()> {
281 // Safety: TODO
282 unsafe {
283 let span = tracing::trace_span!("WASM");
284
285 span.in_scope(|| {
286 let exit = enter_wasm(store);
287 let res = crate::wasm::trap_handler::catch_traps(store, |caller| {
288 tracing::trace!("calling VMFuncRef array call");
289 let success = func_ref.array_call(
290 VMOpaqueContext::from_vmcontext(caller),
291 NonNull::from(params_and_results),
292 );
293 tracing::trace!(success, "returned from VMFuncRef array call");
294 });
295 exit_wasm(store, exit);
296
297 match res {
298 Ok(()) => Ok(()),
299 Err(Trap { reason, backtrace }) => {
300 let error = match reason {
301 TrapReason::User(err) => err,
302 TrapReason::Jit {
303 pc,
304 faulting_addr,
305 trap,
306 } => {
307 let mut err: anyhow::Error = trap.into();
308 if let Some(fault) =
309 faulting_addr.and_then(|addr| store.wasm_fault(pc, addr))
310 {
311 err = err.context(fault);
312 }
313 err
314 }
315 TrapReason::Wasm(trap_code) => trap_code.into(),
316 };
317
318 if let Some(bt) = backtrace {
319 tracing::debug!("TODO properly format wasm backtrace {bt:?}");
320
321 // let bt = WasmBacktrace::from_captured(store, bt, pc);
322 // if !bt.wasm_trace.is_empty() {
323 // error = error.context(bt);
324 // }
325 }
326
327 Err(error)
328 }
329 }
330 })
331 }
332}
333
334/// This function is called to register state within `Store` whenever
335/// WebAssembly is entered within the `Store`.
336///
337/// This function sets up various limits such as:
338///
339/// * The stack limit. This is what ensures that we limit the stack space
340/// allocated by WebAssembly code and it's relative to the initial stack
341/// pointer that called into wasm.
342///
343/// This function may fail if the stack limit can't be set because an
344/// interrupt already happened.
345fn enter_wasm(store: &mut StoreOpaque) -> Option<VirtualAddress> {
346 // If this is a recursive call, e.g. our stack limit is already set, then
347 // we may be able to skip this function.
348 //
349 // // For synchronous stores there's nothing else to do because all wasm calls
350 // // happen synchronously and on the same stack. This means that the previous
351 // // stack limit will suffice for the next recursive call.
352 // //
353 // // For asynchronous stores then each call happens on a separate native
354 // // stack. This means that the previous stack limit is no longer relevant
355 // // because we're on a separate stack.
356 // if unsafe { *store.vm_store_context().stack_limit.get() } != VirtualAddress::MAX
357 // && !store.async_support()
358 // {
359 // return None;
360 // }
361
362 // Ignore this stack pointer business on miri since we can't execute wasm
363 // anyway and the concept of a stack pointer on miri is a bit nebulous
364 // regardless.
365 if cfg!(miri) {
366 return None;
367 }
368
369 // When Cranelift has support for the host then we might be running native
370 // compiled code meaning we need to read the actual stack pointer. If
371 // Cranelift can't be used though then we're guaranteed to be running pulley
372 // in which case this stack pointer isn't actually used as Pulley has custom
373 // mechanisms for stack overflow.
374 let stack_pointer = arch::get_stack_pointer();
375
376 // Determine the stack pointer where, after which, any wasm code will
377 // immediately trap. This is checked on the entry to all wasm functions.
378 //
379 // Note that this isn't 100% precise. We are requested to give wasm
380 // `max_wasm_stack` bytes, but what we're actually doing is giving wasm
381 // probably a little less than `max_wasm_stack` because we're
382 // calculating the limit relative to this function's approximate stack
383 // pointer. Wasm will be executed on a frame beneath this one (or next
384 // to it). In any case it's expected to be at most a few hundred bytes
385 // of slop one way or another. When wasm is typically given a MB or so
386 // (a million bytes) the slop shouldn't matter too much.
387 //
388 // After we've got the stack limit then we store it into the `stack_limit`
389 // variable.
390 let wasm_stack_limit = VirtualAddress::new(stack_pointer - MAX_WASM_STACK).unwrap();
391
392 // Safety: the VMStoreContext is always properly initialized
393 let prev_stack = unsafe {
394 mem::replace(
395 &mut *store.vm_store_context().stack_limit.get(),
396 wasm_stack_limit,
397 )
398 };
399
400 Some(prev_stack)
401}
402
403fn exit_wasm(store: &mut StoreOpaque, prev_stack: Option<VirtualAddress>) {
404 // If we don't have a previous stack pointer to restore, then there's no
405 // cleanup we need to perform here.
406 let Some(prev_stack) = prev_stack else {
407 return;
408 };
409
410 // Safety: the VMStoreContext is always properly initialized
411 unsafe {
412 *store.vm_store_context().stack_limit.get() = prev_stack;
413 }
414}