1use crate::std_facade::{Arc, String, ToOwned, Vec};
11use core::result::Result;
12use core::{fmt, str, u8, convert::TryInto};
13
14use rand::{self, Rng, RngCore, SeedableRng};
15use rand_chacha::ChaChaRng;
16use rand_xorshift::XorShiftRng;
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub enum RngAlgorithm {
25 XorShift,
34 ChaCha,
39 PassThrough,
52 Recorder,
61 #[allow(missing_docs)]
62 #[doc(hidden)]
63 _NonExhaustive,
64}
65
66impl Default for RngAlgorithm {
67 fn default() -> Self {
68 RngAlgorithm::ChaCha
69 }
70}
71
72impl RngAlgorithm {
73 pub(crate) fn persistence_key(self) -> &'static str {
74 match self {
75 RngAlgorithm::XorShift => "xs",
76 RngAlgorithm::ChaCha => "cc",
77 RngAlgorithm::PassThrough => "pt",
78 RngAlgorithm::Recorder => "rc",
79 RngAlgorithm::_NonExhaustive => unreachable!(),
80 }
81 }
82
83 pub(crate) fn from_persistence_key(k: &str) -> Option<Self> {
84 match k {
85 "xs" => Some(RngAlgorithm::XorShift),
86 "cc" => Some(RngAlgorithm::ChaCha),
87 "pt" => Some(RngAlgorithm::PassThrough),
88 "rc" => Some(RngAlgorithm::Recorder),
89 _ => None,
90 }
91 }
92}
93
94impl str::FromStr for RngAlgorithm {
97 type Err = ();
98 fn from_str(s: &str) -> Result<Self, ()> {
99 RngAlgorithm::from_persistence_key(s).ok_or(())
100 }
101}
102impl fmt::Display for RngAlgorithm {
103 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
104 write!(f, "{}", self.persistence_key())
105 }
106}
107
108#[derive(Clone, Debug)]
110pub struct TestRng {
111 rng: TestRngImpl,
112}
113
114#[derive(Clone, Debug)]
115enum TestRngImpl {
116 XorShift(XorShiftRng),
117 ChaCha(ChaChaRng),
118 PassThrough {
119 off: usize,
120 end: usize,
121 data: Arc<[u8]>,
122 },
123 Recorder {
124 rng: ChaChaRng,
125 record: Vec<u8>,
126 },
127}
128
129impl RngCore for TestRng {
130 fn next_u32(&mut self) -> u32 {
131 match &mut self.rng {
132 &mut TestRngImpl::XorShift(ref mut rng) => rng.next_u32(),
133
134 &mut TestRngImpl::ChaCha(ref mut rng) => rng.next_u32(),
135
136 &mut TestRngImpl::PassThrough { .. } => {
137 let mut buf = [0; 4];
138 self.fill_bytes(&mut buf[..]);
139 u32::from_le_bytes(buf)
140 }
141
142 &mut TestRngImpl::Recorder {
143 ref mut rng,
144 ref mut record,
145 } => {
146 let read = rng.next_u32();
147 record.extend_from_slice(&read.to_le_bytes());
148 read
149 }
150 }
151 }
152
153 fn next_u64(&mut self) -> u64 {
154 match &mut self.rng {
155 &mut TestRngImpl::XorShift(ref mut rng) => rng.next_u64(),
156
157 &mut TestRngImpl::ChaCha(ref mut rng) => rng.next_u64(),
158
159 &mut TestRngImpl::PassThrough { .. } => {
160 let mut buf = [0; 8];
161 self.fill_bytes(&mut buf[..]);
162 u64::from_le_bytes(buf)
163 }
164
165 &mut TestRngImpl::Recorder {
166 ref mut rng,
167 ref mut record,
168 } => {
169 let read = rng.next_u64();
170 record.extend_from_slice(&read.to_le_bytes());
171 read
172 }
173 }
174 }
175
176 fn fill_bytes(&mut self, dest: &mut [u8]) {
177 match &mut self.rng {
178 &mut TestRngImpl::XorShift(ref mut rng) => rng.fill_bytes(dest),
179
180 &mut TestRngImpl::ChaCha(ref mut rng) => rng.fill_bytes(dest),
181
182 &mut TestRngImpl::PassThrough {
183 ref mut off,
184 end,
185 ref data,
186 } => {
187 let bytes_to_copy = dest.len().min(end - *off);
188 dest[..bytes_to_copy]
189 .copy_from_slice(&data[*off..*off + bytes_to_copy]);
190 *off += bytes_to_copy;
191 for i in bytes_to_copy..dest.len() {
192 dest[i] = 0;
193 }
194 }
195
196 &mut TestRngImpl::Recorder {
197 ref mut rng,
198 ref mut record,
199 } => {
200 let res = rng.fill_bytes(dest);
201 record.extend_from_slice(&dest);
202 res
203 }
204 }
205 }
206
207 fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
208 match self.rng {
209 TestRngImpl::XorShift(ref mut rng) => rng.try_fill_bytes(dest),
210
211 TestRngImpl::ChaCha(ref mut rng) => rng.try_fill_bytes(dest),
212
213 TestRngImpl::PassThrough { .. } => {
214 self.fill_bytes(dest);
215 Ok(())
216 }
217
218 TestRngImpl::Recorder {
219 ref mut rng,
220 ref mut record,
221 } => {
222 let res = rng.try_fill_bytes(dest);
223 if res.is_ok() {
224 record.extend_from_slice(&dest);
225 }
226 res
227 }
228 }
229 }
230}
231
232#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
233pub(crate) enum Seed {
234 XorShift([u8; 16]),
235 ChaCha([u8; 32]),
236 PassThrough(Option<(usize, usize)>, Arc<[u8]>),
237 Recorder([u8; 32]),
238}
239
240impl Seed {
241 pub(crate) fn from_bytes(algorithm: RngAlgorithm, seed: &[u8]) -> Self {
242 match algorithm {
243 RngAlgorithm::XorShift => {
244 assert_eq!(16, seed.len(), "XorShift requires a 16-byte seed");
245 let mut buf = [0; 16];
246 buf.copy_from_slice(seed);
247 Seed::XorShift(buf)
248 }
249
250 RngAlgorithm::ChaCha => {
251 assert_eq!(32, seed.len(), "ChaCha requires a 32-byte seed");
252 let mut buf = [0; 32];
253 buf.copy_from_slice(seed);
254 Seed::ChaCha(buf)
255 }
256
257 RngAlgorithm::PassThrough => Seed::PassThrough(None, seed.into()),
258
259 RngAlgorithm::Recorder => {
260 assert_eq!(32, seed.len(), "Recorder requires a 32-byte seed");
261 let mut buf = [0; 32];
262 buf.copy_from_slice(seed);
263 Seed::Recorder(buf)
264 }
265
266 RngAlgorithm::_NonExhaustive => unreachable!(),
267 }
268 }
269
270 pub(crate) fn from_persistence(string: &str) -> Option<Seed> {
271 fn from_base16(dst: &mut [u8], src: &str) -> Option<()> {
272 if dst.len() * 2 != src.len() {
273 return None;
274 }
275
276 for (dst_byte, src_pair) in
277 dst.into_iter().zip(src.as_bytes().chunks(2))
278 {
279 *dst_byte =
280 u8::from_str_radix(str::from_utf8(src_pair).ok()?, 16)
281 .ok()?;
282 }
283
284 Some(())
285 }
286
287 let parts =
288 string.trim().split(char::is_whitespace).collect::<Vec<_>>();
289 RngAlgorithm::from_persistence_key(&parts[0]).and_then(
290 |alg| match alg {
291 RngAlgorithm::XorShift => {
292 if 5 != parts.len() {
293 return None;
294 }
295
296 let mut dwords = [0u32; 4];
297 for (dword, part) in
298 (&mut dwords[..]).into_iter().zip(&parts[1..])
299 {
300 *dword = part.parse().ok()?;
301 }
302
303 let mut seed = [0u8; 16];
304 for (chunk, dword) in seed.chunks_mut(4).zip(dwords) {
305 chunk.copy_from_slice(&dword.to_le_bytes());
306 }
307 Some(Seed::XorShift(seed))
308 }
309
310 RngAlgorithm::ChaCha => {
311 if 2 != parts.len() {
312 return None;
313 }
314
315 let mut seed = [0u8; 32];
316 from_base16(&mut seed, &parts[1])?;
317 Some(Seed::ChaCha(seed))
318 }
319
320 RngAlgorithm::PassThrough => {
321 if 1 == parts.len() {
322 return Some(Seed::PassThrough(None, vec![].into()));
323 }
324
325 if 2 != parts.len() {
326 return None;
327 }
328
329 let mut seed = vec![0u8; parts[1].len() / 2];
330 from_base16(&mut seed, &parts[1])?;
331 Some(Seed::PassThrough(None, seed.into()))
332 }
333
334 RngAlgorithm::Recorder => {
335 if 2 != parts.len() {
336 return None;
337 }
338
339 let mut seed = [0u8; 32];
340 from_base16(&mut seed, &parts[1])?;
341 Some(Seed::Recorder(seed))
342 }
343
344 RngAlgorithm::_NonExhaustive => unreachable!(),
345 },
346 )
347 }
348
349 pub(crate) fn to_persistence(&self) -> String {
350 fn to_base16(dst: &mut String, src: &[u8]) {
351 for byte in src {
352 dst.push_str(&format!("{:02x}", byte));
353 }
354 }
355
356 match *self {
357 Seed::XorShift(ref seed) => {
358 let dwords = [
359 u32::from_le_bytes(seed[0..4].try_into().unwrap()),
360 u32::from_le_bytes(seed[4..8].try_into().unwrap()),
361 u32::from_le_bytes(seed[8..12].try_into().unwrap()),
362 u32::from_le_bytes(seed[12..16].try_into().unwrap()),
363 ];
364 format!(
365 "{} {} {} {} {}",
366 RngAlgorithm::XorShift.persistence_key(),
367 dwords[0],
368 dwords[1],
369 dwords[2],
370 dwords[3]
371 )
372 }
373
374 Seed::ChaCha(ref seed) => {
375 let mut string =
376 RngAlgorithm::ChaCha.persistence_key().to_owned();
377 string.push(' ');
378 to_base16(&mut string, seed);
379 string
380 }
381
382 Seed::PassThrough(bounds, ref data) => {
383 let data =
384 bounds.map_or(&data[..], |(start, end)| &data[start..end]);
385 let mut string =
386 RngAlgorithm::PassThrough.persistence_key().to_owned();
387 string.push(' ');
388 to_base16(&mut string, data);
389 string
390 }
391
392 Seed::Recorder(ref seed) => {
393 let mut string =
394 RngAlgorithm::Recorder.persistence_key().to_owned();
395 string.push(' ');
396 to_base16(&mut string, seed);
397 string
398 }
399 }
400 }
401}
402
403impl TestRng {
404 pub fn from_seed(algorithm: RngAlgorithm, seed: &[u8]) -> Self {
414 TestRng::from_seed_internal(Seed::from_bytes(algorithm, seed))
415 }
416
417 pub fn bytes_used(&self) -> Vec<u8> {
424 match self.rng {
425 TestRngImpl::Recorder { ref record, .. } => record.clone(),
426 _ => panic!("bytes_used() called on non-Recorder RNG"),
427 }
428 }
429
430 pub(crate) fn default_rng(algorithm: RngAlgorithm) -> Self {
432 #[cfg(feature = "std")]
433 {
434 Self {
435 rng: match algorithm {
436 RngAlgorithm::XorShift => {
437 TestRngImpl::XorShift(XorShiftRng::from_entropy())
438 }
439 RngAlgorithm::ChaCha => {
440 TestRngImpl::ChaCha(ChaChaRng::from_entropy())
441 }
442 RngAlgorithm::PassThrough => {
443 panic!("cannot create default instance of PassThrough")
444 }
445 RngAlgorithm::Recorder => TestRngImpl::Recorder {
446 rng: ChaChaRng::from_entropy(),
447 record: Vec::new(),
448 },
449 RngAlgorithm::_NonExhaustive => unreachable!(),
450 },
451 }
452 }
453 #[cfg(all(
454 not(feature = "std"),
455 any(target_arch = "x86", target_arch = "x86_64"),
456 feature = "hardware-rng"
457 ))]
458 {
459 return Self::hardware_rng(algorithm);
460 }
461 #[cfg(not(feature = "std"))]
462 {
463 return Self::deterministic_rng(algorithm);
464 }
465 }
466
467 const SEED_FOR_XOR_SHIFT: [u8; 16] = [
468 0xf4, 0x16, 0x16, 0x48, 0xc3, 0xac, 0x77, 0xac, 0x72, 0x20, 0x0b, 0xea,
469 0x99, 0x67, 0x2d, 0x6d,
470 ];
471
472 const SEED_FOR_CHA_CHA: [u8; 32] = [
473 0xf4, 0x16, 0x16, 0x48, 0xc3, 0xac, 0x77, 0xac, 0x72, 0x20, 0x0b, 0xea,
474 0x99, 0x67, 0x2d, 0x6d, 0xca, 0x9f, 0x76, 0xaf, 0x1b, 0x09, 0x73, 0xa0,
475 0x59, 0x22, 0x6d, 0xc5, 0x46, 0x39, 0x1c, 0x4a,
476 ];
477
478 #[cfg(all(
485 not(feature = "std"),
486 any(target_arch = "x86", target_arch = "x86_64"),
487 feature = "hardware-rng"
488 ))]
489 pub fn hardware_rng(algorithm: RngAlgorithm) -> Self {
490 use x86::random::{rdrand_slice, RdRand};
491
492 Self::from_seed_internal(match algorithm {
493 RngAlgorithm::XorShift => {
494 let mut seed: [u8; 16] = TestRng::SEED_FOR_XOR_SHIFT;
496 unsafe {
497 let r = rdrand_slice(&mut seed);
498 debug_assert!(r, "hardware_rng should only be called on machines with support for rdrand");
499 }
500 Seed::XorShift(seed)
501 }
502 RngAlgorithm::ChaCha => {
503 let mut seed: [u8; 32] = TestRng::SEED_FOR_CHA_CHA;
505 unsafe {
506 let r = rdrand_slice(&mut seed);
507 debug_assert!(r, "hardware_rng should only be called on machines with support for rdrand");
508 }
509 Seed::ChaCha(seed)
510 }
511 RngAlgorithm::PassThrough => {
512 panic!("deterministic RNG not available for PassThrough")
513 }
514 RngAlgorithm::Recorder => {
515 let mut seed: [u8; 32] = TestRng::SEED_FOR_CHA_CHA;
517 unsafe {
518 let r = rdrand_slice(&mut seed);
519 debug_assert!(r, "hardware_rng should only be called on machines with support for rdrand");
520 }
521 Seed::Recorder(seed)
522 }
523 RngAlgorithm::_NonExhaustive => unreachable!(),
524 })
525 }
526
527 pub fn deterministic_rng(algorithm: RngAlgorithm) -> Self {
542 Self::from_seed_internal(match algorithm {
543 RngAlgorithm::XorShift => {
544 Seed::XorShift(TestRng::SEED_FOR_XOR_SHIFT)
545 }
546 RngAlgorithm::ChaCha => Seed::ChaCha(TestRng::SEED_FOR_CHA_CHA),
547 RngAlgorithm::PassThrough => {
548 panic!("deterministic RNG not available for PassThrough")
549 }
550 RngAlgorithm::Recorder => Seed::Recorder(TestRng::SEED_FOR_CHA_CHA),
551 RngAlgorithm::_NonExhaustive => unreachable!(),
552 })
553 }
554
555 pub(crate) fn gen_rng(&mut self) -> Self {
558 Self::from_seed_internal(self.new_rng_seed())
559 }
560
561 pub(crate) fn set_seed(&mut self, seed: Seed) {
563 *self = Self::from_seed_internal(seed);
564 }
565
566 pub(crate) fn gen_get_seed(&mut self) -> Seed {
569 let seed = self.new_rng_seed();
570 self.set_seed(seed.clone());
571 seed
572 }
573
574 pub(crate) fn new_rng_seed(&mut self) -> Seed {
576 match self.rng {
577 TestRngImpl::XorShift(ref mut rng) => {
578 let mut seed = rng.gen::<[u8; 16]>();
579
580 for word in seed.chunks_mut(4) {
584 word[3] ^= 0xde;
585 word[2] ^= 0xad;
586 word[1] ^= 0xbe;
587 word[0] ^= 0xef;
588 }
589
590 Seed::XorShift(seed)
591 }
592
593 TestRngImpl::ChaCha(ref mut rng) => Seed::ChaCha(rng.gen()),
594
595 TestRngImpl::PassThrough {
596 ref mut off,
597 ref mut end,
598 ref data,
599 } => {
600 let len = *end - *off;
601 let child_start = *off + len / 2;
602 let child_end = *off + len;
603 *end = child_start;
604 Seed::PassThrough(
605 Some((child_start, child_end)),
606 Arc::clone(data),
607 )
608 }
609
610 TestRngImpl::Recorder { ref mut rng, .. } => {
611 Seed::Recorder(rng.gen())
612 }
613 }
614 }
615
616 fn from_seed_internal(seed: Seed) -> Self {
618 Self {
619 rng: match seed {
620 Seed::XorShift(seed) => {
621 TestRngImpl::XorShift(XorShiftRng::from_seed(seed))
622 }
623
624 Seed::ChaCha(seed) => {
625 TestRngImpl::ChaCha(ChaChaRng::from_seed(seed))
626 }
627
628 Seed::PassThrough(bounds, data) => {
629 let (start, end) = bounds.unwrap_or((0, data.len()));
630 TestRngImpl::PassThrough {
631 off: start,
632 end,
633 data,
634 }
635 }
636
637 Seed::Recorder(seed) => TestRngImpl::Recorder {
638 rng: ChaChaRng::from_seed(seed),
639 record: Vec::new(),
640 },
641 },
642 }
643 }
644}
645
646#[cfg(test)]
647mod test {
648 use crate::std_facade::Vec;
649
650 use rand::{Rng, RngCore};
651
652 use super::{RngAlgorithm, Seed, TestRng};
653 use crate::arbitrary::any;
654 use crate::strategy::*;
655
656 proptest! {
657 #[test]
658 fn gen_parse_seeds(
659 seed in prop_oneof![
660 any::<[u8;16]>().prop_map(Seed::XorShift),
661 any::<[u8;32]>().prop_map(Seed::ChaCha),
662 any::<Vec<u8>>().prop_map(|data| Seed::PassThrough(None, data.into())),
663 any::<[u8;32]>().prop_map(Seed::Recorder),
664 ])
665 {
666 assert_eq!(seed, Seed::from_persistence(&seed.to_persistence()).unwrap());
667 }
668
669 #[test]
670 fn rngs_dont_clone_self_on_genrng(
671 seed in prop_oneof![
672 any::<[u8;16]>().prop_map(Seed::XorShift),
673 any::<[u8;32]>().prop_map(Seed::ChaCha),
674 Just(()).prop_perturb(|_, mut rng| {
675 let mut buf = vec![0u8; 2048];
676 rng.fill_bytes(&mut buf);
677 Seed::PassThrough(None, buf.into())
678 }),
679 any::<[u8;32]>().prop_map(Seed::Recorder),
680 ])
681 {
682 type Value = [u8;32];
683 let orig = TestRng::from_seed_internal(seed);
684
685 {
686 let mut rng1 = orig.clone();
687 let mut rng2 = rng1.gen_rng();
688 assert_ne!(rng1.gen::<Value>(), rng2.gen::<Value>());
689 }
690
691 {
692 let mut rng1 = orig.clone();
693 let mut rng2 = rng1.gen_rng();
694 let mut rng3 = rng1.gen_rng();
695 let mut rng4 = rng2.gen_rng();
696 let a = rng1.gen::<Value>();
697 let b = rng2.gen::<Value>();
698 let c = rng3.gen::<Value>();
699 let d = rng4.gen::<Value>();
700 assert_ne!(a, b);
701 assert_ne!(a, c);
702 assert_ne!(a, d);
703 assert_ne!(b, c);
704 assert_ne!(b, d);
705 assert_ne!(c, d);
706 }
707 }
708 }
709
710 #[test]
711 fn passthrough_rng_behaves_properly() {
712 let mut rng = TestRng::from_seed(
713 RngAlgorithm::PassThrough,
714 &[
715 0xDE, 0xC0, 0x12, 0x34, 0x56, 0x78, 0xFE, 0xCA, 0xEF, 0xBE,
716 0xAD, 0xDE, 0x01, 0x02, 0x03,
717 ],
718 );
719
720 assert_eq!(0x3412C0DE, rng.next_u32());
721 assert_eq!(0xDEADBEEFCAFE7856, rng.next_u64());
722
723 let mut buf = [0u8; 4];
724 rng.try_fill_bytes(&mut buf[0..4]).unwrap();
725 assert_eq!([1, 2, 3, 0], buf);
726 rng.try_fill_bytes(&mut buf[0..4]).unwrap();
727 assert_eq!([0, 0, 0, 0], buf);
728 }
729}