From d0852366331aab3f9783008efdc575433071257e Mon Sep 17 00:00:00 2001
From: Hinton <hinton@users.noreply.github.com>
Date: Tue, 31 Oct 2023 10:45:59 +0100
Subject: [PATCH] Fix make_user_key which previously didn't stretch the users
 key

---
 crates/bitwarden/src/auth/register.rs     |  6 ++--
 crates/bitwarden/src/client/client.rs     | 37 ++++++++++++++++++++++-
 crates/bitwarden/src/crypto/master_key.rs |  8 +++--
 3 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/crates/bitwarden/src/auth/register.rs b/crates/bitwarden/src/auth/register.rs
index 8336ec9c9..206c1905b 100644
--- a/crates/bitwarden/src/auth/register.rs
+++ b/crates/bitwarden/src/auth/register.rs
@@ -79,7 +79,7 @@ pub(super) fn make_register_keys(
 
 #[cfg_attr(feature = "mobile", derive(uniffi::Record))]
 pub struct RegisterKeyResponse {
-    master_password_hash: String,
-    encrypted_user_key: String,
-    keys: RsaKeyPair,
+    pub master_password_hash: String,
+    pub encrypted_user_key: String,
+    pub keys: RsaKeyPair,
 }
diff --git a/crates/bitwarden/src/client/client.rs b/crates/bitwarden/src/client/client.rs
index f6eb0c954..35e230e29 100644
--- a/crates/bitwarden/src/client/client.rs
+++ b/crates/bitwarden/src/client/client.rs
@@ -288,9 +288,14 @@ impl Client {
 
 #[cfg(test)]
 mod tests {
+    use std::num::NonZeroU32;
+
     use wiremock::{matchers, Mock, ResponseTemplate};
 
-    use crate::{auth::login::AccessTokenLoginRequest, secrets_manager::secrets::*};
+    use crate::{
+        auth::login::AccessTokenLoginRequest, client::kdf::Kdf, mobile::crypto::InitCryptoRequest,
+        secrets_manager::secrets::*, Client,
+    };
 
     #[tokio::test]
     async fn test_access_token_login() {
@@ -379,4 +384,34 @@ mod tests {
         assert_eq!(res.note, "TEST");
         assert_eq!(res.value, "TEST");
     }
+
+    #[cfg(feature = "internal")]
+    #[tokio::test]
+    async fn test_register_initialize_crypto() {
+        let mut client = Client::new(None);
+
+        let email = "test@bitwarden.com";
+        let password = "test123";
+        let kdf = Kdf::PBKDF2 {
+            iterations: NonZeroU32::new(600_000).unwrap(),
+        };
+
+        let register_response = client
+            .auth()
+            .make_register_keys(email.to_owned(), password.to_owned(), kdf.clone())
+            .unwrap();
+
+        client
+            .crypto()
+            .initialize_crypto(InitCryptoRequest {
+                kdf_params: kdf,
+                email: email.to_owned(),
+                password: password.to_owned(),
+                user_key: register_response.encrypted_user_key,
+                private_key: register_response.keys.private.to_string(),
+                organization_keys: Default::default(),
+            })
+            .await
+            .unwrap();
+    }
 }
diff --git a/crates/bitwarden/src/crypto/master_key.rs b/crates/bitwarden/src/crypto/master_key.rs
index 064336edb..4e12ac06c 100644
--- a/crates/bitwarden/src/crypto/master_key.rs
+++ b/crates/bitwarden/src/crypto/master_key.rs
@@ -4,8 +4,8 @@ use rand::Rng;
 use sha2::Digest;
 
 use super::{
-    encrypt_aes256, hkdf_expand, EncString, KeyDecryptable, PbkdfSha256Hmac, SymmetricCryptoKey,
-    UserKey, PBKDF_SHA256_HMAC_OUT_SIZE,
+    encrypt_aes256_hmac, hkdf_expand, EncString, KeyDecryptable, PbkdfSha256Hmac,
+    SymmetricCryptoKey, UserKey, PBKDF_SHA256_HMAC_OUT_SIZE,
 };
 use crate::{client::kdf::Kdf, error::Result, util::BASE64_ENGINE};
 
@@ -44,7 +44,9 @@ impl MasterKey {
         let mut user_key = [0u8; 64];
         rand::thread_rng().fill(&mut user_key);
 
-        let protected = encrypt_aes256(&user_key, self.0.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))