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