|
| 1 | +//! `redact_part(in_ptr,in_len):(out_ptr,out_len)` transports UTF-8 JSON through guest-exported linear `memory`. |
| 2 | +
|
| 3 | +use wasmtime::{Engine, Instance, Linker, Module, Store, StoreLimits, StoreLimitsBuilder}; |
| 4 | + |
| 5 | +use crate::error::RedactionError; |
| 6 | + |
| 7 | +const SCRATCH_OFFSET: usize = 0x0800; |
| 8 | +const GUEST_PAGE_BYTES: usize = 65536; |
| 9 | +/// Cap on the length the guest can declare for its output buffer. Without a |
| 10 | +/// ceiling, a malicious or buggy module can return a near-`i32::MAX` length |
| 11 | +/// and force the host to allocate gigabytes (or OOM) before we'd even hit |
| 12 | +/// the linear-memory read. We bound it to the single-page payload window |
| 13 | +/// the guest is allowed to write into in the first place. |
| 14 | +const MAX_GUEST_OUTPUT_BYTES: usize = GUEST_PAGE_BYTES; |
| 15 | +/// Fuel budget for a single `redact_part` call. Wasmtime decrements this per |
| 16 | +/// instruction executed; a guest that loops indefinitely traps with |
| 17 | +/// `OutOfFuel` instead of blocking the caller thread. The value is sized |
| 18 | +/// for the canonical per-part redact workload (one JSON part, scan-and- |
| 19 | +/// replace); the gateway can lift the cap if a skill genuinely needs more. |
| 20 | +const GUEST_FUEL_PER_CALL: u64 = 10_000_000; |
| 21 | +/// Hard cap on a single store's linear-memory growth. The guest already |
| 22 | +/// only writes into one page worth of scratch, but a buggy module could |
| 23 | +/// allocate more pages internally; bound that to keep one bad guest from |
| 24 | +/// pinning the host's RAM. |
| 25 | +const MAX_STORE_MEMORY_BYTES: usize = 16 * 1024 * 1024; |
| 26 | + |
| 27 | +pub(crate) fn new_engine() -> Result<Engine, RedactionError> { |
| 28 | + let mut config = wasmtime::Config::default(); |
| 29 | + config.wasm_multi_value(true); |
| 30 | + // Bound guest CPU time deterministically. Without fuel, a redact_part |
| 31 | + // guest stuck in an infinite loop would block the synchronous Redactor |
| 32 | + // trait call indefinitely. |
| 33 | + config.consume_fuel(true); |
| 34 | + Engine::new(&config).map_err(|e| RedactionError::WasmEngine(e.to_string())) |
| 35 | +} |
| 36 | + |
| 37 | +/// Per-call store state — wraps StoreLimits so wasmtime can enforce the |
| 38 | +/// memory ceiling on `memory.grow` without the host having to check after |
| 39 | +/// the fact. |
| 40 | +struct StoreState { |
| 41 | + limits: StoreLimits, |
| 42 | +} |
| 43 | + |
| 44 | +pub(crate) fn redact_part_guest(engine: &Engine, module: &Module, payload: &[u8]) -> Result<Vec<u8>, RedactionError> { |
| 45 | + if SCRATCH_OFFSET.saturating_add(payload.len()) > GUEST_PAGE_BYTES { |
| 46 | + return Err(RedactionError::WasmMemory( |
| 47 | + "redaction payload does not fit in one wasm guest page".into(), |
| 48 | + )); |
| 49 | + } |
| 50 | + |
| 51 | + let in_len = i32::try_from(payload.len()) |
| 52 | + .map_err(|_| RedactionError::WasmMemory("payload length does not fit in wasm i32 bounds".into()))?; |
| 53 | + |
| 54 | + let mut store = Store::new( |
| 55 | + engine, |
| 56 | + StoreState { |
| 57 | + limits: StoreLimitsBuilder::new().memory_size(MAX_STORE_MEMORY_BYTES).build(), |
| 58 | + }, |
| 59 | + ); |
| 60 | + store.limiter(|state| &mut state.limits); |
| 61 | + store |
| 62 | + .set_fuel(GUEST_FUEL_PER_CALL) |
| 63 | + .map_err(|e| RedactionError::WasmEngine(format!("set_fuel: {e}")))?; |
| 64 | + let linker: Linker<StoreState> = Linker::new(engine); |
| 65 | + let instance: Instance = linker |
| 66 | + .instantiate(&mut store, module) |
| 67 | + .map_err(|e| RedactionError::WasmInstance(e.to_string()))?; |
| 68 | + |
| 69 | + let memory = instance |
| 70 | + .get_memory(&mut store, "memory") |
| 71 | + .ok_or_else(|| RedactionError::WasmAbi("wasm module must export linear memory named memory".into()))?; |
| 72 | + |
| 73 | + memory |
| 74 | + .write(&mut store, SCRATCH_OFFSET, payload) |
| 75 | + .map_err(|e| RedactionError::WasmMemory(e.to_string()))?; |
| 76 | + |
| 77 | + let redact = instance |
| 78 | + .get_typed_func::<(i32, i32), (i32, i32)>(&mut store, "redact_part") |
| 79 | + .map_err(|_| { |
| 80 | + RedactionError::WasmAbi("wasm module must export redact_part with type (i32,i32)->(i32,i32)".into()) |
| 81 | + })?; |
| 82 | + |
| 83 | + let (out_base, out_len) = redact |
| 84 | + .call(&mut store, (SCRATCH_OFFSET as i32, in_len)) |
| 85 | + .map_err(|e| RedactionError::WasmCall(e.to_string()))?; |
| 86 | + |
| 87 | + let out_base = usize::try_from(out_base) |
| 88 | + .map_err(|_| RedactionError::WasmAbi("wasm redact_part returned negative output pointer".into()))?; |
| 89 | + let out_len = usize::try_from(out_len) |
| 90 | + .map_err(|_| RedactionError::WasmAbi("wasm redact_part returned negative output length".into()))?; |
| 91 | + |
| 92 | + if out_len > MAX_GUEST_OUTPUT_BYTES { |
| 93 | + return Err(RedactionError::WasmAbi(format!( |
| 94 | + "wasm redact_part output length {out_len} exceeds guest cap {MAX_GUEST_OUTPUT_BYTES}" |
| 95 | + ))); |
| 96 | + } |
| 97 | + let memory_size = memory.data_size(&store); |
| 98 | + if out_base.saturating_add(out_len) > memory_size { |
| 99 | + return Err(RedactionError::WasmAbi(format!( |
| 100 | + "wasm redact_part output [base={out_base}, len={out_len}) exceeds linear memory size {memory_size}" |
| 101 | + ))); |
| 102 | + } |
| 103 | + |
| 104 | + let mut dst = vec![0u8; out_len]; |
| 105 | + memory |
| 106 | + .read(&store, out_base, &mut dst) |
| 107 | + .map_err(|e| RedactionError::WasmMemory(e.to_string()))?; |
| 108 | + |
| 109 | + Ok(dst) |
| 110 | +} |
| 111 | + |
| 112 | +#[cfg(test)] |
| 113 | +mod tests { |
| 114 | + use super::*; |
| 115 | + use wasmtime::Module; |
| 116 | + |
| 117 | + #[test] |
| 118 | + fn builds_default_compatible_engine() { |
| 119 | + let engine = new_engine().expect("wasmtime engine"); |
| 120 | + drop(engine); |
| 121 | + } |
| 122 | + |
| 123 | + #[test] |
| 124 | + fn identity_guest_round_trips_utf8_payload() { |
| 125 | + let engine = new_engine().unwrap(); |
| 126 | + let wasm = include_bytes!(concat!( |
| 127 | + env!("CARGO_MANIFEST_DIR"), |
| 128 | + "/tests/fixtures/identity_redact_part.wasm" |
| 129 | + )); |
| 130 | + let module = Module::from_binary(&engine, wasm).unwrap(); |
| 131 | + let inp = br#"{"k":42}"#.to_vec(); |
| 132 | + let got = redact_part_guest(&engine, &module, &inp).unwrap(); |
| 133 | + assert_eq!(got, inp); |
| 134 | + } |
| 135 | +} |
0 commit comments