tink-core: allow override of global Rng generator

Author: daviddrysdale

core/src/subtle/random.rs

@@ -19,7 +19,21 @@
 /// Re-export the particular version of the `rand` crate whose types appear in the API.
 pub use rand;
 
+use lazy_static::lazy_static;
 use rand::Rng;
+use std::sync::RwLock;
+
+/// Possible factory functions for producing [`Generator`] instances.
+pub enum Factory {
+    ThreadRng,
+    OsRng,
+    User(fn() -> Box<dyn Generator>),
+}
+
+lazy_static! {
+    /// Global factory that produces [`Generator`] instances.
+    static ref RNG_FACTORY: RwLock<fn() -> Box<dyn Generator>> = RwLock::new(thread_rng_factory);
+}
 
 /// Trait that encapsulates the required traits that a generator instance
 /// must implement.
@@ -28,9 +42,32 @@ pub trait Generator: rand::RngCore + rand::CryptoRng {}
 // Blanket implementation.
 impl<T> Generator for T where T: rand::RngCore + rand::CryptoRng {}
 
+/// Factory function that produces [`rand::rngs::ThreadRng`] instances.
+fn thread_rng_factory() -> Box<dyn Generator> {
+    // Available if `rand` has feature `std` enabled.
+    Box::new(rand::thread_rng())
+}
+
+/// Factory function that produces [`rand::rngs::OsRng`] instances.
+fn os_rng_factory() -> Box<dyn Generator> {
+    // Available if `rand` has feature `getrandom` enabled.
+    Box::new(rand::rngs::OsRng)
+}
+
+/// Set the global factory that produces [`Generator`] instances.
+pub fn set_factory(factory: Factory) {
+    let mut global = RNG_FACTORY.write().unwrap(); // safe: lock
+    *global = match factory {
+        Factory::ThreadRng => thread_rng_factory,
+        Factory::OsRng => os_rng_factory,
+        Factory::User(f) => f,
+    };
+}
+
 /// Return a random number generator suitable for cryptographic operation.
 pub fn rng() -> Box<dyn Generator> {
-    Box::new(rand::thread_rng())
+    let factory = RNG_FACTORY.read().unwrap(); // safe: lock
+    factory()
 }
 
 /// Return a vector of the given `size` filled with random bytes.

tests/tests/core/subtle/random_test.rs

@@ -30,3 +30,44 @@ fn test_random_uint() {
     let v2 = random::get_random_uint32();
     assert_ne!(v1, v2, "Just unlucky?");
 }
+
+#[test]
+fn test_generator_factory() {
+    random::set_factory(random::Factory::User(test_factory));
+    let mut rng = random::rng();
+    assert_eq!(rng.next_u32(), 4);
+    assert_eq!(rng.next_u32(), 4);
+    assert_eq!(rng.next_u32(), 4);
+    random::set_factory(random::Factory::ThreadRng);
+    let mut rng = random::rng();
+    assert_ne!(rng.next_u64(), rng.next_u64());
+}
+
+fn test_factory() -> Box<dyn random::Generator> {
+    Box::new(BogusGenerator {})
+}
+
+struct BogusGenerator;
+
+impl random::rand::RngCore for BogusGenerator {
+    fn next_u32(&mut self) -> u32 {
+        4 // chosen by fair dice roll. guaranteed to be random.
+    }
+
+    fn next_u64(&mut self) -> u64 {
+        self.next_u32() as u64
+    }
+
+    fn fill_bytes(&mut self, dest: &mut [u8]) {
+        for b in dest {
+            *b = 4;
+        }
+    }
+
+    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), random::rand::Error> {
+        self.fill_bytes(dest);
+        Ok(())
+    }
+}
+
+impl random::rand::CryptoRng for BogusGenerator {}