fix(users): Add max wrong attempts for two-fa (#6247)

crates/api_models/src/events/user.rs

@@ -18,8 +18,9 @@ use crate::user::{
     ResetPasswordRequest, RotatePasswordRequest, SendVerifyEmailRequest, SignUpRequest,
     SignUpWithMerchantIdRequest, SsoSignInRequest, SwitchMerchantRequest,
     SwitchOrganizationRequest, SwitchProfileRequest, TokenResponse, TwoFactorAuthStatusResponse,
-    UpdateUserAccountDetailsRequest, UpdateUserAuthenticationMethodRequest, UserFromEmailRequest,
+    TwoFactorStatus, UpdateUserAccountDetailsRequest, UpdateUserAuthenticationMethodRequest,
-    UserMerchantCreate, VerifyEmailRequest, VerifyRecoveryCodeRequest, VerifyTotpRequest,
+    UserFromEmailRequest, UserMerchantCreate, VerifyEmailRequest, VerifyRecoveryCodeRequest,
+    VerifyTotpRequest,
 };
 
 #[cfg(feature = "recon")]
@@ -62,6 +63,7 @@ common_utils::impl_api_event_type!(
         GetUserRoleDetailsResponseV2,
         TokenResponse,
         TwoFactorAuthStatusResponse,
+        TwoFactorStatus,
         UserFromEmailRequest,
         BeginTotpResponse,
         VerifyRecoveryCodeRequest,

crates/api_models/src/user.rs

@@ -196,6 +196,23 @@ pub struct TwoFactorAuthStatusResponse {
     pub recovery_code: bool,
 }
 
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
+pub struct TwoFactorAuthAttempts {
+    pub is_completed: bool,
+    pub remaining_attempts: u8,
+}
+
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
+pub struct TwoFactorAuthStatusResponseWithAttempts {
+    pub totp: TwoFactorAuthAttempts,
+    pub recovery_code: TwoFactorAuthAttempts,
+}
+
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
+pub struct TwoFactorStatus {
+    pub status: Option<TwoFactorAuthStatusResponseWithAttempts>,
+}
+
 #[derive(Debug, serde::Deserialize, serde::Serialize)]
 pub struct UserFromEmailRequest {
     pub token: Secret<String>,

crates/router/src/consts/user.rs

@@ -15,6 +15,10 @@ pub const TOTP_DIGITS: usize = 6;
 pub const TOTP_VALIDITY_DURATION_IN_SECONDS: u64 = 30;
 /// Number of totps allowed as network delay. 1 would mean one totp before current totp and one totp after are valids.
 pub const TOTP_TOLERANCE: u8 = 1;
+/// Number of maximum attempts user has for totp
+pub const TOTP_MAX_ATTEMPTS: u8 = 4;
+/// Number of maximum attempts user has for recovery code
+pub const RECOVERY_CODE_MAX_ATTEMPTS: u8 = 4;
 
 pub const MAX_PASSWORD_LENGTH: usize = 70;
 pub const MIN_PASSWORD_LENGTH: usize = 8;
@@ -23,6 +27,10 @@ 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 = 15 * 60; // 15 minutes
+pub const REDIS_TOTP_ATTEMPTS_PREFIX: &str = "TOTP_ATTEMPTS_";
+pub const REDIS_RECOVERY_CODE_ATTEMPTS_PREFIX: &str = "RC_ATTEMPTS_";
+pub const REDIS_TOTP_ATTEMPTS_TTL_IN_SECS: i64 = 5 * 60; // 5 mins
+pub const REDIS_RECOVERY_CODE_ATTEMPTS_TTL_IN_SECS: i64 = 10 * 60; // 10 mins
 
 pub const REDIS_SSO_PREFIX: &str = "SSO_";
 pub const REDIS_SSO_TTL: i64 = 5 * 60; // 5 minutes

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

@@ -90,6 +90,10 @@ pub enum UserErrors {
     SSOFailed,
     #[error("profile_id missing in JWT")]
     JwtProfileIdMissing,
+    #[error("Maximum attempts reached for TOTP")]
+    MaxTotpAttemptsReached,
+    #[error("Maximum attempts reached for Recovery Code")]
+    MaxRecoveryCodeAttemptsReached,
 }
 
 impl common_utils::errors::ErrorSwitch<api_models::errors::types::ApiErrorResponse> for UserErrors {
@@ -229,6 +233,12 @@ impl common_utils::errors::ErrorSwitch<api_models::errors::types::ApiErrorRespon
             Self::JwtProfileIdMissing => {
                 AER::Unauthorized(ApiError::new(sub_code, 47, self.get_error_message(), None))
             }
+            Self::MaxTotpAttemptsReached => {
+                AER::BadRequest(ApiError::new(sub_code, 48, self.get_error_message(), None))
+            }
+            Self::MaxRecoveryCodeAttemptsReached => {
+                AER::BadRequest(ApiError::new(sub_code, 49, self.get_error_message(), None))
+            }
         }
     }
 }
@@ -269,6 +279,8 @@ impl UserErrors {
             Self::InvalidTotp => "Invalid TOTP",
             Self::TotpRequired => "TOTP required",
             Self::InvalidRecoveryCode => "Invalid Recovery Code",
+            Self::MaxTotpAttemptsReached => "Maximum attempts reached for TOTP",
+            Self::MaxRecoveryCodeAttemptsReached => "Maximum attempts reached for 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

@@ -1686,6 +1686,13 @@ pub async fn verify_totp(
         return Err(UserErrors::TotpNotSetup.into());
     }
 
+    let user_totp_attempts =
+        tfa_utils::get_totp_attempts_from_redis(&state, &user_token.user_id).await?;
+
+    if user_totp_attempts >= consts::user::TOTP_MAX_ATTEMPTS {
+        return Err(UserErrors::MaxTotpAttemptsReached.into());
+    }
+
     let user_totp_secret = user_from_db
         .decrypt_and_get_totp_secret(&state)
         .await?
@@ -1702,6 +1709,13 @@ pub async fn verify_totp(
         .change_context(UserErrors::InternalServerError)?
         != req.totp.expose()
     {
+        let _ = tfa_utils::insert_totp_attempts_in_redis(
+            &state,
+            &user_token.user_id,
+            user_totp_attempts + 1,
+        )
+        .await
+        .inspect_err(|error| logger::error!(?error));
         return Err(UserErrors::InvalidTotp.into());
     }
 
@@ -1856,15 +1870,31 @@ pub async fn verify_recovery_code(
         return Err(UserErrors::TwoFactorAuthNotSetup.into());
     }
 
+    let user_recovery_code_attempts =
+        tfa_utils::get_recovery_code_attempts_from_redis(&state, &user_token.user_id).await?;
+
+    if user_recovery_code_attempts >= consts::user::RECOVERY_CODE_MAX_ATTEMPTS {
+        return Err(UserErrors::MaxRecoveryCodeAttemptsReached.into());
+    }
+
     let mut recovery_codes = user_from_db
         .get_recovery_codes()
         .ok_or(UserErrors::InternalServerError)?;
 
-    let matching_index = utils::user::password::get_index_for_correct_recovery_code(
+    let Some(matching_index) = utils::user::password::get_index_for_correct_recovery_code(
         &req.recovery_code,
         &recovery_codes,
     )?
-    .ok_or(UserErrors::InvalidRecoveryCode)?;
+    else {
+        let _ = tfa_utils::insert_recovery_code_attempts_in_redis(
+            &state,
+            &user_token.user_id,
+            user_recovery_code_attempts + 1,
+        )
+        .await
+        .inspect_err(|error| logger::error!(?error));
+        return Err(UserErrors::InvalidRecoveryCode.into());
+    };
 
     tfa_utils::insert_recovery_code_in_redis(&state, user_from_db.get_user_id()).await?;
     let _ = recovery_codes.remove(matching_index);
@@ -1924,10 +1954,17 @@ pub async fn terminate_two_factor_auth(
         }
     }
 
-    let current_flow = domain::CurrentFlow::new(user_token, domain::SPTFlow::TOTP.into())?;
+    let current_flow = domain::CurrentFlow::new(user_token.clone(), 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_attempts_from_redis(&state, &user_token.user_id)
+        .await
+        .inspect_err(|error| logger::error!(?error));
+    let _ = tfa_utils::delete_recovery_code_attempts_from_redis(&state, &user_token.user_id)
+        .await
+        .inspect_err(|error| logger::error!(?error));
+
     auth::cookies::set_cookie_response(
         user_api::TokenResponse {
             token: token.clone(),
@@ -1950,6 +1987,40 @@ pub async fn check_two_factor_auth_status(
     ))
 }
 
+pub async fn check_two_factor_auth_status_with_attempts(
+    state: SessionState,
+    user_token: auth::UserIdFromAuth,
+) -> UserResponse<user_api::TwoFactorStatus> {
+    let user_from_db: domain::UserFromStorage = state
+        .global_store
+        .find_user_by_id(&user_token.user_id)
+        .await
+        .change_context(UserErrors::InternalServerError)?
+        .into();
+    if user_from_db.get_totp_status() == TotpStatus::NotSet {
+        return Ok(ApplicationResponse::Json(user_api::TwoFactorStatus {
+            status: None,
+        }));
+    };
+
+    let totp = user_api::TwoFactorAuthAttempts {
+        is_completed: tfa_utils::check_totp_in_redis(&state, &user_token.user_id).await?,
+        remaining_attempts: consts::user::TOTP_MAX_ATTEMPTS
+            - tfa_utils::get_totp_attempts_from_redis(&state, &user_token.user_id).await?,
+    };
+    let recovery_code = user_api::TwoFactorAuthAttempts {
+        is_completed: tfa_utils::check_recovery_code_in_redis(&state, &user_token.user_id).await?,
+        remaining_attempts: consts::user::RECOVERY_CODE_MAX_ATTEMPTS
+            - tfa_utils::get_recovery_code_attempts_from_redis(&state, &user_token.user_id).await?,
+    };
+    Ok(ApplicationResponse::Json(user_api::TwoFactorStatus {
+        status: Some(user_api::TwoFactorAuthStatusResponseWithAttempts {
+            totp,
+            recovery_code,
+        }),
+    }))
+}
+
 pub async fn create_user_authentication_method(
     state: SessionState,
     req: user_api::CreateUserAuthenticationMethodRequest,

crates/router/src/routes/app.rs

@@ -1853,7 +1853,12 @@ impl User {
         // Two factor auth routes
         route = route.service(
             web::scope("/2fa")
+                // TODO: to be deprecated
                 .service(web::resource("").route(web::get().to(user::check_two_factor_auth_status)))
+                .service(
+                    web::resource("/v2")
+                        .route(web::get().to(user::check_two_factor_auth_status_with_attempts)),
+                )
                 .service(
                     web::scope("/totp")
                         .service(web::resource("/begin").route(web::get().to(user::totp_begin)))

crates/router/src/routes/user.rs

@@ -682,6 +682,23 @@ pub async fn check_two_factor_auth_status(
     .await
 }
 
+pub async fn check_two_factor_auth_status_with_attempts(
+    state: web::Data<AppState>,
+    req: HttpRequest,
+) -> HttpResponse {
+    let flow = Flow::TwoFactorAuthStatus;
+    Box::pin(api::server_wrap(
+        flow,
+        state.clone(),
+        &req,
+        (),
+        |state, user, _, _| user_core::check_two_factor_auth_status_with_attempts(state, user),
+        &auth::SinglePurposeOrLoginTokenAuth(TokenPurpose::TOTP),
+        api_locking::LockAction::NotApplicable,
+    ))
+    .await
+}
+
 pub async fn get_sso_auth_url(
     state: web::Data<AppState>,
     req: HttpRequest,

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

@@ -139,3 +139,90 @@ pub async fn delete_recovery_code_from_redis(
         .change_context(UserErrors::InternalServerError)
         .map(|_| ())
 }
+
+fn get_totp_attempts_key(user_id: &str) -> String {
+    format!("{}{}", consts::user::REDIS_TOTP_ATTEMPTS_PREFIX, user_id)
+}
+fn get_recovery_code_attempts_key(user_id: &str) -> String {
+    format!(
+        "{}{}",
+        consts::user::REDIS_RECOVERY_CODE_ATTEMPTS_PREFIX,
+        user_id
+    )
+}
+
+pub async fn insert_totp_attempts_in_redis(
+    state: &SessionState,
+    user_id: &str,
+    user_totp_attempts: u8,
+) -> UserResult<()> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .set_key_with_expiry(
+            &get_totp_attempts_key(user_id),
+            user_totp_attempts,
+            consts::user::REDIS_TOTP_ATTEMPTS_TTL_IN_SECS,
+        )
+        .await
+        .change_context(UserErrors::InternalServerError)
+}
+pub async fn get_totp_attempts_from_redis(state: &SessionState, user_id: &str) -> UserResult<u8> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .get_key::<Option<u8>>(&get_totp_attempts_key(user_id))
+        .await
+        .change_context(UserErrors::InternalServerError)
+        .map(|v| v.unwrap_or(0))
+}
+
+pub async fn insert_recovery_code_attempts_in_redis(
+    state: &SessionState,
+    user_id: &str,
+    user_recovery_code_attempts: u8,
+) -> UserResult<()> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .set_key_with_expiry(
+            &get_recovery_code_attempts_key(user_id),
+            user_recovery_code_attempts,
+            consts::user::REDIS_RECOVERY_CODE_ATTEMPTS_TTL_IN_SECS,
+        )
+        .await
+        .change_context(UserErrors::InternalServerError)
+}
+
+pub async fn get_recovery_code_attempts_from_redis(
+    state: &SessionState,
+    user_id: &str,
+) -> UserResult<u8> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .get_key::<Option<u8>>(&get_recovery_code_attempts_key(user_id))
+        .await
+        .change_context(UserErrors::InternalServerError)
+        .map(|v| v.unwrap_or(0))
+}
+
+pub async fn delete_totp_attempts_from_redis(
+    state: &SessionState,
+    user_id: &str,
+) -> UserResult<()> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .delete_key(&get_totp_attempts_key(user_id))
+        .await
+        .change_context(UserErrors::InternalServerError)
+        .map(|_| ())
+}
+
+pub async fn delete_recovery_code_attempts_from_redis(
+    state: &SessionState,
+    user_id: &str,
+) -> UserResult<()> {
+    let redis_conn = super::get_redis_connection(state)?;
+    redis_conn
+        .delete_key(&get_recovery_code_attempts_key(user_id))
+        .await
+        .change_context(UserErrors::InternalServerError)
+        .map(|_| ())
+}