wasi_virt_layer/wasi/
args.rs

1use crate::__private::wasip1::*;
2#[cfg(target_os = "wasi")]
3use const_for::const_for;
4use const_struct::*;
5
6use crate::memory::WasmAccess;
7
8/// @embedded or @dynamic
9/// Whether to import JavaScript runtime args from vfs,
10/// args is automatically imported even if you are not using it,
11/// so that you can block it
12///
13/// @embedded or @dynamic
14/// Whether to use embedded or dynamic args.
15/// @embedded if using embedded args.
16/// @dynamic if using dynamic args.
17/// @embedded is faster and small than @dynamic.
18///
19/// ```rust
20/// // @embedded
21/// import_wasm!(test_wasm);
22///
23/// use const_struct::*;
24/// use wasi_virt_layer::prelude::*;
25/// #[const_struct]
26/// const VIRTUAL_ARGS: VirtualArgsEmbeddedState = VirtualArgsEmbeddedState {
27///     args: &["command", "arg1", "arg2"],
28/// };
29/// plug_args!(@embedded, VirtualArgsTy, test_wasm);
30/// ```
31///
32/// ```rust
33/// // @dynamic
34/// import_wasm!(test_wasm);
35///
36/// use std::sync::{LazyLock, Mutex};
37/// use wasi_virt_layer::prelude::*;
38///
39/// struct VirtualArgsState {
40///    args: Vec<String>,
41/// }
42/// impl<'a> VirtualArgs<'a> for VirtualArgsState {
43///    type Str = String;
44///
45///   fn get_args(&mut self) -> &[Self::Str] {
46///       &self.args
47///   }
48/// }
49/// static VIRTUAL_ARGS: LazyLock<Mutex<VirtualArgsState>> = LazyLock::new(|| {
50///    let mut args = Vec::<String>::new();
51///   args.push("command".into());
52///   args.push("arg1".into());
53///   Mutex::new(VirtualArgsState { args })
54/// });
55/// plug_args!(@dynamic, &mut VIRTUAL_ARGS.lock().unwrap(), test_wasm);
56/// ```
57/// Plugs the command-line arguments ecosystem.
58#[macro_export]
59macro_rules! plug_args {
60    (@embedded, $ty:ty, $($wasm:ident),* $(,)?) => {
61        $crate::__as_t!(@through, $($wasm),* => $crate::plug_args, @inner, @embedded, $ty);
62    };
63
64    (@dynamic, $state:expr, $($wasm:ident),* $(,)?) => {
65        $crate::__as_t!(@through, $($wasm),* => $crate::plug_args, @inner, @dynamic, $state);
66    };
67
68    (@inner, @embedded, $ty:ty, $($wasm:ident),*) => {
69        const _: () = {
70            type __TYPE = $ty;
71        };
72
73        $crate::__private::paste::paste! {
74            $(
75                #[unsafe(no_mangle)]
76                #[cfg(target_os = "wasi")]
77                pub unsafe extern "C" fn [<__wasip1_vfs_ $wasm _args_sizes_get>](
78                    args_count: *mut $crate::__private::wasip1::Size,
79                    args_buf_size: *mut $crate::__private::wasip1::Size,
80                ) -> $crate::__private::wasip1::Errno {
81                    $crate::__as_t!(@as_t, $wasm);
82                    $crate::__private::inner::args::args_sizes_get_embedded_inner::<$ty, T>(args_count, args_buf_size)
83                }
84
85                #[cfg(target_os = "wasi")]
86                #[unsafe(no_mangle)]
87                pub unsafe extern "C" fn [<__wasip1_vfs_ $wasm _args_get>](
88                    args: *mut *const u8,
89                    args_buf: *mut u8,
90                ) -> $crate::__private::wasip1::Errno {
91                    $crate::__as_t!(@as_t, $wasm);
92                    $crate::__private::inner::args::args_get_embedded_inner::<$ty, T>(args, args_buf)
93                }
94            )*
95        }
96    };
97
98    (@dynamic, $state:expr, $($wasm:ident),* $(,)?) => {
99        $crate::__as_t!(@through, $($wasm),* => $crate::plug_args, @inner, @dynamic, $state);
100    };
101
102    (@inner, @embedded, $ty:ty, $($wasm:ident),*) => {
103        const _: () = {
104            type __TYPE = $ty;
105        };
106
107        $crate::__private::paste::paste! {
108            $(
109                #[unsafe(no_mangle)]
110                #[cfg(target_os = "wasi")]
111                pub unsafe extern "C" fn [<__wasip1_vfs_ $wasm _args_sizes_get>](
112                    args_count: *mut $crate::__private::wasip1::Size,
113                    args_buf_size: *mut $crate::__private::wasip1::Size,
114                ) -> $crate::__private::wasip1::Errno {
115                    $crate::__as_t!(@as_t, $wasm);
116                    $crate::__private::inner::args::args_sizes_get_const_inner::<$ty, T>(args_count, args_buf_size)
117                }
118
119                #[cfg(target_os = "wasi")]
120                #[unsafe(no_mangle)]
121                pub unsafe extern "C" fn [<__wasip1_vfs_ $wasm _args_get>](
122                    args: *mut *const u8,
123                    args_buf: *mut u8,
124                ) -> $crate::__private::wasip1::Errno {
125                    $crate::__as_t!(@as_t, $wasm);
126                    $crate::__private::inner::args::args_get_const_inner::<$ty, T>(args, args_buf)
127                }
128            )*
129        }
130    };
131
132    (@inner, @dynamic, $state:expr, $($wasm:ident),*) => {
133        $crate::__private::paste::paste! {
134            $(
135                #[cfg(target_os = "wasi")]
136                #[unsafe(no_mangle)]
137                pub unsafe extern "C" fn [<__wasip1_vfs_ $wasm _args_sizes_get>](
138                    args_count: *mut $crate::__private::wasip1::Size,
139                    args_buf_size: *mut $crate::__private::wasip1::Size,
140                ) -> $crate::__private::wasip1::Errno {
141                    $crate::__as_t!(@as_t, $wasm);
142                    let state = $state;
143                    $crate::__private::inner::args::args_sizes_get_inner::<T>(state, args_count, args_buf_size)
144                }
145
146                #[cfg(target_os = "wasi")]
147                #[unsafe(no_mangle)]
148                pub unsafe extern "C" fn [<__wasip1_vfs_ $wasm _args_get>](
149                    args: *mut *const u8,
150                    args_buf: *mut u8,
151                ) -> $crate::__private::wasip1::Errno {
152                    $crate::__as_t!(@as_t, $wasm);
153                    let state = $state;
154                    $crate::__private::inner::args::args_get_inner::<T>(state, args, args_buf)
155                }
156            )*
157        }
158    };
159}
160
161/// A structure holding constant command-line arguments.
162#[const_struct]
163pub struct VirtualArgsEmbeddedState {
164    /// The command-line arguments.
165    pub args: &'static [&'static str],
166}
167
168/// Inner function for retrieving constant command-line argument sizes.
169#[inline]
170#[cfg(target_os = "wasi")]
171pub fn args_sizes_get_embedded_inner<
172    T: PrimitiveTraits<DATATYPE = VirtualArgsEmbeddedState>,
173    Wasm: WasmAccess,
174>(
175    args_count: *mut Size,
176    args_buf_size: *mut Size,
177) -> Errno {
178    const fn inner<T: PrimitiveTraits<DATATYPE = VirtualArgsEmbeddedState>>() -> (Size, Size) {
179        let mut size = 0;
180        let mut count = 0;
181        const_for!(i in 0..T::__DATA.args.len() => {
182            let len = T::__DATA.args[i].len() + 1; // +1 for null terminator
183            size += len;
184            count += 1;
185        });
186
187        (size, count)
188    }
189
190    Wasm::store_le(args_buf_size, inner::<T>().0);
191    Wasm::store_le(args_count, inner::<T>().1);
192    ERRNO_SUCCESS
193}
194
195/// Inner function for retrieving constant command-line arguments.
196#[inline]
197#[cfg(target_os = "wasi")]
198pub fn args_get_embedded_inner<
199    T: PrimitiveTraits<DATATYPE = VirtualArgsEmbeddedState>,
200    Wasm: WasmAccess,
201>(
202    args: *mut *const u8,
203    args_buf: *mut u8,
204) -> Errno {
205    let mut args = args;
206    let mut args_buf = args_buf;
207
208    const_for!(i in 0..T::__DATA.args.len() => {
209        Wasm::store_le(args, args_buf as *const u8);
210
211        Wasm::memcpy(args_buf, T::__DATA.args[i].as_bytes());
212        Wasm::store_le(unsafe { args_buf.add(T::__DATA.args[i].len()) }, 0u8);
213
214        args = unsafe { args.add(1) };
215        args_buf = unsafe { args_buf.add(T::__DATA.args[i].len() + 1) };
216    });
217
218    ERRNO_SUCCESS
219}
220
221/// Inner function for retrieving constant command-line argument sizes.
222#[inline]
223#[cfg(target_os = "wasi")]
224pub fn args_sizes_get_const_inner<
225    T: PrimitiveTraits<DATATYPE = VirtualArgsEmbeddedState>,
226    Wasm: WasmAccess,
227>(
228    args_count: *mut Size,
229    args_buf_size: *mut Size,
230) -> Errno {
231    const fn inner<T: PrimitiveTraits<DATATYPE = VirtualArgsEmbeddedState>>() -> (Size, Size) {
232        let mut size = 0;
233        let mut count = 0;
234        const_for!(i in 0..T::__DATA.args.len() => {
235            let len = T::__DATA.args[i].len() + 1; // +1 for null terminator
236            size += len;
237            count += 1;
238        });
239
240        (size, count)
241    }
242
243    Wasm::store_le(args_buf_size, inner::<T>().0);
244    Wasm::store_le(args_count, inner::<T>().1);
245    ERRNO_SUCCESS
246}
247
248/// Inner function for retrieving constant command-line arguments.
249#[inline]
250#[cfg(target_os = "wasi")]
251pub fn args_get_const_inner<
252    T: PrimitiveTraits<DATATYPE = VirtualArgsEmbeddedState>,
253    Wasm: WasmAccess,
254>(
255    args: *mut *const u8,
256    args_buf: *mut u8,
257) -> Errno {
258    let mut args = args;
259    let mut args_buf = args_buf;
260
261    const_for!(i in 0..T::__DATA.args.len() => {
262        Wasm::store_le(args, args_buf as *const u8);
263
264        Wasm::memcpy(args_buf, T::__DATA.args[i].as_bytes());
265        Wasm::store_le(unsafe { args_buf.add(T::__DATA.args[i].len()) }, 0u8);
266
267        args = unsafe { args.add(1) };
268        args_buf = unsafe { args_buf.add(T::__DATA.args[i].len() + 1) };
269    });
270
271    ERRNO_SUCCESS
272}
273
274/// Trait for a virtual command-line arguments implementation.
275pub trait VirtualArgs<'a> {
276    /// The type used for argument strings.
277    type Str: AsRef<str>;
278
279    /// Returns the command-line arguments.
280    fn get_args(&'a mut self) -> &'a [Self::Str];
281
282    /// Returns the sizes of the command-line arguments.
283    fn args_sizes_get(&'a mut self) -> (Size, Size) {
284        let args = self.get_args();
285        let mut size = 0;
286        let mut count = 0;
287        for arg in args {
288            let len = arg.as_ref().len() + 1; // +1 for null terminator
289            size += len;
290            count += 1;
291        }
292
293        (size, count)
294    }
295
296    /// Copies the command-line arguments into the provided buffers.
297    fn args_get<Wasm: WasmAccess>(&'a mut self, args: *mut *const u8, args_buf: *mut u8) -> Errno {
298        let mut args = args;
299        let mut args_buf = args_buf;
300
301        for arg in self.get_args() {
302            Wasm::store_le(args, args_buf as *const u8);
303
304            Wasm::memcpy(args_buf, arg.as_ref().as_bytes());
305            Wasm::store_le(unsafe { args_buf.add(arg.as_ref().len()) as *mut u8 }, 0u8);
306
307            args = unsafe { args.add(1) };
308            args_buf = unsafe { args_buf.add(arg.as_ref().len() + 1) };
309        }
310
311        ERRNO_SUCCESS
312    }
313}
314
315impl<'a, T: core::ops::DerefMut<Target = U>, U: VirtualArgs<'a> + 'a> VirtualArgs<'a> for T {
316    type Str = U::Str;
317
318    fn get_args(&'a mut self) -> &'a [Self::Str] {
319        self.deref_mut().get_args()
320    }
321
322    fn args_sizes_get(&'a mut self) -> (Size, Size) {
323        self.deref_mut().args_sizes_get()
324    }
325}
326
327/// Inner function for retrieving dynamic command-line argument sizes.
328#[cfg(target_os = "wasi")]
329pub fn args_sizes_get_inner<'a, Wasm: WasmAccess>(
330    state: &'a mut impl VirtualArgs<'a>,
331    args_count: *mut Size,
332    args_buf_size: *mut Size,
333) -> Errno {
334    let (size, count) = state.args_sizes_get();
335
336    Wasm::store_le(args_buf_size, size);
337    Wasm::store_le(args_count, count);
338
339    ERRNO_SUCCESS
340}
341
342/// Inner function for retrieving dynamic command-line arguments.
343#[inline]
344#[cfg(target_os = "wasi")]
345pub fn args_get_inner<'a, Wasm: WasmAccess>(
346    state: &'a mut impl VirtualArgs<'a>,
347    args: *mut *const u8,
348    args_buf: *mut u8,
349) -> Errno {
350    state.args_get::<Wasm>(args, args_buf)
351}
wasi_virt_layer/wasi/args.rs

wasi_virt_layer/wasi/
args.rs
