From 77d0f75c2a6629d4a604b296b71d7035f34f8eaf Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Tue, 31 Oct 2023 11:34:10 +0100
Subject: [PATCH] Add some unit tests

---
 crates/bitwarden/src/crypto/master_key.rs | 73 +++++++++++++++++++----
 crates/bitwarden/src/crypto/user_key.rs   |  2 +-
 2 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/crates/bitwarden/src/crypto/master_key.rs b/crates/bitwarden/src/crypto/master_key.rs
index 4e12ac06c..93847555c 100644
--- a/crates/bitwarden/src/crypto/master_key.rs
+++ b/crates/bitwarden/src/crypto/master_key.rs
@@ -41,15 +41,7 @@ impl MasterKey {
     }
 
     pub(crate) fn make_user_key(&self) -> Result<(UserKey, EncString)> {
-        let mut user_key = [0u8; 64];
-        rand::thread_rng().fill(&mut user_key);
-
-        let stretched_key = stretch_master_key(self)?;
-        let protected =
-            encrypt_aes256_hmac(&user_key, stretched_key.mac_key.unwrap(), stretched_key.key)?;
-
-        let u: &[u8] = &user_key;
-        Ok((UserKey::new(SymmetricCryptoKey::try_from(u)?), protected))
+        make_user_key(rand::thread_rng(), self)
     }
 
     pub(crate) fn decrypt_user_key(&self, user_key: EncString) -> Result<SymmetricCryptoKey> {
@@ -60,6 +52,22 @@ impl MasterKey {
     }
 }
 
+/// Generate a new random user key and encrypt it with the master key.
+fn make_user_key(
+    mut rng: impl rand::RngCore,
+    master_key: &MasterKey,
+) -> Result<(UserKey, EncString)> {
+    let mut user_key = [0u8; 64];
+    rng.fill(&mut user_key);
+
+    let stretched_key = stretch_master_key(master_key)?;
+    let protected =
+        encrypt_aes256_hmac(&user_key, stretched_key.mac_key.unwrap(), stretched_key.key)?;
+
+    let u: &[u8] = &user_key;
+    Ok((UserKey::new(SymmetricCryptoKey::try_from(u)?), protected))
+}
+
 /// Derive a generic key from a secret and salt using the provided KDF.
 fn derive_key(secret: &[u8], salt: &[u8], kdf: &Kdf) -> Result<SymmetricCryptoKey> {
     let hash = match kdf {
@@ -114,7 +122,9 @@ fn stretch_master_key(master_key: &MasterKey) -> Result<SymmetricCryptoKey> {
 mod tests {
     use std::num::NonZeroU32;
 
-    use super::{stretch_master_key, HashPurpose, MasterKey};
+    use rand::rngs::mock::StepRng;
+
+    use super::{make_user_key, stretch_master_key, HashPurpose, MasterKey};
     use crate::{client::kdf::Kdf, crypto::SymmetricCryptoKey};
 
     #[test]
@@ -227,4 +237,47 @@ mod tests {
                 .unwrap(),
         );
     }
+
+    #[test]
+    fn test_make_user_key() {
+        let mut rng = StepRng::new(2, 1);
+
+        let master_key = MasterKey(SymmetricCryptoKey {
+            key: [
+                31, 79, 104, 226, 150, 71, 177, 90, 194, 80, 172, 209, 17, 129, 132, 81, 138, 167,
+                69, 167, 254, 149, 2, 27, 39, 197, 64, 42, 22, 195, 86, 75,
+            ]
+            .into(),
+            mac_key: None,
+        });
+
+        let (user_key, protected) = make_user_key(&mut rng, &master_key).unwrap();
+
+        assert_eq!(
+            user_key.0.key.as_slice(),
+            [
+                2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0,
+                0, 0, 0, 0
+            ]
+        );
+        assert_eq!(
+            user_key.0.mac_key.unwrap().as_slice(),
+            [
+                6, 0, 0, 0, 0, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0, 0, 9, 0, 0, 0,
+                0, 0, 0, 0
+            ]
+        );
+
+        // Ensure we can decrypt the key and get back the same key
+        let decrypted = master_key.decrypt_user_key(protected).unwrap();
+
+        assert_eq!(
+            decrypted.key, user_key.0.key,
+            "Decrypted key doesn't match user key"
+        );
+        assert_eq!(
+            decrypted.mac_key, user_key.0.mac_key,
+            "Decrypted key doesn't match user key"
+        );
+    }
 }
diff --git a/crates/bitwarden/src/crypto/user_key.rs b/crates/bitwarden/src/crypto/user_key.rs
index 0fe560665..7f5bae413 100644
--- a/crates/bitwarden/src/crypto/user_key.rs
+++ b/crates/bitwarden/src/crypto/user_key.rs
@@ -6,7 +6,7 @@ use crate::{
     error::Result,
 };
 
-pub(crate) struct UserKey(SymmetricCryptoKey);
+pub(crate) struct UserKey(pub(super) SymmetricCryptoKey);
 
 impl UserKey {
     pub(crate) fn new(key: SymmetricCryptoKey) -> Self {