znippy-common/src/plugins/wasm/wasm_loader.rs
//! WASM plugin loader with host-provided parallel decompression services.
//!
//! Host functions give WASM plugins access to native multi-core decompressors
//! (ljar-rs, lbzip2-rs, lgz-rs) without the plugins needing threading support.
use crate::plugin::{ExtensionRow, ExtensionValue};
#[cfg(feature = "wasm-plugins")]
use crate::plugin::{ArchiveTypePlugin, Col, Claim, HandlerSpec};
use std::collections::HashMap;
#[cfg(feature = "wasm-plugins")]
use wasmtime::*;
// ─── Host State (shared decompression services) ──────────────────────
/// Codec identifiers for host_decompress
#[repr(u32)]
pub enum HostCodec {
Deflate = 0,
Gzip = 1, // lgz-rs (parallel)
Bzip2 = 2, // lbzip2-rs (workers per core)
Zstd = 3,
}
/// Archive format identifiers for host_archive_open
#[repr(u32)]
pub enum ArchiveFormat {
Jar = 0, // ljar-rs (parallel JAR/ZIP)
TarGz = 1, // tar + lgz
TarBz2 = 2, // tar + lbzip2
}
/// An opened archive handle — holds only the entries that matched the open-time filter.
struct OpenArchive {
entries: Vec<ArchiveEntry>,
}
struct ArchiveEntry {
name: String,
data: Vec<u8>,
}
/// Host state accessible from WASM via host functions
struct HostState {
/// Currently open archives (handle → entries)
archives: HashMap<u32, OpenArchive>,
next_handle: u32,
}
impl HostState {
fn new() -> Self {
Self { archives: HashMap::new(), next_handle: 1 }
}
}
// ─── WASM Plugin ─────────────────────────────────────────────────────
#[cfg(feature = "wasm-plugins")]
pub struct WasmPlugin {
/// Name, type_id, extensions and columns — read once at load. Leaked to
/// `'static` so the module's discovery record is the same `HandlerSpec` a
/// native handler has (one per loaded module, a few hundred bytes).
spec: &'static HandlerSpec,
engine: Engine,
module: Module,
}
#[cfg(feature = "wasm-plugins")]
impl WasmPlugin {
/// Load a WASM plugin from file
pub fn load(wasm_path: &str, name: &str, type_id: i8) -> anyhow::Result<Self> {
let engine = Engine::default();
let module = Module::from_file(&engine, wasm_path)?;
let leak = |s: String| -> &'static str { Box::leak(s.into_boxed_str()) };
let extensions: Vec<&'static str> =
Self::fetch_extensions(&engine, &module).into_iter().map(leak).collect();
let fields: Vec<(&'static str, Col)> =
Self::fetch_schema(&engine, &module).into_iter().map(|(n, c)| (leak(n), c)).collect();
let spec = Box::leak(Box::new(HandlerSpec {
name: leak(name.to_string()),
aliases: &[],
type_id,
ecosystem: "wasm plugin",
extensions: Box::leak(extensions.into_boxed_slice()),
claim: Claim::Extensions,
description: leak(format!("wasm plugin loaded from {wasm_path}")),
fields: Box::leak(fields.into_boxed_slice()),
commands: &[],
}));
Ok(Self { spec, engine, module })
}
/// Instantiate the module, call a no-arg export that returns a pointer to a UTF-8 string
/// in the result buffer, and read it back. Returns None if the export is missing/fails.
fn call_string_export(engine: &Engine, module: &Module, export: &str) -> Option<String> {
let mut store = Store::new(engine, HostState::new());
let mut linker = Linker::new(engine);
register_host_functions(&mut linker).ok()?;
let instance = linker.instantiate(&mut store, module).ok()?;
let memory = instance.get_memory(&mut store, "memory")?;
let func = instance.get_typed_func::<(), u32>(&mut store, export).ok()?;
let ptr = func.call(&mut store, ()).ok()?;
let result_len = instance.get_typed_func::<(), u32>(&mut store, "result_len").ok()?;
let len = result_len.call(&mut store, ()).ok()? as usize;
let mut buf = vec![0u8; len];
memory.read(&store, ptr as usize, &mut buf).ok()?;
String::from_utf8(buf).ok()
}
/// Call `plugin_extensions` once and cache the result. Empty if the plugin omits it —
/// such a plugin then matches nothing (a loud-enough signal that it's misconfigured).
fn fetch_extensions(engine: &Engine, module: &Module) -> Vec<String> {
match Self::call_string_export(engine, module, "plugin_extensions") {
Some(s) => s.split(',').map(|e| e.trim().to_string()).filter(|e| !e.is_empty()).collect(),
None => Vec::new(),
}
}
/// Call `plugin_schema` once and parse `name:type` pairs into columns.
/// Empty if the plugin omits it — its extracted metadata then has no columns to land in.
fn fetch_schema(engine: &Engine, module: &Module) -> Vec<(String, Col)> {
match Self::call_string_export(engine, module, "plugin_schema") {
Some(s) => parse_schema_str(&s),
None => Vec::new(),
}
}
fn call_extract(&self, path: &str, data: &[u8]) -> Option<String> {
let host_state = HostState::new();
let mut store = Store::new(&self.engine, host_state);
// Build linker with host functions
let mut linker = Linker::new(&self.engine);
register_host_functions(&mut linker).ok()?;
let instance = linker.instantiate(&mut store, &self.module).ok()?;
let memory = instance.get_memory(&mut store, "memory")?;
// Allocate + write path
let alloc = instance.get_typed_func::<u32, u32>(&mut store, "alloc").ok()?;
let path_ptr = alloc.call(&mut store, path.len() as u32).ok()?;
memory.write(&mut store, path_ptr as usize, path.as_bytes()).ok()?;
// Allocate + write data
let data_ptr = alloc.call(&mut store, data.len() as u32).ok()?;
memory.write(&mut store, data_ptr as usize, data).ok()?;
// Call extract
let extract = instance.get_typed_func::<(u32, u32, u32, u32), u32>(&mut store, "extract").ok()?;
let result_ptr = extract.call(&mut store, (path_ptr, path.len() as u32, data_ptr, data.len() as u32)).ok()?;
// Read result
let result_len = instance.get_typed_func::<(), u32>(&mut store, "result_len").ok()?;
let len = result_len.call(&mut store, ()).ok()? as usize;
let mut buf = vec![0u8; len];
memory.read(&store, result_ptr as usize, &mut buf).ok()?;
String::from_utf8(buf).ok()
}
}
#[cfg(feature = "wasm-plugins")]
impl ArchiveTypePlugin for WasmPlugin {
fn spec(&self) -> &HandlerSpec {
self.spec
}
fn extract_metadata(&self, path: &str, data: &[u8]) -> Option<ExtensionRow> {
let json = self.call_extract(path, data)?;
if json == "null" {
return None;
}
parse_json_to_row(&json)
}
}
// ─── Host Functions (native decompression services) ──────────────────
#[cfg(feature = "wasm-plugins")]
fn register_host_functions(linker: &mut Linker<HostState>) -> anyhow::Result<()> {
// host_decompress(data_ptr, data_len, codec) -> result_ptr
// Writes decompressed bytes into WASM memory, returns (ptr << 32 | len)
linker.func_wrap(
"env",
"host_decompress",
|mut caller: Caller<'_, HostState>, data_ptr: u32, data_len: u32, codec: u32| -> u64 {
let memory = match caller.get_export("memory") {
Some(Extern::Memory(m)) => m,
_ => return 0,
};
// Read input data from WASM memory
let mut input = vec![0u8; data_len as usize];
if memory.read(&caller, data_ptr as usize, &mut input).is_err() {
return 0;
}
// Decompress using native multi-core implementations
let decompressed = match codec {
0 => { // Deflate
miniz_oxide_decompress(&input)
}
1 => { // Gzip (lgz-rs when available)
gzip_decompress(&input)
}
2 => { // Bzip2 (lbzip2-rs when available)
bzip2_decompress(&input)
}
3 => { // Zstd
zstd_decompress(&input)
}
_ => return 0,
};
let decompressed = match decompressed {
Some(d) => d,
None => return 0,
};
// Allocate in WASM memory and write result
let alloc = match caller.get_export("alloc") {
Some(Extern::Func(f)) => f,
_ => return 0,
};
let mut results = [Val::I32(0)];
if alloc.call(&mut caller, &[Val::I32(decompressed.len() as i32)], &mut results).is_err() {
return 0;
}
let out_ptr = results[0].unwrap_i32() as u32;
let memory = match caller.get_export("memory") {
Some(Extern::Memory(m)) => m,
_ => return 0,
};
if memory.write(&mut caller, out_ptr as usize, &decompressed).is_err() {
return 0;
}
// Pack ptr and len into u64: high 32 = ptr, low 32 = len
((out_ptr as u64) << 32) | (decompressed.len() as u64)
},
)?;
// host_archive_open(data_ptr, data_len, format, filter_ptr, filter_len) -> handle
// filter is a substring; only matching entries are decompressed and stored.
// Returns 0 if the archive is invalid or no entries match.
linker.func_wrap(
"env",
"host_archive_open",
|mut caller: Caller<'_, HostState>, data_ptr: u32, data_len: u32, format: u32, filter_ptr: u32, filter_len: u32| -> u32 {
let memory = match caller.get_export("memory") {
Some(Extern::Memory(m)) => m,
_ => return 0,
};
let mut input = vec![0u8; data_len as usize];
if memory.read(&caller, data_ptr as usize, &mut input).is_err() {
return 0;
}
let mut filter_buf = vec![0u8; filter_len as usize];
if memory.read(&caller, filter_ptr as usize, &mut filter_buf).is_err() {
return 0;
}
let filter = match std::str::from_utf8(&filter_buf) {
Ok(s) => s.to_string(),
Err(_) => return 0,
};
let entries = match format {
0 => jar_filtered_entries(&input, &filter),
1 => tar_gz_list_entries(&input),
2 => tar_bz2_list_entries(&input),
_ => return 0,
};
let entries = match entries {
Some(e) if !e.is_empty() => e,
_ => return 0,
};
let state = caller.data_mut();
let handle = state.next_handle;
state.next_handle += 1;
state.archives.insert(handle, OpenArchive { entries });
handle
},
)?;
// host_archive_list(handle) -> result_ptr (JSON array of names written to WASM mem)
linker.func_wrap(
"env",
"host_archive_list",
|mut caller: Caller<'_, HostState>, handle: u32| -> u64 {
let names: Vec<String> = {
let state = caller.data();
match state.archives.get(&handle) {
Some(archive) => archive.entries.iter().map(|e| e.name.clone()).collect(),
None => return 0,
}
};
let json = format!("[{}]",
names.iter().map(|n| format!("\"{}\"", n)).collect::<Vec<_>>().join(",")
);
write_to_wasm(&mut caller, json.as_bytes())
},
)?;
// host_archive_entry(handle, name_ptr, name_len) -> result_ptr (raw bytes)
linker.func_wrap(
"env",
"host_archive_entry",
|mut caller: Caller<'_, HostState>, handle: u32, name_ptr: u32, name_len: u32| -> u64 {
let memory = match caller.get_export("memory") {
Some(Extern::Memory(m)) => m,
_ => return 0,
};
let mut name_buf = vec![0u8; name_len as usize];
if memory.read(&caller, name_ptr as usize, &mut name_buf).is_err() {
return 0;
}
let name = match std::str::from_utf8(&name_buf) {
Ok(s) => s.to_string(),
Err(_) => return 0,
};
let data: Option<Vec<u8>> = {
let state = caller.data();
state.archives.get(&handle)
.and_then(|a| a.entries.iter().find(|e| e.name.contains(&name)))
.map(|e| e.data.clone())
};
match data {
Some(d) => write_to_wasm(&mut caller, &d),
None => 0,
}
},
)?;
// host_archive_close(handle)
linker.func_wrap(
"env",
"host_archive_close",
|mut caller: Caller<'_, HostState>, handle: u32| {
caller.data_mut().archives.remove(&handle);
},
)?;
Ok(())
}
// ─── Write helper ────────────────────────────────────────────────────
#[cfg(feature = "wasm-plugins")]
fn write_to_wasm(caller: &mut Caller<'_, HostState>, data: &[u8]) -> u64 {
let alloc = match caller.get_export("alloc") {
Some(Extern::Func(f)) => f,
_ => return 0,
};
let mut results = [Val::I32(0)];
if alloc.call(&mut *caller, &[Val::I32(data.len() as i32)], &mut results).is_err() {
return 0;
}
let out_ptr = results[0].unwrap_i32() as u32;
let memory = match caller.get_export("memory") {
Some(Extern::Memory(m)) => m,
_ => return 0,
};
if memory.write(&mut *caller, out_ptr as usize, data).is_err() {
return 0;
}
((out_ptr as u64) << 32) | (data.len() as u64)
}
// ─── Native decompression backends ──────────────────────────────────
// Each one uses the parallel native library when available,
// falls back to single-threaded otherwise.
/// Hard cap on a single host decompression output (M5 decompression-bomb guard).
/// WASM plugin memory is 32-bit addressed, so anything approaching this is already
/// unusable; the capped codecs `Err` out instead of OOMing the host on a malicious
/// archive that expands far beyond this bound.
#[cfg(feature = "host-decompressors")]
const MAX_HOST_DECOMPRESS: usize = 4 * 1024 * 1024 * 1024; // 4 GiB
fn miniz_oxide_decompress(data: &[u8]) -> Option<Vec<u8>> {
miniz_oxide::inflate::decompress_to_vec(data).ok()
}
fn gzip_decompress(data: &[u8]) -> Option<Vec<u8>> {
#[cfg(feature = "host-decompressors")]
{
lgz::decompress_gz_capped(data, MAX_HOST_DECOMPRESS).ok()
}
#[cfg(not(feature = "host-decompressors"))]
{
miniz_oxide::inflate::decompress_to_vec_zlib(data).ok()
}
}
fn bzip2_decompress(data: &[u8]) -> Option<Vec<u8>> {
#[cfg(feature = "host-decompressors")]
{
lbzip2::stream::decompress_capped(data, MAX_HOST_DECOMPRESS).ok()
}
#[cfg(not(feature = "host-decompressors"))]
{
None
}
}
fn zstd_decompress(data: &[u8]) -> Option<Vec<u8>> {
crate::codec::decompress_frame(data).ok()
}
fn tar_gz_list_entries(data: &[u8]) -> Option<Vec<ArchiveEntry>> {
#[cfg(feature = "host-decompressors")]
{
let decompressed = lgz::decompress_gz_capped(data, MAX_HOST_DECOMPRESS).ok()?;
tar_entries_from_bytes(&decompressed)
}
#[cfg(not(feature = "host-decompressors"))]
{
None
}
}
fn tar_bz2_list_entries(data: &[u8]) -> Option<Vec<ArchiveEntry>> {
#[cfg(feature = "host-decompressors")]
{
let decompressed = lbzip2::stream::decompress_capped(data, MAX_HOST_DECOMPRESS).ok()?;
tar_entries_from_bytes(&decompressed)
}
#[cfg(not(feature = "host-decompressors"))]
{
None
}
}
#[cfg(feature = "host-decompressors")]
fn tar_entries_from_bytes(tar_data: &[u8]) -> Option<Vec<ArchiveEntry>> {
use std::io::Read;
let mut archive = tar::Archive::new(tar_data);
let mut entries = Vec::new();
for entry in archive.entries().ok()? {
let mut entry = entry.ok()?;
if entry.header().entry_type().is_file() {
let name = entry.path().ok()?.to_string_lossy().into_owned();
let mut data = Vec::new();
entry.read_to_end(&mut data).ok()?;
entries.push(ArchiveEntry { name, data });
}
}
Some(entries)
}
/// Decompress only entries whose name contains `filter`.
/// With host-decompressors: filtered at the central-directory level via ljar — non-matching
/// entries are never read or decompressed.
/// Without host-decompressors: minimal single-threaded parser that also skips non-matching entries.
fn jar_filtered_entries(data: &[u8], filter: &str) -> Option<Vec<ArchiveEntry>> {
#[cfg(feature = "host-decompressors")]
{
ljar::decompress_jar_filter(data, filter).ok().map(|entries| {
entries.into_iter().map(|e| ArchiveEntry { name: e.name, data: e.data }).collect()
})
}
#[cfg(not(feature = "host-decompressors"))]
{
minimal_jar_filtered_entries(data, filter)
}
}
/// Minimal single-threaded JAR parser that decompresses only entries matching `filter`.
#[cfg(not(feature = "host-decompressors"))]
fn minimal_jar_filtered_entries(data: &[u8], filter: &str) -> Option<Vec<ArchiveEntry>> {
const EOCD_SIG: u32 = 0x06054b50;
const CD_SIG: u32 = 0x02014b50;
const LOCAL_SIG: u32 = 0x04034b50;
let start = data.len().saturating_sub(65557);
let mut eocd_pos = None;
for i in (start..data.len().saturating_sub(21)).rev() {
if u32::from_le_bytes(data[i..i+4].try_into().ok()?) == EOCD_SIG {
eocd_pos = Some(i);
break;
}
}
let eocd_pos = eocd_pos?;
let cd_offset = u32::from_le_bytes(data[eocd_pos+16..eocd_pos+20].try_into().ok()?) as usize;
let cd_entries = u16::from_le_bytes(data[eocd_pos+10..eocd_pos+12].try_into().ok()?) as usize;
let mut entries = Vec::new();
let mut pos = cd_offset;
for _ in 0..cd_entries {
if pos + 46 > data.len() { break; }
if u32::from_le_bytes(data[pos..pos+4].try_into().ok()?) != CD_SIG { break; }
let compression = u16::from_le_bytes(data[pos+10..pos+12].try_into().ok()?) ;
let comp_size = u32::from_le_bytes(data[pos+20..pos+24].try_into().ok()?) as usize;
let name_len = u16::from_le_bytes(data[pos+28..pos+30].try_into().ok()?) as usize;
let extra_len = u16::from_le_bytes(data[pos+30..pos+32].try_into().ok()?) as usize;
let comment_len = u16::from_le_bytes(data[pos+32..pos+34].try_into().ok()?) as usize;
let local_offset = u32::from_le_bytes(data[pos+42..pos+46].try_into().ok()?) as usize;
let name = std::str::from_utf8(&data[pos+46..pos+46+name_len]).unwrap_or("").to_string();
// Skip non-matching and directory entries without touching their data.
if !name.ends_with('/') && name.contains(filter) {
if local_offset + 30 <= data.len() {
let lsig = u32::from_le_bytes(data[local_offset..local_offset+4].try_into().ok()?);
if lsig == LOCAL_SIG {
let ln = u16::from_le_bytes(data[local_offset+26..local_offset+28].try_into().ok()?) as usize;
let le = u16::from_le_bytes(data[local_offset+28..local_offset+30].try_into().ok()?) as usize;
let ds = local_offset + 30 + ln + le;
if ds + comp_size <= data.len() {
let raw = &data[ds..ds + comp_size];
let entry_data = match compression {
0 => raw.to_vec(),
8 => miniz_oxide::inflate::decompress_to_vec(raw).unwrap_or_default(),
_ => Vec::new(),
};
entries.push(ArchiveEntry { name, data: entry_data });
}
}
}
}
pos += 46 + name_len + extra_len + comment_len;
}
Some(entries)
}
// ─── JSON parser ─────────────────────────────────────────────────────
/// Parse a `plugin_schema` description (`name:type,name:type`) into columns.
/// Type tags: `u32` is UInt32; `utf8` and every other tag is Utf8 (the writer has
/// no Int8 arm, so `i8` is written as text). All columns nullable.
#[cfg(feature = "wasm-plugins")]
fn parse_schema_str(s: &str) -> Vec<(String, Col)> {
s.split(',')
.filter_map(|pair| {
let mut it = pair.splitn(2, ':');
let name = it.next()?.trim();
if name.is_empty() {
return None;
}
let col = match it.next().unwrap_or("utf8").trim() {
"u32" => Col::U32,
_ => Col::Utf8,
};
Some((name.to_string(), col))
})
.collect()
}
fn parse_json_to_row(json: &str) -> Option<ExtensionRow> {
let json = json.trim().trim_start_matches('{').trim_end_matches('}');
let mut fields = HashMap::new();
for pair in json.split(',') {
let mut kv = pair.splitn(2, ':');
let key = kv.next()?.trim().trim_matches('"').to_string();
let val = kv.next()?.trim().trim_matches('"').to_string();
fields.insert(key, ExtensionValue::Str(val));
}
Some(ExtensionRow { fields })
}