use std::sync::Arc;
use anyhow::Result;
use wasmtime::component::ComponentExportIndex;
use wasmtime::{Engine, Store};
use crate::contract::{
self, FilterPreV01, FilterPreV02, FilterPreV03, FilterV01, FilterV02, FilterV03,
request_body_decision_from_v01, request_body_decision_from_v02, request_body_decision_from_v03,
request_decision_from_v01, request_decision_from_v02, request_decision_from_v03,
request_to_v01, request_to_v02, response_decision_from_v01, response_decision_from_v02,
response_decision_from_v03, response_to_v01, response_to_v02,
};
use crate::errors::InvalidGuestOutput;
#[cfg(feature = "outbound-http")]
use crate::outbound_http;
#[cfg(feature = "outbound-tcp")]
use crate::outbound_tcp;
use crate::quota::KvQuota;
use crate::state::{HostState, HostStateInit};
use crate::{
Bucket, HttpRequest, HttpResponse, KvBackend, LogLine, RequestBodyDecision, RequestDecision,
ResponseDecision,
};
pub(crate) trait FilterRuntime: Send + Sync {
type Instance: Send;
fn instantiate_initialized(&self) -> Result<Self::Instance>;
fn begin_request(&self, instance: &mut Self::Instance);
fn set_request_deadline(&self, instance: &mut Self::Instance);
fn take_logs(&self, instance: &mut Self::Instance) -> Vec<LogLine>;
fn take_logs_final(&self, instance: &mut Self::Instance) -> Vec<LogLine>;
}
pub(crate) enum FilterPreBinding {
V01(FilterPreV01<HostState>),
V02(FilterPreV02<HostState>),
V03(FilterPreV03<HostState>),
}
pub(crate) struct WasmtimeRuntime {
pub(crate) engine: Engine,
pub(crate) kv: Arc<dyn KvBackend>,
pub(crate) kv_prefix: String,
pub(crate) pre: FilterPreBinding,
pub(crate) body_export: Option<ComponentExportIndex>,
pub(crate) init_deadline_ms: u64,
pub(crate) request_deadline_ms: u64,
pub(crate) max_memory_bytes: u64,
pub(crate) ratelimit_bucket: Option<Bucket>,
pub(crate) kv_quota: Arc<KvQuota>,
pub(crate) config: Arc<std::collections::BTreeMap<String, String>>,
#[cfg(feature = "outbound-http")]
pub(crate) outbound: Option<outbound_http::OutboundState>,
#[cfg(feature = "outbound-tcp")]
pub(crate) outbound_tcp: Option<outbound_tcp::OutboundTcpState>,
#[cfg(any(feature = "outbound-http", feature = "outbound-tcp"))]
pub(crate) rt: Option<Arc<tokio::runtime::Runtime>>,
#[cfg(feature = "fat-guest")]
pub(crate) wasi_minimal: bool,
}
impl WasmtimeRuntime {
pub(crate) fn drive<F: std::future::Future>(&self, fut: F) -> F::Output {
#[cfg(any(feature = "outbound-http", feature = "outbound-tcp"))]
if let Some(rt) = &self.rt {
return rt.block_on(fut);
}
pollster::block_on(fut)
}
pub(crate) fn drive_call<T>(
&self,
fut: impl std::future::Future<Output = wasmtime::Result<T>>,
) -> wasmtime::Result<T> {
#[cfg(feature = "outbound-tcp")]
if let (Some(rt), Some(state)) = (&self.rt, &self.outbound_tcp) {
return block_on_with_deadline(rt, state.io_deadline(), fut);
}
self.drive(fut)
}
#[cfg(feature = "outbound-tcp")]
pub(crate) fn tcp_guard(&self) -> outbound_tcp::TcpGuard {
match &self.outbound_tcp {
Some(state) => state.guard(),
None => outbound_tcp::TcpGuard::deny_all(),
}
}
pub(crate) fn call_on_request(
&self,
inst: &mut WasmtimeInstance,
req: &HttpRequest,
) -> wasmtime::Result<RequestDecision> {
match &mut inst.filter {
BoundFilter::V03(filter) => {
let raw = self.drive_call(filter.call_on_request(&mut inst.store, req))?;
request_decision_from_v03(raw).ok_or_else(invalid_guest_header_error)
}
BoundFilter::V02(filter) => {
let guest_req = request_to_v02(req);
let raw = self.drive_call(filter.call_on_request(&mut inst.store, &guest_req))?;
request_decision_from_v02(raw).ok_or_else(invalid_guest_header_error)
}
BoundFilter::V01(filter) => {
let guest_req = request_to_v01(req);
let raw = self.drive_call(filter.call_on_request(&mut inst.store, &guest_req))?;
request_decision_from_v01(raw).ok_or_else(invalid_guest_header_error)
}
}
}
pub(crate) fn call_on_response(
&self,
inst: &mut WasmtimeInstance,
req: &HttpRequest,
resp: &HttpResponse,
) -> wasmtime::Result<ResponseDecision> {
match &mut inst.filter {
BoundFilter::V03(filter) => {
let raw = self.drive_call(filter.call_on_response(&mut inst.store, req, resp))?;
response_decision_from_v03(raw).ok_or_else(invalid_guest_header_error)
}
BoundFilter::V02(filter) => {
let guest_resp = response_to_v02(resp);
let raw = self.drive_call(filter.call_on_response(&mut inst.store, &guest_resp))?;
response_decision_from_v02(raw).ok_or_else(invalid_guest_header_error)
}
BoundFilter::V01(filter) => {
let guest_resp = response_to_v01(resp);
let raw = self.drive_call(filter.call_on_response(&mut inst.store, &guest_resp))?;
response_decision_from_v01(raw).ok_or_else(invalid_guest_header_error)
}
}
}
pub(crate) fn call_body_hook(
&self,
inst: &mut WasmtimeInstance,
body: &[u8],
) -> wasmtime::Result<RequestBodyDecision> {
let Some(func) = inst.body_func else {
return Err(wasmtime::Error::msg(
"on-request-body called on a filter without a body export",
));
};
match &inst.filter {
BoundFilter::V03(_) => {
use contract::types_v03::RequestBodyDecision as Raw;
let func = func.typed::<(&[u8],), (Raw,)>(&inst.store)?;
let (decision,) = self.drive_call(func.call_async(&mut inst.store, (body,)))?;
request_body_decision_from_v03(decision).ok_or_else(invalid_guest_header_error)
}
BoundFilter::V02(_) => {
use contract::types_v02::RequestBodyDecision as Raw;
let func = func.typed::<(&[u8],), (Raw,)>(&inst.store)?;
let (decision,) = self.drive_call(func.call_async(&mut inst.store, (body,)))?;
request_body_decision_from_v02(decision).ok_or_else(invalid_guest_header_error)
}
BoundFilter::V01(_) => {
use contract::types_v01::RequestBodyDecision as Raw;
let func = func.typed::<(&[u8],), (Raw,)>(&inst.store)?;
let (decision,) = self.drive_call(func.call_async(&mut inst.store, (body,)))?;
request_body_decision_from_v01(decision).ok_or_else(invalid_guest_header_error)
}
}
}
#[cfg(feature = "outbound-http")]
pub(crate) fn outbound_hooks(&self) -> outbound_http::PlectoHttpHooks {
match &self.outbound {
Some(state) => state.hooks(),
None => outbound_http::PlectoHttpHooks::deny_all(),
}
}
}
fn invalid_guest_header_error() -> wasmtime::Error {
wasmtime::Error::new(InvalidGuestOutput)
}
impl FilterRuntime for WasmtimeRuntime {
type Instance = WasmtimeInstance;
fn instantiate_initialized(&self) -> Result<WasmtimeInstance> {
let mut store = Store::new(
&self.engine,
HostState::new(
HostStateInit {
kv: self.kv.clone(),
kv_prefix: self.kv_prefix.clone(),
max_memory_bytes: self.max_memory_bytes,
ratelimit_bucket: self.ratelimit_bucket,
quota: self.kv_quota.clone(),
config: self.config.clone(),
#[cfg(feature = "fat-guest")]
wasi_minimal: self.wasi_minimal,
},
#[cfg(feature = "outbound-http")]
self.outbound_hooks(),
#[cfg(feature = "outbound-tcp")]
self.tcp_guard(),
),
);
store.limiter(|s| &mut s.limits);
store.set_epoch_deadline(self.init_deadline_ms);
let (filter, instance) = match &self.pre {
FilterPreBinding::V03(pre) => {
let instance = self.drive_call(pre.instance_pre().instantiate_async(&mut store))?;
let filter = FilterV03::new(&mut store, &instance)?;
self.drive_call(filter.call_init(&mut store))?;
(BoundFilter::V03(filter), instance)
}
FilterPreBinding::V02(pre) => {
let instance = self.drive_call(pre.instance_pre().instantiate_async(&mut store))?;
let filter = FilterV02::new(&mut store, &instance)?;
self.drive_call(filter.call_init(&mut store))?;
(BoundFilter::V02(filter), instance)
}
FilterPreBinding::V01(pre) => {
let instance = self.drive_call(pre.instance_pre().instantiate_async(&mut store))?;
let filter = FilterV01::new(&mut store, &instance)?;
self.drive_call(filter.call_init(&mut store))?;
(BoundFilter::V01(filter), instance)
}
};
let body_func = match &self.body_export {
Some(idx) => {
let func = instance.get_func(&mut store, idx).ok_or_else(|| {
anyhow::anyhow!("on-request-body export index did not resolve to a function")
})?;
match &filter {
BoundFilter::V03(_) => {
func.typed::<(&[u8],), (contract::types_v03::RequestBodyDecision,)>(
&store,
)?;
}
BoundFilter::V02(_) => {
func.typed::<(&[u8],), (contract::types_v02::RequestBodyDecision,)>(
&store,
)?;
}
BoundFilter::V01(_) => {
func.typed::<(&[u8],), (contract::types_v01::RequestBodyDecision,)>(
&store,
)?;
}
}
Some(func)
}
None => None,
};
Ok(WasmtimeInstance {
store,
filter,
body_func,
})
}
fn begin_request(&self, instance: &mut WasmtimeInstance) {
instance.store.data_mut().begin_request();
}
fn set_request_deadline(&self, instance: &mut WasmtimeInstance) {
instance.store.set_epoch_deadline(self.request_deadline_ms);
}
fn take_logs(&self, instance: &mut WasmtimeInstance) -> Vec<LogLine> {
instance.store.data_mut().take_logs()
}
fn take_logs_final(&self, instance: &mut WasmtimeInstance) -> Vec<LogLine> {
instance.store.data_mut().take_logs_final()
}
}
pub(crate) enum BoundFilter {
V01(FilterV01),
V02(FilterV02),
V03(FilterV03),
}
pub(crate) struct WasmtimeInstance {
pub(crate) store: Store<HostState>,
pub(crate) filter: BoundFilter,
pub(crate) body_func: Option<wasmtime::component::Func>,
}
#[cfg(feature = "outbound-tcp")]
pub(crate) fn block_on_with_deadline<T>(
rt: &tokio::runtime::Runtime,
deadline: std::time::Duration,
fut: impl std::future::Future<Output = wasmtime::Result<T>>,
) -> wasmtime::Result<T> {
rt.block_on(async {
match tokio::time::timeout(deadline, fut).await {
Ok(result) => result,
Err(_) => Err(wasmtime::Error::msg(format!(
"outbound-tcp io deadline ({deadline:?}) exceeded; the hook call was cancelled \
fail-closed and the instance is discarded"
))),
}
})
}
#[cfg(all(test, feature = "outbound-tcp"))]
mod deadline_tests {
use super::block_on_with_deadline;
use std::time::Duration;
fn rt() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_current_thread()
.enable_time()
.build()
.unwrap()
}
#[test]
fn a_prompt_call_completes_within_the_deadline() {
let out = block_on_with_deadline(&rt(), Duration::from_secs(1), async { Ok(42) });
assert_eq!(out.unwrap(), 42);
}
#[test]
fn a_hanging_call_is_cancelled_fail_closed() {
let out = block_on_with_deadline(&rt(), Duration::from_millis(20), async {
std::future::pending::<wasmtime::Result<()>>().await
});
let err = out.expect_err("a hung call must be cancelled");
assert!(err.to_string().contains("io deadline"));
}
}