541 lines
21 KiB
Rust
Executable File
541 lines
21 KiB
Rust
Executable File
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 unique_type_id::UniqueTypeId;
|
|
use wasmer::{
|
|
AsStoreMut, AsStoreRef, Cranelift, 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};
|
|
|
|
/// A WebAssembly script resolver implements the dynamic scripts functionality with WebAssembly.
|
|
pub struct WebAssemblyScriptResolver {
|
|
/// The global state storage of WASM.
|
|
_store: *mut Store,
|
|
/// The WASM modules we have loaded.
|
|
modules: Vec<Module>,
|
|
/// Our currently loaded WASM instances. Empty until finalize() is called, after which the loaded modules get turned
|
|
/// into actual instances.
|
|
instances: Vec<Instance>,
|
|
|
|
/// This is the WASM function to load a script.
|
|
load_script_fn: Option<TypedFunction<(u8, ExternRef<StringKey>), u32>>,
|
|
|
|
/// The data for use in the scripting function calls.
|
|
environment_data: Arc<WebAssemblyEnvironmentData>,
|
|
}
|
|
|
|
/// 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<WebAssemblyScriptResolver> {
|
|
let compiler = 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 {
|
|
_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);
|
|
}
|
|
|
|
/// Initialise all the data we need.
|
|
pub fn finalize(&mut self) {
|
|
let mut imports = Imports::new();
|
|
//let mut exports = Exports::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);
|
|
//imports.register("env", exports);
|
|
|
|
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::<StringKey>(&"load_script".into()) {
|
|
self.load_script_fn = Some(f.typed(&self.store_ref()).unwrap())
|
|
}
|
|
if let Some(f) = exported_functions.get::<StringKey>(&"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<WebAssemblyEnvironmentData> {
|
|
&self.environment_data
|
|
}
|
|
}
|
|
|
|
impl ScriptResolver for WebAssemblyScriptResolver {
|
|
fn load_script(
|
|
&self,
|
|
owner: ScriptOwnerData,
|
|
category: ScriptCategory,
|
|
script_key: &StringKey,
|
|
) -> PkmnResult<Option<Arc<dyn Script>>> {
|
|
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: &Item) -> PkmnResult<Option<Arc<dyn ItemScript>>> {
|
|
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<Vec<*const u8>>,
|
|
/// 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<HashMap<ExternRefLookupKey, u32>>,
|
|
/// 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<HashSet<ExternRefLookupKey>>,
|
|
|
|
/// Additional security for data slices passed to WASM.
|
|
extern_vec_ref_lookup: RwLock<HashMap<u32, Vec<u32>>>,
|
|
|
|
/// The memory inside of the WASM container.
|
|
memory: RwLock<Option<Memory>>,
|
|
|
|
/// This is a map of all the functions that WASM gives us.
|
|
exported_functions: RwLock<HashMap<StringKey, Function>>,
|
|
|
|
/// 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<Option<TypedFunction<(u32, u32), u32>>>,
|
|
|
|
/// An allocator for quick short lifetime allocations within WASM.
|
|
temp_allocator: RwLock<Option<TempWasmAllocator>>,
|
|
|
|
/// 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<HashMap<ScriptCapabilitiesKey, Arc<HashSet<WebAssemblyScriptCapabilities>>>>,
|
|
|
|
/// A weak reference to ourselves.
|
|
self_arc: RwLock<Option<Weak<Self>>>,
|
|
|
|
/// A lookup from WASM memory pointer to their actual script wrappers.
|
|
loaded_scripts: RwLock<HashMap<u32, Weak<WebAssemblyScript>>>,
|
|
}
|
|
|
|
/// A quick lookup so we can find the extern ref of the value.
|
|
#[derive(Clone, Eq, PartialEq, Hash)]
|
|
struct ExternRefLookupKey {
|
|
/// The raw pointer to the data
|
|
pub ptr: *const u8,
|
|
/// Whether or not the reference is a Vec
|
|
pub is_vec: bool,
|
|
/// The unique identifier of the type.
|
|
pub t: u64,
|
|
}
|
|
|
|
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<T>(&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<StringKey, Function>> {
|
|
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<T>(&self) -> (*mut u8, u32) {
|
|
let wasm_ptr = self
|
|
.allocate_mem_fn
|
|
.read()
|
|
.as_ref()
|
|
.unwrap()
|
|
.call(&mut self.store_mut(), size_of::<T>() as u32, align_of::<T>() as u32)
|
|
.unwrap();
|
|
unsafe { (self.memory().offset(wasm_ptr as isize), 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<T>(&self, value: T) -> AllocatedObject<T> {
|
|
self.temp_allocator.read().as_ref().unwrap().alloc::<T>(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<T: UniqueTypeId<u64> + ?Sized>(&self, value: &T) -> u32 {
|
|
self.get_extern_ref_from_ptr(value as *const T as *const u8, T::id().0, 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<T: UniqueTypeId<u64>>(&self, value: &[T]) -> u32 {
|
|
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_ptr(value as *const [T] as *const u8, T::id().0, 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: u32, index: u32) -> u32 {
|
|
let r = self.extern_vec_ref_lookup.read();
|
|
let v = r.get(&extern_vec_ref).unwrap();
|
|
v[index as usize]
|
|
}
|
|
|
|
/// 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_ptr(&self, ptr: *const u8, type_id: u64, is_vec: bool) -> u32 {
|
|
if let Some(v) = self.extern_ref_pointers_lookup.read().get(&ExternRefLookupKey {
|
|
ptr,
|
|
is_vec,
|
|
t: type_id,
|
|
}) {
|
|
return *v as u32;
|
|
}
|
|
let index = {
|
|
let mut extern_ref_guard = self.extern_ref_pointers.write();
|
|
extern_ref_guard.push(ptr);
|
|
extern_ref_guard.len() as u32
|
|
};
|
|
self.extern_ref_pointers_lookup.write().insert(
|
|
ExternRefLookupKey {
|
|
ptr,
|
|
is_vec,
|
|
t: type_id,
|
|
},
|
|
index,
|
|
);
|
|
self.extern_ref_type_lookup.write().insert(ExternRefLookupKey {
|
|
ptr,
|
|
is_vec,
|
|
t: type_id,
|
|
});
|
|
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<T: UniqueTypeId<u64>>(&self, index: u32) -> &T {
|
|
let read_guard = self.extern_ref_pointers.read();
|
|
let ptr = read_guard.get((index - 1) as usize).unwrap();
|
|
if self
|
|
.extern_ref_type_lookup
|
|
.read()
|
|
.get(&ExternRefLookupKey {
|
|
ptr: *ptr,
|
|
is_vec: false,
|
|
t: T::id().0,
|
|
})
|
|
.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::<T>()
|
|
);
|
|
}
|
|
|
|
unsafe { (*ptr as *const T).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<Arc<WebAssemblyScript>> {
|
|
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<Option<Arc<dyn Script>>> {
|
|
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::<StringKey>(&"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<WebAssemblyEnvironmentData>,
|
|
}
|
|
|
|
impl WebAssemblyEnv {
|
|
/// Instantiates a new Environment with the requested data.
|
|
pub fn new(data: Weak<WebAssemblyEnvironmentData>) -> Self {
|
|
Self { data }
|
|
}
|
|
|
|
/// Get the actual data belonging to the current context.
|
|
pub fn data(&self) -> Arc<WebAssemblyEnvironmentData> {
|
|
self.data.upgrade().unwrap()
|
|
}
|
|
}
|
|
|
|
unsafe impl Sync for WebAssemblyEnv {}
|
|
|
|
unsafe impl Send for WebAssemblyEnv {}
|
|
|
|
unsafe impl Sync for WebAssemblyEnvironmentData {}
|
|
|
|
unsafe impl Send for WebAssemblyEnvironmentData {}
|