From 87b1244a20be640fe4140fa92cc428093995beb0 Mon Sep 17 00:00:00 2001
From: Hasini Samarathunga <hasinis@wso2.com>
Date: Tue, 7 Jan 2025 14:40:05 +0530
Subject: [PATCH] Set correct user store domain in shared token revoke flow

---
 .../wso2/carbon/identity/oauth/OAuthUtil.java | 42 +++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth/OAuthUtil.java b/components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth/OAuthUtil.java
index 4cf88072dc..fa5578918c 100755
--- a/components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth/OAuthUtil.java
+++ b/components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth/OAuthUtil.java
@@ -70,6 +70,7 @@
 import org.wso2.carbon.identity.role.v2.mgt.core.model.Role;
 import org.wso2.carbon.idp.mgt.IdentityProviderManagementException;
 import org.wso2.carbon.user.api.Tenant;
+import org.wso2.carbon.user.api.UserRealm;
 import org.wso2.carbon.user.core.UserStoreException;
 import org.wso2.carbon.user.core.UserStoreManager;
 import org.wso2.carbon.user.core.common.AbstractUserStoreManager;
@@ -747,6 +748,10 @@ private static AuthenticatedUser buildAuthenticatedUser(UserStoreManager userSto
                 authenticatedUser.setUserResidentOrganization(managedOrg);
                 authenticatedUser.setAccessingOrganization(accessingOrg);
 
+                Optional<String> parentUserStoreDomain = getUserStoreDomainOfParentUser(
+                        userId, accessingOrg, tenantDomain);
+                parentUserStoreDomain.ifPresent(authenticatedUser::setUserStoreDomain);
+
                 // SSO login user shared flow.
                 if (!OAuthComponentServiceHolder.getInstance().getOrganizationManager()
                         .isPrimaryOrganization(managedOrg)) {
@@ -1334,4 +1339,41 @@ private static String readServerConfigurationPvtKeyJWTReuse() {
         }
         return tokenEPAllowReusePvtKeyJwtTenantConfig;
     }
+
+    /**
+     * Retrieves the user store domain of the parent user for a shared user in a specific organization.
+     *
+     * @param userId         ID of the shared user.
+     * @param accessingOrgId ID of the shared user's organization.
+     * @param tenantDomain   Tenant domain of the shared user.
+     * @return Optional containing the parent user's user store domain, or empty if not found.
+     * @throws OrganizationManagementException If an error occurs retrieving user association.
+     * @throws UserStoreException              If an error occurs retrieving the user store domain.
+     */
+    private static Optional<String> getUserStoreDomainOfParentUser(String userId, String accessingOrgId,
+                                                                   String tenantDomain)
+            throws OrganizationManagementException, UserStoreException {
+
+        String parentUserId = OAuthComponentServiceHolder.getInstance().getOrganizationUserSharingService()
+                .getUserAssociation(userId, accessingOrgId)
+                .getAssociatedUserId();
+
+        if (parentUserId == null) {
+            return Optional.empty();
+        }
+        try {
+            int tenantId = IdentityTenantUtil.getTenantId(tenantDomain);
+            UserRealm userRealm = OAuthComponentServiceHolder.getInstance()
+                    .getRealmService()
+                    .getTenantUserRealm(tenantId);
+            UserStoreManager userStoreManager = (AbstractUserStoreManager) userRealm.getUserStoreManager();
+
+            return Optional.ofNullable(((AbstractUserStoreManager) userStoreManager)
+                    .getUser(parentUserId, null)
+                    .getUserStoreDomain());
+        } catch (org.wso2.carbon.user.api.UserStoreException e) {
+            throw new UserStoreException("Failed to retrieve the user store domain for the parent user with ID: "
+                    + parentUserId + " in tenant domain: " + tenantDomain, e);
+        }
+    }
 }