wasmtime_wasi/
random.rs

1use rand::{Rng, SeedableRng as _, TryRng};
2use std::convert::Infallible;
3use wasmtime::component::HasData;
4
5/// A helper struct which implements [`HasData`] for the `wasi:random` APIs.
6///
7/// This can be useful when directly calling `add_to_linker` functions directly,
8/// such as [`wasmtime_wasi::p2::bindings::random::random::add_to_linker`] as
9/// the `D` type parameter. See [`HasData`] for more information about the type
10/// parameter's purpose.
11///
12/// When using this type you can skip the [`WasiRandomView`] trait, for
13/// example.
14///
15/// [`wasmtime_wasi::p2::bindings::random::random::add_to_linker`]: crate::p2::bindings::random::random::add_to_linker
16///
17/// # Examples
18///
19/// ```
20/// use wasmtime::component::Linker;
21/// use wasmtime::{Engine, Result};
22/// use wasmtime_wasi::random::*;
23///
24/// struct MyStoreState {
25///     random: WasiRandomCtx,
26/// }
27///
28/// fn main() -> Result<()> {
29///     let engine = Engine::default();
30///     let mut linker = Linker::new(&engine);
31///
32///     wasmtime_wasi::p2::bindings::random::random::add_to_linker::<MyStoreState, WasiRandom>(
33///         &mut linker,
34///         |state| &mut state.random,
35///     )?;
36///     Ok(())
37/// }
38/// ```
39pub struct WasiRandom;
40
41impl HasData for WasiRandom {
42    type Data<'a> = &'a mut WasiRandomCtx;
43}
44
45/// Default largest length accepted by wasi 0.2 `get-random-bytes` and
46/// `get-insecure-random-bytes` methods. This constant must match docs in
47/// cli-flags crate.
48pub const DEFAULT_MAX_SIZE: u64 = 64 << 20;
49
50pub struct WasiRandomCtx {
51    pub(crate) random: Box<dyn Rng + Send>,
52    pub(crate) insecure_random: Box<dyn Rng + Send>,
53    pub(crate) insecure_random_seed: u128,
54    pub(crate) max_size: u64,
55}
56
57impl Default for WasiRandomCtx {
58    fn default() -> Self {
59        // For the insecure random API, use `SmallRng`, which is fast. It's
60        // also insecure, but that's the deal here.
61        let insecure_random = Box::new(rand::rngs::SmallRng::from_rng(&mut rand::rng()));
62        // For the insecure random seed, use a `u128` generated from
63        // `rand::random()`, so that it's not guessable from the
64        // insecure_random API.
65        let insecure_random_seed = rand::random::<u128>();
66        let max_size = DEFAULT_MAX_SIZE;
67        Self {
68            random: thread_rng(),
69            insecure_random,
70            insecure_random_seed,
71            max_size,
72        }
73    }
74}
75
76pub trait WasiRandomView: Send {
77    fn random(&mut self) -> &mut WasiRandomCtx;
78}
79
80impl WasiRandomView for WasiRandomCtx {
81    fn random(&mut self) -> &mut WasiRandomCtx {
82        self
83    }
84}
85
86/// Implement `insecure-random` using a deterministic cycle of bytes.
87pub struct Deterministic {
88    cycle: std::iter::Cycle<std::vec::IntoIter<u8>>,
89}
90
91impl Deterministic {
92    pub fn new(bytes: Vec<u8>) -> Self {
93        Deterministic {
94            cycle: bytes.into_iter().cycle(),
95        }
96    }
97}
98
99impl TryRng for Deterministic {
100    type Error = Infallible;
101    fn try_next_u32(&mut self) -> Result<u32, Infallible> {
102        let b0 = self.cycle.next().expect("infinite sequence");
103        let b1 = self.cycle.next().expect("infinite sequence");
104        let b2 = self.cycle.next().expect("infinite sequence");
105        let b3 = self.cycle.next().expect("infinite sequence");
106        Ok(((b0 as u32) << 24) + ((b1 as u32) << 16) + ((b2 as u32) << 8) + (b3 as u32))
107    }
108    fn try_next_u64(&mut self) -> Result<u64, Infallible> {
109        let w0 = self.next_u32();
110        let w1 = self.next_u32();
111        Ok(((w0 as u64) << 32) + (w1 as u64))
112    }
113    fn try_fill_bytes(&mut self, buf: &mut [u8]) -> Result<(), Infallible> {
114        for b in buf.iter_mut() {
115            *b = self.cycle.next().expect("infinite sequence");
116        }
117        Ok(())
118    }
119}
120
121#[cfg(test)]
122mod test {
123    use super::*;
124    #[test]
125    fn deterministic() {
126        let mut det = Deterministic::new(vec![1, 2, 3, 4]);
127        let mut buf = vec![0; 1024];
128        det.try_fill_bytes(&mut buf).expect("get randomness");
129        for (ix, b) in buf.iter().enumerate() {
130            assert_eq!(*b, (ix % 4) as u8 + 1)
131        }
132    }
133}
134
135pub fn thread_rng() -> Box<dyn Rng + Send> {
136    let mut rng = rand::rng();
137    Box::new(rand::rngs::StdRng::from_rng(&mut rng))
138}
wasmtime_wasi/random.rs

wasmtime_wasi/
random.rs
