feat(auth): Add profile_id in AuthenticationData (#5492)

crates/router/src/core/user.rs

@@ -138,7 +138,9 @@ pub async fn signup(
         .await?;
     utils::user_role::set_role_permissions_in_cache_by_user_role(&state, &user_role).await;
 
-    let token = utils::user::generate_jwt_auth_token(&state, &user_from_db, &user_role).await?;
+    let token =
+        utils::user::generate_jwt_auth_token_without_profile(&state, &user_from_db, &user_role)
+            .await?;
     let response =
         utils::user::get_dashboard_entry_response(&state, user_from_db, user_role, token.clone())?;
 
@@ -894,6 +896,7 @@ async fn handle_new_user_invitation(
             merchant_id: user_from_token.merchant_id.clone(),
             org_id: user_from_token.org_id.clone(),
             role_id: request.role_id.clone(),
+            profile_id: None,
         };
 
         let set_metadata_request = SetMetaDataRequest::IsChangePasswordRequired;
@@ -1036,8 +1039,12 @@ pub async fn accept_invite_from_email(
         .change_context(UserErrors::InternalServerError)?
         .into();
 
-    let token =
+    let token = utils::user::generate_jwt_auth_token_without_profile(
-        utils::user::generate_jwt_auth_token(&state, &user_from_db, &update_status_result).await?;
+        &state,
+        &user_from_db,
+        &update_status_result,
+    )
+    .await?;
     utils::user_role::set_role_permissions_in_cache_by_user_role(&state, &update_status_result)
         .await;
 
@@ -1263,6 +1270,7 @@ pub async fn switch_merchant_id(
             request.merchant_id.clone(),
             org_id.clone(),
             user_from_token.role_id.clone(),
+            None,
         )
         .await?;
 
@@ -1295,7 +1303,8 @@ pub async fn switch_merchant_id(
             .ok_or(report!(UserErrors::InvalidRoleOperation))
             .attach_printable("User doesn't have access to switch")?;
 
-        let token = utils::user::generate_jwt_auth_token(&state, &user, user_role).await?;
+        let token =
+            utils::user::generate_jwt_auth_token_without_profile(&state, &user, user_role).await?;
         utils::user_role::set_role_permissions_in_cache_by_user_role(&state, user_role).await;
 
         (token, user_role.role_id.clone())

crates/router/src/core/user_role.rs

@@ -169,7 +169,9 @@ pub async fn transfer_org_ownership(
 
     utils::user_role::set_role_permissions_in_cache_by_user_role(&state, &user_role).await;
 
-    let token = utils::user::generate_jwt_auth_token(&state, &user_from_db, &user_role).await?;
+    let token =
+        utils::user::generate_jwt_auth_token_without_profile(&state, &user_from_db, &user_role)
+            .await?;
     let response =
         utils::user::get_dashboard_entry_response(&state, user_from_db, user_role, token.clone())?;
 
@@ -246,7 +248,9 @@ pub async fn merchant_select(
 
         utils::user_role::set_role_permissions_in_cache_by_user_role(&state, &user_role).await;
 
-        let token = utils::user::generate_jwt_auth_token(&state, &user_from_db, &user_role).await?;
+        let token =
+            utils::user::generate_jwt_auth_token_without_profile(&state, &user_from_db, &user_role)
+                .await?;
         let response = utils::user::get_dashboard_entry_response(
             &state,
             user_from_db,

crates/router/src/db/merchant_account.rs

@@ -273,6 +273,7 @@ impl MerchantAccountInterface for Store {
                 .change_context(errors::StorageError::DecryptionError)?,
 
             key_store,
+            profile_id: None,
         })
     }
 

crates/router/src/services/authentication.rs

@@ -54,6 +54,14 @@ mod detached;
 pub struct AuthenticationData {
     pub merchant_account: domain::MerchantAccount,
     pub key_store: domain::MerchantKeyStore,
+    pub profile_id: Option<String>,
+}
+
+#[derive(Clone)]
+pub struct AuthenticationDataWithMultipleProfiles {
+    pub merchant_account: domain::MerchantAccount,
+    pub key_store: domain::MerchantKeyStore,
+    pub profile_id_list: Option<Vec<String>>,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Serialize)]
@@ -178,6 +186,7 @@ pub struct AuthToken {
     pub role_id: String,
     pub exp: u64,
     pub org_id: id_type::OrganizationId,
+    pub profile_id: Option<String>,
 }
 
 #[cfg(feature = "olap")]
@@ -188,6 +197,7 @@ impl AuthToken {
         role_id: String,
         settings: &Settings,
         org_id: id_type::OrganizationId,
+        profile_id: Option<String>,
     ) -> UserResult<String> {
         let exp_duration = std::time::Duration::from_secs(consts::JWT_TOKEN_TIME_IN_SECS);
         let exp = jwt::generate_exp(exp_duration)?.as_secs();
@@ -197,6 +207,7 @@ impl AuthToken {
             role_id,
             exp,
             org_id,
+            profile_id,
         };
         jwt::generate_jwt(&token_payload, settings).await
     }
@@ -208,6 +219,7 @@ pub struct UserFromToken {
     pub merchant_id: id_type::MerchantId,
     pub role_id: String,
     pub org_id: id_type::OrganizationId,
+    pub profile_id: Option<String>,
 }
 
 pub struct UserIdFromAuth {
@@ -376,6 +388,7 @@ where
         let auth = AuthenticationData {
             merchant_account: merchant,
             key_store,
+            profile_id: None,
         };
         Ok((
             auth.clone(),
@@ -512,6 +525,7 @@ where
     let auth = AuthenticationData {
         merchant_account: merchant,
         key_store,
+        profile_id: None,
     };
 
     Ok(auth)
@@ -735,6 +749,7 @@ where
         let auth = AuthenticationData {
             merchant_account: merchant,
             key_store,
+            profile_id: None,
         };
         Ok((
             auth.clone(),
@@ -852,6 +867,61 @@ where
                 merchant_id: payload.merchant_id.clone(),
                 org_id: payload.org_id,
                 role_id: payload.role_id,
+                profile_id: payload.profile_id,
+            },
+            AuthenticationType::MerchantJwt {
+                merchant_id: payload.merchant_id,
+                user_id: Some(payload.user_id),
+            },
+        ))
+    }
+}
+
+#[cfg(feature = "olap")]
+#[async_trait]
+impl<A> AuthenticateAndFetch<AuthenticationDataWithMultipleProfiles, A> for JWTAuth
+where
+    A: SessionStateInfo + Sync,
+{
+    async fn authenticate_and_fetch(
+        &self,
+        request_headers: &HeaderMap,
+        state: &A,
+    ) -> RouterResult<(AuthenticationDataWithMultipleProfiles, AuthenticationType)> {
+        let payload = parse_jwt_payload::<A, AuthToken>(request_headers, state).await?;
+        if payload.check_in_blacklist(state).await? {
+            return Err(errors::ApiErrorResponse::InvalidJwtToken.into());
+        }
+
+        let permissions = authorization::get_permissions(state, &payload).await?;
+        authorization::check_authorization(&self.0, &permissions)?;
+        let key_manager_state = &(&state.session_state()).into();
+        let key_store = state
+            .store()
+            .get_merchant_key_store_by_merchant_id(
+                key_manager_state,
+                &payload.merchant_id,
+                &state.store().get_master_key().to_vec().into(),
+            )
+            .await
+            .change_context(errors::ApiErrorResponse::InvalidJwtToken)
+            .attach_printable("Failed to fetch merchant key store for the merchant id")?;
+
+        let merchant = state
+            .store()
+            .find_merchant_account_by_merchant_id(
+                key_manager_state,
+                &payload.merchant_id,
+                &key_store,
+            )
+            .await
+            .change_context(errors::ApiErrorResponse::InvalidJwtToken)?;
+
+        Ok((
+            AuthenticationDataWithMultipleProfiles {
+                key_store,
+                merchant_account: merchant,
+                profile_id_list: None,
             },
             AuthenticationType::MerchantJwt {
                 merchant_id: payload.merchant_id,
@@ -1020,6 +1090,7 @@ where
         let auth = AuthenticationData {
             merchant_account: merchant,
             key_store,
+            profile_id: payload.profile_id,
         };
         Ok((
             auth.clone(),
@@ -1075,6 +1146,7 @@ where
         let auth = AuthenticationData {
             merchant_account: merchant,
             key_store,
+            profile_id: payload.profile_id,
         };
         Ok((
             (auth.clone(), payload.user_id.clone()),
@@ -1110,6 +1182,7 @@ where
                 merchant_id: payload.merchant_id.clone(),
                 org_id: payload.org_id,
                 role_id: payload.role_id,
+                profile_id: payload.profile_id,
             },
             AuthenticationType::MerchantJwt {
                 merchant_id: payload.merchant_id,
@@ -1175,6 +1248,7 @@ where
         let auth = AuthenticationData {
             merchant_account: merchant,
             key_store,
+            profile_id: payload.profile_id,
         };
         Ok((
             auth.clone(),

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

@@ -1144,8 +1144,12 @@ impl SignInWithSingleRoleStrategy {
         self,
         state: &SessionState,
     ) -> UserResult<user_api::SignInResponse> {
-        let token =
+        let token = utils::user::generate_jwt_auth_token_without_profile(
-            utils::user::generate_jwt_auth_token(state, &self.user, &self.user_role).await?;
+            state,
+            &self.user,
+            &self.user_role,
+        )
+        .await?;
         utils::user_role::set_role_permissions_in_cache_by_user_role(state, &self.user_role).await;
 
         let dashboard_entry_response =

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

@@ -101,7 +101,7 @@ impl JWTFlow {
         Ok(true)
     }
 
-    pub async fn generate_jwt(
+    pub async fn generate_jwt_without_profile(
         self,
         state: &SessionState,
         next_flow: &NextFlow,
@@ -119,6 +119,7 @@ impl JWTFlow {
                 .org_id
                 .clone()
                 .ok_or(report!(UserErrors::InternalServerError))?,
+            None,
         )
         .await
         .map(|token| token.into())
@@ -293,7 +294,9 @@ impl NextFlow {
                 utils::user_role::set_role_permissions_in_cache_by_user_role(state, &user_role)
                     .await;
 
-                jwt_flow.generate_jwt(state, self, &user_role).await
+                jwt_flow
+                    .generate_jwt_without_profile(state, self, &user_role)
+                    .await
             }
         }
     }
@@ -313,7 +316,9 @@ impl NextFlow {
                 utils::user_role::set_role_permissions_in_cache_by_user_role(state, user_role)
                     .await;
 
-                jwt_flow.generate_jwt(state, self, user_role).await
+                jwt_flow
+                    .generate_jwt_without_profile(state, self, user_role)
+                    .await
             }
         }
     }

crates/router/src/utils/user.rs

@@ -81,7 +81,7 @@ impl UserFromToken {
     }
 }
 
-pub async fn generate_jwt_auth_token(
+pub async fn generate_jwt_auth_token_without_profile(
     state: &SessionState,
     user: &UserFromStorage,
     user_role: &UserRole,
@@ -102,6 +102,7 @@ pub async fn generate_jwt_auth_token(
             .ok_or(report!(UserErrors::InternalServerError))
             .attach_printable("org_id not found for user_role")?
             .clone(),
+        None,
     )
     .await?;
     Ok(Secret::new(token))
@@ -113,6 +114,7 @@ pub async fn generate_jwt_auth_token_with_custom_role_attributes(
     merchant_id: id_type::MerchantId,
     org_id: id_type::OrganizationId,
     role_id: String,
+    profile_id: Option<String>,
 ) -> UserResult<Secret<String>> {
     let token = AuthToken::new_token(
         user.get_user_id().to_string(),
@@ -120,6 +122,7 @@ pub async fn generate_jwt_auth_token_with_custom_role_attributes(
         role_id,
         &state.conf,
         org_id,
+        profile_id,
     )
     .await?;
     Ok(Secret::new(token))