wasmtime_wasi/random.rs
1use crate::{NamedId, WasiCtxNamedView};
2use rand::{Rng, SeedableRng as _, TryRng};
3use std::convert::Infallible;
4use std::marker;
5use wasmtime::component::HasData;
6
7/// A helper struct which implements [`HasData`] for the `wasi:random` APIs.
8///
9/// This can be useful when directly calling `add_to_linker` functions directly,
10/// such as [`wasmtime_wasi::p2::bindings::random::random::add_to_linker`] as
11/// the `D` type parameter. See [`HasData`] for more information about the type
12/// parameter's purpose.
13///
14/// When using this type you can skip the [`WasiRandomView`] trait, for
15/// example.
16///
17/// [`wasmtime_wasi::p2::bindings::random::random::add_to_linker`]: crate::p2::bindings::random::random::add_to_linker
18///
19/// # Examples
20///
21/// ```
22/// use wasmtime::component::Linker;
23/// use wasmtime::{Engine, Result};
24/// use wasmtime_wasi::random::*;
25///
26/// struct MyStoreState {
27/// random: WasiRandomCtx,
28/// }
29///
30/// fn main() -> Result<()> {
31/// let engine = Engine::default();
32/// let mut linker = Linker::new(&engine);
33///
34/// wasmtime_wasi::p2::bindings::random::random::add_to_linker::<MyStoreState, WasiRandom>(
35/// &mut linker,
36/// |state| &mut state.random,
37/// )?;
38/// Ok(())
39/// }
40/// ```
41pub struct WasiRandom;
42
43impl HasData for WasiRandom {
44 type Data<'a> = &'a mut WasiRandomCtx;
45}
46
47/// Default largest length accepted by wasi 0.2 `get-random-bytes` and
48/// `get-insecure-random-bytes` methods. This constant must match docs in
49/// cli-flags crate.
50pub const DEFAULT_MAX_SIZE: u64 = 64 << 20;
51
52pub struct WasiRandomCtx {
53 pub(crate) random: Box<dyn Rng + Send>,
54 pub(crate) insecure_random: Box<dyn Rng + Send>,
55 pub(crate) insecure_random_seed: u128,
56 pub(crate) max_size: u64,
57}
58
59impl Default for WasiRandomCtx {
60 fn default() -> Self {
61 // For the insecure random API, use `SmallRng`, which is fast. It's
62 // also insecure, but that's the deal here.
63 let insecure_random = Box::new(rand::rngs::SmallRng::from_rng(&mut rand::rng()));
64 // For the insecure random seed, use a `u128` generated from
65 // `rand::random()`, so that it's not guessable from the
66 // insecure_random API.
67 let insecure_random_seed = rand::random::<u128>();
68 let max_size = DEFAULT_MAX_SIZE;
69 Self {
70 random: thread_rng(),
71 insecure_random,
72 insecure_random_seed,
73 max_size,
74 }
75 }
76}
77
78pub trait WasiRandomView: Send {
79 fn random(&mut self) -> &mut WasiRandomCtx;
80}
81
82impl WasiRandomView for WasiRandomCtx {
83 fn random(&mut self) -> &mut WasiRandomCtx {
84 self
85 }
86}
87
88/// Implement `insecure-random` using a deterministic cycle of bytes.
89pub struct Deterministic {
90 cycle: std::iter::Cycle<std::vec::IntoIter<u8>>,
91}
92
93impl Deterministic {
94 pub fn new(bytes: Vec<u8>) -> Self {
95 Deterministic {
96 cycle: bytes.into_iter().cycle(),
97 }
98 }
99}
100
101impl TryRng for Deterministic {
102 type Error = Infallible;
103 fn try_next_u32(&mut self) -> Result<u32, Infallible> {
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 Ok(((b0 as u32) << 24) + ((b1 as u32) << 16) + ((b2 as u32) << 8) + (b3 as u32))
109 }
110 fn try_next_u64(&mut self) -> Result<u64, Infallible> {
111 let w0 = self.next_u32();
112 let w1 = self.next_u32();
113 Ok(((w0 as u64) << 32) + (w1 as u64))
114 }
115 fn try_fill_bytes(&mut self, buf: &mut [u8]) -> Result<(), Infallible> {
116 for b in buf.iter_mut() {
117 *b = self.cycle.next().expect("infinite sequence");
118 }
119 Ok(())
120 }
121}
122
123#[cfg(test)]
124mod test {
125 use super::*;
126 #[test]
127 fn deterministic() {
128 let mut det = Deterministic::new(vec![1, 2, 3, 4]);
129 let mut buf = vec![0; 1024];
130 det.try_fill_bytes(&mut buf).expect("get randomness");
131 for (ix, b) in buf.iter().enumerate() {
132 assert_eq!(*b, (ix % 4) as u8 + 1)
133 }
134 }
135}
136
137pub fn thread_rng() -> Box<dyn Rng + Send> {
138 let mut rng = rand::rng();
139 Box::new(rand::rngs::StdRng::from_rng(&mut rng))
140}
141
142/// A helper struct which implements [`HasData`] for the `wasi:random` APIs
143/// when used in combination with named imports.
144///
145/// This structure is similar in purpose to [`WasiRandom`] and is used
146/// when using the [`named_imports`] module for `wasi:random`. This structure
147/// serves as the `D` type parameter for `add_to_linker` functions.
148///
149/// [`named_imports`]: crate::p3::bindings::named_imports::wasi::random
150///
151/// # Meaning of the `T` parameter
152///
153/// Here the `T` must be something that implements [`WasiRandomNamedView`]. The
154/// corresponding `Data` for this type is [`WasiCtxNamedView`] which internally
155/// will contain `&mut T`.
156///
157/// Effectively you're going to implement [`WasiRandomNamedView`] for something in
158/// your embedding, and that's the `T` you'll fill in here.
159///
160/// # Examples
161///
162/// ```
163/// use wasmtime::component::{Linker, Component};
164/// use wasmtime::{Engine, Result};
165/// use wasmtime_wasi::{NamedId, WasiCtxNamedView};
166/// use wasmtime_wasi::random::*;
167/// use wasmtime_wasi::p2::bindings::named_imports;
168/// use std::collections::HashMap;
169///
170/// struct MyStoreState {
171/// states: HashMap<NamedId, WasiRandomCtx>,
172/// }
173///
174/// fn main() -> Result<()> {
175/// let engine = Engine::default();
176/// let mut linker = Linker::new(&engine);
177/// let component = Component::new(&engine, "(component)")?;
178/// let mut name_map = HashMap::new();
179///
180/// named_imports::wasi::random::random::add_to_linker::<MyStoreState, WasiRandomNamed<MyStoreState>>(
181/// &mut linker,
182/// &component,
183/// |name| {
184/// let len = name_map.len();
185/// Ok(NamedId(*name_map.entry(name.to_string()).or_insert(len)))
186/// },
187/// |state| WasiCtxNamedView(state),
188/// )?;
189/// Ok(())
190/// }
191///
192/// impl WasiRandomNamedView for MyStoreState {
193/// fn random(&mut self, id: NamedId) -> &mut WasiRandomCtx {
194/// self.states.get_mut(&id).expect("state for id")
195/// }
196/// }
197/// ```
198pub struct WasiRandomNamed<T>(marker::PhantomData<fn() -> T>);
199
200impl<T> HasData for WasiRandomNamed<T>
201where
202 T: WasiRandomNamedView,
203{
204 type Data<'a> = WasiCtxNamedView<'a, T>;
205}
206
207/// A trait used to look up a specific `wasi:random` context for a named
208/// import.
209///
210/// This trait is used in conjunction with the [`named_imports`] bindings
211/// generated for all WASI interfaces. The purpose of this trait is for
212/// embedders to define how a [`NamedId`] maps to a particular `wasi:random`
213/// context, here returned as [`WasiRandomCtx`]. Embedders are responsible
214/// for assigning meaning to [`NamedId`] values themselves. These IDs are
215/// assigned when [`add_named_to_linker`] is called, for example, as the
216/// `lookup` argument to that function.
217///
218/// When using [`add_named_to_linker`] it's sufficient to implement this trait
219/// for the `T` in `Store<T>`. You can also instead implement the
220/// [`WasiNamedView`] trait for `T` which implies an implementation of this
221/// trait.
222///
223/// When using `add_to_linker` in the generated `bindings::named_imports`
224/// module then values implementing this live within the `T` of `Store<T>`, and
225/// be temporarily referenced in [`WasiCtxNamedView`] where internally that'll
226/// hold `WasiCtxNamedView(&mut your_type)`.
227///
228/// [`named_imports`]: crate::p3::bindings::named_imports
229/// [`add_named_to_linker`]: crate::p3::random::add_named_to_linker
230/// [`WasiNamedView`]: crate::WasiNamedView
231///
232/// # Examples
233///
234/// ```
235/// use wasmtime::component::{Linker, Component};
236/// use wasmtime::{Engine, Result};
237/// use wasmtime_wasi::{NamedId, WasiCtxNamedView};
238/// use wasmtime_wasi::random::*;
239/// use std::collections::HashMap;
240///
241/// struct MyStoreState {
242/// states: HashMap<NamedId, WasiRandomCtx>,
243/// }
244///
245/// fn main() -> Result<()> {
246/// let engine = Engine::default();
247/// let mut linker = Linker::new(&engine);
248/// let component = Component::new(&engine, "(component)")?;
249/// let mut name_map = HashMap::new();
250///
251/// wasmtime_wasi::p3::random::add_named_to_linker::<MyStoreState>(
252/// &mut linker,
253/// &component,
254/// |_, name| {
255/// let len = name_map.len();
256/// Ok(NamedId(*name_map.entry(name.to_string()).or_insert(len)))
257/// },
258/// )?;
259/// Ok(())
260/// }
261///
262/// impl WasiRandomNamedView for MyStoreState {
263/// fn random(&mut self, id: NamedId) -> &mut WasiRandomCtx {
264/// self.states.get_mut(&id).expect("state for id")
265/// }
266/// }
267/// ```
268pub trait WasiRandomNamedView: Send + 'static {
269 /// Looks up the [`WasiRandomCtx`] for the given [`NamedId`].
270 ///
271 /// This method will resolve the `id` specified to a specific random
272 /// context that is available to be used. Note that this method is
273 /// specifically infallible meaning that a random context must be returned
274 /// and this cannot generate a trap or panic or similar.
275 ///
276 /// Embedders are responsible for allocating [`NamedId`] and assigning
277 /// meaning to ids. When a `Linker` is populated embedders will have the
278 /// ability to generate a `NamedId` for all imports found, and then that
279 /// embedder-allocated id is then passed back here when the corresponding
280 /// imported function is invoked.
281 fn random(&mut self, id: NamedId) -> &mut WasiRandomCtx;
282}