Skip to main content

forest/utils/rand/
mod.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use rand::{CryptoRng, Rng, RngCore, SeedableRng as _};
5
6/// A wrapper of [`uuid::Builder::from_random_bytes`] that uses [`forest_rng`] internally
7pub fn new_uuid_v4() -> uuid::Uuid {
8    let mut random_bytes = uuid::Bytes::default();
9    forest_rng().fill(&mut random_bytes);
10    uuid::Builder::from_random_bytes(random_bytes).into_uuid()
11}
12
13/// A wrapper of [`rand::thread_rng`] that can be overridden by reproducible seeded
14/// [`rand_chacha::ChaChaRng`] via `FOREST_TEST_RNG_FIXED_SEED` environment variable.
15/// This is required for reproducible test cases for normally non-deterministic methods.
16pub fn forest_rng() -> impl Rng + CryptoRng {
17    forest_rng_internal(ForestRngMode::ThreadRng)
18}
19
20/// A wrapper of [`rand::rngs::OsRng`] that can be overridden by reproducible seeded
21/// [`rand_chacha::ChaChaRng`] via `FOREST_TEST_RNG_FIXED_SEED` environment variable.
22/// This is required for reproducible test cases for normally non-deterministic methods.
23pub fn forest_os_rng() -> impl Rng + CryptoRng {
24    forest_rng_internal(ForestRngMode::OsRng)
25}
26
27pub const FIXED_RNG_SEED_ENV: &str = "FOREST_TEST_RNG_FIXED_SEED";
28
29enum ForestRngMode {
30    ThreadRng,
31    OsRng,
32}
33
34fn forest_rng_internal(mode: ForestRngMode) -> impl Rng + CryptoRng {
35    const ENV: &str = FIXED_RNG_SEED_ENV;
36    if let Ok(v) = std::env::var(ENV) {
37        if let Ok(seed) = v.parse() {
38            #[cfg(not(test))]
39            tracing::warn!("[security] using test RNG with fixed seed {seed} set by {ENV}");
40            return Either::Left(rand_chacha::ChaChaRng::seed_from_u64(seed));
41        } else {
42            tracing::warn!("invalid u64 seed set by {ENV}: {v}. Falling back to the default RNG.");
43        }
44    }
45    match mode {
46        #[allow(clippy::disallowed_methods)]
47        ForestRngMode::ThreadRng => Either::Right(Either::Left(rand::thread_rng())),
48        #[allow(clippy::disallowed_types)]
49        ForestRngMode::OsRng => Either::Right(Either::Right(rand::rngs::OsRng)),
50    }
51}
52
53enum Either<A, B> {
54    Left(A),
55    Right(B),
56}
57
58impl<A, B> RngCore for Either<A, B>
59where
60    A: RngCore,
61    B: RngCore,
62{
63    fn next_u32(&mut self) -> u32 {
64        match self {
65            Self::Left(i) => i.next_u32(),
66            Self::Right(i) => i.next_u32(),
67        }
68    }
69
70    fn next_u64(&mut self) -> u64 {
71        match self {
72            Self::Left(i) => i.next_u64(),
73            Self::Right(i) => i.next_u64(),
74        }
75    }
76
77    fn fill_bytes(&mut self, dst: &mut [u8]) {
78        match self {
79            Self::Left(i) => i.fill_bytes(dst),
80            Self::Right(i) => i.fill_bytes(dst),
81        }
82    }
83
84    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
85        match self {
86            Self::Left(i) => i.try_fill_bytes(dest),
87            Self::Right(i) => i.try_fill_bytes(dest),
88        }
89    }
90}
91
92impl<A, B> CryptoRng for Either<A, B>
93where
94    A: RngCore,
95    B: RngCore,
96{
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use serial_test::serial;
103
104    #[test]
105    #[serial]
106    fn test_fixed_seed_env() {
107        unsafe { std::env::set_var(FIXED_RNG_SEED_ENV, "0") };
108
109        let mut a = [0; 1024];
110        let mut b = [0; 1024];
111
112        forest_rng().fill(&mut a);
113        forest_rng().fill(&mut b);
114        assert_eq!(a, b);
115
116        forest_os_rng().fill(&mut a);
117        forest_os_rng().fill(&mut b);
118        assert_eq!(a, b);
119
120        unsafe { std::env::remove_var(FIXED_RNG_SEED_ENV) };
121    }
122
123    #[test]
124    #[serial]
125    fn test_thread_rng() {
126        unsafe { std::env::remove_var(FIXED_RNG_SEED_ENV) };
127        let mut a = [0; 1024];
128        forest_rng().fill(&mut a);
129        let mut b = [0; 1024];
130        forest_rng().fill(&mut b);
131        assert_ne!(a, b);
132    }
133
134    #[test]
135    #[serial]
136    fn test_os_rng() {
137        unsafe { std::env::remove_var(FIXED_RNG_SEED_ENV) };
138        let mut a = [0; 1024];
139        forest_os_rng().fill(&mut a);
140        let mut b = [0; 1024];
141        forest_os_rng().fill(&mut b);
142        assert_ne!(a, b);
143    }
144}