use std::collections::BTreeMap;
use serde::Deserialize;
use crate::error::ExtismError;
use crate::host_fns::HostFnRegistry;
const DEFAULT_TIMEOUT_MS: u64 = 30_000;
const DEFAULT_MEMORY_MAX_PAGES: u32 = 16_384;
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ExtismPluginManifest {
pub id: String,
pub version: String,
#[serde(default, rename = "abi-extism")]
pub abi_extism: Option<String>,
#[serde(default)]
pub capabilities: Vec<uni_plugin::ManifestCapability>,
#[serde(default)]
pub determinism: Option<String>,
#[serde(default)]
pub description: Option<String>,
#[serde(default)]
pub fuel_per_call: Option<u64>,
#[serde(default)]
pub memory_max_pages: Option<u32>,
#[serde(default)]
pub timeout_ms: Option<u64>,
}
impl ExtismPluginManifest {
#[must_use]
pub fn declared_capability_set(&self) -> uni_plugin::CapabilitySet {
uni_plugin::CapabilitySet::from_manifest(self.capabilities.iter().cloned())
}
}
#[derive(Debug, Clone)]
pub struct PreparedExtismPlugin {
pub manifest: ExtismPluginManifest,
pub effective: uni_plugin::CapabilitySet,
pub allowed_host_fns: Vec<String>,
pub denied_capabilities: Vec<String>,
}
#[derive(Default)]
pub struct ExtismLoader {
host_fns: HostFnRegistry,
runtime_fns: BTreeMap<String, extism::Function>,
kms: Option<std::sync::Arc<dyn uni_plugin::KmsProvider>>,
secrets: Option<std::sync::Arc<uni_plugin::secrets::SecretStore>>,
http: Option<std::sync::Arc<dyn uni_plugin::HttpEgress>>,
}
impl std::fmt::Debug for ExtismLoader {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ExtismLoader")
.field("host_fns", &self.host_fns)
.field("runtime_fn_count", &self.runtime_fns.len())
.finish()
}
}
impl ExtismLoader {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn host_fns_mut(&mut self) -> &mut HostFnRegistry {
&mut self.host_fns
}
#[must_use]
pub fn host_fns(&self) -> &HostFnRegistry {
&self.host_fns
}
pub fn register_host_function(
&mut self,
spec: crate::host_fns::HostFnSpec,
function: extism::Function,
) {
let name = spec.name.clone();
self.host_fns.register(spec);
self.runtime_fns.insert(name, function);
}
#[must_use]
pub fn runtime_fn_count(&self) -> usize {
self.runtime_fns.len()
}
fn allowed_host_fn_names(&self, caps: &uni_plugin::CapabilitySet) -> Vec<String> {
self.host_fns
.iter()
.filter(|spec| match &spec.required_capability {
None => true,
Some(req) => caps.contains_variant(req),
})
.map(|s| s.name.clone())
.collect()
}
#[must_use]
pub fn with_kms(mut self, kms: std::sync::Arc<dyn uni_plugin::KmsProvider>) -> Self {
self.kms = Some(kms);
self
}
#[must_use]
pub fn with_secret_store(
mut self,
store: std::sync::Arc<uni_plugin::secrets::SecretStore>,
) -> Self {
self.secrets = Some(store);
self
}
#[must_use]
pub fn with_http(mut self, http: std::sync::Arc<dyn uni_plugin::HttpEgress>) -> Self {
self.http = Some(http);
self
}
fn runtime_fns_for_load(
&self,
prepared: &PreparedExtismPlugin,
) -> BTreeMap<String, extism::Function> {
let mut fns = self.runtime_fns.clone();
let ctx = crate::host_svc::HostSvcCtx {
effective: prepared.effective.clone(),
kms: self.kms.clone(),
secrets: self.secrets.clone(),
http: self.http.clone(),
};
for name in &prepared.allowed_host_fns {
if fns.contains_key(name) {
continue;
}
if let Some(function) = crate::host_svc::build_service_fn(name, &ctx) {
fns.insert(name.clone(), function);
}
}
fns
}
pub fn prepare(
&self,
manifest_json: &[u8],
grants: &uni_plugin::CapabilitySet,
) -> Result<PreparedExtismPlugin, ExtismError> {
let manifest = crate::exports::parse_manifest_json(manifest_json)?;
Ok(self.prepare_parsed(manifest, grants))
}
#[must_use]
pub fn prepare_parsed(
&self,
manifest: ExtismPluginManifest,
grants: &uni_plugin::CapabilitySet,
) -> PreparedExtismPlugin {
let declared = manifest.declared_capability_set();
let effective = declared.intersect(grants);
let denied: Vec<String> = declared
.iter()
.filter(|c| !effective.contains_variant(c))
.map(|c| format!("{c:?}"))
.collect();
let allowed = self.allowed_host_fn_names(&effective);
PreparedExtismPlugin {
manifest,
effective,
allowed_host_fns: allowed,
denied_capabilities: denied,
}
}
pub fn build_plugin(
&self,
bytes: &[u8],
prepared: &PreparedExtismPlugin,
) -> Result<extism::Plugin, ExtismError> {
build_plugin_from_parts(bytes, prepared, &self.runtime_fns_for_load(prepared))
}
pub fn load(
&self,
bytes: &[u8],
host_grants: &uni_plugin::CapabilitySet,
registrar: &mut uni_plugin::PluginRegistrar<'_>,
) -> Result<LoadOutcome, ExtismError> {
let bootstrap_allowed = self.allowed_host_fn_names(host_grants);
let bootstrap_prepared = PreparedExtismPlugin {
manifest: ExtismPluginManifest {
id: String::new(),
version: String::new(),
abi_extism: None,
capabilities: Vec::new(),
determinism: None,
description: None,
fuel_per_call: None,
memory_max_pages: None,
timeout_ms: None,
},
effective: host_grants.clone(),
allowed_host_fns: bootstrap_allowed,
denied_capabilities: Vec::new(),
};
let mut bootstrap_plugin = self.build_plugin(bytes, &bootstrap_prepared)?;
let parsed_manifest = crate::exports::read_manifest_export(&mut bootstrap_plugin)?;
drop(bootstrap_plugin);
registrar.set_plugin_id(uni_plugin::PluginId::new(parsed_manifest.id.clone()));
let prepared = self.prepare_parsed(parsed_manifest, host_grants);
let pool = build_pool(bytes, &prepared, &self.runtime_fns_for_load(&prepared))?;
let mut leased = crate::pool::PooledInstance::acquire(std::sync::Arc::clone(&pool))?;
let registration = crate::exports::read_register_export(leased.get_mut())?;
drop(leased);
let mut scalars_registered: Vec<String> = Vec::new();
let mut aggregates_registered: Vec<String> = Vec::new();
let mut procedures_registered: Vec<String> = Vec::new();
for entry in registration.entries {
match entry {
crate::exports::RegistrationEntry::Scalar { qname, signature } => {
let parsed_qname = uni_plugin::QName::parse(&qname).map_err(|e| {
ExtismError::OutputDecode(format!("invalid qname `{qname}`: {e}"))
})?;
let sig = crate::wire_translate::wire_fn_sig_to_internal(&signature)?;
let adapter = std::sync::Arc::new(crate::adapter::ExtismScalarFn::new(
std::sync::Arc::clone(&pool),
parsed_qname.clone(),
sig.clone(),
));
registrar
.scalar_fn(parsed_qname, sig, adapter)
.map_err(|e| {
ExtismError::Internal(format!("registrar.scalar_fn `{qname}`: {e}"))
})?;
scalars_registered.push(qname);
}
crate::exports::RegistrationEntry::Aggregate {
qname,
signature,
state,
} => {
let parsed_qname = uni_plugin::QName::parse(&qname).map_err(|e| {
ExtismError::OutputDecode(format!("invalid qname `{qname}`: {e}"))
})?;
let sig = crate::wire_translate::wire_agg_sig_to_internal(&signature, &state)?;
let adapter =
std::sync::Arc::new(crate::adapter_aggregate::ExtismAggregateFn::new(
std::sync::Arc::clone(&pool),
parsed_qname.clone(),
sig.clone(),
));
registrar
.aggregate_fn(parsed_qname, sig, adapter)
.map_err(|e| {
ExtismError::Internal(format!("registrar.aggregate_fn `{qname}`: {e}"))
})?;
aggregates_registered.push(qname);
}
crate::exports::RegistrationEntry::Procedure {
qname,
args,
yields,
mode,
} => {
let parsed_qname = uni_plugin::QName::parse(&qname).map_err(|e| {
ExtismError::OutputDecode(format!("invalid qname `{qname}`: {e}"))
})?;
let sig =
crate::wire_translate::wire_proc_sig_to_internal(&args, &yields, &mode)?;
let adapter =
std::sync::Arc::new(crate::adapter_procedure::ExtismProcedure::new(
std::sync::Arc::clone(&pool),
parsed_qname.clone(),
sig.clone(),
));
registrar
.procedure(parsed_qname, sig, adapter)
.map_err(|e| {
ExtismError::Internal(format!("registrar.procedure `{qname}`: {e}"))
})?;
procedures_registered.push(qname);
}
}
}
Ok(LoadOutcome {
plugin_id: prepared.manifest.id.clone(),
version: prepared.manifest.version.clone(),
effective_capabilities: prepared
.effective
.iter()
.map(|c| format!("{c:?}"))
.collect(),
denied_capabilities: prepared.denied_capabilities,
scalars_registered,
aggregates_registered,
procedures_registered,
pool,
})
}
}
fn build_plugin_from_parts(
bytes: &[u8],
prepared: &PreparedExtismPlugin,
runtime_fns: &BTreeMap<String, extism::Function>,
) -> Result<extism::Plugin, ExtismError> {
let manifest = build_extism_manifest(bytes, &prepared.manifest);
let mut builder = extism::PluginBuilder::new(manifest).with_wasi(true);
if let Some(fuel) = prepared.manifest.fuel_per_call {
builder = builder.with_fuel_limit(fuel);
}
let mut selected: Vec<extism::Function> = Vec::with_capacity(prepared.allowed_host_fns.len());
for fn_name in &prepared.allowed_host_fns {
let function = runtime_fns.get(fn_name).ok_or_else(|| {
ExtismError::Internal(format!(
"allowed host fn `{fn_name}` missing from runtime_fns; \
registry-state bug — every spec.name should have a Function"
))
})?;
selected.push(function.clone());
}
builder = builder.with_functions(selected);
builder
.build()
.map_err(|e| ExtismError::Instantiate(e.to_string()))
}
fn build_extism_manifest(bytes: &[u8], plugin_manifest: &ExtismPluginManifest) -> extism::Manifest {
let pages = plugin_manifest
.memory_max_pages
.unwrap_or(DEFAULT_MEMORY_MAX_PAGES);
let ms = plugin_manifest.timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS);
extism::Manifest::new([extism::Wasm::data(bytes.to_vec())])
.with_memory_max(pages)
.with_timeout(std::time::Duration::from_millis(ms))
}
fn build_pool(
bytes: &[u8],
prepared: &PreparedExtismPlugin,
runtime_fns: &BTreeMap<String, extism::Function>,
) -> Result<std::sync::Arc<crate::pool::ExtismInstancePool<extism::Plugin>>, ExtismError> {
let bytes_owned: std::sync::Arc<Vec<u8>> = std::sync::Arc::new(bytes.to_vec());
let prepared_owned: std::sync::Arc<PreparedExtismPlugin> =
std::sync::Arc::new(prepared.clone());
let runtime_fns_owned: std::sync::Arc<BTreeMap<String, extism::Function>> =
std::sync::Arc::new(runtime_fns.clone());
let factory = {
let bytes = std::sync::Arc::clone(&bytes_owned);
let prepared = std::sync::Arc::clone(&prepared_owned);
let runtime_fns = std::sync::Arc::clone(&runtime_fns_owned);
move || build_plugin_from_parts(&bytes, &prepared, &runtime_fns)
};
let pool = crate::pool::ExtismInstancePool::new(crate::pool::PoolConfig::default(), factory)?;
Ok(std::sync::Arc::new(pool))
}
pub struct LoadOutcome {
pub plugin_id: String,
pub version: String,
pub effective_capabilities: Vec<String>,
pub denied_capabilities: Vec<String>,
pub scalars_registered: Vec<String>,
pub aggregates_registered: Vec<String>,
pub procedures_registered: Vec<String>,
pub pool: std::sync::Arc<crate::pool::ExtismInstancePool<extism::Plugin>>,
}
impl std::fmt::Debug for LoadOutcome {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LoadOutcome")
.field("plugin_id", &self.plugin_id)
.field("version", &self.version)
.field("effective_capabilities", &self.effective_capabilities)
.field("denied_capabilities", &self.denied_capabilities)
.field("scalars_registered", &self.scalars_registered)
.field("aggregates_registered", &self.aggregates_registered)
.field("procedures_registered", &self.procedures_registered)
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::host_fns::HostFnSpec;
use uni_plugin::{Capability, CapabilitySet};
fn manifest_json(caps: &[&str]) -> String {
let caps_json: Vec<String> = caps.iter().map(|c| format!("\"{c}\"")).collect();
format!(
r#"{{ "id": "ai.example.test", "version": "1.0.0", "capabilities": [{}] }}"#,
caps_json.join(", ")
)
}
#[test]
fn loader_constructs_with_empty_host_fns() {
let l = ExtismLoader::new();
assert!(l.host_fns().is_empty());
}
fn fs_cap() -> Capability {
Capability::Filesystem {
read: vec![],
write: vec![],
}
}
#[test]
fn loader_accepts_host_fn_registrations() {
let mut l = ExtismLoader::new();
l.host_fns_mut().register(HostFnSpec {
name: "host_fs_read".to_owned(),
required_capability: Some(fs_cap()),
docs: "Read file.".to_owned(),
});
assert_eq!(l.host_fns().len(), 1);
}
#[test]
fn prepare_parses_minimal_manifest() {
let l = ExtismLoader::new();
let json = manifest_json(&[]);
let prep = l.prepare(json.as_bytes(), &CapabilitySet::new()).unwrap();
assert_eq!(prep.manifest.id, "ai.example.test");
assert_eq!(prep.manifest.version, "1.0.0");
assert!(prep.effective.is_empty());
assert!(prep.denied_capabilities.is_empty());
assert!(prep.allowed_host_fns.is_empty());
}
#[test]
fn prepare_intersects_declared_and_granted_capabilities() {
let l = ExtismLoader::new();
let json = manifest_json(&["filesystem", "network", "kms"]);
let grants = CapabilitySet::from_iter_of([fs_cap(), Capability::Network { allow: vec![] }]);
let prep = l.prepare(json.as_bytes(), &grants).unwrap();
assert_eq!(prep.effective.len(), 2);
assert!(prep.effective.contains_variant(&fs_cap()));
assert!(
prep.effective
.contains_variant(&Capability::Network { allow: vec![] })
);
assert!(
!prep
.effective
.contains_variant(&Capability::Kms { key_ids: vec![] })
);
}
#[test]
fn prepare_filters_host_fns_through_effective_capabilities() {
let mut l = ExtismLoader::new();
l.host_fns_mut().register(HostFnSpec {
name: "host_fs_read".to_owned(),
required_capability: Some(fs_cap()),
docs: "Read file.".to_owned(),
});
l.host_fns_mut().register(HostFnSpec {
name: "host_net_http_get".to_owned(),
required_capability: Some(Capability::Network { allow: vec![] }),
docs: "HTTP GET.".to_owned(),
});
l.host_fns_mut().register(HostFnSpec {
name: "host_log".to_owned(),
required_capability: None, docs: "Log a message.".to_owned(),
});
let json = manifest_json(&["filesystem"]);
let prep = l
.prepare(json.as_bytes(), &CapabilitySet::from_iter_of([fs_cap()]))
.unwrap();
assert_eq!(prep.allowed_host_fns.len(), 2);
assert!(prep.allowed_host_fns.iter().any(|n| n == "host_log"));
assert!(prep.allowed_host_fns.iter().any(|n| n == "host_fs_read"));
assert!(
!prep
.allowed_host_fns
.iter()
.any(|n| n == "host_net_http_get")
);
}
#[test]
fn prepare_rejects_malformed_manifest() {
let l = ExtismLoader::new();
let err = l.prepare(b"not json", &CapabilitySet::new()).unwrap_err();
assert!(matches!(err, ExtismError::ManifestInvalid(_)));
}
#[test]
fn build_plugin_rejects_garbage_bytes_as_instantiate_error() {
let l = ExtismLoader::new();
let prep = l
.prepare(
b"{\"id\":\"a.b\",\"version\":\"0.0.0\"}",
&CapabilitySet::new(),
)
.unwrap();
let err = l.build_plugin(b"not real wasm", &prep).unwrap_err();
assert!(
matches!(err, ExtismError::Instantiate(_)),
"expected Instantiate(_), got: {err:?}"
);
}
#[test]
fn undeclared_limits_get_host_defaults() {
let l = ExtismLoader::new();
let json = manifest_json(&[]);
let prep = l.prepare(json.as_bytes(), &CapabilitySet::new()).unwrap();
assert_eq!(prep.manifest.memory_max_pages, None);
assert_eq!(prep.manifest.timeout_ms, None);
let m = build_extism_manifest(b"\0asm", &prep.manifest);
assert_eq!(
m.memory.max_pages,
Some(DEFAULT_MEMORY_MAX_PAGES),
"undeclared memory cap must fall back to the host default"
);
assert_eq!(
m.timeout_ms,
Some(DEFAULT_TIMEOUT_MS),
"undeclared timeout must fall back to the host default"
);
}
#[test]
fn declared_limits_are_honored() {
let l = ExtismLoader::new();
let json = r#"{ "id": "ai.example.test", "version": "1.0.0", "capabilities": [], "memory_max_pages": 4, "timeout_ms": 500 }"#;
let prep = l.prepare(json.as_bytes(), &CapabilitySet::new()).unwrap();
let m = build_extism_manifest(b"\0asm", &prep.manifest);
assert_eq!(m.memory.max_pages, Some(4));
assert_eq!(m.timeout_ms, Some(500));
}
}