feat(users): Add redis in Begin and Verify TOTP and create a new API that updates TOTP (#4765)

crates/api_models/src/user.rs

@@ -255,12 +255,11 @@ pub struct BeginTotpResponse {
 pub struct TotpSecret {
     pub secret: Secret<String>,
     pub totp_url: Secret<String>,
-    pub recovery_codes: Vec<Secret<String>>,
 }
 
 #[derive(Debug, serde::Deserialize, serde::Serialize)]
 pub struct VerifyTotpRequest {
-    pub totp: Option<Secret<String>>,
+    pub totp: Secret<String>,
 }
 
 #[derive(Debug, serde::Deserialize, serde::Serialize)]

crates/router/src/consts/user.rs

@@ -13,5 +13,7 @@ pub const TOTP_TOLERANCE: u8 = 1;
 pub const MAX_PASSWORD_LENGTH: usize = 70;
 pub const MIN_PASSWORD_LENGTH: usize = 8;
 
-pub const TOTP_PREFIX: &str = "TOTP_";
+pub const REDIS_TOTP_PREFIX: &str = "TOTP_";
 pub const REDIS_RECOVERY_CODE_PREFIX: &str = "RC_";
+pub const REDIS_TOTP_SECRET_PREFIX: &str = "TOTP_SEC_";
+pub const REDIS_TOTP_SECRET_TTL_IN_SECS: i64 = 5 * 60; // 5 minutes

crates/router/src/core/errors/user.rs

@@ -78,6 +78,8 @@ pub enum UserErrors {
     TwoFactorAuthRequired,
     #[error("TwoFactorAuthNotSetup")]
     TwoFactorAuthNotSetup,
+    #[error("TOTP secret not found")]
+    TotpSecretNotFound,
 }
 
 impl common_utils::errors::ErrorSwitch<api_models::errors::types::ApiErrorResponse> for UserErrors {
@@ -199,6 +201,9 @@ impl common_utils::errors::ErrorSwitch<api_models::errors::types::ApiErrorRespon
             Self::TwoFactorAuthNotSetup => {
                 AER::BadRequest(ApiError::new(sub_code, 41, self.get_error_message(), None))
             }
+            Self::TotpSecretNotFound => {
+                AER::BadRequest(ApiError::new(sub_code, 42, self.get_error_message(), None))
+            }
         }
     }
 }
@@ -241,6 +246,7 @@ impl UserErrors {
             Self::InvalidRecoveryCode => "Invalid Recovery Code",
             Self::TwoFactorAuthRequired => "Two factor auth required",
             Self::TwoFactorAuthNotSetup => "Two factor auth not setup",
+            Self::TotpSecretNotFound => "TOTP secret not found",
         }
     }
 }

crates/router/src/core/user.rs

@@ -1633,41 +1633,14 @@ pub async fn begin_totp(
     }
 
     let totp = tfa_utils::generate_default_totp(user_from_db.get_email(), None)?;
-    let recovery_codes = domain::RecoveryCodes::generate_new();
+    let secret = totp.get_secret_base32().into();
 
-    let key_store = user_from_db.get_or_create_key_store(&state).await?;
-
-    state
-        .store
-        .update_user_by_user_id(
-            user_from_db.get_user_id(),
-            storage_user::UserUpdate::TotpUpdate {
-                totp_status: Some(TotpStatus::InProgress),
-                totp_secret: Some(
-                    // TODO: Impl conversion trait for User and move this there
-                    domain::types::encrypt::<String, masking::WithType>(
-                        totp.get_secret_base32().into(),
-                        key_store.key.peek(),
-                    )
-                    .await
-                    .change_context(UserErrors::InternalServerError)?
-                    .into(),
-                ),
-                totp_recovery_codes: Some(
-                    recovery_codes
-                        .get_hashed()
-                        .change_context(UserErrors::InternalServerError)?,
-                ),
-            },
-        )
-        .await
-        .change_context(UserErrors::InternalServerError)?;
+    tfa_utils::insert_totp_secret_in_redis(&state, &user_token.user_id, &secret).await?;
 
     Ok(ApplicationResponse::Json(user_api::BeginTotpResponse {
         secret: Some(user_api::TotpSecret {
-            secret: totp.get_secret_base32().into(),
+            secret,
             totp_url: totp.get_url().into(),
-            recovery_codes: recovery_codes.into_inner(),
         }),
     }))
 }
@@ -1684,8 +1657,7 @@ pub async fn verify_totp(
         .change_context(UserErrors::InternalServerError)?
         .into();
 
-    if let Some(user_totp) = req.totp {
-        if user_from_db.get_totp_status() == TotpStatus::NotSet {
+    if user_from_db.get_totp_status() != TotpStatus::Set {
         return Err(UserErrors::TotpNotSetup.into());
     }
 
@@ -1694,44 +1666,84 @@ pub async fn verify_totp(
         .await?
         .ok_or(UserErrors::InternalServerError)?;
 
-        let totp =
-            tfa_utils::generate_default_totp(user_from_db.get_email(), Some(user_totp_secret))?;
+    let totp = tfa_utils::generate_default_totp(user_from_db.get_email(), Some(user_totp_secret))?;
 
     if totp
         .generate_current()
         .change_context(UserErrors::InternalServerError)?
-            != user_totp.expose()
+        != req.totp.expose()
     {
         return Err(UserErrors::InvalidTotp.into());
     }
 
-        if user_from_db.get_totp_status() == TotpStatus::InProgress {
+    tfa_utils::insert_totp_in_redis(&state, &user_token.user_id).await?;
+
+    Ok(ApplicationResponse::StatusOk)
+}
+
+pub async fn update_totp(
+    state: AppState,
+    user_token: auth::UserFromSinglePurposeToken,
+    req: user_api::VerifyTotpRequest,
+) -> UserResponse<()> {
+    let user_from_db: domain::UserFromStorage = state
+        .store
+        .find_user_by_id(&user_token.user_id)
+        .await
+        .change_context(UserErrors::InternalServerError)?
+        .into();
+
+    let new_totp_secret = tfa_utils::get_totp_secret_from_redis(&state, &user_token.user_id)
+        .await?
+        .ok_or(UserErrors::TotpSecretNotFound)?;
+
+    let totp = tfa_utils::generate_default_totp(user_from_db.get_email(), Some(new_totp_secret))?;
+
+    if totp
+        .generate_current()
+        .change_context(UserErrors::InternalServerError)?
+        != req.totp.expose()
+    {
+        return Err(UserErrors::InvalidTotp.into());
+    }
+
+    let key_store = user_from_db.get_or_create_key_store(&state).await?;
+
     state
         .store
         .update_user_by_user_id(
-                    user_from_db.get_user_id(),
+            &user_token.user_id,
             storage_user::UserUpdate::TotpUpdate {
-                        totp_status: Some(TotpStatus::Set),
-                        totp_secret: None,
+                totp_status: None,
+                totp_secret: Some(
+                    // TODO: Impl conversion trait for User and move this there
+                    domain::types::encrypt::<String, masking::WithType>(
+                        totp.get_secret_base32().into(),
+                        key_store.key.peek(),
+                    )
+                    .await
+                    .change_context(UserErrors::InternalServerError)?
+                    .into(),
+                ),
+
                 totp_recovery_codes: None,
             },
         )
         .await
         .change_context(UserErrors::InternalServerError)?;
-        }
-    }
 
-    let current_flow = domain::CurrentFlow::new(user_token.origin, domain::SPTFlow::TOTP.into())?;
-    let next_flow = current_flow.next(user_from_db, &state).await?;
-    let token = next_flow.get_token(&state).await?;
+    let _ = tfa_utils::delete_totp_secret_from_redis(&state, &user_token.user_id)
+        .await
+        .map_err(|e| logger::error!(?e));
 
-    auth::cookies::set_cookie_response(
-        user_api::TokenResponse {
-            token: token.clone(),
-            token_type: next_flow.get_flow().into(),
-        },
-        token,
-    )
+    // This is not the main task of this API, so we don't throw error if this fails.
+    // Any following API which requires TOTP will throw error if TOTP is not set in redis
+    // and FE will ask user to enter TOTP again
+    let _ = tfa_utils::insert_totp_in_redis(&state, &user_token.user_id)
+        .await
+        .map_err(|e| logger::error!(?e));
+
+    Ok(ApplicationResponse::StatusOk)
 }
 
 pub async fn generate_recovery_codes(

crates/router/src/routes/app.rs

@@ -1209,17 +1209,33 @@ impl User {
                 web::resource("/data")
                     .route(web::get().to(get_multiple_dashboard_metadata))
                     .route(web::post().to(set_dashboard_metadata)),
-            )
-            .service(web::resource("/totp/begin").route(web::get().to(totp_begin)))
-            .service(web::resource("/totp/verify").route(web::post().to(totp_verify)))
-            .service(
-                web::resource("/2fa/terminate").route(web::get().to(terminate_two_factor_auth)),
             );
 
+        // Two factor auth routes
         route = route.service(
+            web::scope("/2fa")
+                .service(
+                    web::scope("/totp")
+                        .service(web::resource("/begin").route(web::get().to(totp_begin)))
+                        .service(
+                            web::resource("/verify")
+                                .route(web::post().to(totp_verify))
+                                .route(web::put().to(totp_update)),
+                        ),
+                )
+                .service(
                     web::scope("/recovery_code")
-                .service(web::resource("/verify").route(web::post().to(verify_recovery_code)))
-                .service(web::resource("/generate").route(web::post().to(generate_recovery_codes))),
+                        .service(
+                            web::resource("/verify").route(web::post().to(verify_recovery_code)),
+                        )
+                        .service(
+                            web::resource("/generate")
+                                .route(web::get().to(generate_recovery_codes)),
+                        ),
+                )
+                .service(
+                    web::resource("/terminate").route(web::get().to(terminate_two_factor_auth)),
+                ),
         );
 
         #[cfg(feature = "email")]

crates/router/src/routes/lock_utils.rs

@@ -215,9 +215,11 @@ impl From<Flow> for ApiIdentifier {
             | Flow::UpdateUserAccountDetails
             | Flow::TotpBegin
             | Flow::TotpVerify
+            | Flow::TotpUpdate
             | Flow::RecoveryCodeVerify
             | Flow::RecoveryCodesGenerate
             | Flow::TerminateTwoFactorAuth => Self::User,
+
             Flow::ListRoles
             | Flow::GetRole
             | Flow::GetRoleFromToken

crates/router/src/routes/user.rs

@@ -684,6 +684,24 @@ pub async fn verify_recovery_code(
     .await
 }
 
+pub async fn totp_update(
+    state: web::Data<AppState>,
+    req: HttpRequest,
+    json_payload: web::Json<user_api::VerifyTotpRequest>,
+) -> HttpResponse {
+    let flow = Flow::TotpUpdate;
+    Box::pin(api::server_wrap(
+        flow,
+        state.clone(),
+        &req,
+        json_payload.into_inner(),
+        |state, user, req_body, _| user_core::update_totp(state, user, req_body),
+        &auth::SinglePurposeJWTAuth(TokenPurpose::TOTP),
+        api_locking::LockAction::NotApplicable,
+    ))
+    .await
+}
+
 pub async fn generate_recovery_codes(state: web::Data<AppState>, req: HttpRequest) -> HttpResponse {
     let flow = Flow::RecoveryCodesGenerate;
     Box::pin(api::server_wrap(

crates/router/src/utils/user.rs

@@ -1,9 +1,10 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 use api_models::user as user_api;
 use common_utils::errors::CustomResult;
 use diesel_models::{enums::UserStatus, user_role::UserRole};
 use error_stack::ResultExt;
+use redis_interface::RedisConnectionPool;
 
 use crate::{
     core::errors::{StorageError, UserErrors, UserResult},
@@ -191,3 +192,11 @@ pub fn get_token_from_signin_response(resp: &user_api::SignInResponse) -> maskin
         user_api::SignInResponse::MerchantSelect(data) => data.token.clone(),
     }
 }
+
+pub fn get_redis_connection(state: &AppState) -> UserResult<Arc<RedisConnectionPool>> {
+    state
+        .store
+        .get_redis_conn()
+        .change_context(UserErrors::InternalServerError)
+        .attach_printable("Failed to get redis connection")
+}

crates/router/src/utils/user/two_factor_auth.rs

@@ -1,9 +1,6 @@
-use std::sync::Arc;
-
 use common_utils::pii;
 use error_stack::ResultExt;
-use masking::ExposeInterface;
-use redis_interface::RedisConnectionPool;
+use masking::{ExposeInterface, PeekInterface};
 use totp_rs::{Algorithm, TOTP};
 
 use crate::{
@@ -35,8 +32,8 @@ pub fn generate_default_totp(
 }
 
 pub async fn check_totp_in_redis(state: &AppState, user_id: &str) -> UserResult<bool> {
-    let redis_conn = get_redis_connection(state)?;
-    let key = format!("{}{}", consts::user::TOTP_PREFIX, user_id);
+    let redis_conn = super::get_redis_connection(state)?;
+    let key = format!("{}{}", consts::user::REDIS_TOTP_PREFIX, user_id);
     redis_conn
         .exists::<()>(&key)
         .await
@@ -44,7 +41,7 @@ pub async fn check_totp_in_redis(state: &AppState, user_id: &str) -> UserResult<
 }
 
 pub async fn check_recovery_code_in_redis(state: &AppState, user_id: &str) -> UserResult<bool> {
-    let redis_conn = get_redis_connection(state)?;
+    let redis_conn = super::get_redis_connection(state)?;
     let key = format!("{}{}", consts::user::REDIS_RECOVERY_CODE_PREFIX, user_id);
     redis_conn
         .exists::<()>(&key)
@@ -52,16 +49,62 @@ pub async fn check_recovery_code_in_redis(state: &AppState, user_id: &str) -> Us
         .change_context(UserErrors::InternalServerError)
 }
 
-fn get_redis_connection(state: &AppState) -> UserResult<Arc<RedisConnectionPool>> {
-    state
-        .store
-        .get_redis_conn()
+pub async fn insert_totp_in_redis(state: &AppState, user_id: &str) -> UserResult<()> {
+    let redis_conn = super::get_redis_connection(state)?;
+    let key = format!("{}{}", consts::user::REDIS_TOTP_PREFIX, user_id);
+    redis_conn
+        .set_key_with_expiry(
+            key.as_str(),
+            common_utils::date_time::now_unix_timestamp(),
+            state.conf.user.two_factor_auth_expiry_in_secs,
+        )
+        .await
         .change_context(UserErrors::InternalServerError)
-        .attach_printable("Failed to get redis connection")
+}
+
+pub async fn insert_totp_secret_in_redis(
+    state: &AppState,
+    user_id: &str,
+    secret: &masking::Secret<String>,
+) -> UserResult<()> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .set_key_with_expiry(
+            &get_totp_secret_key(user_id),
+            secret.peek(),
+            consts::user::REDIS_TOTP_SECRET_TTL_IN_SECS,
+        )
+        .await
+        .change_context(UserErrors::InternalServerError)
+}
+
+pub async fn get_totp_secret_from_redis(
+    state: &AppState,
+    user_id: &str,
+) -> UserResult<Option<masking::Secret<String>>> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .get_key::<Option<String>>(&get_totp_secret_key(user_id))
+        .await
+        .change_context(UserErrors::InternalServerError)
+        .map(|secret| secret.map(Into::into))
+}
+
+pub async fn delete_totp_secret_from_redis(state: &AppState, user_id: &str) -> UserResult<()> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .delete_key(&get_totp_secret_key(user_id))
+        .await
+        .change_context(UserErrors::InternalServerError)
+        .map(|_| ())
+}
+
+fn get_totp_secret_key(user_id: &str) -> String {
+    format!("{}{}", consts::user::REDIS_TOTP_SECRET_PREFIX, user_id)
 }
 
 pub async fn insert_recovery_code_in_redis(state: &AppState, user_id: &str) -> UserResult<()> {
-    let redis_conn = get_redis_connection(state)?;
+    let redis_conn = super::get_redis_connection(state)?;
     let key = format!("{}{}", consts::user::REDIS_RECOVERY_CODE_PREFIX, user_id);
     redis_conn
         .set_key_with_expiry(

crates/router_env/src/logger/types.rs

@@ -406,6 +406,8 @@ pub enum Flow {
     TotpBegin,
     /// Verify TOTP
     TotpVerify,
+    /// Update TOTP secret
+    TotpUpdate,
     /// Verify Access Code
     RecoveryCodeVerify,
     /// Generate or Regenerate recovery codes