ThreadRng / EntropyRng improvements

src/rngs/entropy.rs

@@ -8,6 +8,10 @@
 
 //! Entropy generator, or wrapper around external generators
 
+use core::fmt;
+use std::sync::{Once, Mutex, MutexGuard, ONCE_INIT};
+use std::marker::PhantomData;
+
 use rand_core::{RngCore, CryptoRng, Error, ErrorKind, impls};
 #[allow(unused)]
 use rngs;
@@ -28,6 +32,8 @@ use rngs;
 /// jitter); for better performance it is common to seed a local PRNG from
 /// external entropy then primarily use the local PRNG ([`thread_rng`] is
 /// provided as a convenient, local, automatically-seeded CSPRNG).
+/// 
+/// `EntropyRng` instances are explicitly not `Send` or `Sync`.
 ///
 /// # Panics
 ///
@@ -46,6 +52,7 @@ use rngs;
 #[derive(Debug)]
 pub struct EntropyRng {
     source: Source,
+    _not_send: PhantomData<*mut ()>,    // enforce !Send
 }
 
 #[derive(Debug)]
@@ -63,7 +70,7 @@ impl EntropyRng {
     /// those are done on first use. This is done to make `new` infallible,
     /// and `try_fill_bytes` the only place to report errors.
     pub fn new() -> Self {
-        EntropyRng { source: Source::None }
+        EntropyRng { source: Source::None, _not_send: Default::default() }
     }
 }
 
@@ -89,65 +96,37 @@ impl RngCore for EntropyRng {
 
     fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
         let mut reported_error = None;
-
-        if let Source::Os(ref mut os_rng) = self.source {
-            match os_rng.fill(dest) {
-                Ok(()) => return Ok(()),
-                Err(err) => {
-                    warn!("EntropyRng: OsRng failed \
-                          [trying other entropy sources]: {}", err);
-                    reported_error = Some(err);
-                },
-            }
-        } else if Os::is_supported() {
-            match Os::new_and_fill(dest) {
-                Ok(os_rng) => {
-                    debug!("EntropyRng: using OsRng");
-                    self.source = Source::Os(os_rng);
-                    return Ok(());
-                },
-                Err(err) => { reported_error = reported_error.or(Some(err)) },
+
+        macro_rules! try_source {
+            ($source:ident) => {
+                if let Source::$source(ref mut inner) = self.source {
+                    match inner.fill(dest) {
+                        Ok(()) => return Ok(()),
+                        Err(err) => {
+                            warn!("EntropyRng: source {} failed: {}", $source::name(), err);
+                            reported_error = Some(err);
+                        },
+                    }
+                } else if $source::is_available() {
+                    match $source::new().and_then(|mut inner|
+                        inner.fill(dest).and(Ok(inner)))
+                    {
+                        Ok(inner) => {
+                            debug!("EntropyRng: using source {}", $source::name());
+                            self.source = Source::$source(inner);
+                            return Ok(());
+                        },
+                        Err(err) => {
+                            reported_error = reported_error.or(Some(err));
+                        },
+                    }
+                }
             }
         }
 
-        if let Source::Custom(ref mut rng) = self.source {
-            match rng.fill(dest) {
-                Ok(()) => return Ok(()),
-                Err(err) => {
-                    warn!("EntropyRng: custom entropy source failed \
-                          [trying other entropy sources]: {}", err);
-                    reported_error = Some(err);
-                },
-            }
-        } else if Custom::is_supported() {
-            match Custom::new_and_fill(dest) {
-                Ok(custom) => {
-                    debug!("EntropyRng: using custom entropy source");
-                    self.source = Source::Custom(custom);
-                    return Ok(());
-                },
-                Err(err) => { reported_error = reported_error.or(Some(err)) },
-            }
-        }
-
-        if let Source::Jitter(ref mut jitter_rng) = self.source {
-            match jitter_rng.fill(dest) {
-                Ok(()) => return Ok(()),
-                Err(err) => {
-                    warn!("EntropyRng: JitterRng failed: {}", err);
-                    reported_error = Some(err);
-                },
-            }
-        } else if Jitter::is_supported() {
-            match Jitter::new_and_fill(dest) {
-                Ok(jitter_rng) => {
-                    debug!("EntropyRng: using JitterRng");
-                    self.source = Source::Jitter(jitter_rng);
-                    return Ok(());
-                },
-                Err(err) => { reported_error = reported_error.or(Some(err)) },
-            }
-        }
+        try_source!(Os);
+        try_source!(Custom);
+        try_source!(Jitter);
 
         if let Some(err) = reported_error {
             Err(Error::with_cause(ErrorKind::Unavailable,
@@ -165,12 +144,17 @@ impl CryptoRng for EntropyRng {}
 
 
 trait EntropySource {
-    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error>
-        where Self: Sized;
-
+    /// Name of this source
+    fn name() -> &'static str;
+
+    /// Is this source available?
+    fn is_available() -> bool;
+
+    /// Create an instance
+    fn new() -> Result<Self, Error> where Self: Sized;
+
+    /// Fill `dest` with random data from the entropy source
     fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error>;
-
-    fn is_supported() -> bool { true }
 }
 
 #[allow(unused)]
@@ -179,15 +163,16 @@ struct NoSource;
 
 #[allow(unused)]
 impl EntropySource for NoSource {
-    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
+    fn name() -> &'static str { unreachable!() }
+    fn is_available() -> bool { false }
+
+    fn new() -> Result<Self, Error> {
         Err(Error::new(ErrorKind::Unavailable, "Source not supported"))
     }
 
     fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
         unreachable!()
     }
-
-    fn is_supported() -> bool { false }
 }
 
 
@@ -229,10 +214,12 @@ pub struct Os(rngs::OsRng);
               all(target_arch = "wasm32", feature = "wasm-bindgen"),
 )))]
 impl EntropySource for Os {
-    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
-        let mut rng = rngs::OsRng::new()?;
-        rng.try_fill_bytes(dest)?;
-        Ok(Os(rng))
+    fn name() -> &'static str { "OsRng" }
+
+    fn is_available() -> bool { true }
+
+    fn new() -> Result<Self, Error> {
+        Ok(Os(rngs::OsRng::new()?))
     }
 
     fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
@@ -259,8 +246,122 @@ impl EntropySource for Os {
 ))))]
 type Os = NoSource;
 
+#[derive(Clone)]
+struct Custom {
+    source: &'static CustomEntropySource,
+    param: u64,
+    _not_send: PhantomData<*mut ()>,    // enforce !Send
+}
+
+impl fmt::Debug for Custom {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Custom {{ ... }}")
+    }
+}
+
+#[test]
+fn test_size() {
+    use core::mem::size_of;
+    println!("Size of Custom: {}", size_of::<Custom>());
+    assert!(size_of::<Custom>() <= size_of::<rngs::JitterRng>());
+}
 
-type Custom = NoSource;
+/// Properties of a custom entropy source.
+pub trait CustomEntropySource {
+    /// Name of this source
+    fn name(&self) -> &'static str;
+
+    /// Is this source available?
+    /// 
+    /// The default implementation returns `true`.
+    fn is_available(&self) -> bool { true }
+
+    /// Prepare the entropy source for use.
+    /// 
+    /// This is always called before `fill` on each thread. This may be called
+    /// multiple times.
+    /// 
+    /// A `u64` parameter may be returned, which is passed to `fill` when
+    /// called, and is considered `!Send` (i.e. is never passed to a different
+    /// thread).
+    fn init(&self) -> Result<u64, Error>;
+
+    /// Fill `dest` with random data from the entropy source.
+    /// 
+    /// The `u64` parameter from `init` is passed.
+    fn fill(&self, param: &mut u64, dest: &mut [u8]) -> Result<(), Error>;
+}
+
+struct CustomNoSource;
+impl CustomEntropySource for CustomNoSource {
+    fn name(&self) -> &'static str {
+        "no source"
+    }
+
+    fn is_available(&self) -> bool { false }
+
+    fn init(&self) -> Result<u64, Error> {
+        unreachable!()
+    }
+
+    fn fill(&self, _: &mut u64, _: &mut [u8]) -> Result<(), Error> {
+        unreachable!()
+    }
+}
+
+// TODO: remove outer Option when `Mutex::new(&...)` is a constant expression
+static mut CUSTOM_SOURCE: Option<Mutex<&CustomEntropySource>> = None;
+static CUSTOM_SOURCE_ONCE: Once = ONCE_INIT;
+
+fn access_custom_entropy() -> MutexGuard<'static, &'static CustomEntropySource> {
+    CUSTOM_SOURCE_ONCE.call_once(|| {
+        unsafe { CUSTOM_SOURCE = Some(Mutex::new(&CustomNoSource)) }
+    });
+    let mutex = unsafe { CUSTOM_SOURCE.as_ref().unwrap() };
+    mutex.lock().unwrap()
+}
+
+fn get_custom_entropy() -> &'static CustomEntropySource {
+    *access_custom_entropy()
+}
+
+/// Specify a custom entropy source.
+/// 
+/// This must be a static reference to an object implementing the
+/// `CustomEntropySource` trait.
+pub fn set_custom_entropy(source: &'static CustomEntropySource) {
+    let mut guard = access_custom_entropy();
+    *guard = source;
+}
+
+impl EntropySource for Custom {
+    fn name() -> &'static str {
+        get_custom_entropy().name()
+    }
+
+    /// Is this source available?
+    /// 
+    /// The default implementation returns `true`.
+    fn is_available() -> bool {
+        get_custom_entropy().is_available()
+    }
+
+    /// Create an instance
+    fn new() -> Result<Self, Error> where Self: Sized {
+        let source = get_custom_entropy();
+        let param = source.init()?;
+        Ok(Custom {
+            source,
+            param,
+            _not_send: Default::default(),
+        })
+    }
+
+    /// Fill `dest` with random data from the entropy source
+    fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
+        self.source.fill(&mut self.param, dest)
+    }
+}
 
 
 #[cfg(not(target_arch = "wasm32"))]
@@ -269,10 +370,12 @@ pub struct Jitter(rngs::JitterRng);
 
 #[cfg(not(target_arch = "wasm32"))]
 impl EntropySource for Jitter {
-    fn new_and_fill(dest: &mut [u8]) -> Result<Self, Error> {
-        let mut rng = rngs::JitterRng::new()?;
-        rng.try_fill_bytes(dest)?;
-        Ok(Jitter(rng))
+    fn name() -> &'static str { "JitterRng" }
+
+    fn is_available() -> bool { true }
+
+    fn new() -> Result<Self, Error> {
+        Ok(Jitter(rngs::JitterRng::new()?))
     }
 
     fn fill(&mut self, dest: &mut [u8]) -> Result<(), Error> {
@@ -294,4 +397,21 @@ mod test {
         let n = (rng.next_u32() ^ rng.next_u32()).count_ones();
         assert!(n >= 2);    // p(failure) approx 1e-7
     }
+
+    #[test]
+    fn test_custom_entropy() {
+        struct FakeEntropy;
+        impl CustomEntropySource for FakeEntropy {
+            fn name(&self) -> &'static str { "fake entropy" }
+            fn init(&self) -> Result<u64, Error> { Ok(0) }
+            fn fill(&self, _: &mut u64, dest: &mut [u8]) -> Result<(), Error> {
+                for x in dest { *x = 0 }
+                Ok(())
+            }
+        }
+        set_custom_entropy(&FakeEntropy);
+        let mut entropy = EntropyRng::new();
+        // we can't properly test this because we can't disable `OsRng`
+        assert!(entropy.next_u64() != 1);
+    }
 }

src/rngs/mod.rs

@@ -175,7 +175,7 @@ mod std;
 
 
 pub use self::jitter::{JitterRng, TimerError};
-#[cfg(feature="std")] pub use self::entropy::EntropyRng;
+#[cfg(feature="std")] pub use self::entropy::{EntropyRng, CustomEntropySource, set_custom_entropy};
 
 pub use self::small::SmallRng;
 pub use self::std::StdRng;