lunatic_wasi_api/
lib.rs

1use anyhow::Result;
2use lunatic_common_api::{get_memory, IntoTrap};
3use lunatic_process::state::ProcessState;
4use lunatic_stdout_capture::StdoutCapture;
5use wasmtime::{Caller, Linker};
6use wasmtime_wasi::{ambient_authority, Dir, WasiCtx, WasiCtxBuilder};
7
8/// Create a `WasiCtx` from configuration settings.
9pub fn build_wasi(
10    args: Option<&Vec<String>>,
11    envs: Option<&Vec<(String, String)>>,
12    dirs: &[String],
13) -> Result<WasiCtx> {
14    let mut wasi = WasiCtxBuilder::new().inherit_stdio();
15    if let Some(envs) = envs {
16        wasi = wasi.envs(envs)?;
17    }
18    if let Some(args) = args {
19        wasi = wasi.args(args)?;
20    }
21    for preopen_dir_path in dirs {
22        let preopen_dir = Dir::open_ambient_dir(preopen_dir_path, ambient_authority())?;
23        wasi = wasi.preopened_dir(preopen_dir, preopen_dir_path)?;
24    }
25    Ok(wasi.build())
26}
27
28pub trait LunaticWasiConfigCtx {
29    fn add_environment_variable(&mut self, key: String, value: String);
30    fn add_command_line_argument(&mut self, argument: String);
31    fn preopen_dir(&mut self, dir: String);
32}
33
34pub trait LunaticWasiCtx {
35    fn wasi(&self) -> &WasiCtx;
36    fn wasi_mut(&mut self) -> &mut WasiCtx;
37    fn set_stdout(&mut self, stdout: StdoutCapture);
38    fn get_stdout(&self) -> Option<&StdoutCapture>;
39    fn set_stderr(&mut self, stderr: StdoutCapture);
40    fn get_stderr(&self) -> Option<&StdoutCapture>;
41}
42
43// Register WASI APIs to the linker
44pub fn register<T>(linker: &mut Linker<T>) -> Result<()>
45where
46    T: ProcessState + LunaticWasiCtx + Send + 'static,
47    T::Config: LunaticWasiConfigCtx,
48{
49    // Register all wasi host functions
50    wasmtime_wasi::sync::snapshots::preview_1::add_wasi_snapshot_preview1_to_linker(
51        linker,
52        |ctx| ctx.wasi_mut(),
53    )?;
54
55    // Register host functions to configure wasi
56    linker.func_wrap(
57        "lunatic::wasi",
58        "config_add_environment_variable",
59        add_environment_variable,
60    )?;
61    linker.func_wrap(
62        "lunatic::wasi",
63        "config_add_command_line_argument",
64        add_command_line_argument,
65    )?;
66    linker.func_wrap("lunatic::wasi", "config_preopen_dir", preopen_dir)?;
67
68    Ok(())
69}
70
71// Adds environment variable to a configuration.
72//
73// Traps:
74// * If the config ID doesn't exist.
75// * If the key or value string is not a valid utf8 string.
76// * If any of the memory slices falls outside the memory.
77fn add_environment_variable<T>(
78    mut caller: Caller<T>,
79    config_id: u64,
80    key_ptr: u32,
81    key_len: u32,
82    value_ptr: u32,
83    value_len: u32,
84) -> Result<()>
85where
86    T: ProcessState,
87    T::Config: LunaticWasiConfigCtx,
88{
89    let memory = get_memory(&mut caller)?;
90    let key_str = memory
91        .data(&caller)
92        .get(key_ptr as usize..(key_ptr + key_len) as usize)
93        .or_trap("lunatic::wasi::config_add_environment_variable")?;
94    let key = std::str::from_utf8(key_str)
95        .or_trap("lunatic::wasi::config_add_environment_variable")?
96        .to_string();
97    let value_str = memory
98        .data(&caller)
99        .get(value_ptr as usize..(value_ptr + value_len) as usize)
100        .or_trap("lunatic::wasi::config_add_environment_variable")?;
101    let value = std::str::from_utf8(value_str)
102        .or_trap("lunatic::wasi::config_add_environment_variable")?
103        .to_string();
104
105    caller
106        .data_mut()
107        .config_resources_mut()
108        .get_mut(config_id)
109        .or_trap("lunatic::wasi::config_set_max_memory: Config ID doesn't exist")?
110        .add_environment_variable(key, value);
111    Ok(())
112}
113
114// Adds command line argument to a configuration.
115//
116// Traps:
117// * If the config ID doesn't exist.
118// * If the argument string is not a valid utf8 string.
119// * If any of the memory slices falls outside the memory.
120fn add_command_line_argument<T>(
121    mut caller: Caller<T>,
122    config_id: u64,
123    argument_ptr: u32,
124    argument_len: u32,
125) -> Result<()>
126where
127    T: ProcessState,
128    T::Config: LunaticWasiConfigCtx,
129{
130    let memory = get_memory(&mut caller)?;
131    let argument_str = memory
132        .data(&caller)
133        .get(argument_ptr as usize..(argument_ptr + argument_len) as usize)
134        .or_trap("lunatic::wasi::add_command_line_argument")?;
135    let argument = std::str::from_utf8(argument_str)
136        .or_trap("lunatic::wasi::add_command_line_argument")?
137        .to_string();
138
139    caller
140        .data_mut()
141        .config_resources_mut()
142        .get_mut(config_id)
143        .or_trap("lunatic::wasi::add_command_line_argument: Config ID doesn't exist")?
144        .add_command_line_argument(argument);
145    Ok(())
146}
147
148// Mark a directory as preopened in the configuration.
149//
150// Traps:
151// * If the config ID doesn't exist.
152// * If the directory string is not a valid utf8 string.
153// * If any of the memory slices falls outside the memory.
154fn preopen_dir<T>(mut caller: Caller<T>, config_id: u64, dir_ptr: u32, dir_len: u32) -> Result<()>
155where
156    T: ProcessState,
157    T::Config: LunaticWasiConfigCtx,
158{
159    let memory = get_memory(&mut caller)?;
160    let dir_str = memory
161        .data(&caller)
162        .get(dir_ptr as usize..(dir_ptr + dir_len) as usize)
163        .or_trap("lunatic::wasi::preopen_dir")?;
164    let dir = std::str::from_utf8(dir_str)
165        .or_trap("lunatic::wasi::preopen_dir")?
166        .to_string();
167
168    caller
169        .data_mut()
170        .config_resources_mut()
171        .get_mut(config_id)
172        .or_trap("lunatic::wasi::preopen_dir: Config ID doesn't exist")?
173        .preopen_dir(dir);
174    Ok(())
175}
lunatic_wasi_api/lib.rs

lunatic_wasi_api/
lib.rs
