Add Unix fork protection

Author: pitdicker

src/rngs/adapter/reseeding.rs

@@ -16,6 +16,9 @@ use core::mem::size_of;
 use rand_core::{RngCore, CryptoRng, SeedableRng, Error, ErrorKind};
 use rand_core::block::{BlockRngCore, BlockRng};
 
+#[cfg(all(feature="std", unix))]
+extern crate libc;
+
 /// A wrapper around any PRNG which reseeds the underlying PRNG after it has
 /// generated a certain number of random bytes.
 ///
@@ -126,6 +129,7 @@ struct ReseedingCore<R, Rsdr> {
     reseeder: Rsdr,
     threshold: i64,
     bytes_until_reseed: i64,
+    fork_counter: u64,
 }
 
 impl<R, Rsdr> BlockRngCore for ReseedingCore<R, Rsdr>
@@ -136,8 +140,9 @@ where R: BlockRngCore + SeedableRng,
     type Results = <R as BlockRngCore>::Results;
 
     fn generate(&mut self, results: &mut Self::Results) {
-        if self.bytes_until_reseed <= 0 {
-            // We get better performance by not calling only `auto_reseed` here
+        if self.bytes_until_reseed <= 0 ||
+           self.fork_counter < get_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);
@@ -161,11 +166,14 @@ where R: BlockRngCore + SeedableRng,
     /// * `reseeder`: the RNG to use for reseeding.
     pub fn new(rng: R, threshold: u64, reseeder: Rsdr) -> Self {
         assert!(threshold <= ::core::i64::MAX as u64);
+        register_fork_handler();
+
         ReseedingCore {
             inner: rng,
             reseeder,
             threshold: threshold as i64,
             bytes_until_reseed: threshold as i64,
+            fork_counter: 0,
         }
     }
 
@@ -181,9 +189,15 @@ where R: BlockRngCore + SeedableRng,
     fn reseed_and_generate(&mut self,
                            results: &mut <Self as BlockRngCore>::Results)
     {
-        trace!("Reseeding RNG after {} generated bytes",
-               self.threshold - self.bytes_until_reseed);
-        let threshold = if let Err(e) = self.reseed()  {
+        let fork_counter = get_fork_counter();
+        if self.fork_counter < fork_counter {
+            warn!("Fork detected, reseeding RNG");
+        } else {
+            trace!("Reseeding RNG after {} generated bytes",
+                   self.threshold - self.bytes_until_reseed);
+        }
+
+        let threshold = if let Err(e) = self.reseed() {
             let delay = match e.kind {
                 ErrorKind::Transient => 0,
                 kind @ _ if kind.should_retry() => self.threshold >> 8,
@@ -193,11 +207,13 @@ where R: BlockRngCore + SeedableRng,
                     error from source: {}", delay, e);
             delay
         } else {
-            self.threshold
+            let num_bytes =
+                results.as_ref().len() * size_of::<<R as BlockRngCore>::Item>();
+            self.fork_counter = fork_counter;
+            self.threshold - num_bytes as i64
         };
-        
-        let num_bytes = results.as_ref().len() * size_of::<<R as BlockRngCore>::Item>();
-        self.bytes_until_reseed = threshold - num_bytes as i64;
+
+        self.bytes_until_reseed = threshold;
         self.inner.generate(results);
     }
 }
@@ -212,6 +228,7 @@ where R: BlockRngCore + SeedableRng + Clone,
             reseeder: self.reseeder.clone(),
             threshold: self.threshold,
             bytes_until_reseed: 0, // reseed clone on first use
+            fork_counter: self.fork_counter,
         }
     }
 }
@@ -220,6 +237,40 @@ impl<R, Rsdr> CryptoRng for ReseedingCore<R, Rsdr>
 where R: BlockRngCore + SeedableRng + CryptoRng,
       Rsdr: RngCore + CryptoRng {}
 
+
+// Fork protection
+//
+// We implement fork protection on Unix using `pthread_atfork`.
+// When the process is forked, we increment `RESEEDING_RNG_FORK_COUNTER`.
+// Every `ReseedingRng` stores the last known value of the static in
+// `fork_counter`. If the cached `fork_counter` is less than
+// `RESEEDING_RNG_FORK_COUNTER`, it is time to reseed this RNG.
+//
+// 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))]
+static mut RESEEDING_RNG_FORK_COUNTER: u64 = 0;
+
+#[cfg(all(feature="std", unix))]
+extern fn forkhandler() {
+    unsafe { RESEEDING_RNG_FORK_COUNTER += 1; }
+}
+
+#[cfg(all(feature="std", unix))]
+fn register_fork_handler() {
+    unsafe {
+        libc::pthread_atfork(None, None, Some(forkhandler));
+    }
+}
+#[cfg(not(all(feature="std", unix)))]
+fn register_fork_handler() {}
+
+#[cfg(all(feature="std", unix))]
+fn get_fork_counter() -> u64 { unsafe { RESEEDING_RNG_FORK_COUNTER } }
+#[cfg(not(all(feature="std", unix)))]
+fn get_fork_counter() -> u64 { 0 }
+
+
 #[cfg(test)]
 mod test {
     use {Rng, SeedableRng};