Skip to main content

coil_wasm/
engine.rs

1use std::time::Instant;
2
3use wasmtime::{Caller, Config, Engine, Linker, Module, Store, StoreLimits, StoreLimitsBuilder};
4
5use crate::error::WasmModelError;
6use crate::grants::HostCapabilityGrant;
7use crate::ids::ExtensionPointKind;
8use crate::invocation::{ExecutionReceipt, HostCall, InvocationOutcome, WasmExecutionSession};
9use crate::output::TypedExecutionOutput;
10
11#[derive(Debug, Clone)]
12pub struct WasmEngine {
13    engine: Engine,
14}
15
16impl Default for WasmEngine {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl WasmEngine {
23    pub fn new() -> Self {
24        let mut config = Config::new();
25        config.consume_fuel(true);
26        let engine = Engine::new(&config).expect("static wasmtime configuration must be valid");
27        Self { engine }
28    }
29
30    pub fn compile_module(&self, bytes: &[u8]) -> Result<CompiledWasmModule, WasmModelError> {
31        let module =
32            Module::new(&self.engine, bytes).map_err(|error| WasmModelError::EngineCompile {
33                reason: error.to_string(),
34            })?;
35        Ok(CompiledWasmModule { module })
36    }
37
38    pub fn execute_session(
39        &self,
40        module: &CompiledWasmModule,
41        session: WasmExecutionSession,
42        export: &str,
43    ) -> Result<ExecutionReceipt, WasmModelError> {
44        module.execute(&self.engine, session, export)
45    }
46}
47
48#[derive(Debug, Clone)]
49pub struct CompiledWasmModule {
50    module: Module,
51}
52
53#[derive(Debug)]
54struct EngineHostState {
55    session: WasmExecutionSession,
56    grants: Vec<HostCapabilityGrant>,
57    last_error: Option<WasmModelError>,
58    limits: StoreLimits,
59}
60
61impl EngineHostState {
62    fn new(session: WasmExecutionSession) -> Self {
63        let limits = StoreLimitsBuilder::new()
64            .memory_size(session.plan().limits.max_memory_bytes as usize)
65            .instances(1)
66            .tables(4)
67            .memories(1)
68            .build();
69
70        Self {
71            grants: session.grant_slots(),
72            session,
73            last_error: None,
74            limits,
75        }
76    }
77
78    fn record_slot_call(&mut self, slot: i32, metric: i64) -> Result<i32, WasmModelError> {
79        if slot < 0 {
80            return Err(WasmModelError::InvalidHostCapabilitySlot {
81                handler_id: self.session.plan().handler_id.to_string(),
82                slot,
83            });
84        }
85
86        let grant = self.grants.get(slot as usize).ok_or_else(|| {
87            WasmModelError::InvalidHostCapabilitySlot {
88                handler_id: self.session.plan().handler_id.to_string(),
89                slot,
90            }
91        })?;
92        let call = host_call_for_grant(&self.session, grant, metric)?;
93        let _ = self.session.execute_host_call(call)?;
94        Ok(0)
95    }
96}
97
98impl CompiledWasmModule {
99    pub fn execute(
100        &self,
101        engine: &Engine,
102        session: WasmExecutionSession,
103        export: &str,
104    ) -> Result<ExecutionReceipt, WasmModelError> {
105        let export = crate::validation::validate_token("export", export.to_string())?;
106        let mut store = Store::new(engine, EngineHostState::new(session));
107        store.limiter(|state| &mut state.limits);
108
109        let fuel_budget = store
110            .data()
111            .session
112            .plan()
113            .limits
114            .max_runtime
115            .as_millis()
116            .max(1) as u64
117            * 10_000;
118        store
119            .set_fuel(fuel_budget)
120            .map_err(|error| WasmModelError::EngineInstantiate {
121                handler_id: store.data().session.plan().handler_id.to_string(),
122                reason: error.to_string(),
123            })?;
124
125        let mut linker = Linker::new(engine);
126        linker
127            .func_wrap(
128                "coil",
129                "host_call",
130                |mut caller: Caller<'_, EngineHostState>,
131                 slot: i32,
132                 metric: i64|
133                 -> Result<i32, wasmtime::Error> {
134                    let state = caller.data_mut();
135                    match state.record_slot_call(slot, metric) {
136                        Ok(result) => Ok(result),
137                        Err(error) => {
138                            state.last_error = Some(error);
139                            Err(wasmtime::Error::msg("coil host call failed"))
140                        }
141                    }
142                },
143            )
144            .map_err(|error| WasmModelError::EngineInstantiate {
145                handler_id: store.data().session.plan().handler_id.to_string(),
146                reason: error.to_string(),
147            })?;
148
149        let start = Instant::now();
150        let instance = linker
151            .instantiate(&mut store, &self.module)
152            .map_err(|error| WasmModelError::EngineInstantiate {
153                handler_id: store.data().session.plan().handler_id.to_string(),
154                reason: error.to_string(),
155            })?;
156        let handler_id = store.data().session.plan().handler_id.to_string();
157        let function = instance
158            .get_typed_func::<(), i32>(&mut store, &export)
159            .map_err(|_| WasmModelError::EngineExportMissing {
160                handler_id: handler_id.clone(),
161                export: export.clone(),
162            })?;
163        let outcome_code = function.call(&mut store, ()).map_err(|error| {
164            if let Some(host_error) = store.data().last_error.clone() {
165                host_error
166            } else {
167                WasmModelError::EngineTrap {
168                    handler_id: handler_id.clone(),
169                    reason: error.to_string(),
170                }
171            }
172        })?;
173        let runtime = start.elapsed();
174        let point = store.data().session.plan().point;
175        let typed_output = read_typed_output(&mut store, &instance, &handler_id, point)?;
176
177        let state = store.into_data();
178        if let Some(host_error) = state.last_error {
179            return Err(host_error);
180        }
181
182        let outcome = InvocationOutcome::from_engine_code(
183            outcome_code,
184            state.session.plan().handler_id.to_string(),
185        )?;
186        state.session.finish(runtime, outcome, typed_output)
187    }
188}
189
190fn read_typed_output(
191    store: &mut Store<EngineHostState>,
192    instance: &wasmtime::Instance,
193    handler_id: &str,
194    point: ExtensionPointKind,
195) -> Result<Option<TypedExecutionOutput>, WasmModelError> {
196    let Some(export) = instance.get_func(&mut *store, TypedExecutionOutput::ABI_EXPORT) else {
197        return Ok(None);
198    };
199    let func = export.typed::<(), i64>(&mut *store).map_err(|error| {
200        WasmModelError::EngineInstantiate {
201            handler_id: handler_id.to_string(),
202            reason: format!(
203                "typed return export `{}` has unexpected signature: {error}",
204                TypedExecutionOutput::ABI_EXPORT
205            ),
206        }
207    })?;
208    let packed = func
209        .call(&mut *store, ())
210        .map_err(|error| WasmModelError::EngineTrap {
211            handler_id: handler_id.to_string(),
212            reason: error.to_string(),
213        })?;
214    let packed = packed as u64;
215    let ptr = (packed & 0xffff_ffff) as u32 as usize;
216    let len = (packed >> 32) as u32 as usize;
217
218    let memory = instance.get_memory(&mut *store, "memory").ok_or_else(|| {
219        WasmModelError::InvalidTypedReturn {
220            reason: format!(
221                "typed return export `{}` is present but no `memory` export exists",
222                TypedExecutionOutput::ABI_EXPORT
223            ),
224        }
225    })?;
226
227    let memory_size = memory.data_size(&mut *store);
228    let end = ptr
229        .checked_add(len)
230        .ok_or_else(|| WasmModelError::InvalidTypedReturn {
231            reason: "typed return payload length overflows host address space".to_string(),
232        })?;
233    if end > memory_size {
234        return Err(WasmModelError::InvalidTypedReturn {
235            reason: format!(
236                "typed return payload pointer/length `{ptr}..{end}` exceeds guest memory size `{memory_size}`"
237            ),
238        });
239    }
240
241    let mut bytes = vec![0u8; len];
242    memory.read(&mut *store, ptr, &mut bytes).map_err(|error| {
243        WasmModelError::InvalidTypedReturn {
244            reason: format!("failed to read typed return payload: {error}"),
245        }
246    })?;
247
248    TypedExecutionOutput::decode_for_point(&bytes, point).map(Some)
249}
250
251fn host_call_for_grant(
252    session: &WasmExecutionSession,
253    grant: &HostCapabilityGrant,
254    metric: i64,
255) -> Result<HostCall, WasmModelError> {
256    let metric = u64::try_from(metric).map_err(|_| WasmModelError::InvalidHostCallMetric {
257        handler_id: session.plan().handler_id.to_string(),
258        metric,
259    })?;
260
261    Ok(match grant {
262        HostCapabilityGrant::DataRead { resource } => HostCall::DataRead {
263            resource: resource.clone(),
264        },
265        HostCapabilityGrant::DataWrite { resource } => HostCall::DataWrite {
266            resource: resource.clone(),
267        },
268        HostCapabilityGrant::AuthCheck => HostCall::AuthCheck,
269        HostCapabilityGrant::AuthList => HostCall::AuthList,
270        HostCapabilityGrant::AuthLookup => HostCall::AuthLookup,
271        HostCapabilityGrant::AuthTupleWrite => HostCall::AuthTupleWrite,
272        HostCapabilityGrant::StorageRead { class } => HostCall::StorageRead { class: *class },
273        HostCapabilityGrant::StorageWrite { class } => HostCall::StorageWrite {
274            class: *class,
275            bytes: metric,
276        },
277        HostCapabilityGrant::RenderFragment { slot } => {
278            HostCall::RenderFragment { slot: slot.clone() }
279        }
280        HostCapabilityGrant::MetadataWrite { kind } => HostCall::MetadataWrite { kind: *kind },
281        HostCapabilityGrant::CacheHintWrite => HostCall::CacheHintWrite,
282        HostCapabilityGrant::OutboundHttp { integration } => HostCall::OutboundHttp {
283            integration: integration.clone(),
284            response_bytes: metric,
285        },
286        HostCapabilityGrant::SecretRead { secret } => HostCall::SecretRead {
287            secret: secret.clone(),
288        },
289        HostCapabilityGrant::EnqueueJob { queue } => HostCall::EnqueueJob {
290            queue: queue.clone(),
291        },
292    })
293}