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

strate_wasm/
main.rs
