Switch kdf type to enum

5 files changed, 105 insertions, 16 deletions

diff --git a/src/actions.rs b/src/actions.rs
index 2c57405..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,7 +18,7 @@ pub async fn login(
     password: crate::locked::Password,
     two_factor_token: Option<&str>,
     two_factor_provider: Option<crate::api::TwoFactorProviderType>,
-) -> Result<(String, String, u32, u32, Option<u32>, Option<u32>, String)> {
+) -> Result<(String, String, KdfType, u32, Option<u32>, Option<u32>, String)> {
     let (client, config) = api_client_async().await?;
     let (kdf, iterations, memory, parallelism) = client.prelogin(email).await?;
 
@@ -40,7 +40,7 @@ pub async fn login(
 pub fn unlock<S: std::hash::BuildHasher>(
     email: &str,
     password: &crate::locked::Password,
-    kdf: u32,
+    kdf: KdfType,
     iterations: u32,
     memory: Option<u32>,
     parallelism: Option<u32>,
diff --git a/src/api.rs b/src/api.rs
index 6c8545c..85a0ce5 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -170,7 +170,7 @@ struct PreloginReq {
 #[derive(serde::Deserialize, Debug)]
 struct PreloginRes {
     #[serde(rename = "Kdf", alias = "kdf")]
-    kdf: u32,
+    kdf: KdfType,
     #[serde(rename = "KdfIterations", alias = "kdfIterations")]
     kdf_iterations: u32,
     #[serde(rename = "KdfMemory", alias = "kdfMemory")]
@@ -634,7 +634,7 @@ impl Client {
         }
     }
 
-    pub async fn prelogin(&self, email: &str) -> Result<(u32, u32, Option<u32>, Option<u32>)> {
+    pub async fn prelogin(&self, email: &str) -> Result<(KdfType, u32, Option<u32>, Option<u32>)> {
         let prelogin = PreloginReq {
             email: email.to_string(),
         };
@@ -1218,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 4b3267f..56a98e9 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,
@@ -217,7 +218,7 @@ async fn two_factor(
     email: &str,
     password: rbw::locked::Password,
     provider: rbw::api::TwoFactorProviderType,
-) -> anyhow::Result<(String, String, u32, u32, Option<u32>, Option<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 {
@@ -295,7 +296,7 @@ async fn login_success(
     state: std::sync::Arc<tokio::sync::RwLock<crate::agent::State>>,
     access_token: String,
     refresh_token: String,
-    kdf: u32,
+    kdf: KdfType,
     iterations: u32,
     memory: Option<u32>,
     parallelism: Option<u32>,
diff --git a/src/db.rs b/src/db.rs
index 8781459..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,7 @@ pub struct Db {
     pub access_token: Option<String>,
     pub refresh_token: Option<String>,
 
-    pub kdf: Option<u32>,
+    pub kdf: Option<KdfType>,
     pub iterations: Option<u32>,
     pub memory: Option<u32>,
     pub parallelism: Option<u32>,
diff --git a/src/identity.rs b/src/identity.rs
index d63dc9b..3637d75 100644
--- a/src/identity.rs
+++ b/src/identity.rs
@@ -1,4 +1,4 @@
-use crate::prelude::*;
+use crate::{prelude::*, api::KdfType};
 use sha2::Digest;
 extern crate argon2;
 use argon2::{Config, ThreadMode, Variant, Version};
@@ -13,7 +13,7 @@ impl Identity {
     pub fn new(
         email: &str,
         password: &crate::locked::Password,
-        kdf: u32,
+        kdf: KdfType,
         iterations: u32,
         memory: Option<u32>,
         parallelism: Option<u32>,
@@ -27,7 +27,7 @@ impl Identity {
         let enc_key = &mut keys.data_mut()[0..32];
 
         match kdf {
-            0 => {
+            KdfType::Pbkdf2 => {
                 pbkdf2::pbkdf2::<hmac::Hmac<sha2::Sha256>>(
                     password.password(),
                     email.as_bytes(),
@@ -37,7 +37,7 @@ impl Identity {
                 .map_err(|_| Error::Pbkdf2)?;
             }
 
-            1 => {
+            KdfType::Argon2id => {
                 let mut hasher = sha2::Sha256::new();
                 hasher.update(email.as_bytes());
                 let salt = hasher.finalize();
@@ -56,9 +56,6 @@ impl Identity {
                 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();