From 9817b4ffc19c2e92cb646975521966b2fda10ab4 Mon Sep 17 00:00:00 2001
From: Diggory Hardy <git@dhardy.name>
Date: Sun, 20 Oct 2019 12:06:54 +0100
Subject: [PATCH] Add rand_core::RngCore implementation to chacha20

---
 chacha20/Cargo.toml    |  2 ++
 chacha20/src/cipher.rs |  2 +-
 chacha20/src/lib.rs    |  8 +++++
 chacha20/src/rng.rs    | 71 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 chacha20/src/rng.rs

diff --git a/chacha20/Cargo.toml b/chacha20/Cargo.toml
index a30ac950..e7df3ee9 100644
--- a/chacha20/Cargo.toml
+++ b/chacha20/Cargo.toml
@@ -19,6 +19,7 @@ travis-ci = { repository = "RustCrypto/stream-ciphers" }
 byteorder = { version = "1", default-features = false }
 stream-cipher = "0.3"
 salsa20-core = { version = "0.2", path = "../salsa20-core" }
+rand_core = { version = "0.5", optional = true }
 
 [dev-dependencies]
 stream-cipher = { version = "0.3", features = ["dev"] }
@@ -30,6 +31,7 @@ default = ["xchacha20"]
 legacy = []
 xchacha20 = []
 zeroize = ["salsa20-core/zeroize"]
+rng = ["rand_core"]
 
 [[bench]]
 name = "chacha20"
diff --git a/chacha20/src/cipher.rs b/chacha20/src/cipher.rs
index 131efd43..58bf241f 100644
--- a/chacha20/src/cipher.rs
+++ b/chacha20/src/cipher.rs
@@ -4,7 +4,7 @@ use byteorder::{ByteOrder, LE};
 use salsa20_core::{SalsaFamilyCipher, IV_WORDS, KEY_WORDS, STATE_WORDS};
 
 /// ChaCha20 core cipher functionality
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) struct Cipher {
     /// Secret key
     key: [u32; KEY_WORDS],
diff --git a/chacha20/src/lib.rs b/chacha20/src/lib.rs
index 266fa95e..a624f57d 100644
--- a/chacha20/src/lib.rs
+++ b/chacha20/src/lib.rs
@@ -60,6 +60,8 @@ extern crate salsa20_core;
 // TODO: replace with `u32::from_le_bytes`/`to_le_bytes` in libcore (1.32+)
 #[cfg(feature = "xchacha20")]
 extern crate byteorder;
+#[cfg(feature = "rand_core")]
+extern crate rand_core;
 
 mod block;
 pub(crate) mod cipher;
@@ -68,6 +70,9 @@ mod legacy;
 #[cfg(feature = "xchacha20")]
 mod xchacha20;
 
+#[cfg(feature = "rng")]
+mod rng;
+
 use self::cipher::Cipher;
 use salsa20_core::Ctr;
 use stream_cipher::generic_array::{
@@ -81,6 +86,9 @@ pub use self::legacy::ChaCha20Legacy;
 #[cfg(feature = "xchacha20")]
 pub use self::xchacha20::XChaCha20;
 
+#[cfg(feature = "rng")]
+pub use rng::{ChaCha20Rng, ChaCha20RngCore};
+
 /// Maximum number of blocks that can be encrypted with ChaCha20 before the
 /// counter overflows.
 pub const MAX_BLOCKS: usize = core::u32::MAX as usize;
diff --git a/chacha20/src/rng.rs b/chacha20/src/rng.rs
new file mode 100644
index 00000000..8e2129fa
--- /dev/null
+++ b/chacha20/src/rng.rs
@@ -0,0 +1,71 @@
+//! Block RNG based on rand_core::BlockRng
+
+use rand_core::{RngCore, SeedableRng, Error};
+use rand_core::block::{BlockRng, BlockRngCore};
+use salsa20_core::{SalsaFamilyCipher, IV_BYTES, KEY_BYTES, STATE_WORDS};
+
+use crate::cipher::Cipher;
+
+/// Random number generator over the ChaCha20 stream cipher.
+#[derive(Clone, Debug)]
+pub struct ChaCha20Rng(BlockRng<ChaCha20RngCore>);
+
+impl SeedableRng for ChaCha20Rng {
+    type Seed = [u8; KEY_BYTES];
+    #[inline]
+    fn from_seed(seed: Self::Seed) -> Self {
+        let core = ChaCha20RngCore::from_seed(seed);
+        Self(BlockRng::new(core))
+    }
+}
+
+impl RngCore for ChaCha20Rng {
+    #[inline]
+    fn next_u32(&mut self) -> u32 {
+        self.0.next_u32()
+    }
+    #[inline]
+    fn next_u64(&mut self) -> u64 {
+        self.0.next_u64()
+    }
+    #[inline]
+    fn fill_bytes(&mut self, bytes: &mut [u8]) {
+        self.0.fill_bytes(bytes)
+    }
+    #[inline]
+    fn try_fill_bytes(&mut self, bytes: &mut [u8]) -> Result<(), Error> {
+        self.0.try_fill_bytes(bytes)
+    }
+}
+
+
+/// Core of the [`ChaCha20Rng`] random number generator, for use with
+/// [`rand_core::block::BlockRng`].
+#[derive(Clone, Debug)]
+pub struct ChaCha20RngCore {
+    cipher: Cipher,
+    counter: u64,
+}
+
+impl SeedableRng for ChaCha20RngCore {
+    type Seed = [u8; KEY_BYTES];
+    #[inline]
+    fn from_seed(seed: Self::Seed) -> Self {
+        let iv = [0; IV_BYTES];
+        let cipher = Cipher::new(&seed, &iv, 0);
+        Self {
+            cipher,
+            counter: 0,
+        }
+    }
+}
+
+impl BlockRngCore for ChaCha20RngCore {
+    type Item = u32;
+    type Results = [u32; STATE_WORDS];
+    fn generate(&mut self, results: &mut Self::Results) {
+        let out = self.cipher.block(self.counter);
+        self.counter += 1;
+        *results = out;
+    }
+}