1 #![cfg(all(not(target_os = "windows"), not(miri)))]
2 use std::{
3     alloc::{GlobalAlloc, Layout, System},
4     ops::Range,
5     ptr::NonNull,
6     sync::Arc,
7 };
8 use wasmtime::error::Context as _;
9 use wasmtime::*;
10 
11 fn align_up(v: usize, align: usize) -> usize {
12     return (v + (align - 1)) & (!(align - 1));
13 }
14 
15 struct CustomStack {
16     base: NonNull<u8>,
17     len: usize,
18 }
19 unsafe impl Send for CustomStack {}
20 unsafe impl Sync for CustomStack {}
21 impl CustomStack {
22     fn new(base: NonNull<u8>, len: usize) -> Self {
23         CustomStack { base, len }
24     }
25 }
26 unsafe impl StackMemory for CustomStack {
27     fn top(&self) -> *mut u8 {
28         unsafe { self.base.as_ptr().add(self.len) }
29     }
30     fn range(&self) -> Range<usize> {
31         let base = self.base.as_ptr() as usize;
32         base..base + self.len
33     }
34     fn guard_range(&self) -> Range<*mut u8> {
35         std::ptr::null_mut()..std::ptr::null_mut()
36     }
37 }
38 
39 // A creator that allocates stacks on the heap instead of mmap'ing.
40 struct CustomStackCreator {
41     memory: NonNull<u8>,
42     size: usize,
43     layout: Layout,
44 }
45 
46 unsafe impl Send for CustomStackCreator {}
47 unsafe impl Sync for CustomStackCreator {}
48 impl CustomStackCreator {
49     fn new() -> Result<Self> {
50         // 1MB
51         const MINIMUM_STACK_SIZE: usize = 1 * 1_024 * 1_024;
52         let page_size = rustix::param::page_size();
53         let size = align_up(MINIMUM_STACK_SIZE, page_size);
54         // Add an extra page for the guard page
55         let layout = Layout::from_size_align(size + page_size, page_size)
56             .context("unable to compute stack layout")?;
57         let memory = unsafe {
58             let mem = System.alloc(layout);
59             let notnull = NonNull::new(mem);
60             if let Some(mem) = notnull {
61                 // It's required that stack memory is zeroed for wasmtime
62                 libc::memset(mem.as_ptr().cast(), 0, layout.size());
63                 // Mark guard page as protected
64                 rustix::mm::mprotect(
65                     mem.as_ptr().cast(),
66                     page_size,
67                     rustix::mm::MprotectFlags::empty(),
68                 )?;
69             }
70             notnull
71         }
72         .context("unable to allocate stack memory")?;
73         Ok(CustomStackCreator {
74             memory,
75             size,
76             layout,
77         })
78     }
79     fn range(&self) -> Range<usize> {
80         let page_size = rustix::param::page_size();
81         let base = unsafe { self.memory.as_ptr().add(page_size) as usize };
82         base..base + self.size
83     }
84 }
85 impl Drop for CustomStackCreator {
86     fn drop(&mut self) {
87         let page_size = rustix::param::page_size();
88         unsafe {
89             // Unprotect the guard page as the allocator could reuse it.
90             rustix::mm::mprotect(
91                 self.memory.as_ptr().cast(),
92                 page_size,
93                 rustix::mm::MprotectFlags::READ | rustix::mm::MprotectFlags::WRITE,
94             )
95             .unwrap();
96             System.dealloc(self.memory.as_ptr(), self.layout);
97         }
98     }
99 }
100 unsafe impl StackCreator for CustomStackCreator {
101     fn new_stack(&self, size: usize, zeroed: bool) -> Result<Box<dyn StackMemory>> {
102         if zeroed {
103             bail!("CustomStackCreator does not support stack zeroing");
104         }
105         if size != self.size {
106             bail!("must use the size we allocated for this stack memory creator");
107         }
108         let page_size = rustix::param::page_size();
109         // skip over the page size
110         let base_ptr = unsafe { self.memory.as_ptr().add(page_size) };
111         let base = NonNull::new(base_ptr).context("unable to compute stack base")?;
112         Ok(Box::new(CustomStack::new(base, self.size)))
113     }
114 }
115 
116 fn config() -> (Store<()>, Arc<CustomStackCreator>) {
117     let stack_creator = Arc::new(CustomStackCreator::new().unwrap());
118     let mut config = Config::new();
119     config
120         .async_support(true)
121         .max_wasm_stack(stack_creator.size / 2)
122         .async_stack_size(stack_creator.size)
123         .with_host_stack(stack_creator.clone());
124     (
125         Store::new(&Engine::new(&config).unwrap(), ()),
126         stack_creator,
127     )
128 }
129 
130 #[tokio::test]
131 #[cfg_attr(asan, ignore)]
132 async fn called_on_custom_heap_stack() -> Result<()> {
133     let (mut store, stack_creator) = config();
134     let module = Module::new(
135         store.engine(),
136         r#"
137             (module
138                 (import "host" "callback" (func $callback (result i64)))
139                 (func $f (result i64) (call $callback))
140                 (export "f" (func $f))
141             )
142         "#,
143     )?;
144 
145     let ty = FuncType::new(store.engine(), [], [ValType::I64]);
146     let host_func = Func::new(&mut store, ty, move |_caller, _params, results| {
147         let foo = 42;
148         // output an address on the stack
149         results[0] = Val::I64((&foo as *const i32) as usize as i64);
150         Ok(())
151     });
152     let export = wasmtime::Extern::Func(host_func);
153     let instance = Instance::new_async(&mut store, &module, &[export]).await?;
154     let mut results = [Val::I64(0)];
155     instance
156         .get_func(&mut store, "f")
157         .context("missing function export")?
158         .call_async(&mut store, &[], &mut results)
159         .await?;
160     // Make sure the stack address we wrote was within our custom stack range
161     let stack_address = results[0].i64().unwrap() as usize;
162     assert!(stack_creator.range().contains(&stack_address));
163     Ok(())
164 }
165