xref: /wasmtime-44.0.1/crates/wasi/src/random.rs (revision bda02c19)
1 use cap_rand::{Rng as _, RngCore, SeedableRng as _};
2 use 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 /// ```
38 pub struct WasiRandom;
39 
40 impl HasData for WasiRandom {
41     type Data<'a> = &'a mut WasiRandomCtx;
42 }
43 
44 pub struct WasiRandomCtx {
45     pub(crate) random: Box<dyn RngCore + Send>,
46     pub(crate) insecure_random: Box<dyn RngCore + Send>,
47     pub(crate) insecure_random_seed: u128,
48 }
49 
50 impl Default for WasiRandomCtx {
51     fn default() -> Self {
52         // For the insecure random API, use `SmallRng`, which is fast. It's
53         // also insecure, but that's the deal here.
54         let insecure_random = Box::new(
55             cap_rand::rngs::SmallRng::from_rng(cap_rand::thread_rng(cap_rand::ambient_authority()))
56                 .unwrap(),
57         );
58         // For the insecure random seed, use a `u128` generated from
59         // `thread_rng()`, so that it's not guessable from the insecure_random
60         // API.
61         let insecure_random_seed =
62             cap_rand::thread_rng(cap_rand::ambient_authority()).r#gen::<u128>();
63         Self {
64             random: thread_rng(),
65             insecure_random,
66             insecure_random_seed,
67         }
68     }
69 }
70 
71 pub trait WasiRandomView: Send {
72     fn random(&mut self) -> &mut WasiRandomCtx;
73 }
74 
75 impl WasiRandomView for WasiRandomCtx {
76     fn random(&mut self) -> &mut WasiRandomCtx {
77         self
78     }
79 }
80 
81 /// Implement `insecure-random` using a deterministic cycle of bytes.
82 pub struct Deterministic {
83     cycle: std::iter::Cycle<std::vec::IntoIter<u8>>,
84 }
85 
86 impl Deterministic {
87     pub fn new(bytes: Vec<u8>) -> Self {
88         Deterministic {
89             cycle: bytes.into_iter().cycle(),
90         }
91     }
92 }
93 
94 impl RngCore for Deterministic {
95     fn next_u32(&mut self) -> u32 {
96         let b0 = self.cycle.next().expect("infinite sequence");
97         let b1 = self.cycle.next().expect("infinite sequence");
98         let b2 = self.cycle.next().expect("infinite sequence");
99         let b3 = self.cycle.next().expect("infinite sequence");
100         ((b0 as u32) << 24) + ((b1 as u32) << 16) + ((b2 as u32) << 8) + (b3 as u32)
101     }
102     fn next_u64(&mut self) -> u64 {
103         let w0 = self.next_u32();
104         let w1 = self.next_u32();
105         ((w0 as u64) << 32) + (w1 as u64)
106     }
107     fn fill_bytes(&mut self, buf: &mut [u8]) {
108         for b in buf.iter_mut() {
109             *b = self.cycle.next().expect("infinite sequence");
110         }
111     }
112     fn try_fill_bytes(&mut self, buf: &mut [u8]) -> Result<(), cap_rand::Error> {
113         self.fill_bytes(buf);
114         Ok(())
115     }
116 }
117 
118 #[cfg(test)]
119 mod test {
120     use super::*;
121     #[test]
122     fn deterministic() {
123         let mut det = Deterministic::new(vec![1, 2, 3, 4]);
124         let mut buf = vec![0; 1024];
125         det.try_fill_bytes(&mut buf).expect("get randomness");
126         for (ix, b) in buf.iter().enumerate() {
127             assert_eq!(*b, (ix % 4) as u8 + 1)
128         }
129     }
130 }
131 
132 pub fn thread_rng() -> Box<dyn RngCore + Send> {
133     use cap_rand::{Rng, SeedableRng};
134     let mut rng = cap_rand::thread_rng(cap_rand::ambient_authority());
135     Box::new(cap_rand::rngs::StdRng::from_seed(rng.r#gen()))
136 }
137