PkmnLib_rs/src/script_implementations/wasm/script_resolver.rs

298 lines
12 KiB
Rust

use std::fmt::{Debug, Formatter};
use std::ops::DerefMut;
use std::sync::Arc;
use hashbrown::{HashMap, HashSet};
use parking_lot::lock_api::{MappedRwLockReadGuard, RwLockReadGuard};
use parking_lot::{RawRwLock, RwLock};
use unique_type_id::{TypeId, UniqueTypeId};
use wasmer::{
Cranelift, Exports, Extern, Features, Function, ImportObject, Instance, Memory, Module, NativeFunc, Store,
Universal, Value, WasmerEnv,
};
use crate::dynamic_data::{ItemScript, Script, 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::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: 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 a map of all the functions that WASM gives us.
exported_functions: HashMap<StringKey, Function>,
/// This is the WASM function to load a script.
load_script_fn: Option<NativeFunc<(u8, ExternRef<StringKey>), u32>>,
/// 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, HashSet<WebAssemblyScriptCapabilities>>>,
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 config = Cranelift::default();
let mut features = Features::new();
features.multi_value = true;
features.reference_types = true;
let universal = Universal::new(config).features(features);
let engine = universal.engine();
let store = Store::new(&engine);
let s = Self {
store,
modules: Default::default(),
instances: Default::default(),
exported_functions: Default::default(),
load_script_fn: None,
script_capabilities: Default::default(),
environment_data: Arc::new(Default::default()),
};
Box::new(s)
}
/// Load a compiled WASM module.
pub fn load_wasm_from_bytes(&mut self, bytes: &[u8]) {
// FIXME: Error handling
let module = Module::new(&self.store, bytes).unwrap();
self.modules.push(module);
}
/// Initialise all the data we need.
pub fn finalize(&mut self) {
let mut imports = ImportObject::new();
let mut exports = Exports::new();
let env = WebAssemblyEnv {
resolver: self.environment_data.clone(),
};
register_webassembly_funcs(&mut exports, &self.store, env);
imports.register("env", exports);
for module in &self.modules {
let instance = Instance::new(module, &imports).unwrap();
let exports = &instance.exports;
for export in exports.iter() {
match export.1 {
Extern::Function(f) => {
self.exported_functions.insert(export.0.as_str().into(), f.clone());
}
Extern::Memory(m) => {
self.environment_data.memory.write().insert(m.clone());
}
_ => {}
}
}
if let Some(m) = &self.environment_data.memory.read().as_ref() {
m.grow(32).unwrap();
}
if let Some(f) = self.exported_functions.get(&"load_script".into()) {
self.load_script_fn = Some(f.native().unwrap())
}
if let Some(f) = self.exported_functions.get(&"allocate_mem".into()) {
self.environment_data
.allocate_mem_fn
.write()
.insert(f.native().unwrap());
}
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: *const u8,
category: ScriptCategory,
script_key: &StringKey,
) -> PkmnResult<Option<Arc<dyn Script>>> {
let script = self
.load_script_fn
.as_ref()
.unwrap()
.call(category as u8, ExternRef::new_with_resolver(self, script_key))
.unwrap();
if script == 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.get(&"get_script_capabilities".into()) {
let res = get_cap.call(&[Value::I32(script as i32)]).unwrap();
let ptr = (self.environment_data.memory.read().as_ref().unwrap().data_ptr()
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(), capabilities);
}
let read_guard = self.script_capabilities.read();
let capabilities = read_guard.get(&key).unwrap();
Ok(Some(Arc::new(WebAssemblyScript::new(
owner as *mut u8,
script,
capabilities as *const HashSet<WebAssemblyScriptCapabilities>
as *mut HashSet<WebAssemblyScriptCapabilities>,
self as *const WebAssemblyScriptResolver as *mut WebAssemblyScriptResolver,
script_key.clone(),
))))
}
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")
}
}
/// This data is what is passed to every function that requires access to the global runtime context.
#[derive(Default)]
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<*const u8, 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<HashMap<*const u8, TypeId<u64>>>,
/// The memory inside of the WASM container.
memory: RwLock<Option<Memory>>,
/// This is the WASM function to allocate memory inside the WASM container.
allocate_mem_fn: RwLock<Option<NativeFunc<(u32, u32), u32>>>,
}
impl WebAssemblyEnvironmentData {
/// This returns the memory of the WASM container.
pub fn memory(&self) -> MappedRwLockReadGuard<'_, RawRwLock, Memory> {
RwLockReadGuard::map(self.memory.read(), |a| a.as_ref().unwrap())
}
/// 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) -> (*const u8, u32) {
let wasm_ptr = self.allocate_mem_fn.read().as_ref().unwrap().call(size, align).unwrap();
unsafe {
(
self.memory
.read()
.as_ref()
.unwrap()
.data_ptr()
.offset(wasm_ptr as isize),
wasm_ptr,
)
}
}
/// 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_ref_index<T: UniqueTypeId<u64>>(&self, value: &T) -> u32 {
let ptr = value as *const T as *const u8;
if let Some(v) = self.extern_ref_pointers_lookup.read().get(&ptr) {
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(ptr, index);
self.extern_ref_type_lookup.write().insert(ptr, T::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();
let expected_type_id = &self.extern_ref_type_lookup.read()[ptr];
if expected_type_id.0 != T::id().0 {
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 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.
pub resolver: Arc<WebAssemblyEnvironmentData>,
}
impl WebAssemblyEnv {
/// Get the WebAssemblyScriptResolver belonging to the current context.
pub fn resolver(&self) -> &Arc<WebAssemblyEnvironmentData> {
&self.resolver
}
}
unsafe impl Sync for WebAssemblyEnv {}
unsafe impl Send for WebAssemblyEnv {}
impl WasmerEnv for WebAssemblyEnv {}