strate_wasm/
main.rs

1#![no_std]
2#![no_main]
3#![feature(alloc_error_handler)]
4
5extern crate alloc;
6
7use alloc::{format, string::String, vec, vec::Vec};
8use core::{alloc::Layout, panic::PanicInfo};
9use strat9_syscall::{call, data::IpcMessage};
10use talc::*;
11use wasmi::{
12    Caller, Config, Engine, Instance, Linker, Module, Store, StoreLimits, StoreLimitsBuilder,
13};
14
15// ---------------------------------------------------------------------------
16// MEMORY MANAGEMENT (TALC + BRK)
17// ---------------------------------------------------------------------------
18
19struct BrkGrower;
20
21impl OomHandler for BrkGrower {
22    /// Implements handle oom.
23    fn handle_oom(talc: &mut Talc<Self>, layout: Layout) -> Result<(), ()> {
24        let current_brk = call::brk(0).map_err(|_| ())?;
25        let growth = (layout.size().max(layout.align()) + 65536) & !4095;
26        let new_brk = call::brk(current_brk + growth).map_err(|_| ())?;
27        let actual = new_brk.saturating_sub(current_brk);
28        if actual == 0 {
29            return Err(());
30        }
31        unsafe {
32            talc.claim(Span::from_base_size(current_brk as *mut u8, actual))
33                .map_err(|_| ())?;
34        }
35        Ok(())
36    }
37}
38
39#[global_allocator]
40static ALLOCATOR: Talck<spin::Mutex<()>, BrkGrower> = Talck::new(Talc::new(BrkGrower));
41
42#[alloc_error_handler]
43/// Implements alloc error.
44fn alloc_error(_layout: Layout) -> ! {
45    let _ = call::debug_log(b"[strate-wasm] Fatal: OOM\n");
46    call::exit(12);
47}
48
49// ---------------------------------------------------------------------------
50// UTILITIES
51// ---------------------------------------------------------------------------
52
53/// Implements debug log.
54fn debug_log(msg: &str) {
55    let _ = call::debug_log(msg.as_bytes());
56}
57
58#[panic_handler]
59/// Implements panic.
60fn panic(_info: &PanicInfo) -> ! {
61    debug_log("[strate-wasm] PANIC!\n");
62    call::exit(1);
63}
64
65/// Implements extract string.
66fn extract_string(payload: &[u8], offset: usize, len: usize) -> String {
67    if len == 0 || offset + len > payload.len() {
68        return String::from("invalid");
69    }
70    let slice = &payload[offset..(offset + len)];
71    let mut s = String::with_capacity(len);
72    for &b in slice {
73        if b.is_ascii() && !b.is_ascii_control() {
74            s.push(b as char);
75        } else {
76            s.push('?');
77        }
78    }
79    s
80}
81
82/// Implements send response.
83fn send_response(original_sender: u64, status: u32) {
84    let mut msg = IpcMessage::new(REPLY_MSG_TYPE);
85    msg.sender = original_sender;
86    msg.payload[0..4].copy_from_slice(&status.to_le_bytes());
87    let _ = call::ipc_reply(&msg);
88}
89
90/// Implements bind runtime alias.
91fn bind_runtime_alias(port_h: usize, label: &str) {
92    let path = format!("/srv/strate-wasm/{}", label);
93    let _ = call::ipc_bind_port(port_h, path.as_bytes());
94}
95
96/// Reads leb u32.
97fn read_leb_u32(bytes: &[u8], mut idx: usize) -> Option<(u32, usize)> {
98    let mut result: u32 = 0;
99    let mut shift = 0u32;
100    for _ in 0..5 {
101        let b = *bytes.get(idx)?;
102        idx += 1;
103        result |= ((b & 0x7f) as u32) << shift;
104        if (b & 0x80) == 0 {
105            return Some((result, idx));
106        }
107        shift += 7;
108    }
109    None
110}
111
112/// Implements wasm effective len.
113fn wasm_effective_len(bytes: &[u8]) -> Option<usize> {
114    if bytes.len() < 8 {
115        return None;
116    }
117    if bytes[0] != 0x00
118        || bytes[1] != b'a'
119        || bytes[2] != b's'
120        || bytes[3] != b'm'
121        || bytes[4] != 0x01
122        || bytes[5] != 0x00
123        || bytes[6] != 0x00
124        || bytes[7] != 0x00
125    {
126        return None;
127    }
128    let mut i = 8usize;
129    while i < bytes.len() {
130        let id = bytes[i];
131        if id > 12 {
132            return Some(i);
133        }
134        i += 1;
135        let (section_len, next) = read_leb_u32(bytes, i)?;
136        i = next;
137        let end = i.checked_add(section_len as usize)?;
138        if end > bytes.len() {
139            return None;
140        }
141        i = end;
142    }
143    Some(i)
144}
145
146// ---------------------------------------------------------------------------
147// WASI / HOST FUNCTIONS
148// ---------------------------------------------------------------------------
149
150struct HostState {
151    limits: StoreLimits,
152    exit_code: Option<u32>,
153}
154
155/// Implements wasi fd write.
156fn wasi_fd_write(
157    mut caller: Caller<'_, HostState>,
158    fd: u32,
159    iovs_ptr: u32,
160    iovs_len: u32,
161    nwritten_ptr: u32,
162) -> u32 {
163    let Some(memory) = caller.get_export("memory").and_then(|e| e.into_memory()) else {
164        return 1;
165    };
166    if fd != 1 && fd != 2 {
167        return 8; // WASI EBADF
168    }
169    let mut total_written: u32 = 0;
170    for i in 0..iovs_len {
171        let base_addr = (iovs_ptr + i * 8) as usize;
172        let mut iov_desc = [0u8; 8];
173        if memory.read(&caller, base_addr, &mut iov_desc).is_err() {
174            break;
175        }
176        let buf_ptr =
177            u32::from_le_bytes([iov_desc[0], iov_desc[1], iov_desc[2], iov_desc[3]]) as usize;
178        let buf_len =
179            u32::from_le_bytes([iov_desc[4], iov_desc[5], iov_desc[6], iov_desc[7]]) as usize;
180        let mut buffer = vec![0u8; buf_len];
181        if memory.read(&caller, buf_ptr, &mut buffer).is_err() {
182            break;
183        }
184        match call::write(fd as usize, &buffer) {
185            Ok(n) => total_written += n as u32,
186            Err(_) => break,
187        }
188    }
189    let _ = memory.write(
190        &mut caller,
191        nwritten_ptr as usize,
192        &total_written.to_le_bytes(),
193    );
194    0
195}
196
197/// Implements wasi environ get.
198fn wasi_environ_get(_caller: Caller<'_, HostState>, _environ: u32, _environ_buf: u32) -> u32 {
199    0
200}
201/// Implements wasi environ sizes get.
202fn wasi_environ_sizes_get(mut caller: Caller<'_, HostState>, count_ptr: u32, size_ptr: u32) -> u32 {
203    let Some(memory) = caller.get_export("memory").and_then(|e| e.into_memory()) else {
204        return 1;
205    };
206    let _ = memory.write(&mut caller, count_ptr as usize, &0u32.to_le_bytes());
207    let _ = memory.write(&mut caller, size_ptr as usize, &0u32.to_le_bytes());
208    0
209}
210/// Implements wasi proc exit.
211fn wasi_proc_exit(mut caller: Caller<'_, HostState>, code: u32) -> Result<(), wasmi::Error> {
212    caller.data_mut().exit_code = Some(code);
213    Err(wasmi::Error::new("proc_exit"))
214}
215
216// ---------------------------------------------------------------------------
217// IPC PROTOCOL
218// ---------------------------------------------------------------------------
219
220const OP_WASM_LOAD_PATH: u32 = 0x100;
221const OP_WASM_RUN_MAIN: u32 = 0x102;
222const OP_BOOTSTRAP: u32 = 0x10;
223const REPLY_MSG_TYPE: u32 = 0x80;
224
225const RESP_OK: u32 = 0x0;
226const RESP_ERR_OPEN: u32 = 0x10;
227const RESP_ERR_READ: u32 = 0x11;
228const RESP_ERR_MODULE: u32 = 0x12;
229const RESP_ERR_INSTANTIATE: u32 = 0x13;
230const RESP_ERR_START: u32 = 0x14;
231const RESP_ERR_NO_INSTANCE: u32 = 0x20;
232const RESP_ERR_NO_ENTRY: u32 = 0x21;
233const RESP_ERR_RUN: u32 = 0x22;
234const RESP_ERR_TOO_LARGE: u32 = 0x23;
235const RESP_ERR_UNSUPPORTED: u32 = 0x24;
236
237const MAX_WASM_SIZE: usize = 16 * 1024 * 1024;
238
239// ---------------------------------------------------------------------------
240// MAIN
241// ---------------------------------------------------------------------------
242
243#[unsafe(no_mangle)]
244/// Implements start.
245pub unsafe extern "C" fn _start() -> ! {
246    let mut config = Config::default();
247    config.consume_fuel(true);
248
249    let engine = Engine::new(&config);
250    let mut linker = Linker::new(&engine);
251    let limits = StoreLimitsBuilder::new()
252        .memory_size(16 * 1024 * 1024)
253        .instances(1)
254        .build();
255    let mut store = Store::new(
256        &engine,
257        HostState {
258            limits,
259            exit_code: None,
260        },
261    );
262    store.limiter(|state| &mut state.limits);
263
264    let _ = linker.define(
265        "wasi_snapshot_preview1",
266        "fd_write",
267        wasmi::Func::wrap(&mut store, wasi_fd_write),
268    );
269    let _ = linker.define(
270        "wasi_snapshot_preview1",
271        "proc_exit",
272        wasmi::Func::wrap(&mut store, wasi_proc_exit),
273    );
274    let _ = linker.define(
275        "wasi_snapshot_preview1",
276        "environ_get",
277        wasmi::Func::wrap(&mut store, wasi_environ_get),
278    );
279    let _ = linker.define(
280        "wasi_snapshot_preview1",
281        "environ_sizes_get",
282        wasmi::Func::wrap(&mut store, wasi_environ_sizes_get),
283    );
284
285    let mut label = String::from("default");
286    let mut queued_msg: Option<IpcMessage> = None;
287
288    // We need a port to bind to
289    let port_h = call::ipc_create_port(0).expect("failed to create port");
290    let _ = call::ipc_bind_port(port_h, b"/srv/strate-wasm/bootstrap");
291
292    let mut b_msg = IpcMessage::new(0);
293    if let Ok(_) = call::ipc_try_recv(port_h, &mut b_msg) {
294        if b_msg.msg_type == OP_BOOTSTRAP {
295            label = extract_string(&b_msg.payload, 1, b_msg.payload[0] as usize);
296            send_response(b_msg.sender, RESP_OK);
297        } else {
298            queued_msg = Some(b_msg);
299        }
300    }
301
302    bind_runtime_alias(port_h, &label);
303    let final_path = format!("/srv/strate-wasm/{}", label);
304    debug_log("[strate-wasm] running: ");
305    debug_log(&final_path);
306    debug_log("\n");
307
308    let mut current_instance: Option<Instance> = None;
309
310    loop {
311        let msg = if let Some(m) = queued_msg.take() {
312            m
313        } else {
314            let mut m = IpcMessage::new(0);
315            if call::ipc_recv(port_h, &mut m).is_err() {
316                let _ = call::sched_yield();
317                continue;
318            }
319            m
320        };
321        let src = msg.sender;
322        match msg.msg_type {
323            OP_BOOTSTRAP => {
324                let new_label = extract_string(&msg.payload, 1, msg.payload[0] as usize);
325                bind_runtime_alias(port_h, &new_label);
326                send_response(src, RESP_OK);
327            }
328            OP_WASM_LOAD_PATH => {
329                let path = extract_string(&msg.payload, 1, msg.payload[0] as usize);
330                let mut wasm_bytes = Vec::new();
331                if let Ok(fd) = call::openat(0, &path, 0x1, 0) {
332                    let mut chunk = [0u8; 4096];
333                    let mut read_failed = false;
334                    loop {
335                        match call::read(fd as usize, &mut chunk) {
336                            Ok(0) => break,
337                            Ok(n) => {
338                                wasm_bytes.extend_from_slice(&chunk[..n]);
339                                if wasm_bytes.len() > MAX_WASM_SIZE {
340                                    read_failed = true;
341                                    break;
342                                }
343                            }
344                            Err(_) => {
345                                read_failed = true;
346                                break;
347                            }
348                        }
349                    }
350                    let _ = call::close(fd as usize);
351
352                    if read_failed {
353                        if wasm_bytes.len() > MAX_WASM_SIZE {
354                            debug_log("[strate-wasm] load error: too large\n");
355                            send_response(src, RESP_ERR_TOO_LARGE);
356                        } else {
357                            debug_log("[strate-wasm] load error: read\n");
358                            send_response(src, RESP_ERR_READ);
359                        }
360                        continue;
361                    }
362
363                    if wasm_bytes.is_empty() {
364                        debug_log("[strate-wasm] load error: empty\n");
365                        send_response(src, RESP_ERR_READ);
366                        continue;
367                    }
368
369                    let trimmed_len = wasm_bytes
370                        .iter()
371                        .rposition(|&b| b != 0)
372                        .map(|i| i + 1)
373                        .unwrap_or(0);
374                    if trimmed_len == 0 {
375                        debug_log("[strate-wasm] load error: all-zero image\n");
376                        send_response(src, RESP_ERR_MODULE);
377                        continue;
378                    }
379                    let trimmed = &wasm_bytes[..trimmed_len];
380                    let Some(effective_len) = wasm_effective_len(trimmed) else {
381                        debug_log("[strate-wasm] load error: bad magic\n");
382                        send_response(src, RESP_ERR_MODULE);
383                        continue;
384                    };
385                    let module_bytes = &trimmed[..effective_len];
386
387                    match Module::new(&engine, module_bytes) {
388                        Ok(module) => match linker.instantiate(&mut store, &module) {
389                            Ok(pre) => match pre.start(&mut store) {
390                                Ok(inst) => {
391                                    current_instance = Some(inst);
392                                    send_response(src, RESP_OK);
393                                }
394                                Err(_) => {
395                                    debug_log("[strate-wasm] load error: start\n");
396                                    send_response(src, RESP_ERR_START)
397                                }
398                            },
399                            Err(_) => {
400                                debug_log("[strate-wasm] load error: instantiate\n");
401                                send_response(src, RESP_ERR_INSTANTIATE)
402                            }
403                        },
404                        Err(_) => {
405                            let line = format!(
406                                "[strate-wasm] load error: module parse failed (len={})\n",
407                                module_bytes.len()
408                            );
409                            debug_log(&line);
410                            send_response(src, RESP_ERR_MODULE)
411                        }
412                    }
413                } else {
414                    debug_log("[strate-wasm] load error: open\n");
415                    send_response(src, RESP_ERR_OPEN);
416                }
417            }
418
419            OP_WASM_RUN_MAIN => {
420                if let Some(ref instance) = current_instance {
421                    store.data_mut().exit_code = None;
422                    let _ = store.set_fuel(10_000_000);
423                    if let Ok(func) = instance.get_typed_func::<(), ()>(&store, "_start") {
424                        match func.call(&mut store, ()) {
425                            Ok(_) => send_response(src, RESP_OK),
426                            Err(_) => {
427                                if store.data().exit_code.is_some() {
428                                    send_response(src, RESP_OK);
429                                } else {
430                                    debug_log("[strate-wasm] run error: trap\n");
431                                    send_response(src, RESP_ERR_RUN);
432                                }
433                            }
434                        }
435                    } else {
436                        debug_log("[strate-wasm] run error: no _start\n");
437                        send_response(src, RESP_ERR_NO_ENTRY);
438                    }
439                } else {
440                    debug_log("[strate-wasm] run error: no instance\n");
441                    send_response(src, RESP_ERR_NO_INSTANCE);
442                }
443            }
444
445            _ => {
446                send_response(src, RESP_ERR_UNSUPPORTED);
447            }
448        }
449    }
450}
strate_wasm/main.rs

strate_wasm/
main.rs
