[4-1] feat: add 2fa auth route

auth-service/Cargo.toml

@@ -20,6 +20,7 @@ jsonwebtoken = "9.2.0"
 chrono = "0.4.35"
 dotenvy = "0.15.7"
 lazy_static = "1.4.0"
+rand = "0.8.5"
 
 [dev-dependencies]
 reqwest = { version = "0.11.26", default-features = false, features = ["json", "cookies"] }

auth-service/src/domain/data_stores.rs

@@ -1,3 +1,5 @@
+use rand::Rng;
+
 use super::{Email, Password, User};
 
 #[async_trait::async_trait]
@@ -28,3 +30,62 @@ pub enum BannedTokenStoreError {
     InvalidToken,
     UnexpectedError,
 }
+
+
+#[async_trait::async_trait]
+pub trait TwoFACodeStore {
+    async fn add_code(&mut self, email: Email, login_attempt_id: LoginAttemptId, code: TwoFACode) -> Result<(), TwoFACodeStoreError>;
+    async fn remove_code(&mut self, email: &Email) -> Result<(), TwoFACodeStoreError>;
+    async fn get_code(&self, email: &Email) -> Result<(LoginAttemptId, TwoFACode), TwoFACodeStoreError>;
+}
+
+#[derive(Debug, PartialEq)]
+pub enum TwoFACodeStoreError {
+    LoginAttemptIdNotFound,
+    CodeAlreadyExists,
+    UnexpectedError,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct LoginAttemptId(String);
+
+impl LoginAttemptId {
+    pub fn parse(id: String) -> Result<Self, String> {
+        match uuid::Uuid::parse_str(&id) {
+            Ok(val) => Ok(LoginAttemptId(val.to_string())),
+            Err(e) => Err("Failed to parse uuid".to_owned()),
+        }
+    }
+    pub fn to_string(self) -> String {
+        self.0
+    }
+}
+
+impl Default for LoginAttemptId {
+    fn default() -> Self {
+        Self(uuid::Uuid::new_v4().to_string())
+    }
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub struct TwoFACode(String);
+
+impl TwoFACode {
+    pub fn parse(code: String) -> Result<Self, String> {
+        if code.len() != 6 {
+            return Err("Code too short!".to_owned())
+        }
+        match code.parse::<u32>() {
+            Ok(_) => Ok(Self(code)),
+            Err(_) => Err("Failed to parse code".to_owned())
+        }
+    }
+}
+
+impl Default for TwoFACode {
+    fn default() -> Self {
+        let mut rng = rand::thread_rng();
+        let code  = rng.gen_range(100000..999999);
+        Self(code.to_string())
+    }
+}

auth-service/src/main.rs

@@ -1,6 +1,6 @@
 use auth_service::{
-    services::{HashmapBannedTokenStore, HashmapUserStore},
-    store::{AppState, BannedTokenStoreType, UserStoreType},
+    services::{HashmapBannedTokenStore, HashmapTwoFACodeStore, HashmapUserStore},
+    store::{AppState, BannedTokenStoreType, TwoFACodeStoreType, UserStoreType},
     utils::constants::prod,
     Application,
 };
@@ -15,7 +15,10 @@ async fn main() {
     let banned_token_store: BannedTokenStoreType = Arc::new(RwLock::new(HashmapBannedTokenStore {
         tokens: HashMap::new(),
     }));
-    let app_state = AppState::new(user_store, banned_token_store);
+    let two_fa_code_store: TwoFACodeStoreType = Arc::new(RwLock::new(HashmapTwoFACodeStore {
+        codes: HashMap::new(),
+    }));
+    let app_state = AppState::new(user_store, banned_token_store, two_fa_code_store);
 
     let app = Application::build(app_state, prod::APP_ADDRESS)
         .await

auth-service/src/routes/login.rs

@@ -1,9 +1,11 @@
+use std::borrow::Borrow;
+
 use axum::{extract::State, http::StatusCode, response::IntoResponse, Json};
 use axum_extra::extract::CookieJar;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    domain::{AuthAPIError, Email, Password},
+    domain::{AuthAPIError, Email, LoginAttemptId, Password, TwoFACode},
     store::AppState,
     utils::auth::generate_auth_cookie,
 };
@@ -64,26 +66,50 @@ pub async fn login(
     };
 
     match user.requires_2fa {
-        true => handle_2fa(jar).await,
+        true => handle_2fa(&user.email, &state, jar).await,
         false => handle_no_2fa(&user.email, jar).await,
     }
 }
 async fn handle_2fa(
+    email: &Email,
+    state: &AppState,
     jar: CookieJar,
+
 ) -> (
     CookieJar,
     Result<(StatusCode, Json<LoginResponse>), AuthAPIError>,
 ) {
-    (
-        jar,
-        Ok((
-            StatusCode::PARTIAL_CONTENT,
-            Json(LoginResponse::TwoFactorAuth(TwoFactorAuthResponse {
-                message: "2FA required".to_owned(),
-                login_attempt_id: "123456".to_owned(),
-            })),
-        )),
-    )
+    let login_attempt_id = match LoginAttemptId::parse( uuid::Uuid::new_v4().to_string()) {
+        Ok(id) => id,
+        Err(_) => return (jar, Err(AuthAPIError::UnexpectedError))
+    };
+    let two_fa_code = TwoFACode::default();
+
+    match state.two_fa_code_store.write().await.add_code(email.clone(), login_attempt_id.clone(), two_fa_code).await {
+        Err(e) => return (jar, Err(AuthAPIError::UnexpectedError)),
+        Ok(_) => {
+            let auth_cookie = match generate_auth_cookie(email) {
+                Ok(cookie) => cookie,
+                Err(_) => return (jar, Err(AuthAPIError::UnexpectedError)),
+            };
+
+            let updated_jar = jar.add(auth_cookie);
+
+            (
+                updated_jar,
+                Ok((
+                    StatusCode::PARTIAL_CONTENT,
+                    Json(LoginResponse::TwoFactorAuth(TwoFactorAuthResponse {
+                        message: "2FA required".to_owned(),
+                        login_attempt_id: login_attempt_id.to_string(),
+                    })),
+                )),
+            )
+        }
+    }
+
+
+
 }
 
 async fn handle_no_2fa(

auth-service/src/services/hashmap_two_fa_code_store.rs

@@ -0,0 +1,65 @@
+use std::collections::HashMap;
+
+use crate::domain::{
+  data_stores::{LoginAttemptId, TwoFACode, TwoFACodeStore, TwoFACodeStoreError},
+  email::{self, Email}
+};
+
+#[derive(Default)]
+pub struct HashmapTwoFACodeStore {
+  pub codes: HashMap<Email, (LoginAttemptId, TwoFACode)>,
+}
+
+#[async_trait::async_trait]
+impl TwoFACodeStore for HashmapTwoFACodeStore {
+  async fn add_code(&mut self, email: Email, login_attempt_id: LoginAttemptId, code: TwoFACode) -> Result<(), TwoFACodeStoreError> {
+    if self.get_code(&email).await.is_ok() {
+      return Err(TwoFACodeStoreError::CodeAlreadyExists)
+    }
+    self.codes.insert(email.clone(), (login_attempt_id, code));
+
+    if self.get_code(&email).await.is_err() {
+      return Err(TwoFACodeStoreError::UnexpectedError)
+    }
+    Ok(())
+  }
+  async fn remove_code(&mut self, email: &Email) -> Result<(), TwoFACodeStoreError> {
+    match self.codes.remove(&email) {
+      Some(_) => Ok(()),
+      None => Err(TwoFACodeStoreError::LoginAttemptIdNotFound)
+    }
+  }
+  async fn get_code(&self, email: &Email) -> Result<(LoginAttemptId, TwoFACode), TwoFACodeStoreError> {
+    match self.codes.get(email) {
+      Some((login_attempt_id, code)) => Ok((login_attempt_id.clone(), code.clone())),
+      None => Err(TwoFACodeStoreError::LoginAttemptIdNotFound)
+    }
+  }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::domain::Email;
+    use super::*;
+    use uuid::Uuid;
+
+  #[tokio::test]
+  async fn test_add_remove_get_methods() {
+    let email = Email::parse("email@email.com".to_owned()).expect("Failed to create email");
+    let mut store = HashmapTwoFACodeStore {
+      codes: HashMap::new()
+    };
+    let login_attempt_id = LoginAttemptId::parse(Uuid::new_v4().to_string()).expect("Failed to parse uuid");
+    let two_fa_code = TwoFACode::parse("123456".to_owned()).expect("Failed to parse code");
+
+    let inserted = store.add_code(email.clone(), login_attempt_id.clone(), two_fa_code.clone()).await;
+    assert!(inserted.is_ok());
+    assert_eq!(store.codes.is_empty(), false);
+
+    let code = store.get_code(&email).await;
+    assert!(code.is_ok());
+    let (id, code) = code.unwrap();
+    assert_eq!(id, login_attempt_id);
+    assert_eq!(code, two_fa_code);
+  }
+}

auth-service/src/services/mod.rs

@@ -1,5 +1,7 @@
 mod hashmap_banned_token_store;
 mod hashmap_user_store;
+mod hashmap_two_fa_code_store;
 
 pub use hashmap_banned_token_store::*;
 pub use hashmap_user_store::*;
+pub use hashmap_two_fa_code_store::*;

auth-service/src/store/app_state.rs

@@ -2,22 +2,25 @@ use std::sync::Arc;
 
 use tokio::sync::RwLock;
 
-use crate::domain::{BannedTokenStore, UserStore};
+use crate::domain::{BannedTokenStore, TwoFACodeStore, UserStore};
 
 pub type UserStoreType = Arc<RwLock<dyn UserStore + Send + Sync>>;
 pub type BannedTokenStoreType = Arc<RwLock<dyn BannedTokenStore + Send + Sync>>;
+pub type TwoFACodeStoreType = Arc<RwLock<dyn TwoFACodeStore + Send + Sync>>;
 
 #[derive(Clone)]
 pub struct AppState {
     pub user_store: UserStoreType,
     pub banned_tokens_store: BannedTokenStoreType,
+    pub two_fa_code_store: TwoFACodeStoreType,
 }
 
 impl AppState {
-    pub fn new(user_store: UserStoreType, banned_tokens_store: BannedTokenStoreType) -> Self {
+    pub fn new(user_store: UserStoreType, banned_tokens_store: BannedTokenStoreType, two_fa_code_store: TwoFACodeStoreType) -> Self {
         Self {
             user_store,
             banned_tokens_store,
+            two_fa_code_store
         }
     }
 }