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     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) -> Result<Box<dyn StackMemory>> {
102         if size != self.size {
103             bail!("must use the size we allocated for this stack memory creator");
104         }
105         let page_size = rustix::param::page_size();
106         // skip over the page size
107         let base_ptr = unsafe { self.memory.as_ptr().add(page_size) };
108         let base = NonNull::new(base_ptr).context("unable to compute stack base")?;
109         Ok(Box::new(CustomStack::new(base, self.size)))
110     }
111 }
112 
113 fn config() -> (Store<()>, Arc<CustomStackCreator>) {
114     let stack_creator = Arc::new(CustomStackCreator::new().unwrap());
115     let mut config = Config::new();
116     config
117         .async_support(true)
118         .max_wasm_stack(stack_creator.size / 2)
119         .async_stack_size(stack_creator.size)
120         .with_host_stack(stack_creator.clone());
121     (
122         Store::new(&Engine::new(&config).unwrap(), ()),
123         stack_creator,
124     )
125 }
126 
127 #[tokio::test]
128 async fn called_on_custom_heap_stack() -> Result<()> {
129     let (mut store, stack_creator) = config();
130     let module = Module::new(
131         store.engine(),
132         r#"
133             (module
134                 (import "host" "callback" (func $callback (result i64)))
135                 (func $f (result i64) (call $callback))
136                 (export "f" (func $f))
137             )
138         "#,
139     )?;
140 
141     let ty = FuncType::new(store.engine(), [], [ValType::I64]);
142     let host_func = Func::new(&mut store, ty, move |_caller, _params, results| {
143         let foo = 42;
144         // output an address on the stack
145         results[0] = Val::I64((&foo as *const i32) as usize as i64);
146         Ok(())
147     });
148     let export = wasmtime::Extern::Func(host_func);
149     let instance = Instance::new_async(&mut store, &module, &[export]).await?;
150     let mut results = [Val::I64(0)];
151     instance
152         .get_func(&mut store, "f")
153         .context("missing function export")?
154         .call_async(&mut store, &[], &mut results)
155         .await?;
156     // Make sure the stack address we wrote was within our custom stack range
157     let stack_address = results[0].i64().unwrap() as usize;
158     assert!(stack_creator.range().contains(&stack_address));
159     Ok(())
160 }
161