feat(users): Create user_key_store table and begin_totp API (#4577)

Co-authored-by: hyperswitch-bot[bot] <148525504+hyperswitch-bot[bot]@users.noreply.github.com>

Cargo.lock

@@ -1430,6 +1430,12 @@ dependencies = [
  "rustc-demangle",
 ]
 
+[[package]]
+name = "base32"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23ce669cd6c8588f79e15cf450314f9638f967fc5770ff1c7c1deb0925ea7cfa"
+
 [[package]]
 name = "base64"
 version = "0.13.1"
@@ -1559,7 +1565,7 @@ dependencies = [
  "arrayvec",
  "cc",
  "cfg-if 1.0.0",
- "constant_time_eq",
+ "constant_time_eq 0.3.0",
 ]
 
 [[package]]
@@ -2044,6 +2050,12 @@ dependencies = [
  "tiny-keccak",
 ]
 
+[[package]]
+name = "constant_time_eq"
+version = "0.2.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "21a53c0a4d288377e7415b53dcfc3c04da5cdc2cc95c8d5ac178b58f0b861ad6"
+
 [[package]]
 name = "constant_time_eq"
 version = "0.3.0"
@@ -5683,6 +5695,7 @@ dependencies = [
  "thiserror",
  "time",
  "tokio 1.37.0",
+ "totp-rs",
  "tracing-futures",
  "unicode-segmentation",
  "url",
@@ -7473,6 +7486,22 @@ dependencies = [
  "tracing-futures",
 ]
 
+[[package]]
+name = "totp-rs"
+version = "5.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c4ae9724c5888c0417d2396037ed3b60665925624766416e3e342b6ba5dbd3f"
+dependencies = [
+ "base32",
+ "constant_time_eq 0.2.6",
+ "hmac",
+ "rand",
+ "sha1",
+ "sha2",
+ "url",
+ "urlencoding",
+]
+
 [[package]]
 name = "tower"
 version = "0.4.13"

crates/api_models/src/events/user.rs

@@ -10,13 +10,14 @@ use crate::user::{
     dashboard_metadata::{
         GetMetaDataRequest, GetMetaDataResponse, GetMultipleMetaDataPayload, SetMetaDataRequest,
     },
-    AcceptInviteFromEmailRequest, AuthorizeResponse, ChangePasswordRequest, ConnectAccountRequest,
+    AcceptInviteFromEmailRequest, AuthorizeResponse, BeginTotpResponse, ChangePasswordRequest,
-    CreateInternalUserRequest, DashboardEntryResponse, ForgotPasswordRequest,
+    ConnectAccountRequest, CreateInternalUserRequest, DashboardEntryResponse,
-    GetUserDetailsResponse, GetUserRoleDetailsRequest, GetUserRoleDetailsResponse,
+    ForgotPasswordRequest, GetUserDetailsResponse, GetUserRoleDetailsRequest,
-    InviteUserRequest, ListUsersResponse, ReInviteUserRequest, ResetPasswordRequest,
+    GetUserRoleDetailsResponse, InviteUserRequest, ListUsersResponse, ReInviteUserRequest,
-    RotatePasswordRequest, SendVerifyEmailRequest, SignInResponse, SignUpRequest,
+    ResetPasswordRequest, RotatePasswordRequest, SendVerifyEmailRequest, SignInResponse,
-    SignUpWithMerchantIdRequest, SwitchMerchantIdRequest, TokenOrPayloadResponse, TokenResponse,
+    SignUpRequest, SignUpWithMerchantIdRequest, SwitchMerchantIdRequest, TokenOrPayloadResponse,
-    UpdateUserAccountDetailsRequest, UserFromEmailRequest, UserMerchantCreate, VerifyEmailRequest,
+    TokenResponse, UpdateUserAccountDetailsRequest, UserFromEmailRequest, UserMerchantCreate,
+    VerifyEmailRequest,
 };
 
 impl ApiEventMetric for DashboardEntryResponse {
@@ -72,7 +73,8 @@ common_utils::impl_misc_api_event_type!(
     GetUserRoleDetailsRequest,
     GetUserRoleDetailsResponse,
     TokenResponse,
-    UserFromEmailRequest
+    UserFromEmailRequest,
+    BeginTotpResponse
 );
 
 #[cfg(feature = "dummy_connector")]

crates/api_models/src/user.rs

@@ -236,8 +236,19 @@ pub enum TokenOrPayloadResponse<T> {
     Token(TokenResponse),
     Payload(T),
 }
-
 #[derive(Debug, serde::Deserialize, serde::Serialize)]
 pub struct UserFromEmailRequest {
     pub token: Secret<String>,
 }
+
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
+pub struct BeginTotpResponse {
+    pub secret: Option<TotpSecret>,
+}
+
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
+pub struct TotpSecret {
+    pub secret: Secret<String>,
+    pub totp_url: Secret<String>,
+    pub recovery_codes: Vec<Secret<String>>,
+}

crates/diesel_models/src/enums.rs

@@ -18,8 +18,8 @@ pub mod diesel_exports {
         DbRefundStatus as RefundStatus, DbRefundType as RefundType,
         DbRequestIncrementalAuthorization as RequestIncrementalAuthorization,
         DbRoleScope as RoleScope, DbRoutingAlgorithmKind as RoutingAlgorithmKind,
-        DbTransactionType as TransactionType, DbUserStatus as UserStatus,
+        DbTotpStatus as TotpStatus, DbTransactionType as TransactionType,
-        DbWebhookDeliveryAttempt as WebhookDeliveryAttempt,
+        DbUserStatus as UserStatus, DbWebhookDeliveryAttempt as WebhookDeliveryAttempt,
     };
 }
 pub use common_enums::*;
@@ -350,3 +350,26 @@ pub enum DashboardMetadata {
     IsChangePasswordRequired,
     OnboardingSurvey,
 }
+
+#[derive(
+    Clone,
+    Copy,
+    Debug,
+    Default,
+    Eq,
+    PartialEq,
+    serde::Serialize,
+    serde::Deserialize,
+    strum::Display,
+    strum::EnumString,
+    frunk::LabelledGeneric,
+)]
+#[diesel_enum(storage_type = "db_enum")]
+#[serde(rename_all = "snake_case")]
+#[strum(serialize_all = "snake_case")]
+pub enum TotpStatus {
+    Set,
+    InProgress,
+    #[default]
+    NotSet,
+}

crates/diesel_models/src/lib.rs

@@ -44,6 +44,7 @@ pub mod routing_algorithm;
 #[allow(unused_qualifications)]
 pub mod schema;
 pub mod user;
+pub mod user_key_store;
 pub mod user_role;
 
 use diesel_impl::{DieselArray, OptionalDieselArray};

crates/diesel_models/src/query.rs

@@ -36,4 +36,5 @@ pub mod reverse_lookup;
 pub mod role;
 pub mod routing_algorithm;
 pub mod user;
+pub mod user_key_store;
 pub mod user_role;

crates/diesel_models/src/query/user_key_store.rs

@@ -0,0 +1,24 @@
+use diesel::{associations::HasTable, ExpressionMethods};
+
+use super::generics;
+use crate::{
+    schema::user_key_store::dsl,
+    user_key_store::{UserKeyStore, UserKeyStoreNew},
+    PgPooledConn, StorageResult,
+};
+
+impl UserKeyStoreNew {
+    pub async fn insert(self, conn: &PgPooledConn) -> StorageResult<UserKeyStore> {
+        generics::generic_insert(conn, self).await
+    }
+}
+
+impl UserKeyStore {
+    pub async fn find_by_user_id(conn: &PgPooledConn, user_id: &str) -> StorageResult<Self> {
+        generics::generic_find_one::<<Self as HasTable>::Table, _, _>(
+            conn,
+            dsl::user_id.eq(user_id.to_owned()),
+        )
+        .await
+    }
+}

crates/diesel_models/src/schema.rs

@@ -1149,6 +1149,18 @@ diesel::table! {
     }
 }
 
+diesel::table! {
+    use diesel::sql_types::*;
+    use crate::enums::diesel_exports::*;
+
+    user_key_store (user_id) {
+        #[max_length = 64]
+        user_id -> Varchar,
+        key -> Bytea,
+        created_at -> Timestamp,
+    }
+}
+
 diesel::table! {
     use diesel::sql_types::*;
     use crate::enums::diesel_exports::*;
@@ -1192,6 +1204,9 @@ diesel::table! {
         last_modified_at -> Timestamp,
         #[max_length = 64]
         preferred_merchant_id -> Nullable<Varchar>,
+        totp_status -> TotpStatus,
+        totp_secret -> Nullable<Bytea>,
+        totp_recovery_codes -> Nullable<Array<Nullable<Text>>>,
         last_password_modified_at -> Nullable<Timestamp>,
     }
 }
@@ -1232,6 +1247,7 @@ diesel::allow_tables_to_appear_in_same_query!(
     reverse_lookup,
     roles,
     routing_algorithm,
+    user_key_store,
     user_roles,
     users,
 );

crates/diesel_models/src/user.rs

@@ -3,7 +3,9 @@ use diesel::{AsChangeset, Identifiable, Insertable, Queryable};
 use masking::Secret;
 use time::PrimitiveDateTime;
 
-use crate::schema::users;
+use crate::{
+    diesel_impl::OptionalDieselArray, encryption::Encryption, enums::TotpStatus, schema::users,
+};
 
 pub mod dashboard_metadata;
 
@@ -20,6 +22,10 @@ pub struct User {
     pub created_at: PrimitiveDateTime,
     pub last_modified_at: PrimitiveDateTime,
     pub preferred_merchant_id: Option<String>,
+    pub totp_status: TotpStatus,
+    pub totp_secret: Option<Encryption>,
+    #[diesel(deserialize_as = OptionalDieselArray<Secret<String>>)]
+    pub totp_recovery_codes: Option<Vec<Secret<String>>>,
     pub last_password_modified_at: Option<PrimitiveDateTime>,
 }
 
@@ -36,6 +42,9 @@ pub struct UserNew {
     pub created_at: Option<PrimitiveDateTime>,
     pub last_modified_at: Option<PrimitiveDateTime>,
     pub preferred_merchant_id: Option<String>,
+    pub totp_status: TotpStatus,
+    pub totp_secret: Option<Encryption>,
+    pub totp_recovery_codes: Option<Vec<Secret<String>>>,
     pub last_password_modified_at: Option<PrimitiveDateTime>,
 }
 
@@ -47,6 +56,9 @@ pub struct UserUpdateInternal {
     is_verified: Option<bool>,
     last_modified_at: PrimitiveDateTime,
     preferred_merchant_id: Option<String>,
+    totp_status: Option<TotpStatus>,
+    totp_secret: Option<Encryption>,
+    totp_recovery_codes: Option<Vec<Secret<String>>>,
     last_password_modified_at: Option<PrimitiveDateTime>,
 }
 
@@ -58,6 +70,11 @@ pub enum UserUpdate {
         is_verified: Option<bool>,
         preferred_merchant_id: Option<String>,
     },
+    TotpUpdate {
+        totp_status: Option<TotpStatus>,
+        totp_secret: Option<Encryption>,
+        totp_recovery_codes: Option<Vec<Secret<String>>>,
+    },
     PasswordUpdate {
         password: Option<Secret<String>>,
     },
@@ -73,6 +90,9 @@ impl From<UserUpdate> for UserUpdateInternal {
                 is_verified: Some(true),
                 last_modified_at,
                 preferred_merchant_id: None,
+                totp_status: None,
+                totp_secret: None,
+                totp_recovery_codes: None,
                 last_password_modified_at: None,
             },
             UserUpdate::AccountUpdate {
@@ -85,6 +105,24 @@ impl From<UserUpdate> for UserUpdateInternal {
                 is_verified,
                 last_modified_at,
                 preferred_merchant_id,
+                totp_status: None,
+                totp_secret: None,
+                totp_recovery_codes: None,
+                last_password_modified_at: None,
+            },
+            UserUpdate::TotpUpdate {
+                totp_status,
+                totp_secret,
+                totp_recovery_codes,
+            } => Self {
+                name: None,
+                password: None,
+                is_verified: None,
+                last_modified_at,
+                preferred_merchant_id: None,
+                totp_status,
+                totp_secret,
+                totp_recovery_codes,
                 last_password_modified_at: None,
             },
             UserUpdate::PasswordUpdate { password } => Self {
@@ -94,6 +132,9 @@ impl From<UserUpdate> for UserUpdateInternal {
                 last_modified_at,
                 preferred_merchant_id: None,
                 last_password_modified_at: Some(last_modified_at),
+                totp_status: None,
+                totp_secret: None,
+                totp_recovery_codes: None,
             },
         }
     }

crates/diesel_models/src/user_key_store.rs

@@ -0,0 +1,21 @@
+use diesel::{Identifiable, Insertable, Queryable};
+use time::PrimitiveDateTime;
+
+use crate::{encryption::Encryption, schema::user_key_store};
+
+#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Identifiable, Queryable)]
+#[diesel(table_name = user_key_store)]
+#[diesel(primary_key(user_id))]
+pub struct UserKeyStore {
+    pub user_id: String,
+    pub key: Encryption,
+    pub created_at: PrimitiveDateTime,
+}
+
+#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Insertable)]
+#[diesel(table_name = user_key_store)]
+pub struct UserKeyStoreNew {
+    pub user_id: String,
+    pub key: Encryption,
+    pub created_at: PrimitiveDateTime,
+}

crates/router/Cargo.toml

@@ -126,6 +126,7 @@ isocountry = "0.3.2"
 iso_currency = "0.4.4"
 actix-http = "3.6.0"
 events = { version = "0.1.0", path = "../events" }
+totp-rs = { version = "5.5.1", features = ["gen_secret", "otpauth"]}
 
 [build-dependencies]
 router_env = { version = "0.1.0", path = "../router_env", default-features = false }

crates/router/src/consts/user.rs

@@ -1,5 +1,14 @@
 pub const MAX_NAME_LENGTH: usize = 70;
 pub const MAX_COMPANY_NAME_LENGTH: usize = 70;
 pub const BUSINESS_EMAIL: &str = "biz@hyperswitch.io";
+pub const RECOVERY_CODES_COUNT: usize = 8;
+pub const RECOVERY_CODE_LENGTH: usize = 8; // This is without counting the hyphen in between
+pub const TOTP_ISSUER_NAME: &str = "Hyperswitch";
+/// The number of digits composing the auth code.
+pub const TOTP_DIGITS: usize = 6;
+/// Duration in seconds of a step.
+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;
 pub const MAX_PASSWORD_LENGTH: usize = 70;
 pub const MIN_PASSWORD_LENGTH: usize = 8;

crates/router/src/core/user.rs

@@ -1,9 +1,13 @@
 use api_models::user::{self as user_api, InviteMultipleUserResponse};
 #[cfg(feature = "email")]
 use diesel_models::user_role::UserRoleUpdate;
-use diesel_models::{enums::UserStatus, user as storage_user, user_role::UserRoleNew};
+use diesel_models::{
+    enums::{TotpStatus, UserStatus},
+    user as storage_user,
+    user_role::UserRoleNew,
+};
 use error_stack::{report, ResultExt};
-use masking::ExposeInterface;
+use masking::{ExposeInterface, PeekInterface};
 #[cfg(feature = "email")]
 use router_env::env;
 use router_env::logger;
@@ -1581,3 +1585,60 @@ pub async fn user_from_email(
     };
     auth::cookies::set_cookie_response(response, token)
 }
+
+pub async fn begin_totp(
+    state: AppState,
+    user_token: auth::UserFromSinglePurposeToken,
+) -> UserResponse<user_api::BeginTotpResponse> {
+    let user_from_db: domain::UserFromStorage = state
+        .store
+        .find_user_by_id(&user_token.user_id)
+        .await
+        .change_context(UserErrors::InternalServerError)?
+        .into();
+
+    if user_from_db.get_totp_status() == TotpStatus::Set {
+        return Ok(ApplicationResponse::Json(user_api::BeginTotpResponse {
+            secret: None,
+        }));
+    }
+
+    let totp = utils::user::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?;
+
+    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)?;
+
+    Ok(ApplicationResponse::Json(user_api::BeginTotpResponse {
+        secret: Some(user_api::TotpSecret {
+            secret: totp.get_secret_base32().into(),
+            totp_url: totp.get_url().into(),
+            recovery_codes: recovery_codes.into_inner(),
+        }),
+    }))
+}

crates/router/src/db.rs

@@ -33,6 +33,7 @@ pub mod reverse_lookup;
 pub mod role;
 pub mod routing_algorithm;
 pub mod user;
+pub mod user_key_store;
 pub mod user_role;
 
 use diesel_models::{
@@ -118,6 +119,7 @@ pub trait StorageInterface:
     + user::sample_data::BatchSampleDataInterface
     + health_check::HealthCheckDbInterface
     + role::RoleInterface
+    + user_key_store::UserKeyStoreInterface
     + authentication::AuthenticationInterface
     + 'static
 {

crates/router/src/db/kafka_store.rs

@@ -32,6 +32,7 @@ use super::{
     dashboard_metadata::DashboardMetadataInterface,
     role::RoleInterface,
     user::{sample_data::BatchSampleDataInterface, UserInterface},
+    user_key_store::UserKeyStoreInterface,
     user_role::UserRoleInterface,
 };
 #[cfg(feature = "payouts")]
@@ -2743,3 +2744,26 @@ impl RoleInterface for KafkaStore {
         self.diesel_store.list_all_roles(merchant_id, org_id).await
     }
 }
+
+#[async_trait::async_trait]
+impl UserKeyStoreInterface for KafkaStore {
+    async fn insert_user_key_store(
+        &self,
+        user_key_store: domain::UserKeyStore,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<domain::UserKeyStore, errors::StorageError> {
+        self.diesel_store
+            .insert_user_key_store(user_key_store, key)
+            .await
+    }
+
+    async fn get_user_key_store_by_user_id(
+        &self,
+        user_id: &str,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<domain::UserKeyStore, errors::StorageError> {
+        self.diesel_store
+            .get_user_key_store_by_user_id(user_id, key)
+            .await
+    }
+}

crates/router/src/db/user.rs

@@ -162,6 +162,9 @@ impl UserInterface for MockDb {
             created_at: user_data.created_at.unwrap_or(time_now),
             last_modified_at: user_data.created_at.unwrap_or(time_now),
             preferred_merchant_id: user_data.preferred_merchant_id,
+            totp_status: user_data.totp_status,
+            totp_secret: user_data.totp_secret,
+            totp_recovery_codes: user_data.totp_recovery_codes,
             last_password_modified_at: user_data.last_password_modified_at,
         };
         users.push(user.clone());
@@ -229,6 +232,18 @@ impl UserInterface for MockDb {
                             .or(user.preferred_merchant_id.clone()),
                         ..user.to_owned()
                     },
+                    storage::UserUpdate::TotpUpdate {
+                        totp_status,
+                        totp_secret,
+                        totp_recovery_codes,
+                    } => storage::User {
+                        totp_status: totp_status.unwrap_or(user.totp_status),
+                        totp_secret: totp_secret.clone().or(user.totp_secret.clone()),
+                        totp_recovery_codes: totp_recovery_codes
+                            .clone()
+                            .or(user.totp_recovery_codes.clone()),
+                        ..user.to_owned()
+                    },
                     storage::UserUpdate::PasswordUpdate { password } => storage::User {
                         password: password.clone().unwrap_or(user.password.clone()),
                         last_password_modified_at: Some(common_utils::date_time::now()),
@@ -272,6 +287,18 @@ impl UserInterface for MockDb {
                             .or(user.preferred_merchant_id.clone()),
                         ..user.to_owned()
                     },
+                    storage::UserUpdate::TotpUpdate {
+                        totp_status,
+                        totp_secret,
+                        totp_recovery_codes,
+                    } => storage::User {
+                        totp_status: totp_status.unwrap_or(user.totp_status),
+                        totp_secret: totp_secret.clone().or(user.totp_secret.clone()),
+                        totp_recovery_codes: totp_recovery_codes
+                            .clone()
+                            .or(user.totp_recovery_codes.clone()),
+                        ..user.to_owned()
+                    },
                     storage::UserUpdate::PasswordUpdate { password } => storage::User {
                         password: password.clone().unwrap_or(user.password.clone()),
                         last_password_modified_at: Some(common_utils::date_time::now()),

crates/router/src/db/user_key_store.rs

@@ -0,0 +1,121 @@
+use common_utils::errors::CustomResult;
+use error_stack::{report, ResultExt};
+use masking::Secret;
+use router_env::{instrument, tracing};
+use storage_impl::MockDb;
+
+use crate::{
+    connection,
+    core::errors,
+    services::Store,
+    types::domain::{
+        self,
+        behaviour::{Conversion, ReverseConversion},
+    },
+};
+
+#[async_trait::async_trait]
+pub trait UserKeyStoreInterface {
+    async fn insert_user_key_store(
+        &self,
+        user_key_store: domain::UserKeyStore,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<domain::UserKeyStore, errors::StorageError>;
+
+    async fn get_user_key_store_by_user_id(
+        &self,
+        user_id: &str,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<domain::UserKeyStore, errors::StorageError>;
+}
+
+#[async_trait::async_trait]
+impl UserKeyStoreInterface for Store {
+    #[instrument(skip_all)]
+    async fn insert_user_key_store(
+        &self,
+        user_key_store: domain::UserKeyStore,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<domain::UserKeyStore, errors::StorageError> {
+        let conn = connection::pg_connection_write(self).await?;
+        user_key_store
+            .construct_new()
+            .await
+            .change_context(errors::StorageError::EncryptionError)?
+            .insert(&conn)
+            .await
+            .map_err(|error| report!(errors::StorageError::from(error)))?
+            .convert(key)
+            .await
+            .change_context(errors::StorageError::DecryptionError)
+    }
+
+    #[instrument(skip_all)]
+    async fn get_user_key_store_by_user_id(
+        &self,
+        user_id: &str,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<domain::UserKeyStore, errors::StorageError> {
+        let conn = connection::pg_connection_read(self).await?;
+
+        diesel_models::user_key_store::UserKeyStore::find_by_user_id(&conn, user_id)
+            .await
+            .map_err(|error| report!(errors::StorageError::from(error)))?
+            .convert(key)
+            .await
+            .change_context(errors::StorageError::DecryptionError)
+    }
+}
+
+#[async_trait::async_trait]
+impl UserKeyStoreInterface for MockDb {
+    #[instrument(skip_all)]
+    async fn insert_user_key_store(
+        &self,
+        user_key_store: domain::UserKeyStore,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<domain::UserKeyStore, errors::StorageError> {
+        let mut locked_user_key_store = self.user_key_store.lock().await;
+
+        if locked_user_key_store
+            .iter()
+            .any(|user_key| user_key.user_id == user_key_store.user_id)
+        {
+            Err(errors::StorageError::DuplicateValue {
+                entity: "user_key_store",
+                key: Some(user_key_store.user_id.clone()),
+            })?;
+        }
+
+        let user_key_store = Conversion::convert(user_key_store)
+            .await
+            .change_context(errors::StorageError::MockDbError)?;
+        locked_user_key_store.push(user_key_store.clone());
+
+        user_key_store
+            .convert(key)
+            .await
+            .change_context(errors::StorageError::DecryptionError)
+    }
+
+    #[instrument(skip_all)]
+    async fn get_user_key_store_by_user_id(
+        &self,
+        user_id: &str,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<domain::UserKeyStore, errors::StorageError> {
+        self.user_key_store
+            .lock()
+            .await
+            .iter()
+            .find(|user_key_store| user_key_store.user_id == user_id)
+            .cloned()
+            .ok_or(errors::StorageError::ValueNotFound(format!(
+                "No user_key_store is found for user_id={}",
+                user_id
+            )))?
+            .convert(key)
+            .await
+            .change_context(errors::StorageError::DecryptionError)
+    }
+}

crates/router/src/routes/app.rs

@@ -1197,7 +1197,8 @@ 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)));
 
         #[cfg(feature = "email")]
         {

crates/router/src/routes/lock_utils.rs

@@ -210,7 +210,8 @@ impl From<Flow> for ApiIdentifier {
             | Flow::VerifyEmail
             | Flow::AcceptInviteFromEmail
             | Flow::VerifyEmailRequest
-            | Flow::UpdateUserAccountDetails => Self::User,
+            | Flow::UpdateUserAccountDetails
+            | Flow::TotpBegin => Self::User,
 
             Flow::ListRoles
             | Flow::GetRole

crates/router/src/routes/user.rs

@@ -612,3 +612,17 @@ pub async fn user_from_email(
     ))
     .await
 }
+
+pub async fn totp_begin(state: web::Data<AppState>, req: HttpRequest) -> HttpResponse {
+    let flow = Flow::TotpBegin;
+    Box::pin(api::server_wrap(
+        flow,
+        state.clone(),
+        &req,
+        (),
+        |state, user, _, _| user_core::begin_totp(state, user),
+        &auth::SinglePurposeJWTAuth(common_enums::TokenPurpose::TOTP),
+        api_locking::LockAction::NotApplicable,
+    ))
+    .await
+}

crates/router/src/types/domain.rs

@@ -9,6 +9,7 @@ pub mod payments;
 pub mod types;
 #[cfg(feature = "olap")]
 pub mod user;
+pub mod user_key_store;
 
 pub use address::*;
 pub use customer::*;
@@ -19,3 +20,4 @@ pub use merchant_key_store::*;
 pub use payments::*;
 #[cfg(feature = "olap")]
 pub use user::*;
+pub use user_key_store::*;

crates/router/src/types/domain/user.rs

@@ -6,7 +6,7 @@ use api_models::{
 use common_enums::TokenPurpose;
 use common_utils::{errors::CustomResult, pii};
 use diesel_models::{
-    enums::UserStatus,
+    enums::{TotpStatus, UserStatus},
     organization as diesel_org,
     organization::Organization,
     user as storage_user,
@@ -15,6 +15,7 @@ use diesel_models::{
 use error_stack::{report, ResultExt};
 use masking::{ExposeInterface, PeekInterface, Secret};
 use once_cell::sync::Lazy;
+use rand::distributions::{Alphanumeric, DistString};
 use router_env::env;
 use unicode_segmentation::UnicodeSegmentation;
 
@@ -26,7 +27,7 @@ use crate::{
     },
     db::StorageInterface,
     routes::AppState,
-    services::{authentication as auth, authentication::UserFromToken, authorization::info},
+    services::{self, authentication as auth, authentication::UserFromToken, authorization::info},
     types::transformers::ForeignFrom,
     utils::{self, user::password},
 };
@@ -35,6 +36,8 @@ pub mod dashboard_metadata;
 pub mod decision_manager;
 pub use decision_manager::*;
 
+use super::{types as domain_types, UserKeyStore};
+
 #[derive(Clone)]
 pub struct UserName(Secret<String>);
 
@@ -863,6 +866,49 @@ impl UserFromStorage {
                 )
         }
     }
+
+    pub async fn get_or_create_key_store(&self, state: &AppState) -> UserResult<UserKeyStore> {
+        let master_key = state.store.get_master_key();
+        let key_store_result = state
+            .store
+            .get_user_key_store_by_user_id(self.get_user_id(), &master_key.to_vec().into())
+            .await;
+
+        if let Ok(key_store) = key_store_result {
+            Ok(key_store)
+        } else if key_store_result
+            .as_ref()
+            .map_err(|e| e.current_context().is_db_not_found())
+            .err()
+            .unwrap_or(false)
+        {
+            let key = services::generate_aes256_key()
+                .change_context(UserErrors::InternalServerError)
+                .attach_printable("Unable to generate aes 256 key")?;
+
+            let key_store = UserKeyStore {
+                user_id: self.get_user_id().to_string(),
+                key: domain_types::encrypt(key.to_vec().into(), master_key)
+                    .await
+                    .change_context(UserErrors::InternalServerError)?,
+                created_at: common_utils::date_time::now(),
+            };
+            state
+                .store
+                .insert_user_key_store(key_store, &master_key.to_vec().into())
+                .await
+                .change_context(UserErrors::InternalServerError)
+        } else {
+            Err(key_store_result
+                .err()
+                .map(|e| e.change_context(UserErrors::InternalServerError))
+                .unwrap_or(UserErrors::InternalServerError.into()))
+        }
+    }
+
+    pub fn get_totp_status(&self) -> TotpStatus {
+        self.0.totp_status
+    }
 }
 
 impl From<info::ModuleInfo> for user_role_api::ModuleInfo {
@@ -1031,3 +1077,36 @@ impl RoleName {
         self.0
     }
 }
+
+#[derive(serde::Serialize, serde::Deserialize, Debug)]
+pub struct RecoveryCodes(pub Vec<Secret<String>>);
+
+impl RecoveryCodes {
+    pub fn generate_new() -> Self {
+        let mut rand = rand::thread_rng();
+        let recovery_codes = (0..consts::user::RECOVERY_CODES_COUNT)
+            .map(|_| {
+                let code_part_1 =
+                    Alphanumeric.sample_string(&mut rand, consts::user::RECOVERY_CODE_LENGTH / 2);
+                let code_part_2 =
+                    Alphanumeric.sample_string(&mut rand, consts::user::RECOVERY_CODE_LENGTH / 2);
+
+                Secret::new(format!("{}-{}", code_part_1, code_part_2))
+            })
+            .collect::<Vec<_>>();
+
+        Self(recovery_codes)
+    }
+
+    pub fn get_hashed(&self) -> UserResult<Vec<Secret<String>>> {
+        self.0
+            .iter()
+            .cloned()
+            .map(password::generate_password_hash)
+            .collect::<Result<Vec<_>, _>>()
+    }
+
+    pub fn into_inner(self) -> Vec<Secret<String>> {
+        self.0
+    }
+}

crates/router/src/types/domain/user_key_store.rs

@@ -0,0 +1,59 @@
+use common_utils::{
+    crypto::{Encryptable, GcmAes256},
+    date_time,
+};
+use error_stack::ResultExt;
+use masking::{PeekInterface, Secret};
+use time::PrimitiveDateTime;
+
+use crate::{
+    errors::{CustomResult, ValidationError},
+    types::domain::types::TypeEncryption,
+};
+
+#[derive(Clone, Debug, serde::Serialize)]
+pub struct UserKeyStore {
+    pub user_id: String,
+    pub key: Encryptable<Secret<Vec<u8>>>,
+    pub created_at: PrimitiveDateTime,
+}
+
+#[async_trait::async_trait]
+impl super::behaviour::Conversion for UserKeyStore {
+    type DstType = diesel_models::user_key_store::UserKeyStore;
+    type NewDstType = diesel_models::user_key_store::UserKeyStoreNew;
+
+    async fn convert(self) -> CustomResult<Self::DstType, ValidationError> {
+        Ok(diesel_models::user_key_store::UserKeyStore {
+            key: self.key.into(),
+            user_id: self.user_id,
+            created_at: self.created_at,
+        })
+    }
+
+    async fn convert_back(
+        item: Self::DstType,
+        key: &Secret<Vec<u8>>,
+    ) -> CustomResult<Self, ValidationError>
+    where
+        Self: Sized,
+    {
+        Ok(Self {
+            key: Encryptable::decrypt(item.key, key.peek(), GcmAes256)
+                .await
+                .change_context(ValidationError::InvalidValue {
+                    message: "Failed while decrypting customer data".to_string(),
+                })?,
+            user_id: item.user_id,
+            created_at: item.created_at,
+        })
+    }
+
+    async fn construct_new(self) -> CustomResult<Self::NewDstType, ValidationError> {
+        Ok(diesel_models::user_key_store::UserKeyStoreNew {
+            user_id: self.user_id,
+            key: self.key.into(),
+            created_at: date_time::now(),
+        })
+    }
+}

crates/router/src/utils/user.rs

@@ -1,12 +1,14 @@
 use std::collections::HashMap;
 
 use api_models::user as user_api;
-use common_utils::errors::CustomResult;
+use common_utils::{errors::CustomResult, pii};
 use diesel_models::{enums::UserStatus, user_role::UserRole};
 use error_stack::ResultExt;
-use masking::Secret;
+use masking::ExposeInterface;
+use totp_rs::{Algorithm, TOTP};
 
 use crate::{
+    consts,
     core::errors::{StorageError, UserErrors, UserResult},
     routes::AppState,
     services::{
@@ -74,7 +76,7 @@ pub async fn generate_jwt_auth_token(
     state: &AppState,
     user: &UserFromStorage,
     user_role: &UserRole,
-) -> UserResult<Secret<String>> {
+) -> UserResult<masking::Secret<String>> {
     let token = AuthToken::new_token(
         user.get_user_id().to_string(),
         user_role.merchant_id.clone(),
@@ -83,7 +85,7 @@ pub async fn generate_jwt_auth_token(
         user_role.org_id.clone(),
     )
     .await?;
-    Ok(Secret::new(token))
+    Ok(masking::Secret::new(token))
 }
 
 pub async fn generate_jwt_auth_token_with_custom_role_attributes(
@@ -92,7 +94,7 @@ pub async fn generate_jwt_auth_token_with_custom_role_attributes(
     merchant_id: String,
     org_id: String,
     role_id: String,
-) -> UserResult<Secret<String>> {
+) -> UserResult<masking::Secret<String>> {
     let token = AuthToken::new_token(
         user.get_user_id().to_string(),
         merchant_id,
@@ -101,14 +103,14 @@ pub async fn generate_jwt_auth_token_with_custom_role_attributes(
         org_id,
     )
     .await?;
-    Ok(Secret::new(token))
+    Ok(masking::Secret::new(token))
 }
 
 pub fn get_dashboard_entry_response(
     state: &AppState,
     user: UserFromStorage,
     user_role: UserRole,
-    token: Secret<String>,
+    token: masking::Secret<String>,
 ) -> UserResult<user_api::DashboardEntryResponse> {
     let verification_days_left = get_verification_days_left(state, &user)?;
 
@@ -185,9 +187,31 @@ pub async fn get_user_from_db_by_email(
         .map(UserFromStorage::from)
 }
 
-pub fn get_token_from_signin_response(resp: &user_api::SignInResponse) -> Secret<String> {
+pub fn get_token_from_signin_response(resp: &user_api::SignInResponse) -> masking::Secret<String> {
     match resp {
         user_api::SignInResponse::DashboardEntry(data) => data.token.clone(),
         user_api::SignInResponse::MerchantSelect(data) => data.token.clone(),
     }
 }
+
+pub fn generate_default_totp(
+    email: pii::Email,
+    secret: Option<masking::Secret<String>>,
+) -> UserResult<TOTP> {
+    let secret = secret
+        .map(|sec| totp_rs::Secret::Encoded(sec.expose()))
+        .unwrap_or_else(totp_rs::Secret::generate_secret)
+        .to_bytes()
+        .change_context(UserErrors::InternalServerError)?;
+
+    TOTP::new(
+        Algorithm::SHA1,
+        consts::user::TOTP_DIGITS,
+        consts::user::TOTP_TOLERANCE,
+        consts::user::TOTP_VALIDITY_DURATION_IN_SECONDS,
+        secret,
+        Some(consts::user::TOTP_ISSUER_NAME.to_string()),
+        email.expose().expose(),
+    )
+    .change_context(UserErrors::InternalServerError)
+}

crates/router_env/src/logger/types.rs

@@ -396,6 +396,8 @@ pub enum Flow {
     UpdateRole,
     /// User email flow start
     UserFromEmail,
+    /// Begin TOTP
+    TotpBegin,
     /// List initial webhook delivery attempts
     WebhookEventInitialDeliveryAttemptList,
     /// List delivery attempts for a webhook event

crates/storage_impl/src/mock_db.rs

@@ -57,6 +57,7 @@ pub struct MockDb {
     pub payouts: Arc<Mutex<Vec<store::payouts::Payouts>>>,
     pub authentications: Arc<Mutex<Vec<store::authentication::Authentication>>>,
     pub roles: Arc<Mutex<Vec<store::role::Role>>>,
+    pub user_key_store: Arc<Mutex<Vec<store::user_key_store::UserKeyStore>>>,
 }
 
 impl MockDb {
@@ -100,6 +101,7 @@ impl MockDb {
             payouts: Default::default(),
             authentications: Default::default(),
             roles: Default::default(),
+            user_key_store: Default::default(),
         })
     }
 }

migrations/2024-05-06-105026_user_key_store_table/down.sql

@@ -0,0 +1,2 @@
+-- This file should undo anything in `up.sql`
+DROP TABLE IF EXISTS user_key_store;

migrations/2024-05-06-105026_user_key_store_table/up.sql

@@ -0,0 +1,6 @@
+-- Your SQL goes here
+CREATE TABLE IF NOT EXISTS user_key_store (
+    user_id VARCHAR(64) PRIMARY KEY,
+    key bytea NOT NULL,
+    created_at TIMESTAMP NOT NULL DEFAULT now()
+);

migrations/2024-05-07-080628_user_totp/down.sql

@@ -0,0 +1,6 @@
+-- This file should undo anything in `up.sql`
+ALTER TABLE users DROP COLUMN totp_status;
+ALTER TABLE users DROP COLUMN totp_secret;
+ALTER TABLE users DROP COLUMN totp_recovery_codes;
+
+DROP TYPE "TotpStatus";

migrations/2024-05-07-080628_user_totp/up.sql

@@ -0,0 +1,10 @@
+-- Your SQL goes here
+CREATE TYPE "TotpStatus" AS ENUM (
+  'set',
+  'in_progress',
+  'not_set'
+);
+
+ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_status "TotpStatus" DEFAULT 'not_set' NOT NULL;
+ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_secret bytea DEFAULT NULL;
+ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_recovery_codes TEXT[] DEFAULT NULL;