use std::any::Any; use std::fmt::{Debug, Formatter}; use std::mem::{align_of, forget, size_of}; use std::sync::{Arc, Weak}; use hashbrown::{HashMap, HashSet}; use parking_lot::lock_api::RwLockReadGuard; use parking_lot::{RawRwLock, RwLock}; use wasmer::{ AsStoreMut, AsStoreRef, EngineBuilder, Extern, Features, Function, FunctionEnv, Imports, Instance, Memory, Module, Store, StoreMut, StoreRef, TypedFunction, Value, }; use crate::dynamic_data::{ItemScript, Script, ScriptOwnerData, ScriptResolver}; use crate::script_implementations::wasm::export_registry::register_webassembly_funcs; use crate::script_implementations::wasm::extern_ref::ExternRef; use crate::script_implementations::wasm::script::WebAssemblyScript; use crate::script_implementations::wasm::script_function_cache::ScriptFunctionCache; use crate::script_implementations::wasm::temp_wasm_allocator::{AllocatedObject, TempWasmAllocator}; use crate::script_implementations::wasm::WebAssemblyScriptCapabilities; use crate::static_data::Item; use crate::{PkmnResult, ScriptCategory, StringKey, ValueIdentifiable, ValueIdentifier}; /// A WebAssembly script resolver implements the dynamic scripts functionality with WebAssembly. pub struct WebAssemblyScriptResolver { /// A unique identifier so we know what value this is. identifier: ValueIdentifier, /// The global state storage of WASM. _store: *mut Store, /// The WASM modules we have loaded. modules: Vec, /// Our currently loaded WASM instances. Empty until finalize() is called, after which the loaded modules get turned /// into actual instances. instances: Vec, /// This is the WASM function to load a script. load_script_fn: Option), u32>>, /// The data for use in the scripting function calls. environment_data: Arc, } /// This struct allows us to index a hashmap with both a category and name of a script. #[derive(Debug, Clone, Eq, PartialEq, Hash)] struct ScriptCapabilitiesKey { /// The category for the script we're looking for capabilities for. category: ScriptCategory, /// The name of the script we're looking for capabilities for. script_key: StringKey, } impl WebAssemblyScriptResolver { /// Instantiates a new WebAssemblyScriptResolver. pub fn new() -> Box { let compiler = wasmer::Cranelift::default(); let mut features = Features::new(); features.multi_value = true; features.reference_types = true; let engine = EngineBuilder::new(compiler).set_features(Some(features)); let mut store = Box::new(Store::new(engine)); let store_ptr: *mut Store = store.as_mut(); forget(store); let environment = Arc::new(WebAssemblyEnvironmentData::new(store_ptr)); environment.self_arc.write().replace(Arc::downgrade(&environment)); let s = Self { identifier: Default::default(), _store: store_ptr, modules: Default::default(), instances: Default::default(), load_script_fn: None, environment_data: environment, }; Box::new(s) } /// Get an immutable reference to the current WASM Store. fn store_ref(&self) -> StoreRef<'_> { unsafe { self._store.as_ref().unwrap().as_store_ref() } } /// Get a mutable reference to the current WASM Store. fn store_mut(&self) -> StoreMut<'_> { unsafe { self._store.as_mut().unwrap().as_store_mut() } } /// Load a compiled WASM module. pub fn load_wasm_from_bytes(&mut self, bytes: &[u8]) { // FIXME: Error handling let module = Module::new(&self.store_ref(), bytes).unwrap(); self.modules.push(module); } /// Tells the script resolver we're done loading wasm modules, and to finalize the resolver. pub fn finalize(&mut self) { let mut imports = Imports::new(); let env = FunctionEnv::new( &mut self.store_mut(), WebAssemblyEnv::new(Arc::downgrade(&self.environment_data)), ); register_webassembly_funcs(&mut imports, &mut self.store_mut(), &env); for module in &self.modules { for import in module.imports() { if imports.get_export("env", import.name()).is_none() { println!( "\x1b[91mMissing import: \"{}\" with type: {:?} \x1b[0m", import.name(), import.ty() ); } } let instance = Instance::new(&mut self.store_mut(), module, &imports).unwrap(); let exports = &instance.exports; let init_fn = exports.get_extern("_init"); if let Some(Extern::Function(init_fn)) = init_fn { init_fn.call(&mut self.store_mut(), &[]).unwrap(); } let mut exported_functions = self.environment_data.exported_functions.write(); for export in exports.iter() { match export.1 { Extern::Function(f) => { exported_functions.insert(export.0.as_str().into(), f.clone()); } Extern::Memory(m) => { let _ = self.environment_data.memory.write().insert(m.clone()); } _ => {} } } if let Some(m) = &self.environment_data.memory.read().as_ref() { m.grow(&mut self.store_mut(), 32).unwrap(); } if let Some(f) = exported_functions.get::(&"load_script".into()) { self.load_script_fn = Some(f.typed(&self.store_ref()).unwrap()) } if let Some(f) = exported_functions.get::(&"allocate_mem".into()) { let _ = self .environment_data .allocate_mem_fn .write() .insert(f.typed(&self.store_ref()).unwrap()); let temp_memory_slab = self.environment_data.allocate_mem(128, 1); let _ = self .environment_data .temp_allocator .write() .insert(TempWasmAllocator::new(temp_memory_slab.0, temp_memory_slab.1)); } self.instances.push(instance); } } /// Gets the data passed to every function as environment data. pub fn environment_data(&self) -> &Arc { &self.environment_data } } impl ValueIdentifiable for WebAssemblyScriptResolver { fn value_identifier(&self) -> ValueIdentifier { self.identifier } } impl ScriptResolver for WebAssemblyScriptResolver { fn load_script( &self, owner: ScriptOwnerData, category: ScriptCategory, script_key: &StringKey, ) -> PkmnResult>> { let script = self .load_script_fn .as_ref() .unwrap() .call( &mut self.store_mut(), category as u8, ExternRef::new_with_resolver(self, script_key), ) .unwrap(); self.environment_data.setup_script(script, category, script_key, owner) } fn load_item_script(&self, _key: &dyn Item) -> PkmnResult>> { todo!() } } impl Debug for WebAssemblyScriptResolver { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.write_str("WebAssemblyScriptResolver") } } impl Drop for WebAssemblyScriptResolver { fn drop(&mut self) { unsafe { drop(Box::from_raw(self._store)); } } } /// This data is what is passed to every function that requires access to the global runtime context. pub struct WebAssemblyEnvironmentData { /// We currently have a hacky implementation of extern refs while we're waiting for ExternRef support to hit the /// wasm32-unknown-unknown target of Rust. As we don't want to pass raw memory pointers to WASM for security reasons, /// we instead keep track of all the data we've sent to WASM, and pass the ID of that data to WASM. This allows us /// to only operate on data we know WASM owns. We currently store this data in this continuous Vec, and give the index /// of the data as the ID. extern_ref_pointers: RwLock>, /// To make sure we send the same identifier to WASM when we send the same piece of data multiple times, we have a /// backwards lookup on extern_ref_pointers. This allows us to get the index for a given piece of data. extern_ref_pointers_lookup: RwLock>, /// As an added security measure on our extern refs, we keep track of the types of the extern ref data we've sent. /// This prevents illegal arbitrary memory operations, where we expect type X, but the actual type is Y, which would /// allow for modifying memory we might not want to. If we get a type mismatch, we will panic, preventing this. extern_ref_type_lookup: RwLock>, /// Additional security for data slices passed to WASM. extern_vec_ref_lookup: RwLock>>, /// The memory inside of the WASM container. memory: RwLock>, /// This is a map of all the functions that WASM gives us. exported_functions: RwLock>, /// A cache of all script functions for faster calls. script_function_cache: ScriptFunctionCache, /// This is the WASM function to allocate memory inside the WASM container. allocate_mem_fn: RwLock>>, /// An allocator for quick short lifetime allocations within WASM. temp_allocator: RwLock>, /// The WASM store. store: *mut Store, /// Script capabilities tell us which functions are implemented on a given script. This allows us to skip unneeded /// WASM calls. script_capabilities: RwLock>>>, /// A weak reference to ourselves. self_arc: RwLock>>, /// A lookup from WASM memory pointer to their actual script wrappers. loaded_scripts: RwLock>>, } /// A quick lookup so we can find the extern ref of the value. #[allow(clippy::vtable_address_comparisons)] #[derive(Clone, Eq, PartialEq, Hash)] struct ExternRefLookupKey { /// The raw pointer to the data pub ptr: *const dyn Any, /// Whether or not the reference is a Vec pub is_vec: bool, } impl WebAssemblyEnvironmentData { /// Instantiates new Environment data with a given store. pub(crate) fn new(store: *mut Store) -> Self { Self { extern_ref_pointers: Default::default(), extern_ref_pointers_lookup: Default::default(), extern_ref_type_lookup: Default::default(), extern_vec_ref_lookup: Default::default(), memory: Default::default(), exported_functions: Default::default(), script_function_cache: Default::default(), allocate_mem_fn: Default::default(), temp_allocator: Default::default(), store, script_capabilities: Default::default(), self_arc: Default::default(), loaded_scripts: Default::default(), } } /// This returns the memory of the WASM container. pub fn memory(&self) -> *mut u8 { self.memory.read().as_ref().unwrap().view(&self.store_ref()).data_ptr() } /// Return a pointer to something inside the WASM memory. pub fn get_raw_pointer(&self, offset: u32) -> *mut T { unsafe { self.memory .read() .as_ref() .unwrap() .view(&self.store_ref()) .data_ptr() .offset(offset as isize) as *mut T } } /// This returns the functions exported from WASM. pub fn exported_functions(&self) -> RwLockReadGuard<'_, RawRwLock, HashMap> { self.exported_functions.read() } /// pub(super) fn script_function_cache(&self) -> &ScriptFunctionCache { &self.script_function_cache } /// Allocates memory inside the WASM container with a given size and alignment. This memory is /// owned by WASM, and is how we can pass memory references that the host allocated to WASM. /// The return is a tuple containing both the actual pointer to the memory (usable by the host), /// and the WASM offset to the memory (usable by the client). pub fn allocate_mem(&self, size: u32, align: u32) -> (*mut u8, u32) { let wasm_ptr = self .allocate_mem_fn .read() .as_ref() .unwrap() .call(&mut self.store_mut(), size, align) .unwrap(); unsafe { (self.memory().offset(wasm_ptr as isize), wasm_ptr) } } /// Allocates memory inside the WASM container with a given size and alignment. This memory is /// owned by WASM, and is how we can pass memory references that the host allocated to WASM. /// The return is a tuple containing both the actual pointer to the memory (usable by the host), /// and the WASM offset to the memory (usable by the client). pub fn allocate_mem_typed(&self) -> (*mut u8, u32) { let wasm_ptr = self .allocate_mem_fn .read() .as_ref() .unwrap() .call(&mut self.store_mut(), size_of::() as u32, align_of::() as u32) .unwrap(); unsafe { (self.memory().offset(wasm_ptr as isize), wasm_ptr) } } /// Allocates a raw array in wasm, and copy the values from a given slice to it. Returns the /// pointer in wasm memory. pub fn copy_value_vec_to_wasm(&self, v: &[T]) -> u32 { let wasm_ptr = self .allocate_mem_fn .read() .as_ref() .unwrap() .call( &mut self.store_mut(), (size_of::() * v.len()) as u32, align_of::() as u32, ) .unwrap(); unsafe { let raw = self.memory().offset(wasm_ptr as isize) as *mut T; v.as_ptr().copy_to(raw, v.len()); } wasm_ptr } /// Allocate a piece of memory inside WASM with a very short lifespan. This is mainly used for /// rapid allocation of script function parameters, where WASM needs to write to a specific /// pointer. pub(super) fn allocate_temp(&self, value: T) -> AllocatedObject { self.temp_allocator.read().as_ref().unwrap().alloc::(value) } /// Get a numeric value from any given value. This is not a true Extern Ref from WASM, as this /// is not supported by our current WASM platform (Rust). Instead, this is simply a way to not /// have to send arbitrary pointer values back and forth with WASM. Only values WASM can actually /// access can be touched through this, and we ensure the value is the correct type. In the future, /// when extern refs get actually properly implemented at compile time we might want to get rid /// of this code. #[inline(always)] pub fn get_extern_ref_index(&self, value: &dyn Any) -> usize { self.get_extern_ref_from_identifier::(value, false) } /// Get a numeric value from any given value. This is not a true Extern Ref from WASM, as this /// is not supported by our current WASM platform (Rust). Instead, this is simply a way to not /// have to send arbitrary pointer values back and forth with WASM. Only values WASM can actually /// access can be touched through this, and we ensure the value is the correct type. In the future, /// when extern refs get actually properly implemented at compile time we might want to get rid /// of this code. pub fn get_extern_vec_ref_index(&self, value: &Vec) -> usize { let mut vec = Vec::with_capacity(value.len()); for v in value { vec.push(self.get_extern_ref_index::(v)); } let p = self.get_extern_ref_from_identifier::>(value, true); self.extern_vec_ref_lookup.write().insert(p, vec); p } /// Get an extern ref belonging to a vector we have passed to WASM. pub fn get_extern_vec_ref_extern_ref(&self, extern_vec_ref: usize, index: usize) -> usize { let r = self.extern_vec_ref_lookup.read(); let v = r.get(&extern_vec_ref).unwrap(); v[index] } /// Gets the extern ref index belonging to a specific pointer. If none exists, this will create /// a new one. #[inline(always)] fn get_extern_ref_from_identifier(&self, value: &dyn Any, is_vec: bool) -> usize { if let Some(v) = self.extern_ref_pointers_lookup.read().get(&ExternRefLookupKey { ptr: value as *const dyn Any, is_vec, }) { return *v; } let index = { let mut extern_ref_guard = self.extern_ref_pointers.write(); extern_ref_guard.push(value as *const dyn Any); extern_ref_guard.len() }; self.extern_ref_pointers_lookup.write().insert( ExternRefLookupKey { ptr: value as *const dyn Any, is_vec, }, index, ); self.extern_ref_type_lookup.write().insert(ExternRefLookupKey { ptr: value as *const dyn Any, is_vec, }); index } /// Gets a value from the extern ref lookup. This turns an earlier registered index back into /// its proper value, validates its type, and returns the value. pub fn get_extern_ref_value<'a, T: ?Sized>(&self, index: usize) -> &'a dyn Any { let read_guard = self.extern_ref_pointers.read(); let ptr = read_guard.get(index - 1).unwrap(); if self .extern_ref_type_lookup .read() .get(&ExternRefLookupKey { ptr: *ptr, is_vec: false, }) .is_none() { panic!( "Extern ref was accessed with wrong type. Requested type {}, but this was not the type the extern ref was stored with.", std::any::type_name::() ); } unsafe { ptr.as_ref().unwrap() } } /// The WASM store. pub fn store_ref(&self) -> StoreRef<'_> { unsafe { self.store.as_ref().unwrap().as_store_ref() } } /// The mutable WASM store. pub fn store_mut(&self) -> StoreMut<'_> { unsafe { self.store.as_mut().unwrap().as_store_mut() } } /// Find a loaded script based on the pointer in WASM memory. pub(crate) fn get_loaded_script(&self, wasm_ptr: u32) -> Option> { if let Some(script) = self.loaded_scripts.read().get(&wasm_ptr) { script.upgrade() } else { None } } /// Wrap a script pointer in WASM memory into a host managed script. pub fn setup_script( &self, script_ptr: u32, category: ScriptCategory, script_key: &StringKey, owner: ScriptOwnerData, ) -> PkmnResult>> { if script_ptr == 0 { return Ok(None); } let key = ScriptCapabilitiesKey { category, script_key: script_key.clone(), }; if !self.script_capabilities.read().contains_key(&key) { let mut capabilities = HashSet::new(); unsafe { if let Some(get_cap) = self .exported_functions .read() .get::(&"get_script_capabilities".into()) { let res = get_cap .call(&mut self.store_mut(), &[Value::I32(script_ptr as i32)]) .unwrap(); let ptr = (self.memory() as *const WebAssemblyScriptCapabilities).offset(res[0].i32().unwrap() as isize); let length = res[1].i32().unwrap() as usize; for i in 0..length { capabilities.insert(*ptr.add(i)); } } } self.script_capabilities .write() .insert(key.clone(), Arc::new(capabilities)); } let read_guard = self.script_capabilities.read(); let capabilities = read_guard.get(&key).unwrap(); let script = Arc::new(WebAssemblyScript::new( owner, script_ptr, capabilities.clone(), self.self_arc.read().as_ref().unwrap().upgrade().unwrap(), script_key.clone(), )); self.loaded_scripts.write().insert(script_ptr, Arc::downgrade(&script)); Ok(Some(script)) } } /// The runtime environment for script execution. This is passed to most of the host functions being called. #[derive(Clone)] pub(crate) struct WebAssemblyEnv { /// A pointer to the WebAssemblyScriptResolver belonging to the current script environment. data: Weak, } impl WebAssemblyEnv { /// Instantiates a new Environment with the requested data. pub fn new(data: Weak) -> Self { Self { data } } /// Get the actual data belonging to the current context. pub fn data(&self) -> Arc { self.data.upgrade().unwrap() } } unsafe impl Sync for WebAssemblyEnv {} unsafe impl Send for WebAssemblyEnv {} unsafe impl Sync for WebAssemblyEnvironmentData {} unsafe impl Send for WebAssemblyEnvironmentData {}