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         .max_wasm_stack(stack_creator.size / 2)
121         .async_stack_size(stack_creator.size)
122         .with_host_stack(stack_creator.clone());
123     (
124         Store::new(&Engine::new(&config).unwrap(), ()),
125         stack_creator,
126     )
127 }
128 
129 #[tokio::test]
130 #[cfg_attr(asan, ignore)]
131 async fn called_on_custom_heap_stack() -> Result<()> {
132     let (mut store, stack_creator) = config();
133     let module = Module::new(
134         store.engine(),
135         r#"
136             (module
137                 (import "host" "callback" (func $callback (result i64)))
138                 (func $f (result i64) (call $callback))
139                 (export "f" (func $f))
140             )
141         "#,
142     )?;
143 
144     let ty = FuncType::new(store.engine(), [], [ValType::I64]);
145     let host_func = Func::new(&mut store, ty, move |_caller, _params, results| {
146         let foo = 42;
147         // output an address on the stack
148         results[0] = Val::I64((&foo as *const i32) as usize as i64);
149         Ok(())
150     });
151     let export = wasmtime::Extern::Func(host_func);
152     let instance = Instance::new_async(&mut store, &module, &[export]).await?;
153     let mut results = [Val::I64(0)];
154     instance
155         .get_func(&mut store, "f")
156         .context("missing function export")?
157         .call_async(&mut store, &[], &mut results)
158         .await?;
159     // Make sure the stack address we wrote was within our custom stack range
160     let stack_address = results[0].i64().unwrap() as usize;
161     assert!(stack_creator.range().contains(&stack_address));
162     Ok(())
163 }
164