Merge pull request #109 from quexten/feature/argon2

Implement argon2 kdf

8 files changed, 231 insertions, 20 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 4cc30c3..92a042c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -40,6 +40,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "224afbd727c3d6e4b90103ece64b8d1b67fbb1973b1046c2281eed3f3803f800"
 
 [[package]]
+name = "argon2"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "95c2fcf79ad1932ac6269a738109997a83c227c09b75842ae564dc8ede6a861c"
+dependencies = [
+ "base64ct",
+ "blake2",
+ "password-hash",
+]
+
+[[package]]
 name = "arrayvec"
 version = "0.7.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -87,6 +98,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
 [[package]]
+name = "blake2"
+version = "0.10.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe"
+dependencies = [
+ "digest",
+]
+
+[[package]]
 name = "block-buffer"
 version = "0.10.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -857,6 +877,17 @@ dependencies = [
 ]
 
 [[package]]
+name = "password-hash"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "346f04948ba92c43e8469c1ee6736c7563d71012b17d40745260fe106aac2166"
+dependencies = [
+ "base64ct",
+ "rand_core",
+ "subtle",
+]
+
+[[package]]
 name = "pbkdf2"
 version = "0.12.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -999,6 +1030,7 @@ version = "1.6.0"
 dependencies = [
  "aes",
  "anyhow",
+ "argon2",
  "arrayvec",
  "async-trait",
  "base32",
diff --git a/Cargo.toml b/Cargo.toml
index 703a3ba..ae6d641 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ include = ["src/**/*", "bin/**/*", "LICENSE", "README.md", "CHANGELOG.md"]
 [dependencies]
 aes = "0.8.2"
 anyhow = "1.0.69"
+argon2 = "0.5.0"
 arrayvec = "0.7.2"
 async-trait = "0.1.66"
 base32 = "0.4.0"
diff --git a/src/actions.rs b/src/actions.rs
index dc78f7e..7212415 100644
--- a/src/actions.rs
+++ b/src/actions.rs
@@ -1,4 +1,4 @@
-use crate::prelude::*;
+use crate::{prelude::*, api::KdfType};
 
 pub async fn register(
     email: &str,
@@ -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, KdfType, 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: KdfType,
     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..85a0ce5 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: KdfType,
     #[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<(KdfType, 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(
@@ -1212,3 +1218,94 @@ fn classify_login_error(error_res: &ConnectErrorRes, code: u16) -> Error {
     log::warn!("unexpected error received during login: {:?}", error_res);
     Error::RequestFailed { status: code }
 }
+
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum KdfType {
+    Pbkdf2 = 0,
+    Argon2id = 1,
+}
+
+impl<'de> serde::Deserialize<'de> for KdfType {
+    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        struct KdfTypeVisitor;
+        impl<'de> serde::de::Visitor<'de> for KdfTypeVisitor {
+            type Value = KdfType;
+
+            fn expecting(
+                &self,
+                formatter: &mut std::fmt::Formatter,
+            ) -> std::fmt::Result {
+                formatter.write_str("two factor provider id")
+            }
+
+            fn visit_str<E>(
+                self,
+                value: &str,
+            ) -> std::result::Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                value.parse().map_err(serde::de::Error::custom)
+            }
+
+            fn visit_u64<E>(
+                self,
+                value: u64,
+            ) -> std::result::Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                std::convert::TryFrom::try_from(value)
+                    .map_err(serde::de::Error::custom)
+            }
+        }
+
+        deserializer.deserialize_any(KdfTypeVisitor)
+    }
+}
+
+impl std::convert::TryFrom<u64> for KdfType {
+    type Error = Error;
+
+    fn try_from(ty: u64) -> Result<Self> {
+        match ty {
+            0 => Ok(Self::Pbkdf2),
+            1 => Ok(Self::Argon2id),
+            _ => Err(Error::InvalidTwoFactorProvider {
+                ty: format!("{ty}"),
+            }),
+        }
+    }
+}
+
+impl std::str::FromStr for KdfType {
+    type Err = Error;
+
+    fn from_str(ty: &str) -> Result<Self> {
+        match ty {
+            "0" => Ok(Self::Pbkdf2),
+            "1" => Ok(Self::Argon2id),
+            _ => Err(Error::InvalidTwoFactorProvider { ty: ty.to_string() }),
+        }
+    }
+}
+
+impl serde::Serialize for KdfType {
+    fn serialize<S>(
+        &self,
+        serializer: S,
+    ) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        let s = match self {
+            Self::Pbkdf2 => "0",
+            Self::Argon2id => "1",
+        };
+        serializer.serialize_str(s)
+    }
+}
diff --git a/src/bin/rbw-agent/actions.rs b/src/bin/rbw-agent/actions.rs
index 88236ba..361009b 100644
--- a/src/bin/rbw-agent/actions.rs
+++ b/src/bin/rbw-agent/actions.rs
@@ -1,4 +1,5 @@
 use anyhow::Context as _;
+use rbw::api::KdfType;
 
 pub async fn register(
     sock: &mut crate::sock::Sock,
@@ -123,7 +124,10 @@ pub async fn login(
                 Ok((
                     access_token,
                     refresh_token,
+                    kdf,
                     iterations,
+                    memory,
+                    parallelism,
                     protected_key,
                 )) => {
                     login_success(
@@ -131,7 +135,10 @@ pub async fn login(
                         state,
                         access_token,
                         refresh_token,
+                        kdf,
                         iterations,
+                        memory,
+                        parallelism,
                         protected_key,
                         password,
                         db,
@@ -151,7 +158,10 @@ pub async fn login(
                             let (
                                 access_token,
                                 refresh_token,
+                                kdf,
                                 iterations,
+                                memory,
+                                parallelism,
                                 protected_key,
                             ) = two_factor(
                                 tty,
@@ -165,7 +175,10 @@ pub async fn login(
                                 state,
                                 access_token,
                                 refresh_token,
+                                kdf,
                                 iterations,
+                                memory,
+                                parallelism,
                                 protected_key,
                                 password,
                                 db,
@@ -205,7 +218,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, KdfType, u32, Option<u32>, Option<u32>, String)> {
     let mut err_msg = None;
     for i in 1_u8..=3 {
         let err = if i > 1 {
@@ -235,11 +248,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 +296,10 @@ async fn login_success(
     state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
     access_token: String,
     refresh_token: String,
+    kdf: KdfType,
     iterations: u32,
+    memory: Option<u32>,
+    parallelism: Option<u32>,
     protected_key: String,
     password: rbw::locked::Password,
     mut db: rbw::db::Db,
@@ -288,7 +307,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 +327,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 +356,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"
             ));
         };
+
+        let memory= db.memory;
+        let parallelism = db.parallelism;
+
         let Some(protected_key) = db.protected_key
         else {
             return Err(anyhow::anyhow!(
@@ -377,7 +413,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..74be28a 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1,4 +1,4 @@
-use crate::prelude::*;
+use crate::{prelude::*, api::KdfType};
 
 use std::io::{Read as _, Write as _};
 use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
@@ -164,7 +164,10 @@ pub struct Db {
     pub access_token: Option<String>,
     pub refresh_token: Option<String>,
 
+    pub kdf: Option<KdfType>,
     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..9bc435f 100644
--- a/src/identity.rs
+++ b/src/identity.rs
@@ -1,5 +1,6 @@
-use crate::prelude::*;
-
+use crate::{prelude::*, api::KdfType};
+use sha2::Digest;
+use argon2::Argon2;
 pub struct Identity {
     pub email: String,
     pub keys: crate::locked::Keys,
@@ -10,7 +11,10 @@ impl Identity {
     pub fn new(
         email: &str,
         password: &crate::locked::Password,
+        kdf: KdfType,
         iterations: u32,
+        memory: Option<u32>,
+        parallelism: Option<u32>,
     ) -> Result<Self> {
         let iterations = std::num::NonZeroU32::new(iterations)
             .ok_or(Error::Pbkdf2ZeroIterations)?;
@@ -19,14 +23,37 @@ 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 {
+            KdfType::Pbkdf2 => {
+                pbkdf2::pbkdf2::<hmac::Hmac<sha2::Sha256>>(
+                    password.password(),
+                    email.as_bytes(),
+                    iterations.get(),
+                    enc_key,
+                )
+                .map_err(|_| Error::Pbkdf2)?;
+            }
+
+            KdfType::Argon2id => {
+                let mut hasher = sha2::Sha256::new();
+                hasher.update(email.as_bytes());
+                let mut salt = hasher.finalize();
+
+                let mut output_key_material = [0u8];
+                let argon2_config = Argon2::new(
+                    argon2::Algorithm::Argon2id,
+                    argon2::Version::V0x13,
+                    argon2::Params::new(memory.unwrap() * 1024,
+                        iterations.get(),
+                        parallelism.unwrap(),
+                        Some(32)).unwrap());
+                argon2::Argon2::hash_password_into(&argon2_config, password.password(), &mut salt, &mut output_key_material)
+                    .map_err(|_| Error::Argon2)?;
+                enc_key.copy_from_slice(&output_key_material);
+            }
+        };
+        
         let mut hash = crate::locked::Vec::new();
         hash.extend(std::iter::repeat(0).take(32));
         pbkdf2::pbkdf2::<hmac::Hmac<sha2::Sha256>>(