Implement argon2 kdf

8 files changed, 171 insertions, 20 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 4cc30c3..a6afafc 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -40,6 +40,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "224afbd727c3d6e4b90103ece64b8d1b67fbb1973b1046c2281eed3f3803f800"
 
 [[package]]
+name = "arrayref"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4c527152e37cf757a3f78aae5a06fbeefdb07ccc535c980a3208ee3060dd544"
+
+[[package]]
 name = "arrayvec"
 version = "0.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -70,6 +76,12 @@ checksum = "23ce669cd6c8588f79e15cf450314f9638f967fc5770ff1c7c1deb0925ea7cfa"
 
 [[package]]
 name = "base64"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
+
+[[package]]
+name = "base64"
 version = "0.21.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a4a4ddaa51a5bc52a6948f74c06d20aaaddb71924eab79b8c97a8c556e942d6a"
@@ -87,6 +99,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
 [[package]]
+name = "blake2b_simd"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3c2f0dc9a68c6317d884f97cc36cf5a3d20ba14ce404227df55e1af708ab04bc"
+dependencies = [
+ "arrayref",
+ "arrayvec",
+ "constant_time_eq 0.2.5",
+]
+
+[[package]]
 name = "block-buffer"
 version = "0.10.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -207,6 +230,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "520fbf3c07483f94e3e3ca9d0cfd913d7718ef2483d2cfd91c0d9e91474ab913"
 
 [[package]]
+name = "constant_time_eq"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc"
+
+[[package]]
+name = "constant_time_eq"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13418e745008f7349ec7e449155f419a61b92b58a99cc3616942b926825ec76b"
+
+[[package]]
 name = "core-foundation"
 version = "0.9.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -232,6 +267,15 @@ dependencies = [
 ]
 
 [[package]]
+name = "crossbeam-utils"
+version = "0.8.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3c063cd8cc95f5c377ed0d4b49a4b21f632396ff690e8470c29b3359b346984b"
+dependencies = [
+ "cfg-if",
+]
+
+[[package]]
 name = "crypto-common"
 version = "0.1.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1002,7 +1046,7 @@ dependencies = [
  "arrayvec",
  "async-trait",
  "base32",
- "base64",
+ "base64 0.21.0",
  "block-padding",
  "cbc",
  "clap",
@@ -1023,6 +1067,7 @@ dependencies = [
  "region",
  "reqwest",
  "rsa",
+ "rust-argon2",
  "serde",
  "serde_json",
  "serde_path_to_error",
@@ -1095,7 +1140,7 @@ version = "0.11.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "21eed90ec8570952d53b772ecf8f206aa1ec9a3d76b2521c56c42973f2d91ee9"
 dependencies = [
- "base64",
+ "base64 0.21.0",
  "bytes",
  "encoding_rs",
  "futures-core",
@@ -1164,6 +1209,18 @@ dependencies = [
 ]
 
 [[package]]
+name = "rust-argon2"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b50162d19404029c1ceca6f6980fe40d45c8b369f6f44446fa14bb39573b5bb9"
+dependencies = [
+ "base64 0.13.1",
+ "blake2b_simd",
+ "constant_time_eq 0.1.5",
+ "crossbeam-utils",
+]
+
+[[package]]
 name = "rustix"
 version = "0.36.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1207,7 +1264,7 @@ version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d194b56d58803a43635bdc398cd17e383d6f71f9182b9a192c127ca42494a59b"
 dependencies = [
- "base64",
+ "base64 0.21.0",
 ]
 
 [[package]]
diff --git a/Cargo.toml b/Cargo.toml
index 703a3ba..19d11c5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -39,6 +39,7 @@ rand = "0.8.5"
 region = "3.0.0"
 reqwest = { version = "0.11.14", default-features = false, features = ["blocking", "json", "rustls-tls-native-roots"] }
 rsa = "0.8.2"
+rust-argon2 = "1.0.0"
 serde = { version = "1.0.154", features = ["derive"] }
 serde_json = "1.0.94"
 serde_path_to_error = "0.1.10"
diff --git a/src/actions.rs b/src/actions.rs
index dc78f7e..2c57405 100644
--- a/src/actions.rs
+++ b/src/actions.rs
@@ -18,11 +18,12 @@ pub async fn login(
     password: crate::locked::Password,
     two_factor_token: Option<&str>,
     two_factor_provider: Option<crate::api::TwoFactorProviderType>,
-) -> Result<(String, String, u32, String)> {
+) -> Result<(String, String, u32, u32, Option<u32>, Option<u32>, String)> {
     let (client, config) = api_client_async().await?;
-    let iterations = client.prelogin(email).await?;
+    let (kdf, iterations, memory, parallelism) = client.prelogin(email).await?;
+
     let identity =
-        crate::identity::Identity::new(email, &password, iterations)?;
+        crate::identity::Identity::new(email, &password, kdf, iterations, memory, parallelism)?;
     let (access_token, refresh_token, protected_key) = client
         .login(
             email,
@@ -33,13 +34,16 @@ pub async fn login(
         )
         .await?;
 
-    Ok((access_token, refresh_token, iterations, protected_key))
+    Ok((access_token, refresh_token, kdf, iterations, memory, parallelism, protected_key))
 }
 
 pub fn unlock<S: std::hash::BuildHasher>(
     email: &str,
     password: &crate::locked::Password,
+    kdf: u32,
     iterations: u32,
+    memory: Option<u32>,
+    parallelism: Option<u32>,
     protected_key: &str,
     protected_private_key: &str,
     protected_org_keys: &std::collections::HashMap<String, String, S>,
@@ -48,7 +52,7 @@ pub fn unlock<S: std::hash::BuildHasher>(
     std::collections::HashMap<String, crate::locked::Keys>,
 )> {
     let identity =
-        crate::identity::Identity::new(email, password, iterations)?;
+        crate::identity::Identity::new(email, password, kdf, iterations, memory, parallelism)?;
 
     let protected_key =
         crate::cipherstring::CipherString::new(protected_key)?;
diff --git a/src/api.rs b/src/api.rs
index f2c853d..6c8545c 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -169,8 +169,14 @@ struct PreloginReq {
 
 #[derive(serde::Deserialize, Debug)]
 struct PreloginRes {
+    #[serde(rename = "Kdf", alias = "kdf")]
+    kdf: u32,
     #[serde(rename = "KdfIterations", alias = "kdfIterations")]
     kdf_iterations: u32,
+    #[serde(rename = "KdfMemory", alias = "kdfMemory")]
+    kdf_memory: Option<u32>,
+    #[serde(rename = "KdfParallelism", alias = "kdfParallelism")]
+    kdf_parallelism: Option<u32>,
 }
 
 #[derive(serde::Serialize, Debug)]
@@ -628,7 +634,7 @@ impl Client {
         }
     }
 
-    pub async fn prelogin(&self, email: &str) -> Result<u32> {
+    pub async fn prelogin(&self, email: &str) -> Result<(u32, u32, Option<u32>, Option<u32>)> {
         let prelogin = PreloginReq {
             email: email.to_string(),
         };
@@ -640,7 +646,7 @@ impl Client {
             .await
             .map_err(|source| Error::Reqwest { source })?;
         let prelogin_res: PreloginRes = res.json_with_path().await?;
-        Ok(prelogin_res.kdf_iterations)
+        Ok((prelogin_res.kdf, prelogin_res.kdf_iterations, prelogin_res.kdf_memory, prelogin_res.kdf_parallelism))
     }
 
     pub async fn register(
diff --git a/src/bin/rbw-agent/actions.rs b/src/bin/rbw-agent/actions.rs
index 88236ba..4b3267f 100644
--- a/src/bin/rbw-agent/actions.rs
+++ b/src/bin/rbw-agent/actions.rs
@@ -123,7 +123,10 @@ pub async fn login(
                 Ok((
                     access_token,
                     refresh_token,
+                    kdf,
                     iterations,
+                    memory,
+                    parallelism,
                     protected_key,
                 )) => {
                     login_success(
@@ -131,7 +134,10 @@ pub async fn login(
                         state,
                         access_token,
                         refresh_token,
+                        kdf,
                         iterations,
+                        memory,
+                        parallelism,
                         protected_key,
                         password,
                         db,
@@ -151,7 +157,10 @@ pub async fn login(
                             let (
                                 access_token,
                                 refresh_token,
+                                kdf,
                                 iterations,
+                                memory,
+                                parallelism,
                                 protected_key,
                             ) = two_factor(
                                 tty,
@@ -165,7 +174,10 @@ pub async fn login(
                                 state,
                                 access_token,
                                 refresh_token,
+                                kdf,
                                 iterations,
+                                memory,
+                                parallelism,
                                 protected_key,
                                 password,
                                 db,
@@ -205,7 +217,7 @@ async fn two_factor(
     email: &str,
     password: rbw::locked::Password,
     provider: rbw::api::TwoFactorProviderType,
-) -> anyhow::Result<(String, String, u32, String)> {
+) -> anyhow::Result<(String, String, u32, u32, Option<u32>, Option<u32>, String)> {
     let mut err_msg = None;
     for i in 1_u8..=3 {
         let err = if i > 1 {
@@ -235,11 +247,14 @@ async fn two_factor(
         )
         .await
         {
-            Ok((access_token, refresh_token, iterations, protected_key)) => {
+            Ok((access_token, refresh_token, kdf, iterations, memory, parallelism, protected_key)) => {
                 return Ok((
                     access_token,
                     refresh_token,
+                    kdf,
                     iterations,
+                    memory,
+                    parallelism,
                     protected_key,
                 ))
             }
@@ -280,7 +295,10 @@ async fn login_success(
     state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
     access_token: String,
     refresh_token: String,
+    kdf: u32,
     iterations: u32,
+    memory: Option<u32>,
+    parallelism: Option<u32>,
     protected_key: String,
     password: rbw::locked::Password,
     mut db: rbw::db::Db,
@@ -288,7 +306,10 @@ async fn login_success(
 ) -> anyhow::Result<()> {
     db.access_token = Some(access_token.to_string());
     db.refresh_token = Some(refresh_token.to_string());
+    db.kdf = Some(kdf);
     db.iterations = Some(iterations);
+    db.memory = memory;
+    db.parallelism = parallelism;
     db.protected_key = Some(protected_key.to_string());
     save_db(&db).await?;
 
@@ -305,7 +326,10 @@ async fn login_success(
     let res = rbw::actions::unlock(
         &email,
         &password,
+        kdf,
         iterations,
+        memory,
+        parallelism,
         &protected_key,
         &protected_private_key,
         &db.protected_org_keys,
@@ -331,12 +355,23 @@ pub async fn unlock(
     if state.read().await.needs_unlock() {
         let db = load_db().await?;
 
+        let Some(kdf) = db.kdf
+        else {
+            return Err(anyhow::anyhow!(
+                "failed to find kdf type in db"
+            ));
+        };
+
         let Some(iterations) = db.iterations
         else {
             return Err(anyhow::anyhow!(
-                "failed to find number of iterations in db"
+                "failed to find iterations in db"
             ));
         };
+
+        let memory= db.memory;
+        let parallelism = db.parallelism;
+
         let Some(protected_key) = db.protected_key
         else {
             return Err(anyhow::anyhow!(
@@ -377,7 +412,10 @@ pub async fn unlock(
             match rbw::actions::unlock(
                 &email,
                 &password,
+                kdf,
                 iterations,
+                memory,
+                parallelism,
                 &protected_key,
                 &protected_private_key,
                 &db.protected_org_keys,
diff --git a/src/db.rs b/src/db.rs
index 7d09db1..8781459 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -164,7 +164,10 @@ pub struct Db {
     pub access_token: Option<String>,
     pub refresh_token: Option<String>,
 
+    pub kdf: Option<u32>,
     pub iterations: Option<u32>,
+    pub memory: Option<u32>,
+    pub parallelism: Option<u32>,
     pub protected_key: Option<String>,
     pub protected_private_key: Option<String>,
     pub protected_org_keys: std::collections::HashMap<String, String>,
@@ -293,6 +296,7 @@ impl Db {
         self.access_token.is_none()
             || self.refresh_token.is_none()
             || self.iterations.is_none()
+            || self.kdf.is_none()
             || self.protected_key.is_none()
     }
 }
diff --git a/src/error.rs b/src/error.rs
index 6d49e09..9b7261f 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -142,6 +142,9 @@ pub enum Error {
     #[error("failed to run pbkdf2")]
     Pbkdf2,
 
+    #[error("failed to run argon2")]
+    Argon2,
+
     #[error("pinentry cancelled")]
     PinentryCancelled,
 
@@ -224,6 +227,10 @@ pub enum Error {
 
     #[error("error writing to pinentry stdin")]
     WriteStdin { source: tokio::io::Error },
+
+
+    #[error("invalid kdf type: {ty}")]
+    InvalidKdfType { ty: String },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
diff --git a/src/identity.rs b/src/identity.rs
index 8a5dc61..d63dc9b 100644
--- a/src/identity.rs
+++ b/src/identity.rs
@@ -1,4 +1,7 @@
 use crate::prelude::*;
+use sha2::Digest;
+extern crate argon2;
+use argon2::{Config, ThreadMode, Variant, Version};
 
 pub struct Identity {
     pub email: String,
@@ -10,7 +13,10 @@ impl Identity {
     pub fn new(
         email: &str,
         password: &crate::locked::Password,
+        kdf: u32,
         iterations: u32,
+        memory: Option<u32>,
+        parallelism: Option<u32>,
     ) -> Result<Self> {
         let iterations = std::num::NonZeroU32::new(iterations)
             .ok_or(Error::Pbkdf2ZeroIterations)?;
@@ -19,14 +25,42 @@ impl Identity {
         keys.extend(std::iter::repeat(0).take(64));
 
         let enc_key = &mut keys.data_mut()[0..32];
-        pbkdf2::pbkdf2::<hmac::Hmac<sha2::Sha256>>(
-            password.password(),
-            email.as_bytes(),
-            iterations.get(),
-            enc_key,
-        )
-        .map_err(|_| Error::Pbkdf2)?;
 
+        match kdf {
+            0 => {
+                pbkdf2::pbkdf2::<hmac::Hmac<sha2::Sha256>>(
+                    password.password(),
+                    email.as_bytes(),
+                    iterations.get(),
+                    enc_key,
+                )
+                .map_err(|_| Error::Pbkdf2)?;
+            }
+
+            1 => {
+                let mut hasher = sha2::Sha256::new();
+                hasher.update(email.as_bytes());
+                let salt = hasher.finalize();
+
+                let config = Config {
+                    variant: Variant::Argon2id,
+                    version: Version::Version13,
+                    mem_cost: memory.unwrap() * 1024,
+                    time_cost: iterations.get(),
+                    lanes: parallelism.unwrap(),
+                    thread_mode: ThreadMode::Parallel,
+                    secret: &[],
+                    ad: &[],
+                    hash_length: 32
+                };
+                let hash = argon2::hash_raw(password.password(), &salt[..], &config).map_err(|_| Error::Argon2)?;
+                enc_key.copy_from_slice(&hash);
+            }
+            _ => {
+                // todo throw error or switch to enum?
+            }
+        };
+        
         let mut hash = crate::locked::Vec::new();
         hash.extend(std::iter::repeat(0).take(32));
         pbkdf2::pbkdf2::<hmac::Hmac<sha2::Sha256>>(