Use atomic fork counter

src/lib.rs

@@ -251,7 +251,7 @@ extern crate rand_core;
 #[cfg(feature = "log")] #[macro_use] extern crate log;
 #[cfg(not(feature = "log"))] macro_rules! trace { ($($x:tt)*) => () }
 #[cfg(not(feature = "log"))] macro_rules! debug { ($($x:tt)*) => () }
-#[cfg(all(feature="std", not(feature = "log")))] macro_rules! info { ($($x:tt)*) => () }
+#[cfg(not(feature = "log"))] macro_rules! info { ($($x:tt)*) => () }
 #[cfg(not(feature = "log"))] macro_rules! warn { ($($x:tt)*) => () }
 #[cfg(all(feature="std", not(feature = "log")))] macro_rules! error { ($($x:tt)*) => () }
 

src/rngs/adapter/reseeding.rs

@@ -19,6 +19,11 @@ use rand_core::block::{BlockRngCore, BlockRng};
 #[cfg(all(feature="std", unix, not(target_os="emscripten")))]
 extern crate libc;
 
+#[cfg(all(feature="std", unix, not(target_os="emscripten")))]
+use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT,
+                        AtomicBool, ATOMIC_BOOL_INIT, Ordering};
+
+
 /// A wrapper around any PRNG which reseeds the underlying PRNG after it has
 /// generated a certain number of random bytes.
 ///
@@ -129,7 +134,7 @@ struct ReseedingCore<R, Rsdr> {
     reseeder: Rsdr,
     threshold: i64,
     bytes_until_reseed: i64,
-    fork_counter: u64,
+    fork_counter: usize,
 }
 
 impl<R, Rsdr> BlockRngCore for ReseedingCore<R, Rsdr>
@@ -140,12 +145,13 @@ where R: BlockRngCore + SeedableRng,
     type Results = <R as BlockRngCore>::Results;
 
     fn generate(&mut self, results: &mut Self::Results) {
+        let global_fork_counter = get_fork_counter();
         if self.bytes_until_reseed <= 0 ||
-           self.fork_counter < get_fork_counter() {
+           self.is_forked(global_fork_counter) {
             // We get better performance by not calling only `reseed` here
             // and continuing with the rest of the function, but by directly
             // returning from a non-inlined function.
-            return self.reseed_and_generate(results);
+            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;
@@ -185,35 +191,52 @@ where R: BlockRngCore + SeedableRng,
         })
     }
 
+    fn is_forked(&self, global_fork_counter: usize) -> bool {
+        // In theory, on 32-bit platforms, it is possible for
+        // `global_fork_counter` to wrap around after ~4e9 forks.
+        //
+        // This check will detect a fork in the normal case where
+        // `fork_counter < global_fork_counter`, and also when the difference
+        // between both is greater than `isize::MAX` (wrapped around).
+        //
+        // It will still fail to detect a fork if there have been more than
+        // `isize::MAX` forks, without any reseed in between. Seems unlikely
+        // enough.
+        (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)
+                           results: &mut <Self as BlockRngCore>::Results,
+                           global_fork_counter: usize)
     {
-        let fork_counter = get_fork_counter();
-        if self.fork_counter < fork_counter {
-            warn!("Fork detected, reseeding RNG");
+        if self.is_forked(global_fork_counter) {
+            info!("Fork detected, reseeding RNG");
         } else {
+            // FIXME: the amount logged here may not be correct if the intial
+            // `bytes_until_reseed` was set by a delay instead of `threshold`.
             trace!("Reseeding RNG after {} generated bytes",
                    self.threshold - self.bytes_until_reseed);
         }
 
+        let num_bytes =
+            results.as_ref().len() * size_of::<<R as BlockRngCore>::Item>();
+
         let threshold = if let Err(e) = self.reseed() {
             let delay = match e.kind {
-                ErrorKind::Transient => 0,
+                ErrorKind::Transient => num_bytes as i64,
                 kind @ _ if kind.should_retry() => self.threshold >> 8,
                 _ => self.threshold,
             };
             warn!("Reseeding RNG delayed reseeding by {} bytes due to \
-                    error from source: {}", delay, e);
+                   error from source: {}", delay, e);
             delay
         } else {
-            let num_bytes =
-                results.as_ref().len() * size_of::<<R as BlockRngCore>::Item>();
-            self.fork_counter = fork_counter;
-            self.threshold - num_bytes as i64
+            self.fork_counter = global_fork_counter;
+            self.threshold
         };
 
-        self.bytes_until_reseed = threshold;
+        self.bytes_until_reseed = threshold - num_bytes as i64;
         self.inner.generate(results);
     }
 }
@@ -249,30 +272,35 @@ where R: BlockRngCore + SeedableRng + CryptoRng,
 // If reseeding fails, we don't deal with this by setting a delay, but just
 // don't update `fork_counter`, so a reseed is attempted a soon as possible.
 #[cfg(all(feature="std", unix, not(target_os="emscripten")))]
-static mut RESEEDING_RNG_FORK_COUNTER: u64 = ::core::u64::MAX;
+static RESEEDING_RNG_FORK_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
 
 #[cfg(all(feature="std", unix, not(target_os="emscripten")))]
-extern fn forkhandler() {
-    unsafe { RESEEDING_RNG_FORK_COUNTER += 1; }
+fn get_fork_counter() -> usize {
+    RESEEDING_RNG_FORK_COUNTER.load(Ordering::Relaxed)
+}
+#[cfg(not(all(feature="std", unix, not(target_os="emscripten"))))]
+fn get_fork_counter() -> usize { 0 }
+
+
+#[cfg(all(feature="std", unix, not(target_os="emscripten")))]
+static FORK_HANDLER_REGISTERED: AtomicBool = ATOMIC_BOOL_INIT;
+
+#[cfg(all(feature="std", unix, not(target_os="emscripten")))]
+extern fn fork_handler() {
+    // Note: fetch_add is defined to wrap on overflow (which is what we want)
+    RESEEDING_RNG_FORK_COUNTER.fetch_add(1, Ordering::Relaxed);
 }
 
 #[cfg(all(feature="std", unix, not(target_os="emscripten")))]
 fn register_fork_handler() {
-    if get_fork_counter() == ::core::u64::MAX {
-        unsafe {
-            RESEEDING_RNG_FORK_COUNTER = 0;
-            libc::pthread_atfork(None, None, Some(forkhandler));
-        }
+    if FORK_HANDLER_REGISTERED.load(Ordering::Relaxed) == false {
+        libc::pthread_atfork(None, None, Some(fork_handler));
+        FORK_HANDLER_REGISTERED.store(true, Ordering::Relaxed);
     }
 }
 #[cfg(not(all(feature="std", unix, not(target_os="emscripten"))))]
 fn register_fork_handler() {}
 
-#[cfg(all(feature="std", unix, not(target_os="emscripten")))]
-fn get_fork_counter() -> u64 { unsafe { RESEEDING_RNG_FORK_COUNTER } }
-#[cfg(not(all(feature="std", unix, not(target_os="emscripten"))))]
-fn get_fork_counter() -> u64 { 0 }
-
 
 #[cfg(test)]
 mod test {