Merge pull request #2663 from Malith-19/token-fix-for-role-update

Fix token revoke failure in role update flow for sub org.
wso2-extensions · Jan 16, 2025 · f2bb062 · f2bb062
2 parents f11275f + ab83714
commit f2bb062
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 1 deletion.
diff --git a/...rg.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth/OAuthUtil.java b/...rg.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth/OAuthUtil.java
@@ -1039,7 +1039,8 @@ public static boolean revokeTokens(String username, UserStoreManager userStoreMa
                             .getTokenManagementDAO().getAllTimeAuthorizedClientIds(authenticatedUser);
 
                 if (role != null && RoleConstants.ORGANIZATION.equals(role.getAudience())) {
-                    clientIds = filterClientIdsWithOrganizationAudience(new ArrayList<>(clientIds), tenantDomain);
+                    clientIds = filterClientIdsWithOrganizationAudience(new ArrayList<>(clientIds),
+                            authenticatedUser.getTenantDomain());
                 }
 
             } catch (IdentityOAuth2Exception e) {

diff --git a/...so2.carbon.identity.oauth/src/test/java/org/wso2/carbon/identity/oauth/OAuthUtilTest.java b/...so2.carbon.identity.oauth/src/test/java/org/wso2/carbon/identity/oauth/OAuthUtilTest.java
@@ -41,6 +41,7 @@
 import org.wso2.carbon.identity.oauth.internal.OAuthComponentServiceHolder;
 import org.wso2.carbon.identity.oauth2.dao.AccessTokenDAO;
 import org.wso2.carbon.identity.oauth2.dao.OAuthTokenPersistenceFactory;
+import org.wso2.carbon.identity.oauth2.dao.TokenManagementDAO;
 import org.wso2.carbon.identity.oauth2.model.AccessTokenDO;
 import org.wso2.carbon.identity.oauth2.util.OAuth2Util;
 import org.wso2.carbon.identity.organization.management.service.util.OrganizationManagementUtil;
@@ -291,6 +292,80 @@ public void testRevokeTokensForApplicationAudienceRoles() throws Exception {
         assertTrue(result, "Token revocation failed.");
     }
 
+    @Test
+    public void testRevokeTokensForOrganizationAudienceRoles() throws Exception {
+
+        String username = "testUser";
+        String roleId = "testRoleId";
+        String roleName = "testRole";
+        String appId = "testAppId";
+        String clientId = "testClientId";
+        String accessToken = "testAccessToken";
+
+        UserStoreManager userStoreManager = mock(UserStoreManager.class);
+        when(userStoreManager.getTenantId()).thenReturn(-1234);
+        when(userStoreManager.getRealmConfiguration()).thenReturn(mock(RealmConfiguration.class));
+        when(userStoreManager.getRealmConfiguration().getUserStoreProperty(anyString())).thenReturn("PRIMARY");
+
+        when(OrganizationManagementUtil.isOrganization(anyString())).thenReturn(false);
+        when(OAuth2Util.getTenantId(anyString())).thenReturn(-1234);
+
+        OAuthComponentServiceHolder mockOAuthComponentServiceHolder = mock(OAuthComponentServiceHolder.class);
+        when(OAuthComponentServiceHolder.getInstance()).thenReturn(mockOAuthComponentServiceHolder);
+
+        when(mockOAuthComponentServiceHolder.getRoleV2ManagementService()).thenReturn(roleManagementService);
+        RoleBasicInfo roleBasicInfo = new RoleBasicInfo();
+        roleBasicInfo.setId(roleId);
+        roleBasicInfo.setAudience(RoleConstants.ORGANIZATION);
+        roleBasicInfo.setAudienceId(appId);
+        roleBasicInfo.setName(roleName);
+        when(roleManagementService.getRoleBasicInfoById(roleId, MultitenantConstants.SUPER_TENANT_DOMAIN_NAME))
+                .thenReturn(roleBasicInfo);
+
+        when(mockOAuthComponentServiceHolder.getApplicationManagementService())
+                .thenReturn(applicationManagementService);
+        ServiceProvider serviceProvider = new ServiceProvider();
+        InboundAuthenticationConfig inboundAuthenticationConfig = new InboundAuthenticationConfig();
+        InboundAuthenticationRequestConfig[] inboundAuthenticationRequestConfigs =
+                new InboundAuthenticationRequestConfig[1];
+        InboundAuthenticationRequestConfig inboundAuthenticationRequestConfig =
+                new InboundAuthenticationRequestConfig();
+        inboundAuthenticationRequestConfig.setInboundAuthKey(clientId);
+        inboundAuthenticationRequestConfig.setInboundAuthType(ApplicationConstants.StandardInboundProtocols.OAUTH2);
+        inboundAuthenticationRequestConfigs[0] = inboundAuthenticationRequestConfig;
+        inboundAuthenticationConfig.setInboundAuthenticationRequestConfigs(inboundAuthenticationRequestConfigs);
+        serviceProvider.setInboundAuthenticationConfig(inboundAuthenticationConfig);
+        when(applicationManagementService.getApplicationByResourceId(
+                appId, MultitenantConstants.SUPER_TENANT_DOMAIN_NAME)).thenReturn(serviceProvider);
+        when(applicationManagementService.getApplicationResourceIDByInboundKey(anyString(), anyString(), anyString())).
+                thenReturn(appId);
+        when(applicationManagementService.getAllowedAudienceForRoleAssociation(anyString(), anyString())).
+                thenReturn(RoleConstants.ORGANIZATION);
+        OAuthTokenPersistenceFactory mockOAuthTokenPersistenceFactory = mock(OAuthTokenPersistenceFactory.class);
+        when(OAuthTokenPersistenceFactory.getInstance()).thenReturn(mockOAuthTokenPersistenceFactory);
+        AccessTokenDAO mockAccessTokenDAO = mock(AccessTokenDAO.class);
+        when(mockOAuthTokenPersistenceFactory.getAccessTokenDAO()).thenReturn(mockAccessTokenDAO);
+        Set<AccessTokenDO> accessTokens = new HashSet<>();
+        AccessTokenDO accessTokenDO = new AccessTokenDO();
+        accessTokenDO.setAccessToken(accessToken);
+        accessTokenDO.setConsumerKey(clientId);
+        accessTokenDO.setScope(new String[]{"default"});
+        accessTokenDO.setAuthzUser(new AuthenticatedUser());
+        accessTokens.add(accessTokenDO);
+        when(mockAccessTokenDAO.getAccessTokens(anyString(),
+                any(AuthenticatedUser.class), nullable(String.class), anyBoolean())).thenReturn(accessTokens);
+
+        TokenManagementDAO mockTokenManagementDao = mock(TokenManagementDAO.class);
+        when(mockOAuthTokenPersistenceFactory.getTokenManagementDAO()).thenReturn(mockTokenManagementDao);
+        Set<String> clientIds = new HashSet<>();
+        clientIds.add(clientId);
+        when(mockTokenManagementDao.getAllTimeAuthorizedClientIds(any())).thenReturn(clientIds);
+
+        boolean result = OAuthUtil.revokeTokens(username, userStoreManager, roleId);
+        verify(mockAccessTokenDAO, times(1)).revokeAccessTokens(any(), anyBoolean());
+        assertTrue(result, "Token revocation failed.");
+    }
+
     private OAuthCache getOAuthCache(OAuthCacheKey oAuthCacheKey) {