use anyhow::Result;
use wasmtime::component::{Component, Linker, ResourceTable, StreamReader};
use wasmtime::{Config, Engine, Store, StoreLimits, StoreLimitsBuilder};
use wasmtime_wasi::{WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView};
use crate::engine::EpochTicker;
mod bindings {
wasmtime::component::bindgen!({
path: "wit-streaming",
world: "streaming-filter",
exports: { default: async },
});
}
use bindings::exports::plecto::filter_streaming::body_filter::RequestBodyDecision;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StreamingDecision {
Continue,
ShortCircuit { status: u16, body: Vec<u8> },
}
pub struct StreamingLimits {
pub memory_cap: usize,
pub deadline_ms: u64,
}
struct Ctx {
limits: StoreLimits,
table: ResourceTable,
wasi: WasiCtx,
}
impl WasiView for Ctx {
fn ctx(&mut self) -> WasiCtxView<'_> {
WasiCtxView {
ctx: &mut self.wasi,
table: &mut self.table,
}
}
}
fn build_engine() -> Result<Engine> {
let mut config = Config::new();
config.wasm_component_model(true);
config.wasm_component_model_async(true);
config.epoch_interruption(true);
Ok(Engine::new(&config)?)
}
pub struct StreamingFilterRuntime {
engine: Engine,
_ticker: EpochTicker,
linker: Linker<Ctx>,
component: Component,
}
impl StreamingFilterRuntime {
pub fn load(component: &[u8]) -> Result<Self> {
let engine = build_engine()?;
let ticker = EpochTicker::spawn(vec![engine.clone()]);
let component = Component::from_binary(&engine, component)?;
let mut linker: Linker<Ctx> = Linker::new(&engine);
wasmtime_wasi::p2::add_to_linker_proxy_interfaces_async(&mut linker)?;
crate::state::add_cli_runtime::<Ctx>(&mut linker)?;
Ok(Self {
engine,
_ticker: ticker,
linker,
component,
})
}
pub fn run(&self, body: Vec<u8>, limits: &StreamingLimits) -> Result<StreamingDecision> {
let store_limits = StoreLimitsBuilder::new()
.memory_size(limits.memory_cap)
.table_elements(crate::state::MAX_TABLE_ELEMENTS)
.build();
let mut store = Store::new(
&self.engine,
Ctx {
limits: store_limits,
table: ResourceTable::new(),
wasi: WasiCtxBuilder::new().build(),
},
);
store.limiter(|c| &mut c.limits);
store.set_epoch_deadline(limits.deadline_ms);
let instance = pollster::block_on(bindings::StreamingFilter::instantiate_async(
&mut store,
&self.component,
&self.linker,
))?;
let reader = StreamReader::new(&mut store, body)?;
let decision = pollster::block_on(store.run_concurrent(async move |accessor| {
instance
.plecto_filter_streaming_body_filter()
.call_process_body(accessor, reader)
.await
}))??;
Ok(map_decision(decision))
}
}
pub fn run_streaming_body(
component: &[u8],
body: Vec<u8>,
limits: &StreamingLimits,
) -> Result<StreamingDecision> {
StreamingFilterRuntime::load(component)?.run(body, limits)
}
fn map_decision(decision: RequestBodyDecision) -> StreamingDecision {
match decision {
RequestBodyDecision::Continue => StreamingDecision::Continue,
RequestBodyDecision::ShortCircuit(resp) => StreamingDecision::ShortCircuit {
status: resp.status,
body: resp.body,
},
}
}
#[cfg(test)]
mod tests {
use super::*;
use bindings::plecto::filter_streaming::types::HttpResponse as GuestResponse;
#[test]
fn map_decision_lowers_guest_variants() {
struct Case {
name: &'static str,
guest: RequestBodyDecision,
want: StreamingDecision,
}
let cases = vec![
Case {
name: "continue",
guest: RequestBodyDecision::Continue,
want: StreamingDecision::Continue,
},
Case {
name: "short-circuit",
guest: RequestBodyDecision::ShortCircuit(GuestResponse {
status: 403,
headers: Vec::new(),
body: b"denied".to_vec(),
}),
want: StreamingDecision::ShortCircuit {
status: 403,
body: b"denied".to_vec(),
},
},
];
for case in cases {
let got = map_decision(case.guest);
assert_eq!(got, case.want, "case: {}", case.name);
}
}
}