Skip to main content

commonware_utils/
rng.rs

1//! Utilities for random number generation.
2
3use rand::{rngs::StdRng, CryptoRng, RngCore, SeedableRng};
4use std::mem::size_of;
5
6/// Returns a seeded RNG for deterministic testing.
7///
8/// Uses seed 0 by default to ensure reproducible test results.
9pub fn test_rng() -> StdRng {
10    StdRng::seed_from_u64(0)
11}
12
13/// Returns a seeded RNG with a custom seed for deterministic testing.
14///
15/// Use this when you need multiple independent RNG streams in the same test,
16/// or when a helper function needs its own RNG that won't collide with the caller's.
17pub fn test_rng_seeded(seed: u64) -> StdRng {
18    StdRng::seed_from_u64(seed)
19}
20
21/// Applies a SplitMix64-style finalizer to a deterministic input word.
22///
23/// This is useful for cheaply decorrelating derived deterministic seeds while
24/// preserving reproducibility.
25#[inline]
26pub const fn mix64(mut word: u64) -> u64 {
27    word ^= word >> 30;
28    word = word.wrapping_mul(0xbf58_476d_1ce4_e5b9);
29    word ^= word >> 27;
30    word = word.wrapping_mul(0x94d0_49bb_1331_11eb);
31    word ^ (word >> 31)
32}
33
34/// Width of each source window in bytes.
35const BLOCK_BYTES: usize = size_of::<u64>();
36
37/// An RNG that expands a fuzzer byte slice into an infinite deterministic stream.
38///
39/// # Design
40///
41/// `FuzzRng` maps a fuzzer-controlled byte slice to output blocks.
42///
43/// For each block counter `ctr`, it:
44/// 1. Reads a wrapping `u64`-wide window from the input bytes.
45/// 2. Xors in `ctr` and a domain constant.
46/// 3. Applies a SplitMix64-style finalizer.
47///
48/// ```text
49/// input bytes (len = N):
50///   [b0 b1 b2 ... b(N-1)]
51///
52/// block ctr = i:
53///   word_i bytes = [b(i+0)%N, b(i+1)%N, ... b(i+7)%N]
54///   word_i       = big-endian u64 of those bytes
55///   out_i        = mix64(word_i ^ i ^ DOMAIN)
56/// ```
57///
58/// # Why this mapping
59///
60/// Hashing the full input once and then seeding a PRNG makes tiny input changes
61/// look globally unrelated. This adapter avoids that by using a sliding window
62/// keyed by the block counter.
63///
64/// ```text
65/// byte k affects anchors:
66///   i in [k-(BLOCK_BYTES-1), ..., k] (mod N)
67/// ```
68///
69/// # Worked Example
70///
71/// With `N = 4`, input bytes repeat inside each block:
72///
73/// ```text
74/// input: [a b c d]
75///
76/// ctr=0: word bytes [a b c d a b c d]
77/// ctr=1: word bytes [b c d a b c d a]
78/// ctr=2: word bytes [c d a b c d a b]
79/// ...
80/// ```
81///
82/// Even for low-entropy input like `[0 0 0 0]`, output still changes because
83/// `ctr` is mixed into every block before finalization.
84///
85/// `fill_bytes` serves output from cached block bytes so callers get a stable
86/// byte stream regardless of whether they request randomness as `next_u64`,
87/// `next_u32`, or arbitrary byte slices.
88pub struct FuzzRng {
89    bytes: Vec<u8>,
90    ctr: u64,
91    cache: [u8; BLOCK_BYTES],
92    cache_pos: usize,
93}
94
95impl FuzzRng {
96    /// Creates a new `FuzzRng` from a byte buffer.
97    pub const fn new(bytes: Vec<u8>) -> Self {
98        Self {
99            bytes,
100            ctr: 0,
101            cache: [0u8; BLOCK_BYTES],
102            cache_pos: BLOCK_BYTES,
103        }
104    }
105
106    /// Generates the next mixed `u64` block from the fuzz input.
107    ///
108    /// Conceptually:
109    /// 1. Build `word` from a wrapping `BLOCK_BYTES` window anchored at `ctr`.
110    /// 2. Compute `mixed = mix64(word ^ ctr ^ GOLDEN_RATIO)`.
111    /// 3. Increment `ctr`.
112    ///
113    /// This keeps the output deterministic while preserving local mutation
114    /// influence: one input-byte mutation only affects nearby anchor counters.
115    #[inline]
116    fn next_block_u64(&mut self) -> u64 {
117        // Build a wrapping u64-width source word anchored at this block counter.
118        // A single fuzz-byte mutation only impacts nearby anchors.
119        let mut bytes = [0u8; BLOCK_BYTES];
120        if !self.bytes.is_empty() {
121            let len = self.bytes.len() as u64;
122            for (i, byte) in bytes.iter_mut().enumerate() {
123                *byte = self.bytes[(self.ctr.wrapping_add(i as u64) % len) as usize];
124            }
125        }
126        let word = u64::from_be_bytes(bytes);
127
128        // Mix the structured word into a high-quality output block without
129        // hashing the entire seed into an avalanche-style global state.
130        let ctr = self.ctr;
131        self.ctr = self.ctr.wrapping_add(1);
132        mix64(word ^ ctr ^ crate::GOLDEN_RATIO)
133    }
134}
135
136impl RngCore for FuzzRng {
137    fn next_u32(&mut self) -> u32 {
138        let mut buf = [0u8; 4];
139        self.fill_bytes(&mut buf);
140        u32::from_be_bytes(buf)
141    }
142
143    fn next_u64(&mut self) -> u64 {
144        let mut buf = [0u8; BLOCK_BYTES];
145        self.fill_bytes(&mut buf);
146        u64::from_be_bytes(buf)
147    }
148
149    fn fill_bytes(&mut self, dest: &mut [u8]) {
150        let mut written = 0;
151        while written < dest.len() {
152            if self.cache_pos == self.cache.len() {
153                // Cache block bytes so outputs are stable regardless of whether
154                // callers pull randomness as bytes or words:
155                //
156                // next_u64() stream bytes == fill_bytes() stream bytes.
157                self.cache = self.next_block_u64().to_be_bytes();
158                self.cache_pos = 0;
159            }
160
161            let available = self.cache.len() - self.cache_pos;
162            let need = dest.len() - written;
163            let take = available.min(need);
164            dest[written..written + take]
165                .copy_from_slice(&self.cache[self.cache_pos..self.cache_pos + take]);
166            self.cache_pos += take;
167            written += take;
168        }
169    }
170
171    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
172        self.fill_bytes(dest);
173        Ok(())
174    }
175}
176
177// SAFETY: FuzzRng is not cryptographically secure. It implements CryptoRng
178// only because the consensus fuzzer requires CryptoRng-bounded RNG. This type
179// must never be used outside of fuzz/test contexts.
180impl CryptoRng for FuzzRng {}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn test_empty_bytes_not_constant() {
188        let mut rng = FuzzRng::new(vec![]);
189
190        let values: Vec<_> = (0..BLOCK_BYTES).map(|_| rng.next_u64()).collect();
191        assert!(values.windows(2).any(|w| w[0] != w[1]));
192    }
193
194    #[test]
195    fn test_empty_bytes_deterministic() {
196        let mut rng1 = FuzzRng::new(vec![]);
197        let mut rng2 = FuzzRng::new(vec![]);
198
199        for _ in 0..256 {
200            assert_eq!(rng1.next_u64(), rng2.next_u64());
201        }
202    }
203
204    #[test]
205    fn test_all_zero_bytes_not_constant() {
206        let bytes = vec![0; BLOCK_BYTES];
207        let mut rng = FuzzRng::new(bytes);
208        let values: Vec<_> = (0..BLOCK_BYTES).map(|_| rng.next_u64()).collect();
209        assert!(values.windows(2).any(|w| w[0] != w[1]));
210    }
211
212    #[test]
213    fn test_deterministic_with_same_input() {
214        let bytes = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
215
216        let mut rng1 = FuzzRng::new(bytes.clone());
217        let mut rng2 = FuzzRng::new(bytes);
218
219        for _ in 0..1000 {
220            assert_eq!(rng1.next_u64(), rng2.next_u64());
221        }
222    }
223
224    #[test]
225    fn test_short_input_wraparound() {
226        for len in 1..=3 {
227            let bytes = vec![0xAB; len];
228            let mut rng1 = FuzzRng::new(bytes.clone());
229            let mut rng2 = FuzzRng::new(bytes);
230            let out1: Vec<_> = (0..32).map(|_| rng1.next_u64()).collect();
231            let out2: Vec<_> = (0..32).map(|_| rng2.next_u64()).collect();
232            assert_eq!(out1, out2);
233            assert!(out1.windows(2).any(|w| w[0] != w[1]));
234        }
235    }
236
237    #[test]
238    fn test_small_mutation_locality() {
239        let mut base = vec![0u8; 64];
240        for (i, byte) in base.iter_mut().enumerate() {
241            *byte = i as u8;
242        }
243        let mut mutated = base.clone();
244        let mutated_pos = 20usize;
245        mutated[mutated_pos] ^= 0x01;
246
247        let mut rng_a = FuzzRng::new(base);
248        let mut rng_b = FuzzRng::new(mutated);
249
250        let draws = 40usize;
251        let mut diff_indices = Vec::new();
252        for i in 0..draws {
253            if rng_a.next_u64() != rng_b.next_u64() {
254                diff_indices.push(i);
255            }
256        }
257
258        let expected: Vec<usize> = ((mutated_pos - 7)..=mutated_pos).collect();
259        assert_eq!(diff_indices, expected);
260    }
261
262    #[test]
263    fn test_small_mutation_locality_wraparound() {
264        let mut base = vec![0u8; 64];
265        for (i, byte) in base.iter_mut().enumerate() {
266            *byte = i as u8;
267        }
268        let mut mutated = base.clone();
269        let mutated_pos = 2usize;
270        mutated[mutated_pos] ^= 0x01;
271
272        let mut rng_a = FuzzRng::new(base);
273        let mut rng_b = FuzzRng::new(mutated);
274
275        let draws = 64usize;
276        let mut diff_indices = Vec::new();
277        for i in 0..draws {
278            if rng_a.next_u64() != rng_b.next_u64() {
279                diff_indices.push(i);
280            }
281        }
282
283        assert_eq!(diff_indices, vec![0, 1, 2, 59, 60, 61, 62, 63]);
284    }
285
286    #[test]
287    fn test_fill_bytes_shape_stability() {
288        let bytes: Vec<u8> = (0..32u8).collect();
289
290        let mut from_u64_rng = FuzzRng::new(bytes.clone());
291        let mut from_u64 = Vec::with_capacity(128);
292        for _ in 0..16 {
293            from_u64.extend_from_slice(&from_u64_rng.next_u64().to_be_bytes());
294        }
295
296        let mut from_fill_rng = FuzzRng::new(bytes);
297        let mut from_fill = vec![0u8; from_u64.len()];
298        let chunk_sizes = [3usize, 1, 7, 2, 11, 5, 13, 17];
299        let mut offset = 0;
300        let mut idx = 0;
301        while offset < from_fill.len() {
302            let chunk = chunk_sizes[idx % chunk_sizes.len()].min(from_fill.len() - offset);
303            from_fill_rng.fill_bytes(&mut from_fill[offset..offset + chunk]);
304            offset += chunk;
305            idx += 1;
306        }
307        assert_eq!(from_u64, from_fill);
308    }
309
310    #[test]
311    fn test_next_u32_consistency_with_fill_bytes() {
312        let bytes: Vec<u8> = (0..16u8).collect();
313
314        let mut from_u32_rng = FuzzRng::new(bytes.clone());
315        let mut from_u32 = Vec::with_capacity(64);
316        for _ in 0..16 {
317            from_u32.extend_from_slice(&from_u32_rng.next_u32().to_be_bytes());
318        }
319
320        let mut from_fill_rng = FuzzRng::new(bytes);
321        let mut from_fill = vec![0u8; from_u32.len()];
322        from_fill_rng.fill_bytes(&mut from_fill);
323        assert_eq!(from_u32, from_fill);
324    }
325
326    #[test]
327    fn test_try_fill_bytes_consistency_with_fill_bytes() {
328        let bytes: Vec<u8> = (0..16u8).collect();
329
330        let mut fill_rng = FuzzRng::new(bytes.clone());
331        let mut try_fill_rng = FuzzRng::new(bytes);
332
333        let mut fill_out = vec![0u8; 257];
334        fill_rng.fill_bytes(&mut fill_out);
335
336        let mut try_out = vec![0u8; 257];
337        try_fill_rng
338            .try_fill_bytes(&mut try_out)
339            .expect("try_fill_bytes should never fail");
340
341        assert_eq!(fill_out, try_out);
342    }
343
344    #[test]
345    fn test_next_u64_includes_counter_in_mix_input() {
346        // Use a constant source window so any change between blocks comes from
347        // counter mixing, not from different window bytes.
348        let bytes = vec![0xAA; BLOCK_BYTES];
349        let mut rng = FuzzRng::new(bytes.clone());
350
351        let mut source = [0u8; BLOCK_BYTES];
352        source.copy_from_slice(&bytes[..BLOCK_BYTES]);
353        let word = u64::from_be_bytes(source);
354
355        let mix = |mut x: u64| {
356            x ^= x >> 30;
357            x = x.wrapping_mul(0xbf58476d1ce4e5b9);
358            x ^= x >> 27;
359            x = x.wrapping_mul(0x94d049bb133111eb);
360            x ^= x >> 31;
361            x
362        };
363
364        #[allow(clippy::identity_op)]
365        let expected0 = mix(word ^ 0 ^ crate::GOLDEN_RATIO);
366        let expected1 = mix(word ^ 1 ^ crate::GOLDEN_RATIO);
367
368        assert_eq!(rng.next_u64(), expected0);
369        assert_eq!(rng.next_u64(), expected1);
370    }
371
372    #[cfg(feature = "arbitrary")]
373    mod conformance {
374        use super::*;
375        use commonware_conformance::Conformance;
376        use rand::Rng;
377
378        /// Conformance wrapper for FuzzRng that tests output stability.
379        ///
380        /// Derives both the input length and content from a seeded RNG so
381        /// conformance covers variable-length inputs including non-aligned
382        /// lengths that exercise wrapping.
383        struct FuzzRngConformance;
384
385        impl Conformance for FuzzRngConformance {
386            async fn commit(seed: u64) -> Vec<u8> {
387                let mut seed_rng = test_rng_seeded(seed);
388                let len = seed_rng.gen_range(1..=64);
389                let mut input = vec![0u8; len];
390                seed_rng.fill_bytes(&mut input);
391
392                let mut rng = FuzzRng::new(input);
393                const CONFORMANCE_BLOCKS: usize = 32;
394
395                // Generate enough output to exercise wrapping and mixing.
396                let mut output = Vec::with_capacity(CONFORMANCE_BLOCKS * BLOCK_BYTES);
397                for _ in 0..CONFORMANCE_BLOCKS {
398                    output.extend_from_slice(&rng.next_u64().to_be_bytes());
399                }
400                output
401            }
402        }
403
404        commonware_conformance::conformance_tests! {
405            FuzzRngConformance => 1024,
406        }
407    }
408}