Skip to content

Commit

Permalink
feat(users): Create terminate 2fa API (#4731)
Browse files Browse the repository at this point in the history
  • Loading branch information
Riddhiagrawal001 committed May 23, 2024
1 parent ae77373 commit 42e5ef1
Show file tree
Hide file tree
Showing 10 changed files with 115 additions and 7 deletions.
5 changes: 5 additions & 0 deletions crates/api_models/src/user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,11 @@ pub struct TokenOnlyQueryParam {
pub token_only: Option<bool>,
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
pub struct SkipTwoFactorAuthQueryParam {
pub skip_two_factor_auth: Option<bool>,
}

#[derive(Debug, serde::Deserialize, serde::Serialize)]
pub struct TokenResponse {
pub token: Secret<String>,
Expand Down
1 change: 1 addition & 0 deletions crates/router/src/consts/user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,3 +14,4 @@ pub const MAX_PASSWORD_LENGTH: usize = 70;
pub const MIN_PASSWORD_LENGTH: usize = 8;

pub const TOTP_PREFIX: &str = "TOTP_";
pub const REDIS_RECOVERY_CODES_PREFIX: &str = "RC_";
12 changes: 12 additions & 0 deletions crates/router/src/core/errors/user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,10 @@ pub enum UserErrors {
InvalidTotp,
#[error("TotpRequired")]
TotpRequired,
#[error("TwoFactorAuthRequired")]
TwoFactorAuthRequired,
#[error("TwoFactorAuthNotSetup")]
TwoFactorAuthNotSetup,
}

impl common_utils::errors::ErrorSwitch<api_models::errors::types::ApiErrorResponse> for UserErrors {
Expand Down Expand Up @@ -184,6 +188,12 @@ impl common_utils::errors::ErrorSwitch<api_models::errors::types::ApiErrorRespon
Self::TotpRequired => {
AER::BadRequest(ApiError::new(sub_code, 38, self.get_error_message(), None))
}
Self::TwoFactorAuthRequired => {
AER::BadRequest(ApiError::new(sub_code, 39, self.get_error_message(), None))
}
Self::TwoFactorAuthNotSetup => {
AER::BadRequest(ApiError::new(sub_code, 40, self.get_error_message(), None))
}
}
}
}
Expand Down Expand Up @@ -223,6 +233,8 @@ impl UserErrors {
Self::TotpNotSetup => "TOTP not setup",
Self::InvalidTotp => "Invalid TOTP",
Self::TotpRequired => "TOTP required",
Self::TwoFactorAuthRequired => "Two factor auth required",
Self::TwoFactorAuthNotSetup => "Two factor auth not setup",
}
}
}
65 changes: 58 additions & 7 deletions crates/router/src/core/user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,9 @@ use crate::{
routes::{app::ReqState, AppState},
services::{authentication as auth, authorization::roles, ApplicationResponse},
types::{domain, transformers::ForeignInto},
utils,
utils::{self, user::two_factor_auth as tfa_utils},
};

pub mod dashboard_metadata;
#[cfg(feature = "dummy_connector")]
pub mod sample_data;
Expand Down Expand Up @@ -1631,7 +1632,7 @@ pub async fn begin_totp(
}));
}

let totp = utils::user::two_factor_auth::generate_default_totp(user_from_db.get_email(), None)?;
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?;
Expand Down Expand Up @@ -1693,10 +1694,8 @@ pub async fn verify_totp(
.await?
.ok_or(UserErrors::InternalServerError)?;

let totp = utils::user::two_factor_auth::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()
Expand Down Expand Up @@ -1739,7 +1738,7 @@ pub async fn generate_recovery_codes(
state: AppState,
user_token: auth::UserFromSinglePurposeToken,
) -> UserResponse<user_api::RecoveryCodes> {
if !utils::user::two_factor_auth::check_totp_in_redis(&state, &user_token.user_id).await? {
if !tfa_utils::check_totp_in_redis(&state, &user_token.user_id).await? {
return Err(UserErrors::TotpRequired.into());
}

Expand All @@ -1766,3 +1765,55 @@ pub async fn generate_recovery_codes(
recovery_codes: recovery_codes.into_inner(),
}))
}

pub async fn terminate_two_factor_auth(
state: AppState,
user_token: auth::UserFromSinglePurposeToken,
skip_two_factor_auth: bool,
) -> UserResponse<user_api::TokenResponse> {
let user_from_db: domain::UserFromStorage = state
.store
.find_user_by_id(&user_token.user_id)
.await
.change_context(UserErrors::InternalServerError)?
.into();

if !skip_two_factor_auth {
if !tfa_utils::check_totp_in_redis(&state, &user_token.user_id).await?
&& !tfa_utils::check_recovery_code_in_redis(&state, &user_token.user_id).await?
{
return Err(UserErrors::TwoFactorAuthRequired.into());
}

if user_from_db.get_recovery_codes().is_none() {
return Err(UserErrors::TwoFactorAuthNotSetup.into());
}

if user_from_db.get_totp_status() != TotpStatus::Set {
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)?;
}
}

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?;

auth::cookies::set_cookie_response(
user_api::TokenResponse {
token: token.clone(),
token_type: next_flow.get_flow().into(),
},
token,
)
}
3 changes: 3 additions & 0 deletions crates/router/src/routes/app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1215,6 +1215,9 @@ impl User {
.service(
web::resource("/recovery_codes/generate")
.route(web::get().to(generate_recovery_codes)),
)
.service(
web::resource("/2fa/terminate").route(web::get().to(terminate_two_factor_auth)),
);

#[cfg(feature = "email")]
Expand Down
1 change: 1 addition & 0 deletions crates/router/src/routes/lock_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,7 @@ impl From<Flow> for ApiIdentifier {
| Flow::UpdateUserAccountDetails
| Flow::TotpBegin
| Flow::TotpVerify
| Flow::TerminateTwoFactorAuth
| Flow::GenerateRecoveryCodes => Self::User,

Flow::ListRoles
Expand Down
20 changes: 20 additions & 0 deletions crates/router/src/routes/user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -679,3 +679,23 @@ pub async fn generate_recovery_codes(state: web::Data<AppState>, req: HttpReques
))
.await
}

pub async fn terminate_two_factor_auth(
state: web::Data<AppState>,
req: HttpRequest,
query: web::Query<user_api::SkipTwoFactorAuthQueryParam>,
) -> HttpResponse {
let flow = Flow::TerminateTwoFactorAuth;
let skip_two_factor_auth = query.into_inner().skip_two_factor_auth.unwrap_or(false);

Box::pin(api::server_wrap(
flow,
state.clone(),
&req,
(),
|state, user, _, _| user_core::terminate_two_factor_auth(state, user, skip_two_factor_auth),
&auth::SinglePurposeJWTAuth(TokenPurpose::TOTP),
api_locking::LockAction::NotApplicable,
))
.await
}
4 changes: 4 additions & 0 deletions crates/router/src/types/domain/user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -930,6 +930,10 @@ impl UserFromStorage {
self.0.totp_status
}

pub fn get_recovery_codes(&self) -> Option<Vec<Secret<String>>> {
self.0.totp_recovery_codes.clone()
}

pub async fn decrypt_and_get_totp_secret(
&self,
state: &AppState,
Expand Down
9 changes: 9 additions & 0 deletions crates/router/src/utils/user/two_factor_auth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,15 @@ pub async fn check_totp_in_redis(state: &AppState, user_id: &str) -> UserResult<
.change_context(UserErrors::InternalServerError)
}

pub async fn check_recovery_code_in_redis(state: &AppState, user_id: &str) -> UserResult<bool> {
let redis_conn = get_redis_connection(state)?;
let key = format!("{}{}", consts::user::REDIS_RECOVERY_CODES_PREFIX, user_id);
redis_conn
.exists::<()>(&key)
.await
.change_context(UserErrors::InternalServerError)
}

fn get_redis_connection(state: &AppState) -> UserResult<Arc<RedisConnectionPool>> {
state
.store
Expand Down
2 changes: 2 additions & 0 deletions crates/router_env/src/logger/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -408,6 +408,8 @@ pub enum Flow {
TotpVerify,
/// Generate or Regenerate recovery codes
GenerateRecoveryCodes,
// Terminate two factor authentication
TerminateTwoFactorAuth,
/// List initial webhook delivery attempts
WebhookEventInitialDeliveryAttemptList,
/// List delivery attempts for a webhook event
Expand Down

0 comments on commit 42e5ef1

Please sign in to comment.