Bug: feat: Use redis in Begin and Verify TOTP



- Currently begin TOTP is directly replacing the secret in the db. This will cause issues when user didn't complete generating recovery codes. So we should keep the new secret in redis temporarily.
- And Verify TOTP should not complete the TOTP flow as user should complete the generation of recovery codes.
- There should be a new API which verifies the TOTP with the redis's secret and update the DB.
diff --git a/crates/api_models/src/user.rs b/crates/api_models/src/user.rs
index 5423fc830a5..0c8678d2ef8 100644
--- a/crates/api_models/src/user.rs
+++ b/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)]
diff --git a/crates/router/src/consts/user.rs b/crates/router/src/consts/user.rs
index 2f347dc3d3e..3c6cd8b6ccd 100644
--- a/crates/router/src/consts/user.rs
+++ b/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
diff --git a/crates/router/src/core/errors/user.rs b/crates/router/src/core/errors/user.rs
index adcf8794b83..865bc1d4124 100644
--- a/crates/router/src/core/errors/user.rs
+++ b/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",
         }
     }
 }
diff --git a/crates/router/src/core/user.rs b/crates/router/src/core/user.rs
index 4fd9fa865cd..7b6e8ebd365 100644
--- a/crates/router/src/core/user.rs
+++ b/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 key_store = user_from_db.get_or_create_key_store(&state).await?;
+    let secret = totp.get_secret_base32().into();
 
-    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,54 +1657,93 @@ 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 {
-            return Err(UserErrors::TotpNotSetup.into());
-        }
+    if user_from_db.get_totp_status() != TotpStatus::Set {
+        return Err(UserErrors::TotpNotSetup.into());
+    }
 
-        let user_totp_secret = user_from_db
-            .decrypt_and_get_totp_secret(&state)
-            .await?
-            .ok_or(UserErrors::InternalServerError)?;
+    let user_totp_secret = user_from_db
+        .decrypt_and_get_totp_secret(&state)
+        .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()
-        {
-            return Err(UserErrors::InvalidTotp.into());
-        }
+    if totp
+        .generate_current()
+        .change_context(UserErrors::InternalServerError)?
+        != req.totp.expose()
+    {
+        return Err(UserErrors::InvalidTotp.into());
+    }
 
-        if user_from_db.get_totp_status() == TotpStatus::InProgress {
-            state
-                .store
-                .update_user_by_user_id(
-                    user_from_db.get_user_id(),
-                    storage_user::UserUpdate::TotpUpdate {
-                        totp_status: Some(TotpStatus::Set),
-                        totp_secret: None,
-                        totp_recovery_codes: None,
-                    },
-                )
-                .await
-                .change_context(UserErrors::InternalServerError)?;
-        }
+    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 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 key_store = user_from_db.get_or_create_key_store(&state).await?;
 
-    auth::cookies::set_cookie_response(
-        user_api::TokenResponse {
-            token: token.clone(),
-            token_type: next_flow.get_flow().into(),
-        },
-        token,
-    )
+    state
+        .store
+        .update_user_by_user_id(
+            &user_token.user_id,
+            storage_user::UserUpdate::TotpUpdate {
+                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 _ = tfa_utils::delete_totp_secret_from_redis(&state, &user_token.user_id)
+        .await
+        .map_err(|e| logger::error!(?e));
+
+    // 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(
diff --git a/crates/router/src/routes/app.rs b/crates/router/src/routes/app.rs
index bad5ab5a926..b38479cd2b6 100644
--- a/crates/router/src/routes/app.rs
+++ b/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("/recovery_code")
-                .service(web::resource("/verify").route(web::post().to(verify_recovery_code)))
-                .service(web::resource("/generate").route(web::post().to(generate_recovery_codes))),
+            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::get().to(generate_recovery_codes)),
+                        ),
+                )
+                .service(
+                    web::resource("/terminate").route(web::get().to(terminate_two_factor_auth)),
+                ),
         );
 
         #[cfg(feature = "email")]
diff --git a/crates/router/src/routes/lock_utils.rs b/crates/router/src/routes/lock_utils.rs
index 75821bbd2ba..9e53eb35473 100644
--- a/crates/router/src/routes/lock_utils.rs
+++ b/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
diff --git a/crates/router/src/routes/user.rs b/crates/router/src/routes/user.rs
index 6b8015ac3e5..5ba7ec8da25 100644
--- a/crates/router/src/routes/user.rs
+++ b/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(
diff --git a/crates/router/src/utils/user.rs b/crates/router/src/utils/user.rs
index 9c67dbfcab5..db1d0ea2508 100644
--- a/crates/router/src/utils/user.rs
+++ b/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")
+}
diff --git a/crates/router/src/utils/user/two_factor_auth.rs b/crates/router/src/utils/user/two_factor_auth.rs
index e9e3f66005c..6bd693c07f5 100644
--- a/crates/router/src/utils/user/two_factor_auth.rs
+++ b/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(
diff --git a/crates/router_env/src/logger/types.rs b/crates/router_env/src/logger/types.rs
index 35b55b6fb97..1bfc20ff1ca 100644
--- a/crates/router_env/src/logger/types.rs
+++ b/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

## Type of Change
<!-- Put an `x` in the boxes that apply -->

- [ ] Bugfix
- [x] New feature
- [ ] Enhancement
- [ ] Refactoring
- [ ] Dependency updates
- [ ] Documentation
- [ ] CI/CD

## Description
<!-- Describe your changes in detail -->
- Begin TOTP will now insert the TOTP secret in redis instead of DB directly, which allows the old secret to stay if it exists and recovery codes are also removed from Begin TOTP.
- Verify TOTP will now send 200 OK and puts the timestamp in the redis when the TOTP is verified.
- Update TOTP is a new API which will take the TOTP and verifies the TOTP against the redis secret instead of db secret. If the TOTP is correct, it will replace the redis secret in the db.

### Additional Changes

- [x] This PR modifies the API contract
- [ ] This PR modifies the database schema
- [ ] This PR modifies application configuration/environment variables

<!--
Provide links to the files with corresponding changes.

Following are the paths where you can find config files:
1. `config`
2. `crates/router/src/configs`
3. `loadtest/config`
-->


## Motivation and Context
<!--
Why is this change required? What problem does it solve?
If it fixes an open issue, please link to the issue here.

If you don't have an issue, we'd recommend starting with one first so the PR
can focus on the implementation (unless it is an obvious bug or documentation fix
that will have little conversation).
-->
Closes #4764.
Closes #4707.

## How did you test it?
<!--
Did you write an integration/unit/API test to verify the code changes?
Or did you test this change manually (provide relevant screenshots)?
-->
1. TOTP Routes
	1. Begin TOTP
		```curl
		curl --location 'http://localhost:8080/user/2fa/totp/begin' \
		--header 'Authorization: Bearer SPT with Purpose as TOTP'
		```
		```
		{
		    "secret": {
		        "secret": "Q54GEKFSVTP4FU3YHDFMEN5KMGD2P3IN",
		        "totp_url": "otpauth://totp/Hyperswitch:mani.dchandra%40juspay.in?secret=Q54GEKFSVTP4FU3YHDFMEN5KMGD2P3IN&issuer=Hyperswitch"
		    }
		}
		```
	2. Verify TOTP
		```
		curl --location 'http://localhost:8080/user/2fa/totp/verify' \
		--header 'Content-Type: application/json' \
		--header 'Authorization: Bearer SPT with Purpose as TOTP' \
		--data '{
		    "totp": "701394"
		}'
		```
		Response will be 200 OK if the TOTP is correct.
2. Recovery Codes Routes
	1. Verify Recovery Code
		```
		curl --location 'http://localhost:8080/user/2fa/recovery_code/verify' \
		--header 'Content-Type: application/json' \
		--header 'Authorization: Bearer SPT with Purpose as TOTP' \
		--data '{
		    "recovery_code": "m14a-0ni6"
		}'
		```
		Response will be 200 OK if the Recovery code is correct.
	2. Generate Recovery codes
		```
		curl --location 'http://localhost:8080/user/2fa/recovery_code/generate' \
		--header 'Authorization: Bearer SPT with Purpose as TOTP' \
		```
		```
		{
		    "recovery_codes": [
		        "w5fM-NLm0",
		        "TumH-oyXE",
		        "wDIr-zEmu",
		        "HZjm-e16M",
		        "Q4qB-MupL",
		        "5BtN-vRuY",
		        "4vG6-URGf",
		        "Dbq8-n5V1"
		    ]
		}
		
		```

## Checklist
<!-- Put an `x` in the boxes that apply -->

- [x] I formatted the code `cargo +nightly fmt --all`
- [x] I addressed lints thrown by `cargo clippy`
- [x] I reviewed the submitted code
- [ ] I added unit tests for my changes where possible
