feat(users): Send profile_id in JWT and user_info APIs (#5817)

crates/api_models/src/user.rs

@@ -1,6 +1,6 @@
 use std::fmt::Debug;
 
-use common_enums::{PermissionGroup, RoleScope, TokenPurpose};
+use common_enums::{EntityType, PermissionGroup, RoleScope, TokenPurpose};
 use common_utils::{crypto::OptionalEncryptableName, id_type, pii};
 use masking::Secret;
 
@@ -158,6 +158,8 @@ pub struct GetUserDetailsResponse {
     pub org_id: id_type::OrganizationId,
     pub is_two_factor_auth_setup: bool,
     pub recovery_codes_left: Option<usize>,
+    pub profile_id: id_type::ProfileId,
+    pub entity_type: EntityType,
 }
 
 #[derive(Debug, serde::Deserialize, serde::Serialize)]
@@ -185,7 +187,7 @@ pub struct GetUserRoleDetailsResponseV2 {
     pub merchant: Option<NameIdUnit<OptionalEncryptableName, id_type::MerchantId>>,
     pub profile: Option<NameIdUnit<String, id_type::ProfileId>>,
     pub status: UserStatus,
-    pub entity_type: common_enums::EntityType,
+    pub entity_type: EntityType,
 }
 
 #[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]

crates/api_models/src/user_role/role.rs

@@ -63,5 +63,5 @@ pub enum RoleCheckType {
 #[derive(Debug, serde::Serialize, Clone)]
 pub struct MinimalRoleInfo {
     pub role_id: String,
-    pub role_name: String,
+    pub role_name: Option<String>,
 }

crates/router/src/core/user.rs

@@ -100,6 +100,14 @@ pub async fn get_user_details(
 ) -> UserResponse<user_api::GetUserDetailsResponse> {
     let user = user_from_token.get_user_from_db(&state).await?;
     let verification_days_left = utils::user::get_verification_days_left(&state, &user)?;
+    let role_info = roles::RoleInfo::from_role_id(
+        &state,
+        &user_from_token.role_id,
+        &user_from_token.merchant_id,
+        &user_from_token.org_id,
+    )
+    .await
+    .change_context(UserErrors::InternalServerError)?;
 
     Ok(ApplicationResponse::Json(
         user_api::GetUserDetailsResponse {
@@ -112,6 +120,10 @@ pub async fn get_user_details(
             org_id: user_from_token.org_id,
             is_two_factor_auth_setup: user.get_totp_status() == TotpStatus::Set,
             recovery_codes_left: user.get_recovery_codes().map(|codes| codes.len()),
+            profile_id: user_from_token
+                .profile_id
+                .ok_or(UserErrors::JwtProfileIdMissing)?,
+            entity_type: role_info.get_entity_type(),
         },
     ))
 }
@@ -1185,13 +1197,12 @@ pub async fn switch_merchant_id(
             })?
             .organization_id;
 
-        let token = utils::user::generate_jwt_auth_token_with_attributes(
+        let token = utils::user::generate_jwt_auth_token_with_attributes_without_profile(
             &state,
             user_from_token.user_id,
             request.merchant_id.clone(),
             org_id.clone(),
             user_from_token.role_id.clone(),
-            None,
         )
         .await?;
 
@@ -2792,7 +2803,6 @@ pub async fn switch_org_for_user(
         .into());
     }
 
-    let key_manager_state = &(&state).into();
     let role_info = roles::RoleInfo::from_role_id(
         &state,
         &user_from_token.role_id,
@@ -2830,38 +2840,8 @@ pub async fn switch_org_for_user(
             "No user role found for the requested org_id".to_string(),
         ))?;
 
-    let merchant_id = utils::user_role::get_single_merchant_id(&state, &user_role).await?;
+    let (merchant_id, profile_id) =
-
+        utils::user_role::get_single_merchant_id_and_profile_id(&state, &user_role).await?;
-    let profile_id = if let Some(profile_id) = &user_role.profile_id {
-        profile_id.clone()
-    } else {
-        let merchant_key_store = state
-            .store
-            .get_merchant_key_store_by_merchant_id(
-                key_manager_state,
-                &merchant_id,
-                &state.store.get_master_key().to_vec().into(),
-            )
-            .await
-            .change_context(UserErrors::InternalServerError)
-            .attach_printable("Failed to retrieve merchant key store by merchant_id")?;
-
-        state
-            .store
-            .list_business_profile_by_merchant_id(
-                key_manager_state,
-                &merchant_key_store,
-                &merchant_id,
-            )
-            .await
-            .change_context(UserErrors::InternalServerError)
-            .attach_printable("Failed to list business profiles by merchant_id")?
-            .pop()
-            .ok_or(UserErrors::InternalServerError)
-            .attach_printable("No business profile found for the merchant_id")?
-            .get_id()
-            .to_owned()
-    };
 
     let token = utils::user::generate_jwt_auth_token_with_attributes(
         &state,
@@ -2869,7 +2849,7 @@ pub async fn switch_org_for_user(
         merchant_id.clone(),
         request.org_id.clone(),
         user_role.role_id.clone(),
-        Some(profile_id.clone()),
+        profile_id.clone(),
     )
     .await?;
 
@@ -3078,7 +3058,7 @@ pub async fn switch_merchant_for_user_in_org(
         merchant_id.clone(),
         org_id.clone(),
         role_id.clone(),
-        Some(profile_id),
+        profile_id,
     )
     .await?;
 
@@ -3183,7 +3163,7 @@ pub async fn switch_profile_for_user_in_org_and_merchant(
         user_from_token.merchant_id.clone(),
         user_from_token.org_id.clone(),
         role_id.clone(),
-        Some(profile_id),
+        profile_id,
     )
     .await?;
 

crates/router/src/core/user_role.rs

@@ -744,30 +744,6 @@ pub async fn list_users_in_lineage(
         .map(|user| (user.user_id.clone(), user.email))
         .collect::<HashMap<_, _>>();
 
-    let role_info_map =
-        futures::future::try_join_all(user_roles_set.iter().map(|user_role| async {
-            roles::RoleInfo::from_role_id(
-                &state,
-                &user_role.role_id,
-                &user_from_token.merchant_id,
-                &user_from_token.org_id,
-            )
-            .await
-            .map(|role_info| {
-                (
-                    user_role.role_id.clone(),
-                    user_role_api::role::MinimalRoleInfo {
-                        role_id: user_role.role_id.clone(),
-                        role_name: role_info.get_role_name().to_string(),
-                    },
-                )
-            })
-        }))
-        .await
-        .change_context(UserErrors::InternalServerError)?
-        .into_iter()
-        .collect::<HashMap<_, _>>();
-
     let user_role_map = user_roles_set
         .into_iter()
         .fold(HashMap::new(), |mut map, user_role| {
@@ -787,13 +763,11 @@ pub async fn list_users_in_lineage(
                         .ok_or(UserErrors::InternalServerError)?,
                     roles: role_id_vec
                         .into_iter()
-                        .map(|role_id| {
+                        .map(|role_id| user_role_api::role::MinimalRoleInfo {
-                            role_info_map
+                            role_id,
-                                .get(&role_id)
+                            role_name: None,
-                                .cloned()
-                                .ok_or(UserErrors::InternalServerError)
                         })
-                        .collect::<Result<Vec<_>, _>>()?,
+                        .collect(),
                 })
             })
             .collect::<Result<Vec<_>, _>>()?,

crates/router/src/core/user_role/role.rs

@@ -348,7 +348,7 @@ pub async fn list_roles_at_entity_level(
             if check_type && role_info.get_entity_type() == req.entity_type {
                 Some(role_api::MinimalRoleInfo {
                     role_id: role_info.get_role_id().to_string(),
-                    role_name: role_info.get_role_name().to_string(),
+                    role_name: Some(role_info.get_role_name().to_string()),
                 })
             } else {
                 None

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

@@ -5,7 +5,7 @@ use api_models::{
 };
 use common_enums::EntityType;
 use common_utils::{
-    crypto::Encryptable, errors::CustomResult, id_type, new_type::MerchantName, pii, type_name,
+    crypto::Encryptable, id_type, new_type::MerchantName, pii, type_name,
     types::keymanager::Identifier,
 };
 use diesel_models::{
@@ -28,7 +28,7 @@ use crate::{
     consts,
     core::{
         admin,
-        errors::{self, UserErrors, UserResult},
+        errors::{UserErrors, UserResult},
     },
     db::{user_role::InsertUserRolePayload, GlobalStorageInterface},
     routes::SessionState,
@@ -867,22 +867,6 @@ impl UserFromStorage {
         self.0.email.clone()
     }
 
-    pub async fn get_role_from_db(&self, state: SessionState) -> UserResult<UserRole> {
-        state
-            .store
-            .find_user_role_by_user_id(&self.0.user_id, UserRoleVersion::V1)
-            .await
-            .change_context(UserErrors::InternalServerError)
-    }
-
-    pub async fn get_roles_from_db(&self, state: &SessionState) -> UserResult<Vec<UserRole>> {
-        state
-            .store
-            .list_user_roles_by_user_id_and_version(&self.0.user_id, UserRoleVersion::V1)
-            .await
-            .change_context(UserErrors::InternalServerError)
-    }
-
     #[cfg(feature = "email")]
     pub fn get_verification_days_left(&self, state: &SessionState) -> UserResult<Option<i64>> {
         if self.0.is_verified {
@@ -930,21 +914,6 @@ impl UserFromStorage {
         Ok(days_left_for_password_rotate.whole_days() < 0)
     }
 
-    pub async fn get_role_from_db_by_merchant_id(
-        &self,
-        state: &SessionState,
-        merchant_id: &id_type::MerchantId,
-    ) -> CustomResult<UserRole, errors::StorageError> {
-        state
-            .store
-            .find_user_role_by_user_id_merchant_id(
-                self.get_user_id(),
-                merchant_id,
-                UserRoleVersion::V1,
-            )
-            .await
-    }
-
     pub async fn get_or_create_key_store(&self, state: &SessionState) -> UserResult<UserKeyStore> {
         let master_key = state.store.get_master_key();
         let key_manager_state = &state.into();
@@ -1253,7 +1222,7 @@ where
         }
     }
 
-    async fn insert_v1_and_v2_in_db_and_get_v1(
+    async fn insert_v1_and_v2_in_db_and_get_v2(
         state: &SessionState,
         v1_role: UserRoleNew,
         v2_role: UserRoleNew,
@@ -1264,10 +1233,9 @@ where
             .await
             .change_context(UserErrors::InternalServerError)?;
 
-        // Returning v1 role so other code which was not migrated doesn't break
         inserted_roles
             .into_iter()
-            .find(|role| role.version == UserRoleVersion::V1)
+            .find(|role| role.version == UserRoleVersion::V2)
             .ok_or(report!(UserErrors::InternalServerError))
     }
 }
@@ -1323,7 +1291,7 @@ impl NewUserRole<OrganizationLevel> {
             entity_type: EntityType::Organization,
         });
 
-        Self::insert_v1_and_v2_in_db_and_get_v1(state, new_v1_role, new_v2_role).await
+        Self::insert_v1_and_v2_in_db_and_get_v2(state, new_v1_role, new_v2_role).await
     }
 }
 
@@ -1343,7 +1311,7 @@ impl NewUserRole<MerchantLevel> {
             entity_type: EntityType::Merchant,
         });
 
-        Self::insert_v1_and_v2_in_db_and_get_v1(state, new_v1_role, new_v2_role).await
+        Self::insert_v1_and_v2_in_db_and_get_v2(state, new_v1_role, new_v2_role).await
     }
 }
 
@@ -1366,7 +1334,7 @@ impl NewUserRole<InternalLevel> {
             entity_type: EntityType::Internal,
         });
 
-        Self::insert_v1_and_v2_in_db_and_get_v1(state, new_v1_role, new_v2_role).await
+        Self::insert_v1_and_v2_in_db_and_get_v2(state, new_v1_role, new_v2_role).await
     }
 }
 

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

@@ -1,8 +1,5 @@
 use common_enums::TokenPurpose;
-use diesel_models::{
+use diesel_models::{enums::UserStatus, user_role::UserRole};
-    enums::{UserRoleVersion, UserStatus},
-    user_role::UserRole,
-};
 use error_stack::{report, ResultExt};
 use masking::Secret;
 
@@ -67,10 +64,21 @@ impl SPTFlow {
             Self::ForceSetPassword => user
                 .is_password_rotate_required(state)
                 .map(|rotate_required| rotate_required && !path.contains(&TokenPurpose::SSO)),
-            Self::MerchantSelect => user
+            Self::MerchantSelect => Ok(state
-                .get_roles_from_db(state)
+                .store
+                .list_user_roles_by_user_id(ListUserRolesByUserIdPayload {
+                    user_id: user.get_user_id(),
+                    org_id: None,
+                    merchant_id: None,
+                    profile_id: None,
+                    entity_id: None,
+                    version: None,
+                    status: Some(UserStatus::Active),
+                    limit: Some(1),
+                })
                 .await
-                .map(|roles| !roles.iter().any(|role| role.status == UserStatus::Active)),
+                .change_context(UserErrors::InternalServerError)?
+                .is_empty()),
         }
     }
 
@@ -105,15 +113,17 @@ impl JWTFlow {
         Ok(true)
     }
 
-    pub async fn generate_jwt_without_profile(
+    pub async fn generate_jwt(
         self,
         state: &SessionState,
         next_flow: &NextFlow,
         user_role: &UserRole,
     ) -> UserResult<Secret<String>> {
+        let (merchant_id, profile_id) =
+            utils::user_role::get_single_merchant_id_and_profile_id(state, user_role).await?;
         auth::AuthToken::new_token(
             next_flow.user.get_user_id().to_string(),
-            utils::user_role::get_single_merchant_id(state, user_role).await?,
+            merchant_id,
             user_role.role_id.clone(),
             &state.conf,
             user_role
@@ -121,7 +131,7 @@ impl JWTFlow {
                 .clone()
                 .ok_or(report!(UserErrors::InternalServerError))
                 .attach_printable("org_id not found")?,
-            None,
+            Some(profile_id),
         )
         .await
         .map(|token| token.into())
@@ -296,7 +306,7 @@ impl NextFlow {
                         merchant_id: None,
                         profile_id: None,
                         entity_id: None,
-                        version: Some(UserRoleVersion::V1),
+                        version: None,
                         status: Some(UserStatus::Active),
                         limit: Some(1),
                     })
@@ -307,9 +317,7 @@ impl NextFlow {
                 utils::user_role::set_role_permissions_in_cache_by_user_role(state, &user_role)
                     .await;
 
-                jwt_flow
+                jwt_flow.generate_jwt(state, self, &user_role).await
-                    .generate_jwt_without_profile(state, self, &user_role)
-                    .await
             }
         }
     }
@@ -329,9 +337,7 @@ impl NextFlow {
                 utils::user_role::set_role_permissions_in_cache_by_user_role(state, user_role)
                     .await;
 
-                jwt_flow
+                jwt_flow.generate_jwt(state, self, user_role).await
-                    .generate_jwt_without_profile(state, self, user_role)
-                    .await
             }
         }
     }

crates/router/src/utils/user.rs

@@ -108,13 +108,25 @@ pub async fn generate_jwt_auth_token_without_profile(
     Ok(Secret::new(token))
 }
 
+pub async fn generate_jwt_auth_token_with_attributes_without_profile(
+    state: &SessionState,
+    user_id: String,
+    merchant_id: id_type::MerchantId,
+    org_id: id_type::OrganizationId,
+    role_id: String,
+) -> UserResult<Secret<String>> {
+    let token =
+        AuthToken::new_token(user_id, merchant_id, role_id, &state.conf, org_id, None).await?;
+    Ok(Secret::new(token))
+}
+
 pub async fn generate_jwt_auth_token_with_attributes(
     state: &SessionState,
     user_id: String,
     merchant_id: id_type::MerchantId,
     org_id: id_type::OrganizationId,
     role_id: String,
-    profile_id: Option<id_type::ProfileId>,
+    profile_id: id_type::ProfileId,
 ) -> UserResult<Secret<String>> {
     let token = AuthToken::new_token(
         user_id,
@@ -122,7 +134,7 @@ pub async fn generate_jwt_auth_token_with_attributes(
         role_id,
         &state.conf,
         org_id,
-        profile_id,
+        Some(profile_id),
     )
     .await?;
     Ok(Secret::new(token))

crates/router/src/utils/user_role.rs

@@ -358,3 +358,42 @@ pub async fn get_lineage_for_user_id_and_entity_for_accepting_invite(
         }
     }
 }
+
+pub async fn get_single_merchant_id_and_profile_id(
+    state: &SessionState,
+    user_role: &UserRole,
+) -> UserResult<(id_type::MerchantId, id_type::ProfileId)> {
+    let merchant_id = get_single_merchant_id(state, user_role).await?;
+    let (_, entity_type) = user_role
+        .get_entity_id_and_type()
+        .ok_or(UserErrors::InternalServerError)?;
+    let profile_id = match entity_type {
+        EntityType::Organization | EntityType::Merchant | EntityType::Internal => {
+            let key_store = state
+                .store
+                .get_merchant_key_store_by_merchant_id(
+                    &state.into(),
+                    &merchant_id,
+                    &state.store.get_master_key().to_vec().into(),
+                )
+                .await
+                .change_context(UserErrors::InternalServerError)?;
+
+            state
+                .store
+                .list_business_profile_by_merchant_id(&state.into(), &key_store, &merchant_id)
+                .await
+                .change_context(UserErrors::InternalServerError)?
+                .pop()
+                .ok_or(UserErrors::InternalServerError)?
+                .get_id()
+                .to_owned()
+        }
+        EntityType::Profile => user_role
+            .profile_id
+            .clone()
+            .ok_or(UserErrors::InternalServerError)?,
+    };
+
+    Ok((merchant_id, profile_id))
+}