use core::mem::size_of;
use rand_core::block::{BlockRng, BlockRngCore};
use rand_core::{CryptoRng, Error, RngCore, SeedableRng};
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr>(BlockRng<ReseedingCore<R, Rsdr>>)
where
R: BlockRngCore + SeedableRng,
Rsdr: RngCore;
impl<R, Rsdr> ReseedingRng<R, Rsdr>
where
R: BlockRngCore + SeedableRng,
Rsdr: RngCore,
{
pub fn new(rng: R, threshold: u64, reseeder: Rsdr) -> Self {
ReseedingRng(BlockRng::new(ReseedingCore::new(rng, threshold, reseeder)))
}
pub fn reseed(&mut self) -> Result<(), Error> {
self.0.core.reseed()
}
}
impl<R, Rsdr: RngCore> RngCore for ReseedingRng<R, Rsdr>
where
R: BlockRngCore<Item = u32> + SeedableRng,
<R as BlockRngCore>::Results: AsRef<[u32]> + AsMut<[u32]>,
{
#[inline(always)]
fn next_u32(&mut self) -> u32 {
self.0.next_u32()
}
#[inline(always)]
fn next_u64(&mut self) -> u64 {
self.0.next_u64()
}
fn fill_bytes(&mut self, dest: &mut [u8]) {
self.0.fill_bytes(dest)
}
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
self.0.try_fill_bytes(dest)
}
}
impl<R, Rsdr> Clone for ReseedingRng<R, Rsdr>
where
R: BlockRngCore + SeedableRng + Clone,
Rsdr: RngCore + Clone,
{
fn clone(&self) -> ReseedingRng<R, Rsdr> {
ReseedingRng(BlockRng::new(self.0.core.clone()))
}
}
impl<R, Rsdr> CryptoRng for ReseedingRng<R, Rsdr>
where
R: BlockRngCore + SeedableRng + CryptoRng,
Rsdr: RngCore + CryptoRng,
{
}
#[derive(Debug)]
struct ReseedingCore<R, Rsdr> {
inner: R,
reseeder: Rsdr,
threshold: i64,
bytes_until_reseed: i64,
fork_counter: usize,
}
impl<R, Rsdr> BlockRngCore for ReseedingCore<R, Rsdr>
where
R: BlockRngCore + SeedableRng,
Rsdr: RngCore,
{
type Item = <R as BlockRngCore>::Item;
type Results = <R as BlockRngCore>::Results;
fn generate(&mut self, results: &mut Self::Results) {
let global_fork_counter = fork::get_fork_counter();
if self.bytes_until_reseed <= 0 || self.is_forked(global_fork_counter) {
return self.reseed_and_generate(results, global_fork_counter);
}
let num_bytes = results.as_ref().len() * size_of::<Self::Item>();
self.bytes_until_reseed -= num_bytes as i64;
self.inner.generate(results);
}
}
impl<R, Rsdr> ReseedingCore<R, Rsdr>
where
R: BlockRngCore + SeedableRng,
Rsdr: RngCore,
{
fn new(rng: R, threshold: u64, reseeder: Rsdr) -> Self {
use ::core::i64::MAX;
fork::register_fork_handler();
let threshold = if threshold == 0 {
MAX
} else if threshold <= MAX as u64 {
threshold as i64
} else {
MAX
};
ReseedingCore {
inner: rng,
reseeder,
threshold: threshold as i64,
bytes_until_reseed: threshold as i64,
fork_counter: 0,
}
}
fn reseed(&mut self) -> Result<(), Error> {
R::from_rng(&mut self.reseeder).map(|result| {
self.bytes_until_reseed = self.threshold;
self.inner = result
})
}
fn is_forked(&self, global_fork_counter: usize) -> bool {
(self.fork_counter.wrapping_sub(global_fork_counter) as isize) < 0
}
#[inline(never)]
fn reseed_and_generate(
&mut self, results: &mut <Self as BlockRngCore>::Results, global_fork_counter: usize,
) {
#![allow(clippy::if_same_then_else)] if self.is_forked(global_fork_counter) {
info!("Fork detected, reseeding RNG");
} else {
trace!("Reseeding RNG (periodic reseed)");
}
let num_bytes = results.as_ref().len() * size_of::<<R as BlockRngCore>::Item>();
if let Err(e) = self.reseed() {
warn!("Reseeding RNG failed: {}", e);
let _ = e;
}
self.fork_counter = global_fork_counter;
self.bytes_until_reseed = self.threshold - num_bytes as i64;
self.inner.generate(results);
}
}
impl<R, Rsdr> Clone for ReseedingCore<R, Rsdr>
where
R: BlockRngCore + SeedableRng + Clone,
Rsdr: RngCore + Clone,
{
fn clone(&self) -> ReseedingCore<R, Rsdr> {
ReseedingCore {
inner: self.inner.clone(),
reseeder: self.reseeder.clone(),
threshold: self.threshold,
bytes_until_reseed: 0, fork_counter: self.fork_counter,
}
}
}
impl<R, Rsdr> CryptoRng for ReseedingCore<R, Rsdr>
where
R: BlockRngCore + SeedableRng + CryptoRng,
Rsdr: RngCore + CryptoRng,
{
}
#[cfg(all(unix, not(target_os = "emscripten")))]
mod fork {
use core::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Once;
static RESEEDING_RNG_FORK_COUNTER: AtomicUsize = AtomicUsize::new(0);
pub fn get_fork_counter() -> usize {
RESEEDING_RNG_FORK_COUNTER.load(Ordering::Relaxed)
}
extern "C" fn fork_handler() {
RESEEDING_RNG_FORK_COUNTER.fetch_add(1, Ordering::Relaxed);
}
pub fn register_fork_handler() {
static REGISTER: Once = Once::new();
REGISTER.call_once(|| {
let ret = unsafe { libc::pthread_atfork(
Some(fork_handler),
Some(fork_handler),
Some(fork_handler),
) };
if ret != 0 {
panic!("libc::pthread_atfork failed with code {}", ret);
}
});
}
}
#[cfg(not(all(unix, not(target_os = "emscripten"))))]
mod fork {
pub fn get_fork_counter() -> usize {
0
}
pub fn register_fork_handler() {}
}
#[cfg(feature = "std_rng")]
#[cfg(test)]
mod test {
use super::ReseedingRng;
use crate::rngs::mock::StepRng;
use crate::rngs::std::Core;
use crate::{Rng, SeedableRng};
#[test]
fn test_reseeding() {
let mut zero = StepRng::new(0, 0);
let rng = Core::from_rng(&mut zero).unwrap();
let thresh = 1; let mut reseeding = ReseedingRng::new(rng, thresh, zero);
let mut buf = ([0u32; 32], [0u32; 32]);
reseeding.fill(&mut buf.0);
reseeding.fill(&mut buf.1);
let seq = buf;
for _ in 0..10 {
reseeding.fill(&mut buf.0);
reseeding.fill(&mut buf.1);
assert_eq!(buf, seq);
}
}
#[test]
fn test_clone_reseeding() {
#![allow(clippy::redundant_clone)]
let mut zero = StepRng::new(0, 0);
let rng = Core::from_rng(&mut zero).unwrap();
let mut rng1 = ReseedingRng::new(rng, 32 * 4, zero);
let first: u32 = rng1.gen();
for _ in 0..10 {
let _ = rng1.gen::<u32>();
}
let mut rng2 = rng1.clone();
assert_eq!(first, rng2.gen::<u32>());
}
}