From 42e004a23c1cb64d8a8bf386c46e610ef42cbd30 Mon Sep 17 00:00:00 2001
From: Ericson Soares <ericson.ds999@gmail.com>
Date: Thu, 15 Aug 2024 00:46:40 -0300
Subject: [PATCH] More tweaks on crypto crate

---
 crates/crypto/src/cloud/secret_key.rs |  8 ++++----
 crates/crypto/src/erase.rs            |  4 ++--
 crates/crypto/src/error.rs            |  6 +++++-
 crates/crypto/src/lib.rs              |  2 ++
 crates/crypto/src/rng/csprng.rs       | 26 ++++++++++++++++----------
 5 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/crates/crypto/src/cloud/secret_key.rs b/crates/crypto/src/cloud/secret_key.rs
index dcef9e90e..9fbf78b21 100644
--- a/crates/crypto/src/cloud/secret_key.rs
+++ b/crates/crypto/src/cloud/secret_key.rs
@@ -117,7 +117,7 @@ mod tests {
 	#[test]
 	fn one_shot_test() {
 		use super::super::{decrypt::OneShotDecryption, encrypt::OneShotEncryption};
-		let mut rng = CryptoRng::new();
+		let mut rng = CryptoRng::new().unwrap();
 
 		let message = b"Eu queria um apartamento no Guarujah; \
 		Mas o melhor que eu consegui foi um barraco em Itaquah.";
@@ -161,12 +161,12 @@ mod tests {
 		E como aquele ditado que jah dizia; \
 		Pau que nasce torto mija fora da bacia";
 
-		stream_test(&mut CryptoRng::new(), message).await;
+		stream_test(&mut CryptoRng::new().unwrap(), message).await;
 	}
 
 	#[tokio::test]
 	async fn stream_test_big() {
-		let mut rng = CryptoRng::new();
+		let mut rng = CryptoRng::new().unwrap();
 
 		let mut message =
 			vec![0u8; EncryptedBlock::PLAIN_TEXT_SIZE * 10 + EncryptedBlock::PLAIN_TEXT_SIZE / 2];
@@ -178,7 +178,7 @@ mod tests {
 
 	#[tokio::test]
 	async fn stream_test_big_exact() {
-		let mut rng = CryptoRng::new();
+		let mut rng = CryptoRng::new().unwrap();
 
 		let mut message = vec![0u8; EncryptedBlock::PLAIN_TEXT_SIZE * 20];
 
diff --git a/crates/crypto/src/erase.rs b/crates/crypto/src/erase.rs
index d13b7c086..f9a327f61 100644
--- a/crates/crypto/src/erase.rs
+++ b/crates/crypto/src/erase.rs
@@ -26,7 +26,7 @@ pub async fn erase<RW>(stream: &mut RW, size: usize, passes: usize) -> Result<us
 where
 	RW: AsyncReadExt + AsyncWriteExt + AsyncSeekExt + Unpin + Send,
 {
-	let mut rng = CryptoRng::new();
+	let mut rng = CryptoRng::new()?;
 
 	let mut buf = vec![0u8; BLOCK_LEN].into_boxed_slice();
 	let mut end_buf = vec![0u8; size % BLOCK_LEN].into_boxed_slice();
@@ -87,7 +87,7 @@ pub fn erase_sync<RW>(stream: &mut RW, size: usize, passes: usize) -> Result<usi
 where
 	RW: Read + Write + Seek,
 {
-	let mut rng = CryptoRng::new();
+	let mut rng = CryptoRng::new()?;
 
 	let mut buf = vec![0u8; BLOCK_LEN].into_boxed_slice();
 	let mut end_buf = vec![0u8; size % BLOCK_LEN].into_boxed_slice();
diff --git a/crates/crypto/src/error.rs b/crates/crypto/src/error.rs
index efeb5c2ca..f7444b20f 100644
--- a/crates/crypto/src/error.rs
+++ b/crates/crypto/src/error.rs
@@ -5,7 +5,6 @@ use tokio::io;
 /// This enum defines all possible errors that this crate can give
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
-	// crypto errors
 	#[error("Block too big for oneshot encryption: size in bytes = {0}")]
 	BlockTooBig(usize),
 
@@ -16,6 +15,7 @@ pub enum Error {
 	#[error("Decryption error")]
 	Decrypt,
 
+	/// I/O error while encrypting
 	#[error("I/O error while encrypting: {{context: {context}, source: {source}}}")]
 	EncryptIo {
 		context: &'static str,
@@ -29,6 +29,7 @@ pub enum Error {
 		source: io::Error,
 	},
 
+	/// I/O error while erasing a file
 	#[error("I/O error while erasing: {{context: {context}, source: {source}}}")]
 	EraseIo {
 		context: &'static str,
@@ -38,4 +39,7 @@ pub enum Error {
 
 	#[error("hex error: {0}")]
 	Hex(#[from] hex::FromHexError),
+
+	#[error("Entropy source error: {0}")]
+	EntropySource(#[from] rand_core::getrandom::Error),
 }
diff --git a/crates/crypto/src/lib.rs b/crates/crypto/src/lib.rs
index 1115212c6..d5b5fee06 100644
--- a/crates/crypto/src/lib.rs
+++ b/crates/crypto/src/lib.rs
@@ -40,3 +40,5 @@ pub mod rng;
 pub use error::Error;
 pub use protected::Protected;
 pub use rng::CryptoRng;
+
+pub use rand_core::{RngCore, SeedableRng};
diff --git a/crates/crypto/src/rng/csprng.rs b/crates/crypto/src/rng/csprng.rs
index 7ca85eccb..6275aafea 100644
--- a/crates/crypto/src/rng/csprng.rs
+++ b/crates/crypto/src/rng/csprng.rs
@@ -1,3 +1,5 @@
+use crate::Error;
+
 use rand::RngCore;
 use rand_chacha::ChaCha20Rng;
 use rand_core::{impl_try_crypto_rng_from_crypto_rng, SeedableRng};
@@ -14,9 +16,8 @@ impl CryptoRng {
 	/// This creates a new [`ChaCha20Rng`]-backed [`rand::CryptoRng`] from entropy
 	/// (via the [getrandom](https://docs.rs/getrandom) crate).
 	#[inline]
-	#[must_use]
-	pub fn new() -> Self {
-		Self(ChaCha20Rng::from_os_rng())
+	pub fn new() -> Result<Self, Error> {
+		ChaCha20Rng::try_from_os_rng().map(Self).map_err(Into::into)
 	}
 
 	/// Used to generate completely random bytes, with the use of [`ChaCha20Rng`]
@@ -57,10 +58,21 @@ impl RngCore for CryptoRng {
 	}
 }
 
+impl SeedableRng for CryptoRng {
+	type Seed = <ChaCha20Rng as SeedableRng>::Seed;
+
+	fn from_seed(seed: Self::Seed) -> Self {
+		Self(ChaCha20Rng::from_seed(seed))
+	}
+}
+
 impl Zeroize for CryptoRng {
 	#[inline]
 	fn zeroize(&mut self) {
-		self.0 = ChaCha20Rng::from_os_rng();
+		let mut seed = <Self as SeedableRng>::Seed::default();
+		self.0.fill_bytes(&mut seed);
+
+		self.0 = ChaCha20Rng::from_seed(seed);
 	}
 }
 
@@ -68,12 +80,6 @@ impl rand::CryptoRng for CryptoRng {}
 
 impl_try_crypto_rng_from_crypto_rng!(CryptoRng);
 
-impl Default for CryptoRng {
-	fn default() -> Self {
-		Self::new()
-	}
-}
-
 impl Drop for CryptoRng {
 	#[inline]
 	fn drop(&mut self) {