1use candid::CandidType;
11use rand_chacha::{
12 ChaCha20Rng,
13 rand_core::{RngCore, SeedableRng},
14};
15use std::cell::RefCell;
16
17thread_local! {
18 static RNG: RefCell<Option<ChaCha20Rng>> = const { RefCell::new(None) };
19}
20
21#[derive(CandidType, Clone, Debug, Eq, PartialEq)]
31pub enum RngError {
32 RngNotInitialized(String),
33}
34
35impl RngError {
36 fn not_initialized() -> Self {
37 Self::RngNotInitialized("Randomness is not initialized. Please try again later".to_string())
38 }
39}
40
41pub fn seed_from(seed: [u8; 32]) {
47 RNG.with_borrow_mut(|rng| {
48 *rng = Some(ChaCha20Rng::from_seed(seed));
49 });
50}
51
52#[must_use]
54pub fn is_seeded() -> bool {
55 RNG.with_borrow(Option::is_some)
56}
57
58fn with_rng<T>(f: impl FnOnce(&mut ChaCha20Rng) -> T) -> Result<T, RngError> {
59 RNG.with_borrow_mut(|rng| match rng.as_mut() {
60 Some(rand) => Ok(f(rand)),
61 None => Err(RngError::not_initialized()),
62 })
63}
64
65pub fn fill_bytes(dest: &mut [u8]) -> Result<(), RngError> {
71 with_rng(|rand| rand.fill_bytes(dest))
72}
73
74pub fn random_bytes(size: usize) -> Result<Vec<u8>, RngError> {
76 let mut buf = vec![0u8; size];
77 fill_bytes(&mut buf)?;
78 Ok(buf)
79}
80
81pub fn random_hex(size: usize) -> Result<String, RngError> {
83 let bytes = random_bytes(size)?;
84 Ok(hex::encode(bytes))
85}
86
87pub fn next_u8() -> Result<u8, RngError> {
89 Ok((next_u16()? & 0xFF) as u8)
90}
91
92#[allow(clippy::cast_possible_truncation)]
94pub fn next_u16() -> Result<u16, RngError> {
95 with_rng(|rand| rand.next_u32() as u16)
96}
97
98pub fn next_u32() -> Result<u32, RngError> {
100 with_rng(RngCore::next_u32)
101}
102
103pub fn next_u64() -> Result<u64, RngError> {
105 with_rng(RngCore::next_u64)
106}
107
108pub fn next_u128() -> Result<u128, RngError> {
110 with_rng(|rand| {
111 let hi = u128::from(rand.next_u64());
112 let lo = u128::from(rand.next_u64());
113 (hi << 64) | lo
114 })
115}
116
117#[cfg(test)]
122mod tests {
123 use super::*;
124
125 #[test]
126 fn test_unique_u64s() {
127 use std::collections::HashSet;
128
129 seed_from([7; 32]);
130
131 let mut set = HashSet::new();
132 while set.len() < 1000 {
133 let random_value = next_u64().expect("seeded RNG");
134 assert!(set.insert(random_value), "value already in set");
135 }
136 }
137
138 #[test]
139 fn test_rng_reseeding() {
140 seed_from([1; 32]);
141 let first = next_u64().expect("seeded RNG");
142 seed_from([2; 32]);
143 let second = next_u64().expect("seeded RNG");
144
145 assert_ne!(
146 first, second,
147 "RNGs with different seeds unexpectedly produced the same value"
148 );
149 }
150
151 #[test]
152 fn test_determinism_with_fixed_seed() {
153 let seed = [42u8; 32];
154 seed_from(seed);
155
156 let values: Vec<u64> = (0..100).map(|_| next_u64().expect("seeded RNG")).collect();
157
158 seed_from(seed);
159 for value in values {
160 assert_eq!(next_u64().expect("seeded RNG"), value);
161 }
162 }
163
164 #[test]
165 fn test_missing_seed_errors() {
166 RNG.with_borrow_mut(|rng| {
167 *rng = None;
168 });
169
170 assert!(matches!(
171 random_bytes(8),
172 Err(RngError::RngNotInitialized(_))
173 ));
174 }
175
176 #[test]
177 fn test_random_hex_length() {
178 seed_from([9; 32]);
179
180 let value = random_hex(6).expect("seeded RNG");
181 assert_eq!(value.len(), 12);
182 }
183
184 #[test]
187 fn test_bit_entropy() {
188 seed_from([3; 32]);
189
190 let mut bits = 0u64;
191 for _ in 0..100 {
192 bits |= next_u64().expect("seeded RNG");
193 }
194
195 let bit_count = bits.count_ones();
196 assert!(
197 bit_count > 8,
198 "Low entropy: only {bit_count} bits set in 100 samples",
199 );
200 }
201}