Lines
56.54 %
Functions
26.79 %
Branches
100 %
use std::collections::HashMap;
use std::sync::{Arc, Mutex, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};
use nomiscript::SymbolTable;
use wasmtime::{Caller, Engine, Linker, Memory, Module};
pub struct WasmHost {
engine: Engine,
symbol_table: Arc<RwLock<SymbolTable>>,
module_cache: Arc<Mutex<HashMap<Vec<u8>, Module>>>,
}
impl WasmHost {
#[must_use]
pub fn new(engine: Engine, symbol_table: SymbolTable) -> Self {
Self {
engine,
symbol_table: Arc::new(RwLock::new(symbol_table)),
module_cache: Arc::new(Mutex::new(HashMap::new())),
pub fn engine(&self) -> &Engine {
&self.engine
pub fn symbol_table(&self) -> &Arc<RwLock<SymbolTable>> {
&self.symbol_table
pub fn module_cache(&self) -> &Arc<Mutex<HashMap<Vec<u8>, Module>>> {
&self.module_cache
pub fn execution_state(
&self,
input_offset: u32,
output_offset: u32,
strings_offset: u32,
) -> ExecutionState {
ExecutionState {
input_offset,
output_offset,
strings_offset,
output_strings_offset: Arc::new(Mutex::new(0)),
memory: None,
symbol_table: Arc::clone(&self.symbol_table),
pub struct ExecutionState {
pub input_offset: u32,
pub output_offset: u32,
pub strings_offset: u32,
pub output_strings_offset: Arc<Mutex<u32>>,
pub memory: Option<Memory>,
pub symbol_table: Arc<RwLock<SymbolTable>>,
impl ExecutionState {
pub fn new(input_offset: u32, output_offset: u32, strings_offset: u32) -> Self {
symbol_table: Arc::new(RwLock::new(SymbolTable::new())),
pub fn define_host_functions(linker: &mut Linker<ExecutionState>) -> wasmtime::Result<()> {
linker.func_wrap(
"env",
"get_input_offset",
|caller: Caller<ExecutionState>| -> u32 { caller.data().input_offset },
)?;
"get_output_offset",
|caller: Caller<ExecutionState>| -> u32 { caller.data().output_offset },
"get_strings_offset",
|caller: Caller<ExecutionState>| -> u32 { caller.data().strings_offset },
"symbol_resolve",
|caller: Caller<ExecutionState>, _name_ptr: u32, _name_len: u32| {
let _memory = match caller.data().memory {
Some(mem) => mem,
None => return,
};
tracing::debug!(
name_ptr = _name_ptr,
name_len = _name_len,
"symbol_resolve called"
);
},
"write_bytes",
|mut caller: Caller<ExecutionState>, dst: u32, src: u32, len: u32| -> u32 {
let memory = match caller.data().memory {
None => return 0,
let data = memory.data_mut(&mut caller);
let src_start = src as usize;
let src_end = src_start + len as usize;
let dst_start = dst as usize;
if src_end > data.len() || dst_start + len as usize > data.len() {
return 0;
let bytes: Vec<u8> = data[src_start..src_end].to_vec();
data[dst_start..dst_start + len as usize].copy_from_slice(&bytes);
len
"write_string",
|mut caller: Caller<ExecutionState>, ptr: u32, len: u32| -> u32 {
let output_offset = caller.data().output_offset;
let output_strings = caller.data().output_strings_offset.clone();
let src_start = ptr as usize;
if src_end > data.len() {
let mut strings_offset = match output_strings.lock() {
Ok(guard) => guard,
Err(_) => return 0,
let current_offset = *strings_offset;
let dst = output_offset as usize + current_offset as usize;
if dst + len as usize > data.len() {
data[dst..dst + len as usize].copy_from_slice(&bytes);
*strings_offset += len;
current_offset
"log",
|caller: Caller<ExecutionState>, level: u32, msg_ptr: u32, msg_len: u32| {
tracing::debug!(level, msg_ptr, msg_len, "host log called");
let data = memory.data(&caller);
let start = msg_ptr as usize;
let end = start + msg_len as usize;
if end > data.len() {
return;
let msg = match std::str::from_utf8(&data[start..end]) {
Ok(s) => s,
Err(_) => return,
match level {
0 => tracing::debug!("[script] {msg}"),
1 => tracing::info!("[script] {msg}"),
2 => tracing::warn!("[script] {msg}"),
_ => tracing::error!("[script] {msg}"),
linker.func_wrap("env", "get_timestamp", || -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0)
})?;
"generate_uuid",
|mut caller: Caller<ExecutionState>, out_ptr: u32| {
let uuid_bytes = uuid::Uuid::new_v4().into_bytes();
let start = out_ptr as usize;
if start + 16 > data.len() {
data[start..start + 16].copy_from_slice(&uuid_bytes);
"get_input_entities_count",
|caller: Caller<ExecutionState>| -> i32 {
use crate::format::GlobalHeader;
let input_offset = caller.data().input_offset;
let input_start = input_offset as usize;
if input_start + std::mem::size_of::<GlobalHeader>() > data.len() {
if let Some(header) = GlobalHeader::from_bytes(&data[input_start..]) {
header.input_entity_count as i32
} else {
0
Ok(())
#[cfg(test)]
mod tests {
use super::*;
use crate::format::BASE_OFFSET;
#[test]
fn test_execution_state_creation() {
let state = ExecutionState::new(BASE_OFFSET, BASE_OFFSET + 1024, BASE_OFFSET + 512);
assert_eq!(state.input_offset, BASE_OFFSET);
assert_eq!(state.output_offset, BASE_OFFSET + 1024);
assert_eq!(state.strings_offset, BASE_OFFSET + 512);
fn test_wasm_host_creation() {
let host = WasmHost::new(Engine::default(), SymbolTable::new());
assert!(host.module_cache().lock().unwrap().is_empty());