1use std::sync::Arc;
65use std::time::Duration;
66
67use wasmtime::{Engine, Linker, Module, ResourceLimiter, Store};
68
69use crate::config::NanoServiceAction;
70
71const NANO_FUEL: u64 = 1_000_000;
75
76const NANO_MAX_MEMORY: usize = 64 * 1024 * 1024;
80
81#[derive(Clone, Debug)]
87pub struct NanoService {
88 pub name: String,
90 pub module_path: String,
92 pub memory_limit: usize,
94 pub execution_timeout: Duration,
96}
97
98impl NanoService {
99 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#[derive(Clone, Debug, PartialEq, Eq)]
112pub struct NanoServiceResponse {
113 pub status: u16,
115 pub headers: Vec<(String, String)>,
117 pub body: Vec<u8>,
119}
120
121#[derive(Debug)]
127pub enum NanoServiceError {
128 ModuleLoadFailed(String),
130 ExecutionFailed(String),
132 Timeout,
134 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
157struct NanoContext {
159 response_status: u16,
163 response_headers: Vec<(String, String)>,
165 response_body: Vec<u8>,
167 logs: Vec<String>,
169 #[allow(dead_code)]
171 memory_used: usize,
172 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#[derive(Clone)]
214pub struct NanoServiceHandler {
215 engine: Engine,
216 linker: Arc<Linker<NanoContext>>,
217}
218
219impl NanoServiceHandler {
220 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 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 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 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 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 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 Ok(Err(e)) => Err(NanoServiceError::ExecutionFailed(format!(
375 "nano-service task failed: {e}"
376 ))),
377 Err(_) => Err(NanoServiceError::Timeout),
379 }
380 }
381}
382
383pub 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
400fn 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 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 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 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
489fn 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 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
522pub 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
550pub struct DeserializedRequest {
552 pub method: String,
553 pub path: String,
554 pub headers: Vec<(String, String)>,
555 pub body: Vec<u8>,
556}
557
558pub 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 let buf = (1u32).to_be_bytes().to_vec();
641 assert!(deserialize_request(&buf).is_none());
642 }
643}