ayaka_plugin_wasmer/
lib.rs

1//! Wasmer-based plugin backend.
2
3#![warn(missing_docs)]
4
5use ayaka_plugin::*;
6use std::{
7    collections::HashMap,
8    sync::{Arc, Mutex},
9};
10use wasmer::*;
11
12unsafe fn mem_slice<R>(
13    store: &impl AsStoreRef,
14    memory: &Memory,
15    start: i32,
16    len: i32,
17    f: impl FnOnce(&[u8]) -> R,
18) -> R {
19    f(memory
20        .view(store)
21        .data_unchecked()
22        .get_unchecked(start as usize..)
23        .get_unchecked(..len as usize))
24}
25
26unsafe fn mem_slice_mut<R>(
27    store: &impl AsStoreRef,
28    memory: &Memory,
29    start: i32,
30    len: i32,
31    f: impl FnOnce(&mut [u8]) -> R,
32) -> R {
33    f(memory
34        .view(store)
35        .data_unchecked_mut()
36        .get_unchecked_mut(start as usize..)
37        .get_unchecked_mut(..len as usize))
38}
39
40type HostStore = Arc<Mutex<Store>>;
41
42/// A Wasmer [`Instance`].
43pub struct WasmerModule {
44    store: HostStore,
45    instance: Instance,
46    memory: Memory,
47    abi_free: TypedFunction<(i32, i32), ()>,
48    abi_alloc: TypedFunction<i32, i32>,
49}
50
51impl WasmerModule {
52    /// Loads the WASM [`Module`], with some imports.
53    fn new(store: HostStore, instance: Instance) -> Result<Self> {
54        let memory = instance.exports.get_memory(MEMORY_NAME)?.clone();
55        let inner_store = store.lock().unwrap();
56        let abi_free = instance
57            .exports
58            .get_typed_function(&inner_store.as_store_ref(), ABI_FREE_NAME)?;
59        let abi_alloc = instance
60            .exports
61            .get_typed_function(&inner_store.as_store_ref(), ABI_ALLOC_NAME)?;
62        drop(inner_store);
63        Ok(Self {
64            store,
65            instance,
66            memory,
67            abi_free,
68            abi_alloc,
69        })
70    }
71
72    fn call_impl<T>(
73        &self,
74        mut store: StoreMut,
75        name: &str,
76        data: &[u8],
77        f: impl FnOnce(&[u8]) -> Result<T>,
78    ) -> Result<T> {
79        let func = self
80            .instance
81            .exports
82            .get_typed_function::<(i32, i32), u64>(&store, name)?;
83
84        let ptr = self.abi_alloc.call(&mut store, data.len() as i32)?;
85        unsafe {
86            mem_slice_mut(&store, &self.memory, ptr, data.len() as i32, |s| {
87                s.copy_from_slice(data)
88            })
89        };
90
91        let res = func.call(&mut store, data.len() as i32, ptr);
92
93        self.abi_free.call(&mut store, ptr, data.len() as i32)?;
94
95        let res = res?;
96        let (len, res) = ((res >> 32) as i32, (res & 0xFFFFFFFF) as i32);
97
98        let res_data = unsafe { mem_slice(&store, &self.memory, res, len, |s| f(s)) };
99
100        self.abi_free.call(&mut store, res, len)?;
101
102        let res_data = res_data?;
103        Ok(res_data)
104    }
105}
106
107impl RawModule for WasmerModule {
108    type Linker = WasmerLinker;
109
110    type LinkerHandle<'a> = WasmerLinkerHandle<'a>;
111
112    type Func = WasmerFunction;
113
114    fn call<T>(&self, name: &str, data: &[u8], f: impl FnOnce(&[u8]) -> Result<T>) -> Result<T> {
115        self.call_impl(self.store.lock().unwrap().as_store_mut(), name, data, f)
116    }
117}
118
119#[doc(hidden)]
120#[derive(Clone)]
121pub struct RuntimeInstanceData {
122    memory: Option<Memory>,
123    abi_alloc: Option<TypedFunction<i32, i32>>,
124    #[allow(clippy::type_complexity)]
125    func: Arc<dyn (Fn(WasmerLinkerHandle, i32, i32) -> Result<Vec<u8>>) + Send + Sync + 'static>,
126}
127
128impl RuntimeInstanceData {
129    pub fn new(
130        func: impl (Fn(WasmerLinkerHandle, i32, i32) -> Result<Vec<u8>>) + Send + Sync + 'static,
131    ) -> Self {
132        Self {
133            memory: None,
134            abi_alloc: None,
135            func: Arc::new(func),
136        }
137    }
138
139    pub fn set_memory(&mut self, memory: Memory) {
140        self.memory = Some(memory);
141    }
142
143    pub fn memory(&self) -> &Memory {
144        self.memory.as_ref().expect("memory should be set first")
145    }
146
147    pub fn set_abi_alloc(&mut self, func: TypedFunction<i32, i32>) {
148        self.abi_alloc = Some(func);
149    }
150
151    pub fn call(mut this: FunctionEnvMut<Self>, len: i32, ptr: i32) -> u64 {
152        // TODO: should we unwrap here?
153        unsafe {
154            let memory = this.data().memory().clone();
155            let data = {
156                let func = this.data().func.clone();
157                let handle = WasmerLinkerHandle {
158                    store: this.as_store_mut(),
159                    memory: memory.clone(),
160                };
161                (func)(handle, ptr, len).unwrap()
162            };
163            let abi_alloc = this.data().abi_alloc.clone().unwrap();
164            let ptr = abi_alloc.call(&mut this, data.len() as i32).unwrap();
165            mem_slice_mut(&this, &memory, ptr, data.len() as i32, |slice| {
166                slice.copy_from_slice(&data);
167            });
168            ((data.len() as u64) << 32) | (ptr as u64)
169        }
170    }
171}
172
173/// A Wasmer [`Store`] with some imports.
174pub struct WasmerLinker {
175    store: HostStore,
176    imports: HashMap<String, HashMap<String, WasmerFunction>>,
177}
178
179impl WasmerLinker {
180    fn wrap_impl(
181        store: &mut impl AsStoreMut,
182        func: WasmerFunction,
183    ) -> (Function, Option<FunctionEnv<RuntimeInstanceData>>) {
184        let env_data = func.0;
185        let env = FunctionEnv::new(store, env_data);
186        let func = Function::new_typed_with_env(store, &env, RuntimeInstanceData::call);
187        (func, Some(env))
188    }
189}
190
191impl Linker<WasmerModule> for WasmerLinker {
192    type Config = ();
193
194    fn new(_: ()) -> Result<Self> {
195        let store = Store::default();
196        Ok(Self {
197            store: Arc::new(Mutex::new(store)),
198            imports: HashMap::default(),
199        })
200    }
201
202    fn create(&self, binary: &[u8]) -> Result<WasmerModule> {
203        let instance = {
204            let mut store = self.store.lock().unwrap();
205            let module = Module::from_binary(&store.as_store_ref(), binary)?;
206            let mut imports = Imports::default();
207            let mut envs = vec![];
208            for (ns, funcs) in &self.imports {
209                imports.register_namespace(
210                    ns,
211                    funcs.iter().map(|(name, func)| {
212                        (
213                            name.clone(),
214                            Extern::Function({
215                                let (func, env) =
216                                    Self::wrap_impl(&mut store.as_store_mut(), func.clone());
217                                if let Some(env) = env {
218                                    envs.push(env)
219                                }
220                                func
221                            }),
222                        )
223                    }),
224                );
225            }
226            let instance = Instance::new(&mut store.as_store_mut(), &module, &imports)?;
227            let memory = instance.exports.get_memory(MEMORY_NAME)?;
228            let abi_alloc = instance
229                .exports
230                .get_typed_function(&store.as_store_ref(), ABI_ALLOC_NAME)?;
231            let mut store = store.as_store_mut();
232            for env in &envs {
233                let env_mut = env.as_mut(&mut store);
234                env_mut.set_memory(memory.clone());
235                env_mut.set_abi_alloc(abi_alloc.clone());
236            }
237            instance
238        };
239        let host = WasmerModule::new(self.store.clone(), instance)?;
240        Ok(host)
241    }
242
243    fn import(
244        &mut self,
245        ns: impl Into<String>,
246        funcs: HashMap<String, WasmerFunction>,
247    ) -> Result<()> {
248        self.imports.insert(ns.into(), funcs);
249        Ok(())
250    }
251
252    fn wrap_raw(
253        &self,
254        f: impl (Fn(WasmerLinkerHandle, i32, i32) -> Result<Vec<u8>>) + Send + Sync + 'static,
255    ) -> WasmerFunction {
256        WasmerFunction(RuntimeInstanceData::new(f))
257    }
258}
259
260/// Represents a wrapped wasmer function.
261#[derive(Clone)]
262pub struct WasmerFunction(RuntimeInstanceData);
263
264/// A Wasmer [`StoreMut`].
265pub struct WasmerLinkerHandle<'a> {
266    store: StoreMut<'a>,
267    memory: Memory,
268}
269
270impl<'a> LinkerHandle<'a, WasmerModule> for WasmerLinkerHandle<'a> {
271    fn call<T>(
272        &mut self,
273        m: &WasmerModule,
274        name: &str,
275        data: &[u8],
276        f: impl FnOnce(&[u8]) -> Result<T>,
277    ) -> Result<T> {
278        m.call_impl(self.store.as_store_mut(), name, data, f)
279    }
280
281    fn slice<T>(&self, start: i32, len: i32, f: impl FnOnce(&[u8]) -> T) -> T {
282        unsafe { mem_slice(&self.store.as_store_ref(), &self.memory, start, len, f) }
283    }
284
285    fn slice_mut<T>(&mut self, start: i32, len: i32, f: impl FnOnce(&mut [u8]) -> T) -> T {
286        unsafe { mem_slice_mut(&self.store.as_store_ref(), &self.memory, start, len, f) }
287    }
288}
ayaka_plugin_wasmer/lib.rs

ayaka_plugin_wasmer/
lib.rs
