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 /// 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.
47 pub const DEFAULT_MAX_SIZE: u64 = 64 << 20;
48
49 pub 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
56 impl Default for WasiRandomCtx {
default() -> Self57 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
79 pub trait WasiRandomView: Send {
random(&mut self) -> &mut WasiRandomCtx80 fn random(&mut self) -> &mut WasiRandomCtx;
81 }
82
83 impl WasiRandomView for WasiRandomCtx {
random(&mut self) -> &mut WasiRandomCtx84 fn random(&mut self) -> &mut WasiRandomCtx {
85 self
86 }
87 }
88
89 /// Implement `insecure-random` using a deterministic cycle of bytes.
90 pub struct Deterministic {
91 cycle: std::iter::Cycle<std::vec::IntoIter<u8>>,
92 }
93
94 impl Deterministic {
new(bytes: Vec<u8>) -> Self95 pub fn new(bytes: Vec<u8>) -> Self {
96 Deterministic {
97 cycle: bytes.into_iter().cycle(),
98 }
99 }
100 }
101
102 impl RngCore for Deterministic {
next_u32(&mut self) -> u32103 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 }
next_u64(&mut self) -> u64110 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 }
fill_bytes(&mut self, buf: &mut [u8])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 }
try_fill_bytes(&mut self, buf: &mut [u8]) -> Result<(), cap_rand::Error>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)]
127 mod test {
128 use super::*;
129 #[test]
deterministic()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
thread_rng() -> Box<dyn RngCore + Send>140 pub 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 }
145