ReseedingRng: use trait Generator

dhardy · dhardy · commit 31d28f74dbeb · 2025-11-24T08:58:42.000Z
diff --git a/rand_chacha/src/chacha.rs b/rand_chacha/src/chacha.rs
@@ -10,7 +10,7 @@
 
 use crate::guts::ChaCha;
 use core::fmt;
-use rand_core::{le, CryptoRng, RngCore, SeedableRng, Generator, CryptoGenerator};
+use rand_core::{CryptoGenerator, CryptoRng, Generator, RngCore, SeedableRng, le};
 
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -185,8 +185,7 @@ macro_rules! chacha_impl {
             pub fn set_word_pos(&mut self, word_offset: u128) {
                 let block = (word_offset / u128::from(BLOCK_WORDS)) as u64;
                 self.core.state.set_block_pos(block);
-                self
-                    .generate_and_set((word_offset % u128::from(BLOCK_WORDS)) as usize);
+                self.generate_and_set((word_offset % u128::from(BLOCK_WORDS)) as usize);
             }
 
             /// Set the stream number.
diff --git a/src/rngs/reseeding.rs b/src/rngs/reseeding.rs
@@ -10,10 +10,15 @@
 //! A wrapper around another PRNG that reseeds it after it
 //! generates a certain number of random bytes.
 
+use core::fmt;
 use core::mem::size_of_val;
-use rand_core::{CryptoRng, RngCore, SeedableRng, TryCryptoRng, TryRngCore};
 
-/// A wrapper around any PRNG that implements [`RngCore`], that adds the
+use rand_core::le::{self, Word};
+use rand_core::{
+    CryptoGenerator, CryptoRng, Generator, RngCore, SeedableRng, TryCryptoRng, TryRngCore,
+};
+
+/// A wrapper around any PRNG that implements [`Generator`], that adds the
 /// ability to reseed it.
 ///
 /// `ReseedingRng` reseeds the underlying PRNG in the following cases:
@@ -50,34 +55,46 @@ use rand_core::{CryptoRng, RngCore, SeedableRng, TryCryptoRng, TryRngCore};
 /// # Example
 ///
 /// ```
-/// use chacha20::ChaCha20Rng; // Internal part of ChaChaRng that
-///                             // implements BlockRngCore
+/// use chacha20::ChaCha20Core; // Internal part of ChaChaRng that
+///                             // implements Generator
 /// use rand::prelude::*;
 /// use rand::rngs::OsRng;
 /// use rand::rngs::ReseedingRng;
 ///
-/// let mut reseeding_rng = ReseedingRng::<ChaCha20Rng, _>::new(0, OsRng).unwrap();
+/// let mut reseeding_rng = ReseedingRng::<ChaCha20Core, _>::new(0, OsRng).unwrap();
 ///
 /// println!("{}", reseeding_rng.random::<u64>());
 /// ```
 ///
+/// [`Generator`]: rand_core::Generator
 /// [`ReseedingRng::new`]: ReseedingRng::new
 /// [`reseed()`]: ReseedingRng::reseed
-#[derive(Debug)]
 pub struct ReseedingRng<R, Rsdr>
 where
-    R: RngCore + SeedableRng,
+    R: Generator + SeedableRng,
     Rsdr: TryRngCore,
 {
-    prng: R,
-    reseed_rng: Rsdr,
-    threshold: i64,
-    bytes_until_reseed: i64,
+    core: ReseedingCore<R, Rsdr>,
+    buffer: <R as Generator>::Result,
+}
+
+impl<R: fmt::Debug, Rsdr: fmt::Debug> fmt::Debug for ReseedingRng<R, Rsdr>
+where
+    R: Generator + SeedableRng,
+    R::Result: fmt::Debug,
+    Rsdr: TryRngCore,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ReseedingRng")
+            .field("core", &self.core)
+            .field("buffer", &self.buffer)
+            .finish()
+    }
 }
 
-impl<R, Rsdr> ReseedingRng<R, Rsdr>
+impl<W: Word, const N: usize, R, Rsdr> ReseedingRng<R, Rsdr>
 where
-    R: RngCore + SeedableRng,
+    R: Generator<Result = [W; N]> + SeedableRng,
     Rsdr: TryRngCore,
 {
     /// Create a new `ReseedingRng` from an existing PRNG, combined with a RNG
@@ -86,8 +103,93 @@ where
     /// `threshold` sets the number of generated bytes after which to reseed the
     /// PRNG. Set it to zero to never reseed based on the number of generated
     /// values.
+    pub fn new(threshold: u64, reseeder: Rsdr) -> Result<Self, Rsdr::Error> {
+        Ok(ReseedingRng {
+            core: ReseedingCore::new(threshold, reseeder)?,
+            buffer: le::new_buffer(),
+        })
+    }
+
+    /// Immediately reseed the generator
+    ///
+    /// This discards any remaining random data in the cache.
+    pub fn reseed(&mut self) -> Result<(), Rsdr::Error> {
+        self.buffer = le::new_buffer();
+        self.core.reseed()
+    }
+}
+
+impl<const N: usize, R, Rsdr> RngCore for ReseedingRng<R, Rsdr>
+where
+    R: Generator<Result = [u32; N]> + SeedableRng,
+    Rsdr: TryRngCore,
+{
+    #[inline]
+    fn next_u32(&mut self) -> u32 {
+        let Self { core, buffer } = self;
+        // NOTE: this fn always returns one word, thus cannot return u32 for any W: Word
+        le::next_word_via_gen_block(buffer, |block| core.generate(block))
+    }
+
+    #[inline]
+    fn next_u64(&mut self) -> u64 {
+        let Self { core, buffer } = self;
+        // NOTE: this fn is specific to u32 buffers
+        le::next_u64_via_gen_block(buffer, |block| core.generate(block))
+    }
+
     #[inline]
-    pub fn new(threshold: u64, mut reseed_rng: Rsdr) -> Result<Self, Rsdr::Error> {
+    fn fill_bytes(&mut self, dst: &mut [u8]) {
+        let Self { core, buffer } = self;
+        le::fill_bytes_via_gen_block(dst, buffer, |block| core.generate(block));
+    }
+}
+
+impl<const N: usize, R, Rsdr> CryptoRng for ReseedingRng<R, Rsdr>
+where
+    R: Generator<Result = [u32; N]> + SeedableRng + CryptoGenerator,
+    Rsdr: TryCryptoRng,
+{
+}
+
+#[derive(Debug)]
+struct ReseedingCore<R, Rsdr> {
+    inner: R,
+    reseeder: Rsdr,
+    threshold: i64,
+    bytes_until_reseed: i64,
+}
+
+impl<R, Rsdr> Generator for ReseedingCore<R, Rsdr>
+where
+    R: Generator + SeedableRng,
+    Rsdr: TryRngCore,
+{
+    type Result = <R as Generator>::Result;
+
+    fn generate(&mut self, results: &mut Self::Result) {
+        if self.bytes_until_reseed <= 0 {
+            // 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);
+        }
+        let num_bytes = size_of_val(results);
+        self.bytes_until_reseed -= num_bytes as i64;
+        self.inner.generate(results);
+    }
+}
+
+impl<R, Rsdr> ReseedingCore<R, Rsdr>
+where
+    R: Generator + SeedableRng,
+    Rsdr: TryRngCore,
+{
+    /// Create a new `ReseedingCore`.
+    ///
+    /// `threshold` is the maximum number of bytes produced by
+    /// [`Generator::generate`] before attempting reseeding.
+    fn new(threshold: u64, mut reseeder: Rsdr) -> Result<Self, Rsdr::Error> {
         // Because generating more values than `i64::MAX` takes centuries on
         // current hardware, we just clamp to that value.
         // Also we set a threshold of 0, which indicates no limit, to that
@@ -100,68 +202,43 @@ where
             i64::MAX
         };
 
-        R::try_from_rng(&mut reseed_rng).map(|prng| Self {
-            prng,
-            reseed_rng,
+        let inner = R::try_from_rng(&mut reseeder)?;
+
+        Ok(ReseedingCore {
+            inner,
+            reseeder,
             threshold,
             bytes_until_reseed: threshold,
         })
     }
 
-    /// Immediately reseed the generator.
-    ///
-    /// This discards any remaining random data in the cache.
-    #[inline]
-    pub fn reseed(&mut self) -> Result<(), Rsdr::Error> {
-        self.prng = R::try_from_rng(&mut self.reseed_rng)?;
-        Ok(())
-    }
-
-    #[inline]
-    fn reseed_check(&mut self, bytes: usize) {
-        if self.bytes_until_reseed <= 0 {
-            let res = self.reseed();
-            if let Err(_e) = res {
-                warn!("Reseeding RNG failed: {_e}");
-            }
+    /// Reseed the internal PRNG.
+    fn reseed(&mut self) -> Result<(), Rsdr::Error> {
+        R::try_from_rng(&mut self.reseeder).map(|result| {
             self.bytes_until_reseed = self.threshold;
-        }
-        self.bytes_until_reseed -= bytes as i64;
+            self.inner = result
+        })
     }
-}
 
-impl<R, Rsdr> RngCore for ReseedingRng<R, Rsdr>
-where
-    R: RngCore + SeedableRng,
-    Rsdr: TryRngCore,
-{
-    #[inline(always)]
-    fn next_u32(&mut self) -> u32 {
-        self.reseed_check(size_of::<u32>());
-        self.prng.next_u32()
-    }
+    #[inline(never)]
+    fn reseed_and_generate(&mut self, results: &mut <Self as Generator>::Result) {
+        trace!("Reseeding RNG (periodic reseed)");
 
-    #[inline(always)]
-    fn next_u64(&mut self) -> u64 {
-        self.reseed_check(size_of::<u64>());
-        self.prng.next_u64()
-    }
+        let num_bytes = size_of_val(results);
 
-    #[inline(always)]
-    fn fill_bytes(&mut self, dst: &mut [u8]) {
-        // Do not perform `reseed_check` for an empty `dst` to allow the compiler
-        // to infer that we always generate PRNG data after reseeding.
-        if dst.is_empty() {
-            return;
+        if let Err(e) = self.reseed() {
+            warn!("Reseeding RNG failed: {}", e);
+            let _ = e;
         }
-        self.reseed_check(size_of_val(dst));
-        self.prng.fill_bytes(dst)
+
+        self.bytes_until_reseed = self.threshold - num_bytes as i64;
+        self.inner.generate(results);
     }
 }
 
-impl<R, Rsdr> CryptoRng for ReseedingRng<R, Rsdr>
+impl<R, Rsdr> CryptoGenerator for ReseedingCore<R, Rsdr>
 where
-    R: CryptoRng + SeedableRng,
+    R: Generator + SeedableRng + CryptoGenerator,
     Rsdr: TryCryptoRng,
 {
 }
@@ -170,7 +247,7 @@ where
 #[cfg(test)]
 mod test {
     use crate::Rng;
-    use crate::rngs::std::StdRng;
+    use crate::rngs::std::Core;
     use crate::test::const_rng;
 
     use super::ReseedingRng;
@@ -179,7 +256,7 @@ mod test {
     fn test_reseeding() {
         let zero = const_rng(0);
         let thresh = 1; // reseed every time the buffer is exhausted
-        let mut reseeding = ReseedingRng::<StdRng, _>::new(thresh, zero).unwrap();
+        let mut reseeding = ReseedingRng::<Core, _>::new(thresh, zero).unwrap();
 
         // RNG buffer size is [u32; 64]
         // Debug is only implemented up to length 32 so use two arrays
diff --git a/src/rngs/std.rs b/src/rngs/std.rs
@@ -10,6 +10,9 @@
 
 use rand_core::{CryptoRng, RngCore, SeedableRng};
 
+#[cfg(any(test, feature = "os_rng"))]
+pub(crate) use chacha20::ChaCha12Core as Core;
+
 use chacha20::ChaCha12Rng as Rng;
 
 /// A strong, fast (amortized), non-portable RNG
diff --git a/src/rngs/thread.rs b/src/rngs/thread.rs
@@ -13,7 +13,7 @@ use std::fmt;
 use std::rc::Rc;
 use std::thread_local;
 
-use super::{OsError, OsRng, ReseedingRng, std::StdRng};
+use super::{OsError, OsRng, ReseedingRng, std::Core};
 use rand_core::{CryptoRng, RngCore};
 
 // Rationale for using `UnsafeCell` in `ThreadRng`:
@@ -90,7 +90,7 @@ const THREAD_RNG_RESEED_THRESHOLD: u64 = 1024 * 64;
 #[derive(Clone)]
 pub struct ThreadRng {
     // Rc is explicitly !Send and !Sync
-    rng: Rc<UnsafeCell<ReseedingRng<StdRng, OsRng>>>,
+    rng: Rc<UnsafeCell<ReseedingRng<Core, OsRng>>>,
 }
 
 impl ThreadRng {
@@ -115,7 +115,7 @@ impl fmt::Debug for ThreadRng {
 thread_local!(
     // We require Rc<..> to avoid premature freeing when ThreadRng is used
     // within thread-local destructors. See #968.
-    static THREAD_RNG_KEY: Rc<UnsafeCell<ReseedingRng<StdRng, OsRng>>> = {
+    static THREAD_RNG_KEY: Rc<UnsafeCell<ReseedingRng<Core, OsRng>>> = {
         let rng = ReseedingRng::new(THREAD_RNG_RESEED_THRESHOLD,
                                     OsRng).unwrap_or_else(|err|
                 panic!("could not initialize ThreadRng: {}", err));
