dwara_core/dataplane/
nano_service.rs

1//! WASM route handlers (nano-services) -- DW-106.
2//!
3//! A nano-service is a route action that runs a WebAssembly module to
4//! generate the response directly, instead of proxying to an upstream.
5//! The module implements a simple request->response handler ABI over
6//! the existing wasmtime runtime (the `wasm` cargo feature, brought in
7//! by `nano_services`). This is the "function-as-a-route" path: a
8//! self-contained WASM module owns the whole response for a route --
9//! status, headers, and body -- with no backend contact.
10//!
11//! ## Handler ABI
12//!
13//! The module exports:
14//!
15//! - `memory` -- the linear memory the host and module share.
16//! - `alloc(size: i32) -> i32` -- allocate `size` bytes in linear
17//!   memory and return a pointer (a simple bump allocator is fine; the
18//!   host uses it to place the serialized request).
19//! - `handle(req_ptr: i32, req_len: i32) -> i32` -- handle the request
20//!   whose serialized form lives at `[req_ptr, req_ptr+req_len)` in
21//!   linear memory. Returns 0 on success (the response is communicated
22//!   back via the host imports below) and non-zero on error (the route
23//!   answers 502).
24//!
25//! The host provides these imports under the module name `dwara`:
26//!
27//! - `response_status(status: i32)` -- set the HTTP response status.
28//! - `response_header(k_ptr, k_len, v_ptr, v_len)` -- add a response
29//!   header.
30//! - `response_body(ptr, len)` -- set the response body.
31//! - `log(ptr, len)` -- emit a log line (tracing::debug).
32//!
33//! The request is serialized into linear memory in a length-prefixed
34//! binary format (all lengths u32 big-endian):
35//!
36//! ```text
37//! u32 method_len, method bytes
38//! u32 path_len, path bytes
39//! u32 headers_count, [u32 key_len, key, u32 val_len, val]...
40//! u32 body_len, body bytes
41//! ```
42//!
43//! The module parses this to read the request, then calls the host
44//! imports to build the response and returns 0 from `handle`.
45//!
46//! ## Resource limits
47//!
48//! `memory_limit` caps the linear memory the module may allocate
49//! (enforced via wasmtime's `ResourceLimiter`). `execution_timeout_ms`
50//! caps the wall-clock time the `handle` call may run: the call runs on
51//! a blocking pool thread wrapped in a `tokio::time::timeout`, so a
52//! module that exceeds it is interrupted and the route answers 504. A
53//! fuel budget caps CPU work (a module that exhausts it traps and the
54//! route answers 502).
55//!
56//! ## Feature gate
57//!
58//! The module compiles only when the `nano_services` cargo feature is
59//! enabled (which pulls in `wasm` / wasmtime). The config schema is
60//! always present, so configs round-trip without the feature; when the
61//! feature is off the action is accepted but inert (validation warns,
62//! the route returns 502).
63
64use std::sync::Arc;
65use std::time::Duration;
66
67use wasmtime::{Engine, Linker, Module, ResourceLimiter, Store};
68
69use crate::config::NanoServiceAction;
70
71/// The fuel budget for one `handle` call. Generous enough for typical
72/// request inspection and response building; a module that loops
73/// forever traps with an out-of-fuel error (converted to 502).
74const NANO_FUEL: u64 = 1_000_000;
75
76/// The maximum memory the module may allocate, in bytes (sane ceiling
77/// mirrored from the validation bound so a config at the cap still
78/// works).
79const NANO_MAX_MEMORY: usize = 64 * 1024 * 1024;
80
81/// A compiled nano-service: the module path and the resource limits,
82/// resolved from the route's [`NanoServiceAction`]. The WASM bytes are
83/// read and compiled lazily by `NanoServiceHandler::load` so a
84/// missing/broken module is reported at handler construction, not at
85/// config publish time (validation already checks the file exists).
86#[derive(Clone, Debug)]
87pub struct NanoService {
88    /// The route name (used for metrics labeling).
89    pub name: String,
90    /// Filesystem path to the `.wasm` module.
91    pub module_path: String,
92    /// Maximum linear memory in bytes.
93    pub memory_limit: usize,
94    /// Maximum execution time in milliseconds.
95    pub execution_timeout: Duration,
96}
97
98impl NanoService {
99    /// Build from a route's config action.
100    pub fn from_action(name: &str, action: &NanoServiceAction) -> Self {
101        Self {
102            name: name.to_string(),
103            module_path: action.module.clone(),
104            memory_limit: action.memory_limit.min(NANO_MAX_MEMORY),
105            execution_timeout: Duration::from_millis(action.execution_timeout_ms),
106        }
107    }
108}
109
110/// The response a nano-service module produced.
111#[derive(Clone, Debug, PartialEq, Eq)]
112pub struct NanoServiceResponse {
113    /// HTTP status code.
114    pub status: u16,
115    /// Response headers (name, value) pairs.
116    pub headers: Vec<(String, String)>,
117    /// Response body bytes.
118    pub body: Vec<u8>,
119}
120
121/// A nano-service error. Maps to HTTP responses:
122/// - [`NanoServiceError::ModuleLoadFailed`] -> 502
123/// - [`NanoServiceError::ExecutionFailed`] -> 502
124/// - [`NanoServiceError::Timeout`] -> 504
125/// - [`NanoServiceError::MemoryLimitExceeded`] -> 502
126#[derive(Debug)]
127pub enum NanoServiceError {
128    /// The WASM module could not be read or compiled.
129    ModuleLoadFailed(String),
130    /// The `handle` call trapped or returned a non-zero status.
131    ExecutionFailed(String),
132    /// The `handle` call exceeded the execution timeout.
133    Timeout,
134    /// The module tried to allocate past its memory limit.
135    MemoryLimitExceeded(String),
136}
137
138impl std::fmt::Display for NanoServiceError {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        match self {
141            NanoServiceError::ModuleLoadFailed(m) => {
142                write!(f, "nano-service module load failed: {m}")
143            }
144            NanoServiceError::ExecutionFailed(m) => {
145                write!(f, "nano-service execution failed: {m}")
146            }
147            NanoServiceError::Timeout => write!(f, "nano-service execution timed out"),
148            NanoServiceError::MemoryLimitExceeded(m) => {
149                write!(f, "nano-service memory limit exceeded: {m}")
150            }
151        }
152    }
153}
154
155impl std::error::Error for NanoServiceError {}
156
157/// The per-instance context the host imports read from and write to.
158struct NanoContext {
159    /// The HTTP status the module set via `response_status`. Defaults
160    /// to 200 when the module never calls it (a module that only writes
161    /// a body still answers 200).
162    response_status: u16,
163    /// Response headers accumulated via `response_header`.
164    response_headers: Vec<(String, String)>,
165    /// The response body set via `response_body`.
166    response_body: Vec<u8>,
167    /// Log lines emitted via `log`.
168    logs: Vec<String>,
169    /// Memory limiter state: bytes currently in use.
170    #[allow(dead_code)]
171    memory_used: usize,
172    /// Memory cap in bytes.
173    memory_cap: usize,
174}
175
176impl NanoContext {
177    fn new(memory_cap: usize) -> Self {
178        Self {
179            response_status: 200,
180            response_headers: Vec::new(),
181            response_body: Vec::new(),
182            logs: Vec::new(),
183            memory_used: 0,
184            memory_cap,
185        }
186    }
187}
188
189impl ResourceLimiter for NanoContext {
190    fn memory_growing(
191        &mut self,
192        _current: usize,
193        desired: usize,
194        _maximum: Option<usize>,
195    ) -> wasmtime::Result<bool> {
196        Ok(desired <= self.memory_cap)
197    }
198
199    fn table_growing(
200        &mut self,
201        _current: usize,
202        desired: usize,
203        _maximum: Option<usize>,
204    ) -> wasmtime::Result<bool> {
205        Ok(desired <= 10_000)
206    }
207}
208
209/// The nano-service handler: holds a process-wide wasmtime engine and
210/// the linker wired with the `dwara` host imports. Cheap to clone (Arc
211/// internals). One instance per `crate::dataplane::DataPlane` (or per
212/// route set); [`NanoServiceHandler::handle`] is called per request.
213#[derive(Clone)]
214pub struct NanoServiceHandler {
215    engine: Engine,
216    linker: Arc<Linker<NanoContext>>,
217}
218
219impl NanoServiceHandler {
220    /// Create a new handler with the nano-service ABI linker.
221    pub fn new() -> Result<Self, NanoServiceError> {
222        let mut config = wasmtime::Config::new();
223        config.consume_fuel(true);
224        config.strategy(wasmtime::Strategy::Cranelift);
225        let engine = Engine::new(&config)
226            .map_err(|e| NanoServiceError::ModuleLoadFailed(format!("wasmtime engine: {e}")))?;
227
228        let mut linker: Linker<NanoContext> = Linker::new(&engine);
229
230        // dwara::response_status(status: i32)
231        linker
232            .func_wrap(
233                "dwara",
234                "response_status",
235                |mut caller: wasmtime::Caller<NanoContext>, status: i32| {
236                    if status > 0 && status <= 599 {
237                        caller.data_mut().response_status = status as u16;
238                    }
239                },
240            )
241            .map_err(|e| {
242                NanoServiceError::ModuleLoadFailed(format!("linker response_status: {e}"))
243            })?;
244
245        // dwara::response_header(k_ptr, k_len, v_ptr, v_len)
246        linker
247            .func_wrap(
248                "dwara",
249                "response_header",
250                |mut caller: wasmtime::Caller<NanoContext>,
251                 k_ptr: i32,
252                 k_len: i32,
253                 v_ptr: i32,
254                 v_len: i32|
255                 -> i32 {
256                    if k_ptr < 0 || k_len < 0 || v_ptr < 0 || v_len < 0 {
257                        return 1;
258                    }
259                    let memory = match caller.get_export("memory") {
260                        Some(wasmtime::Extern::Memory(m)) => m,
261                        _ => return 1,
262                    };
263                    let data = memory.data(&caller);
264                    let key = match data.get(k_ptr as usize..(k_ptr as usize + k_len as usize)) {
265                        Some(slice) => slice.to_vec(),
266                        None => return 1,
267                    };
268                    let value = match data.get(v_ptr as usize..(v_ptr as usize + v_len as usize)) {
269                        Some(slice) => slice.to_vec(),
270                        None => return 1,
271                    };
272                    let key = String::from_utf8_lossy(&key).into_owned();
273                    let value = String::from_utf8_lossy(&value).into_owned();
274                    caller.data_mut().response_headers.push((key, value));
275                    0
276                },
277            )
278            .map_err(|e| {
279                NanoServiceError::ModuleLoadFailed(format!("linker response_header: {e}"))
280            })?;
281
282        // dwara::response_body(ptr, len)
283        linker
284            .func_wrap(
285                "dwara",
286                "response_body",
287                |mut caller: wasmtime::Caller<NanoContext>, ptr: i32, len: i32| -> i32 {
288                    if ptr < 0 || len < 0 {
289                        return 1;
290                    }
291                    let memory = match caller.get_export("memory") {
292                        Some(wasmtime::Extern::Memory(m)) => m,
293                        _ => return 1,
294                    };
295                    let body = match memory
296                        .data(&caller)
297                        .get(ptr as usize..(ptr as usize + len as usize))
298                    {
299                        Some(slice) => slice.to_vec(),
300                        None => return 1,
301                    };
302                    caller.data_mut().response_body = body;
303                    0
304                },
305            )
306            .map_err(|e| {
307                NanoServiceError::ModuleLoadFailed(format!("linker response_body: {e}"))
308            })?;
309
310        // dwara::log(ptr, len)
311        linker
312            .func_wrap(
313                "dwara",
314                "log",
315                |mut caller: wasmtime::Caller<NanoContext>, ptr: i32, len: i32| -> i32 {
316                    if ptr < 0 || len < 0 {
317                        return 1;
318                    }
319                    let memory = match caller.get_export("memory") {
320                        Some(wasmtime::Extern::Memory(m)) => m,
321                        _ => return 1,
322                    };
323                    let msg = match memory
324                        .data(&caller)
325                        .get(ptr as usize..(ptr as usize + len as usize))
326                    {
327                        Some(slice) => slice.to_vec(),
328                        None => return 1,
329                    };
330                    let msg = String::from_utf8_lossy(&msg).into_owned();
331                    caller.data_mut().logs.push(msg);
332                    0
333                },
334            )
335            .map_err(|e| NanoServiceError::ModuleLoadFailed(format!("linker log: {e}")))?;
336
337        Ok(Self {
338            engine,
339            linker: Arc::new(linker),
340        })
341    }
342
343    /// Run a nano-service module for one request. Reads the module from
344    /// `service.module_path`, instantiates it, writes the serialized
345    /// request into linear memory, calls `handle`, and returns the
346    /// response the module produced. The `handle` call runs on a
347    /// blocking-pool thread bounded by `service.execution_timeout`.
348    pub async fn handle(
349        &self,
350        service: &NanoService,
351        method: &str,
352        path: &str,
353        headers: &[(String, String)],
354        body: &[u8],
355    ) -> Result<NanoServiceResponse, NanoServiceError> {
356        let engine = self.engine.clone();
357        let linker = self.linker.clone();
358        let module_path = service.module_path.clone();
359        let memory_limit = service.memory_limit;
360        let timeout = service.execution_timeout;
361        let req = serialize_request(method, path, headers, body);
362
363        let join = tokio::task::spawn_blocking(move || {
364            run_handle_sync(&engine, &linker, &module_path, memory_limit, &req)
365        });
366
367        match tokio::time::timeout(timeout, join).await {
368            Ok(Ok(Ok(resp))) => Ok(resp),
369            Ok(Ok(Err(e))) => match e {
370                NanoServiceError::MemoryLimitExceeded(_) => Err(e),
371                _ => Err(e),
372            },
373            // The blocking task panicked or was cancelled.
374            Ok(Err(e)) => Err(NanoServiceError::ExecutionFailed(format!(
375                "nano-service task failed: {e}"
376            ))),
377            // The timeout elapsed before the task finished.
378            Err(_) => Err(NanoServiceError::Timeout),
379        }
380    }
381}
382
383/// A process-wide shared nano-service handler. wasmtime recommends a
384/// single `Engine` per process (compiled modules are cached and shared
385/// across threads), so the dataplane reuses one handler for every
386/// nano-service route. The handler is stateless beyond the immutable
387/// engine + linker (both Arc-held), so sharing it across requests and
388/// tests is safe. Lazily initialized on first use; a construction
389/// failure (wasmtime misconfiguration) is reported once and cached as
390/// an error so every subsequent request answers 502 without retrying.
391pub fn shared_handler() -> Result<NanoServiceHandler, NanoServiceError> {
392    static HANDLER: std::sync::OnceLock<Result<NanoServiceHandler, NanoServiceError>> =
393        std::sync::OnceLock::new();
394    match HANDLER.get_or_init(NanoServiceHandler::new).as_ref() {
395        Ok(h) => Ok(h.clone()),
396        Err(e) => Err(NanoServiceError::ModuleLoadFailed(e.to_string())),
397    }
398}
399
400/// The synchronous core: compile, instantiate, write the request, call
401/// `handle`, read the response. Runs on a blocking-pool thread so the
402/// async caller is never blocked on wasm compilation/execution.
403fn run_handle_sync(
404    engine: &Engine,
405    linker: &Linker<NanoContext>,
406    module_path: &str,
407    memory_limit: usize,
408    req: &[u8],
409) -> Result<NanoServiceResponse, NanoServiceError> {
410    let wasm_bytes = std::fs::read(module_path)
411        .map_err(|e| NanoServiceError::ModuleLoadFailed(format!("read {module_path}: {e}")))?;
412    let module = Module::new(engine, &wasm_bytes)
413        .map_err(|e| NanoServiceError::ModuleLoadFailed(format!("compile {module_path}: {e}")))?;
414
415    let memory_cap = memory_limit.max(1);
416    let ctx = NanoContext::new(memory_cap);
417    let mut store = Store::new(engine, ctx);
418    store.limiter(move |ctx| ctx as &mut dyn ResourceLimiter);
419    store
420        .set_fuel(NANO_FUEL)
421        .map_err(|e| NanoServiceError::ModuleLoadFailed(format!("set_fuel: {e}")))?;
422
423    let instance = linker
424        .instantiate(&mut store, &module)
425        .map_err(|e| NanoServiceError::ModuleLoadFailed(format!("instantiate: {e}")))?;
426
427    let memory = instance
428        .get_export(&mut store, "memory")
429        .and_then(|e| e.into_memory())
430        .ok_or_else(|| {
431            NanoServiceError::ModuleLoadFailed("module does not export `memory`".to_string())
432        })?;
433
434    // Place the serialized request into the module's linear memory via
435    // its `alloc` export (a bump allocator). Fall back to growing the
436    // memory if the export is absent.
437    let req_ptr = allocate(&mut store, &instance, &memory, req.len())?;
438    memory
439        .data_mut(&mut store)
440        .get_mut(req_ptr..req_ptr + req.len())
441        .ok_or_else(|| {
442            NanoServiceError::ExecutionFailed("request write out of bounds".to_string())
443        })?
444        .copy_from_slice(req);
445
446    // Call handle(req_ptr, req_len) -> i32.
447    let handle_export = instance
448        .get_export(&mut store, "handle")
449        .and_then(|e| e.into_func())
450        .ok_or_else(|| {
451            NanoServiceError::ModuleLoadFailed("module does not export `handle`".to_string())
452        })?;
453    let typed: wasmtime::TypedFunc<(i32, i32), i32> = handle_export
454        .typed(&store)
455        .map_err(|e| NanoServiceError::ModuleLoadFailed(format!("handle typed: {e}")))?;
456
457    let result = typed
458        .call(&mut store, (req_ptr as i32, req.len() as i32))
459        .map_err(|e| {
460            // Detect fuel exhaustion vs other traps.
461            let fuel_left = store.get_fuel().unwrap_or(0);
462            if fuel_left == 0 {
463                NanoServiceError::ExecutionFailed(format!("handle: fuel exhausted: {e}"))
464            } else {
465                NanoServiceError::ExecutionFailed(format!("handle: {e}"))
466            }
467        })?;
468
469    if result != 0 {
470        return Err(NanoServiceError::ExecutionFailed(format!(
471            "handle returned non-zero status {result}"
472        )));
473    }
474
475    let ctx = store.into_data();
476    for line in &ctx.logs {
477        tracing::debug!(
478            module = %module_path,
479            "nano-service log: {line}"
480        );
481    }
482    Ok(NanoServiceResponse {
483        status: ctx.response_status,
484        headers: ctx.response_headers,
485        body: ctx.response_body,
486    })
487}
488
489/// Allocate `size` bytes in the module's linear memory. Calls the
490/// module's `alloc` export when present; otherwise grows the memory by
491/// one page and uses the new space.
492fn allocate(
493    store: &mut Store<NanoContext>,
494    instance: &wasmtime::Instance,
495    memory: &wasmtime::Memory,
496    size: usize,
497) -> Result<usize, NanoServiceError> {
498    if let Some(export) = instance.get_export(&mut *store, "alloc") {
499        if let Some(func) = export.into_func() {
500            if let Ok(typed) = func.typed::<(i32,), i32>(&mut *store) {
501                let ptr = typed
502                    .call(&mut *store, (size as i32,))
503                    .map_err(|e| NanoServiceError::ExecutionFailed(format!("alloc: {e}")))?;
504                if ptr > 0 {
505                    return Ok(ptr as usize);
506                }
507            }
508        }
509    }
510    // Fallback: grow the memory and use the tail.
511    let current = memory.data_size(&mut *store);
512    let needed = current + size;
513    let pages_needed = (needed - current).div_ceil(64 * 1024);
514    if pages_needed > 0 {
515        memory
516            .grow(&mut *store, pages_needed as u64)
517            .map_err(|e| NanoServiceError::MemoryLimitExceeded(format!("memory.grow: {e}")))?;
518    }
519    Ok(current)
520}
521
522/// Serialize a request into the nano-service wire format (all lengths
523/// u32 big-endian):
524///
525/// ```text
526/// u32 method_len, method
527/// u32 path_len, path
528/// u32 headers_count, [u32 key_len, key, u32 val_len, val]...
529/// u32 body_len, body
530/// ```
531pub fn serialize_request(
532    method: &str,
533    path: &str,
534    headers: &[(String, String)],
535    body: &[u8],
536) -> Vec<u8> {
537    let mut buf = Vec::with_capacity(method.len() + path.len() + body.len() + 32);
538    write_str(&mut buf, method);
539    write_str(&mut buf, path);
540    buf.extend_from_slice(&(headers.len() as u32).to_be_bytes());
541    for (key, value) in headers {
542        write_str(&mut buf, key);
543        write_str(&mut buf, value);
544    }
545    buf.extend_from_slice(&(body.len() as u32).to_be_bytes());
546    buf.extend_from_slice(body);
547    buf
548}
549
550/// The deserialized parts of a nano-service wire-format request.
551pub struct DeserializedRequest {
552    pub method: String,
553    pub path: String,
554    pub headers: Vec<(String, String)>,
555    pub body: Vec<u8>,
556}
557
558/// Deserialize a nano-service wire-format request back into its parts.
559/// Returns `None` on a truncated/corrupt buffer.
560pub fn deserialize_request(buf: &[u8]) -> Option<DeserializedRequest> {
561    let mut pos = 0;
562    let method = read_str(buf, &mut pos)?;
563    let path = read_str(buf, &mut pos)?;
564    let header_count = read_u32(buf, &mut pos)? as usize;
565    let mut headers = Vec::with_capacity(header_count);
566    for _ in 0..header_count {
567        let key = read_str(buf, &mut pos)?;
568        let value = read_str(buf, &mut pos)?;
569        headers.push((key, value));
570    }
571    let body_len = read_u32(buf, &mut pos)? as usize;
572    if pos + body_len > buf.len() {
573        return None;
574    }
575    let body = buf[pos..pos + body_len].to_vec();
576    Some(DeserializedRequest {
577        method,
578        path,
579        headers,
580        body,
581    })
582}
583
584fn write_str(buf: &mut Vec<u8>, s: &str) {
585    buf.extend_from_slice(&(s.len() as u32).to_be_bytes());
586    buf.extend_from_slice(s.as_bytes());
587}
588
589fn read_u32(buf: &[u8], pos: &mut usize) -> Option<u32> {
590    if *pos + 4 > buf.len() {
591        return None;
592    }
593    let v = u32::from_be_bytes(buf[*pos..*pos + 4].try_into().ok()?);
594    *pos += 4;
595    Some(v)
596}
597
598fn read_str(buf: &[u8], pos: &mut usize) -> Option<String> {
599    let len = read_u32(buf, pos)? as usize;
600    if *pos + len > buf.len() {
601        return None;
602    }
603    let s = String::from_utf8(buf[*pos..*pos + len].to_vec()).ok()?;
604    *pos += len;
605    Some(s)
606}
607
608#[cfg(test)]
609mod tests {
610    use super::*;
611
612    #[test]
613    fn request_round_trip() {
614        let headers = vec![
615            ("host".to_string(), "example.com".to_string()),
616            ("x-trace".to_string(), "abc".to_string()),
617        ];
618        let body = b"{\"hello\":\"world\"}";
619        let wire = serialize_request("POST", "/v1/echo", &headers, body);
620        let req = deserialize_request(&wire).expect("round trip");
621        assert_eq!(req.method, "POST");
622        assert_eq!(req.path, "/v1/echo");
623        assert_eq!(req.headers, headers);
624        assert_eq!(req.body, body);
625    }
626
627    #[test]
628    fn request_round_trip_empty() {
629        let wire = serialize_request("GET", "/", &[], &[]);
630        let req = deserialize_request(&wire).expect("round trip");
631        assert_eq!(req.method, "GET");
632        assert_eq!(req.path, "/");
633        assert!(req.headers.is_empty());
634        assert!(req.body.is_empty());
635    }
636
637    #[test]
638    fn deserialize_truncated_returns_none() {
639        // A buffer that claims a 1-byte method but has no method bytes.
640        let buf = (1u32).to_be_bytes().to_vec();
641        assert!(deserialize_request(&buf).is_none());
642    }
643}