Skip to main content

wasmtime_wasi/
random.rs

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