Lines
87.06 %
Functions
17.86 %
Branches
100 %
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use wasmtime::{Engine, Linker, Module, Store};
use crate::error::HookError;
use crate::format::{BASE_OFFSET, GlobalHeader, OUTPUT_HEADER_SIZE, OutputHeader};
use crate::host::{HostState, define_host_functions};
use crate::parser::{OutputParser, ParsedEntity};
const DEFAULT_OUTPUT_SIZE: u32 = 64 * 1024;
const WASM_PAGE_SIZE: u32 = 65536;
pub struct ScriptExecutor {
engine: Engine,
module_cache: Arc<Mutex<HashMap<Vec<u8>, Module>>>,
}
impl Default for ScriptExecutor {
fn default() -> Self {
Self::new()
impl ScriptExecutor {
#[must_use]
pub fn new() -> Self {
Self {
engine: Engine::default(),
module_cache: Arc::new(Mutex::new(HashMap::new())),
pub fn with_engine(engine: Engine) -> Self {
engine,
fn get_or_compile_module(&self, bytecode: &[u8]) -> Result<Module, HookError> {
let cache = self.module_cache.lock()?;
if let Some(module) = cache.get(bytecode) {
return Ok(module.clone());
drop(cache);
let module = Module::new(&self.engine, bytecode)?;
let mut cache = self.module_cache.lock()?;
cache.insert(bytecode.to_vec(), module.clone());
Ok(module)
pub fn execute(
&self,
bytecode: &[u8],
input: &[u8],
output_size: Option<u32>,
) -> Result<Vec<ParsedEntity>, HookError> {
let output_size = output_size.unwrap_or(DEFAULT_OUTPUT_SIZE);
let module = self.get_or_compile_module(bytecode)?;
let header = GlobalHeader::from_bytes(input)
.ok_or_else(|| HookError::Parse("Invalid input header".to_string()))?;
let input_offset = BASE_OFFSET;
let output_offset = input_offset + input.len() as u32;
let strings_offset = header.strings_pool_offset;
let host_state = HostState::new(input_offset, output_offset, strings_offset);
let mut store = Store::new(&self.engine, host_state);
let mut linker = Linker::new(&self.engine);
define_host_functions(&mut linker)?;
let instance = linker.instantiate(&mut store, &module)?;
let memory = instance
.get_memory(&mut store, "memory")
.ok_or(HookError::WASMMem)?;
store.data_mut().memory = Some(memory);
let total_size = input.len() + output_size as usize;
let required_pages = (BASE_OFFSET as usize + total_size).div_ceil(WASM_PAGE_SIZE as usize);
let current_pages = memory.size(&store) as usize;
if required_pages > current_pages {
memory.grow(&mut store, (required_pages - current_pages) as u64)?;
let mem_data = memory.data_mut(&mut store);
let input_start = BASE_OFFSET as usize;
mem_data[input_start..input_start + input.len()].copy_from_slice(input);
let output_start = output_offset as usize;
let input_entity_count = header.input_entity_count;
let output_header = OutputHeader::new(input_entity_count);
mem_data[output_start..output_start + OUTPUT_HEADER_SIZE]
.copy_from_slice(&output_header.to_bytes());
let should_apply = instance
.get_typed_func::<(), i32>(&mut store, "should_apply")
.map_err(|e| HookError::Script(format!("Missing should_apply export: {e}")))?;
let result = should_apply.call(&mut store, ())?;
if result == 0 {
return Ok(Vec::new());
let process = instance
.get_typed_func::<(), ()>(&mut store, "process")
.map_err(|e| HookError::Script(format!("Missing process export: {e}")))?;
process.call(&mut store, ())?;
let mem_data = memory.data(&store);
let output_data = &mem_data[output_start..output_start + output_size as usize];
let output_header = OutputHeader::from_bytes(output_data)
.ok_or_else(|| HookError::Parse("Invalid output header".to_string()))?;
let output_strings_offset = { output_header.strings_offset } as usize;
let parser = OutputParser::new(output_data, output_strings_offset)?;
parser.entities().collect()
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_executor_creation() {
let executor = ScriptExecutor::new();
assert!(executor.module_cache.lock().unwrap().is_empty());