sidevm_env/
args_stack.rs

1use crate::{IntPtr, IntRet, OcallError, Result, VmMemory};
2use log::Level;
3
4const OCALL_N_ARGS: usize = 4;
5
6pub(crate) struct StackedArgs<Args> {
7    args: Args,
8}
9
10impl StackedArgs<()> {
11    pub(crate) const fn empty() -> Self {
12        StackedArgs { args: () }
13    }
14}
15
16impl<A: Nargs> StackedArgs<A> {
17    pub(crate) fn load(mut raw: &[IntPtr]) -> Option<Self> {
18        Some(check_args_length(StackedArgs {
19            args: Nargs::load(&mut raw)?,
20        }))
21    }
22
23    pub(crate) fn dump(self) -> [IntPtr; OCALL_N_ARGS] {
24        let mut ret = [Default::default(); OCALL_N_ARGS];
25        let data = check_args_length(self).args.dump();
26        ret[..data.len()].copy_from_slice(&data);
27        ret
28    }
29}
30
31impl<A, B> StackedArgs<(A, B)> {
32    fn pop(self) -> (A, StackedArgs<B>) {
33        let (a, args) = self.args;
34        (a, StackedArgs { args })
35    }
36}
37
38impl<B> StackedArgs<B> {
39    fn push<A>(self, arg: A) -> StackedArgs<(A, B)> {
40        StackedArgs {
41            args: (arg, self.args),
42        }
43    }
44}
45
46pub(crate) trait Nargs {
47    const N_ARGS: usize;
48    fn load(buf: &mut &[IntPtr]) -> Option<Self>
49    where
50        Self: Sized;
51
52    // Since #![feature(generic_const_exprs)] is not yet stable, we use OCALL_N_ARGS instead of
53    // Self::N_ARGS
54    fn dump(self) -> [IntPtr; OCALL_N_ARGS];
55}
56
57impl Nargs for () {
58    const N_ARGS: usize = 0;
59    fn load(_buf: &mut &[IntPtr]) -> Option<Self> {
60        Some(())
61    }
62    fn dump(self) -> [IntPtr; OCALL_N_ARGS] {
63        Default::default()
64    }
65}
66
67impl Nargs for IntPtr {
68    const N_ARGS: usize = 1;
69    fn load(buf: &mut &[IntPtr]) -> Option<Self> {
70        let me = *buf.get(0)?;
71        *buf = &buf[1..];
72        Some(me)
73    }
74
75    fn dump(self) -> [IntPtr; OCALL_N_ARGS] {
76        let mut ret = [0; OCALL_N_ARGS];
77        ret[0] = self;
78        ret
79    }
80}
81
82impl<A, B> Nargs for (A, B)
83where
84    A: Nargs,
85    B: Nargs,
86{
87    const N_ARGS: usize = A::N_ARGS + B::N_ARGS;
88
89    fn load(buf: &mut &[IntPtr]) -> Option<Self> {
90        let b = B::load(buf)?;
91        let a = A::load(buf)?;
92        Some((a, b))
93    }
94
95    fn dump(self) -> [IntPtr; OCALL_N_ARGS] {
96        let (a, b) = self;
97        let mut buf = [IntPtr::default(); OCALL_N_ARGS];
98        buf[0..B::N_ARGS].copy_from_slice(&b.dump()[0..B::N_ARGS]);
99        buf[B::N_ARGS..Self::N_ARGS].copy_from_slice(&a.dump()[..A::N_ARGS]);
100        buf
101    }
102}
103
104// Since the const evaluation of Rust is not powerful enough yet, we use this trick to statically
105// check the argument types encode output do not exceed the maximum number of arguments.
106pub(crate) trait NotTooManyArgs {
107    const TOO_MANY_ARGUMENTS: ();
108}
109impl<T: Nargs> NotTooManyArgs for T {
110    const TOO_MANY_ARGUMENTS: () = [()][(Self::N_ARGS > OCALL_N_ARGS) as usize];
111}
112
113pub(crate) fn check_args_length<T: Nargs + NotTooManyArgs>(v: StackedArgs<T>) -> StackedArgs<T> {
114    let _ = T::TOO_MANY_ARGUMENTS;
115    v
116}
117
118pub(crate) trait I32Convertible {
119    fn to_i32(&self) -> i32;
120    fn from_i32(i: i32) -> Result<Self>
121    where
122        Self: Sized;
123}
124
125pub(crate) trait ArgEncode<A> {
126    type Encoded;
127
128    fn encode_arg(self, stack: StackedArgs<A>) -> StackedArgs<(Self::Encoded, A)>;
129}
130
131pub(crate) trait ArgDecode<'a, A> {
132    type Encoded;
133    fn decode_arg(
134        stack: StackedArgs<(Self::Encoded, A)>,
135        vm: &'a impl VmMemory,
136    ) -> Result<(Self, StackedArgs<A>)>
137    where
138        Self: Sized;
139}
140
141/// Trait for types that can be encoded to a return value of a ocall.
142pub trait RetEncode {
143    /// Encode the ocall return value into a IntRet
144    fn encode_ret(self) -> IntRet;
145}
146
147pub(crate) trait RetDecode {
148    fn decode_ret(encoded: IntRet) -> Self
149    where
150        Self: Sized;
151}
152
153impl<A> ArgEncode<A> for &[u8] {
154    type Encoded = (IntPtr, IntPtr);
155
156    fn encode_arg(self, stack: StackedArgs<A>) -> StackedArgs<(Self::Encoded, A)> {
157        let ptr = self.as_ptr() as IntPtr;
158        let len = self.len() as IntPtr;
159        stack.push((len, ptr))
160    }
161}
162
163impl<'a, A> ArgDecode<'a, A> for &'a [u8] {
164    type Encoded = (IntPtr, IntPtr);
165
166    fn decode_arg(
167        stack: StackedArgs<(Self::Encoded, A)>,
168        vm: &'a impl VmMemory,
169    ) -> Result<(Self, StackedArgs<A>)>
170    where
171        Self: Sized,
172    {
173        let ((len, ptr), stack) = stack.pop();
174        Ok((vm.slice_from_vm(ptr, len)?, stack))
175    }
176}
177
178impl<A> ArgEncode<A> for &str {
179    type Encoded = (IntPtr, IntPtr);
180
181    fn encode_arg(self, stack: StackedArgs<A>) -> StackedArgs<(Self::Encoded, A)> {
182        let bytes = self.as_bytes();
183        bytes.encode_arg(stack)
184    }
185}
186
187impl<'a, A> ArgDecode<'a, A> for &'a str {
188    type Encoded = (IntPtr, IntPtr);
189
190    fn decode_arg(
191        stack: StackedArgs<(Self::Encoded, A)>,
192        vm: &'a impl VmMemory,
193    ) -> Result<(Self, StackedArgs<A>)>
194    where
195        Self: Sized,
196    {
197        let (bytes, stack): (&[u8], _) = ArgDecode::decode_arg(stack, vm)?;
198        Ok((
199            core::str::from_utf8(bytes).or(Err(OcallError::InvalidEncoding))?,
200            stack,
201        ))
202    }
203}
204
205impl<A> ArgEncode<A> for &mut [u8] {
206    type Encoded = (IntPtr, IntPtr);
207
208    fn encode_arg(self, stack: StackedArgs<A>) -> StackedArgs<(Self::Encoded, A)> {
209        let ptr = self.as_mut_ptr() as IntPtr;
210        let len = self.len() as IntPtr;
211        stack.push((len, ptr))
212    }
213}
214
215impl<'a, A> ArgDecode<'a, A> for &'a mut [u8] {
216    type Encoded = (IntPtr, IntPtr);
217
218    fn decode_arg(
219        stack: StackedArgs<(Self::Encoded, A)>,
220        vm: &'a impl VmMemory,
221    ) -> Result<(Self, StackedArgs<A>)>
222    where
223        Self: Sized,
224    {
225        let ((len, ptr), stack) = stack.pop();
226        Ok((vm.slice_from_vm_mut(ptr, len)?, stack))
227    }
228}
229
230impl<B> StackedArgs<B> {
231    pub(crate) fn push_arg<Arg: ArgEncode<B>>(self, v: Arg) -> StackedArgs<(Arg::Encoded, B)> {
232        v.encode_arg(self)
233    }
234}
235
236impl<A, B> StackedArgs<(A, B)> {
237    pub(crate) fn pop_arg<'a, Arg: ArgDecode<'a, B, Encoded = A>>(
238        self,
239        vm: &'a impl VmMemory,
240    ) -> Result<(Arg, StackedArgs<B>)> {
241        Arg::decode_arg(self, vm)
242    }
243}
244
245macro_rules! impl_codec_i {
246    ($typ: ty) => {
247        impl I32Convertible for $typ {
248            fn to_i32(&self) -> i32 {
249                *self as i32
250            }
251            fn from_i32(i: i32) -> Result<Self> {
252                if i > <$typ>::MAX as i32 || i < (-<$typ>::MAX - 1) as i32 {
253                    Err(OcallError::InvalidEncoding)
254                } else {
255                    Ok(i as Self)
256                }
257            }
258        }
259    };
260}
261impl_codec_i!(i8);
262impl_codec_i!(i16);
263
264macro_rules! impl_codec_u {
265    ($typ: ty) => {
266        impl I32Convertible for $typ {
267            fn to_i32(&self) -> i32 {
268                *self as i32
269            }
270            fn from_i32(i: i32) -> Result<Self> {
271                if i as u32 > <$typ>::MAX as u32 {
272                    Err(OcallError::InvalidEncoding)
273                } else {
274                    Ok(i as Self)
275                }
276            }
277        }
278    };
279}
280impl_codec_u!(u8);
281impl_codec_u!(u16);
282
283macro_rules! impl_codec {
284    ($typ: ty) => {
285        impl I32Convertible for $typ {
286            fn to_i32(&self) -> i32 {
287                *self as i32
288            }
289            fn from_i32(i: i32) -> Result<Self> {
290                Ok(i as Self)
291            }
292        }
293    };
294}
295impl_codec!(i32);
296impl_codec!(u32);
297
298macro_rules! impl_codec64 {
299    ($typ: ty) => {
300        impl<R> ArgEncode<R> for $typ {
301            type Encoded = (IntPtr, IntPtr);
302
303            fn encode_arg(self, stack: StackedArgs<R>) -> StackedArgs<(Self::Encoded, R)> {
304                let low = (self & 0xffffffff) as IntPtr;
305                let high = ((self >> 32) & 0xffffffff) as IntPtr;
306                stack.push((low, high))
307            }
308        }
309
310        impl<'a, R> ArgDecode<'a, R> for $typ {
311            type Encoded = (IntPtr, IntPtr);
312
313            fn decode_arg(
314                stack: StackedArgs<(Self::Encoded, R)>,
315                _vm: &'a impl VmMemory,
316            ) -> Result<(Self, StackedArgs<R>)>
317            where
318                Self: Sized,
319            {
320                let ((low, high), stack) = stack.pop();
321                let high = ((high as Self) << 32);
322                let v = high & (low as Self);
323                Ok((v, stack))
324            }
325        }
326    };
327}
328
329impl_codec64!(i64);
330impl_codec64!(u64);
331
332impl<R, I: I32Convertible> ArgEncode<R> for I {
333    type Encoded = IntPtr;
334
335    fn encode_arg(self, stack: StackedArgs<R>) -> StackedArgs<(Self::Encoded, R)> {
336        stack.push(self.to_i32() as _)
337    }
338}
339
340impl<'a, R, I: I32Convertible> ArgDecode<'a, R> for I {
341    type Encoded = IntPtr;
342
343    fn decode_arg(
344        stack: StackedArgs<(Self::Encoded, R)>,
345        _vm: &'a impl VmMemory,
346    ) -> Result<(Self, StackedArgs<R>)>
347    where
348        Self: Sized,
349    {
350        let (v, stack) = stack.pop();
351        Ok((I::from_i32(v as _)?, stack))
352    }
353}
354
355impl I32Convertible for bool {
356    fn to_i32(&self) -> i32 {
357        *self as i32
358    }
359    fn from_i32(i: i32) -> Result<Self> {
360        match i {
361            0 => Ok(false),
362            1 => Ok(true),
363            _ => Err(OcallError::InvalidEncoding),
364        }
365    }
366}
367
368impl I32Convertible for OcallError {
369    fn to_i32(&self) -> i32 {
370        *self as u8 as i32
371    }
372    fn from_i32(i: i32) -> Result<Self> {
373        let code = u8::from_i32(i)?;
374        OcallError::try_from(code).or(Err(OcallError::InvalidEncoding))
375    }
376}
377
378impl I32Convertible for () {
379    fn to_i32(&self) -> i32 {
380        0
381    }
382    fn from_i32(i: i32) -> Result<()> {
383        if i == 0 {
384            Ok(())
385        } else {
386            Err(OcallError::InvalidEncoding)
387        }
388    }
389}
390
391impl I32Convertible for Level {
392    fn to_i32(&self) -> i32 {
393        match self {
394            Level::Error => 1,
395            Level::Warn => 2,
396            Level::Info => 3,
397            Level::Debug => 4,
398            Level::Trace => 5,
399        }
400    }
401
402    fn from_i32(i: i32) -> Result<Self> {
403        match i {
404            1 => Ok(Level::Error),
405            2 => Ok(Level::Warn),
406            3 => Ok(Level::Info),
407            4 => Ok(Level::Debug),
408            5 => Ok(Level::Trace),
409            _ => Err(OcallError::InvalidEncoding),
410        }
411    }
412}
413
414impl<A, B> RetEncode for Result<A, B>
415where
416    A: I32Convertible,
417    B: I32Convertible,
418{
419    fn encode_ret(self) -> IntRet {
420        let (tp, val) = match self {
421            Ok(v) => (0, v.to_i32()),
422            Err(err) => (1, err.to_i32()),
423        };
424        ((tp as u32 as i64) << 32) | (val as u32 as i64)
425    }
426}
427
428impl<A, B> RetDecode for Result<A, B>
429where
430    A: I32Convertible,
431    B: I32Convertible,
432{
433    fn decode_ret(encoded: IntRet) -> Self {
434        let tp = ((encoded >> 32) & 0xffffffff) as i32;
435        let val = (encoded & 0xffffffff) as i32;
436        if tp == 0 {
437            Ok(A::from_i32(val).expect("Invalid ocall return"))
438        } else {
439            Err(B::from_i32(val).expect("Invalid ocall return"))
440        }
441    }
442}