Merge pull request #240 from sadilchamishka/improve-claim-provider

Improve claim provider according to the organization bound token improvements
wso2-extensions · Oct 17, 2023 · 5343d35 · 5343d35
2 parents 147492e + 870b88b
commit 5343d35
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 22 deletions.
diff --git a/components/org.wso2.carbon.identity.organization.management.claim.provider/pom.xml b/components/org.wso2.carbon.identity.organization.management.claim.provider/pom.xml
@@ -36,6 +36,10 @@
             <groupId>org.apache.felix</groupId>
             <artifactId>org.apache.felix.scr.ds-annotations</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.wso2.carbon.identity.framework</groupId>
+            <artifactId>org.wso2.carbon.identity.application.authentication.framework</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.wso2.carbon.identity.inbound.auth.oauth2</groupId>
             <artifactId>org.wso2.carbon.identity.oauth</artifactId>

diff --git a/...so2/carbon/identity/organization/management/claim/provider/OrganizationClaimProvider.java b/...so2/carbon/identity/organization/management/claim/provider/OrganizationClaimProvider.java
@@ -41,69 +41,92 @@
  */
 public class OrganizationClaimProvider implements ClaimProvider, JWTAccessTokenClaimProvider {
 
-    private static final String ORGANIZATION_ID_ATTRIBUTE = "org_id";
-    private static final String ORGANIZATION_NAME_ATTRIBUTE = "org_name";
+    private static final String AUTHORIZED_ORGANIZATION_ID_ATTRIBUTE = "org_id";
+    private static final String AUTHORIZED_ORGANIZATION_NAME_ATTRIBUTE = "org_name";
+    private static final String USER_RESIDENT_ORGANIZATION_NAME_ATTRIBUTE = "user_org";
 
     @Override
     public Map<String, Object> getAdditionalClaims(OAuthAuthzReqMessageContext oAuthAuthzReqMessageContext,
                                                    OAuth2AuthorizeRespDTO oAuth2AuthorizeRespDTO)
             throws IdentityOAuth2Exception {
 
         String tenantDomain = oAuthAuthzReqMessageContext.getAuthorizationReqDTO().getLoggedInTenantDomain();
-        return getOrganizationInformation(tenantDomain);
+        String organizationId = resolveOrganizationId(tenantDomain);
+        return buildOrganizationInformation(organizationId, organizationId);
     }
 
     @Override
     public Map<String, Object> getAdditionalClaims(OAuthTokenReqMessageContext oAuthTokenReqMessageContext,
                                                    OAuth2AccessTokenRespDTO oAuth2AccessTokenRespDTO)
             throws IdentityOAuth2Exception {
 
-        String tenantDomain = oAuthTokenReqMessageContext.getAuthorizedUser().getTenantDomain();
-        return getOrganizationInformation(tenantDomain);
-    }
-
-    private OrganizationManager getOrganizationManager() {
-
-        return OrganizationClaimProviderServiceComponentHolder.getInstance().getOrganizationManager();
+        String userResidentOrgId = oAuthTokenReqMessageContext.getAuthorizedUser().getUserResidentOrganization();
+        String authorizedOrgId = oAuthTokenReqMessageContext.getAuthorizedUser().getAccessingOrganization();
+        if (StringUtils.isEmpty(authorizedOrgId)) {
+            authorizedOrgId = resolveOrganizationId(oAuthTokenReqMessageContext.getAuthorizedUser().getTenantDomain());
+        }
+        return buildOrganizationInformation(userResidentOrgId, authorizedOrgId);
     }
 
     @Override
     public Map<String, Object> getAdditionalClaims(OAuthAuthzReqMessageContext oAuthAuthzReqMessageContext)
             throws IdentityOAuth2Exception {
 
         String tenantDomain = oAuthAuthzReqMessageContext.getAuthorizationReqDTO().getLoggedInTenantDomain();
-        return getOrganizationInformation(tenantDomain);
+        String organizationId = resolveOrganizationId(tenantDomain);
+        return buildOrganizationInformation(organizationId, organizationId);
     }
 
     @Override
     public Map<String, Object> getAdditionalClaims(OAuthTokenReqMessageContext oAuthTokenReqMessageContext)
             throws IdentityOAuth2Exception {
 
-        String tenantDomain = oAuthTokenReqMessageContext.getAuthorizedUser().getTenantDomain();
-        return getOrganizationInformation(tenantDomain);
+        String userResidentOrgId = oAuthTokenReqMessageContext.getAuthorizedUser().getUserResidentOrganization();
+        String authorizedOrgId = oAuthTokenReqMessageContext.getAuthorizedUser().getAccessingOrganization();
+        // The below condition is not required once console is modeled as B2B app.
+        if (StringUtils.isEmpty(authorizedOrgId)) {
+            authorizedOrgId = resolveOrganizationId(oAuthTokenReqMessageContext.getAuthorizedUser().getTenantDomain());
+        }
+        return buildOrganizationInformation(userResidentOrgId, authorizedOrgId);
     }
 
-    private Map<String, Object> getOrganizationInformation(String tenantDomain) throws IdentityOAuth2Exception {
+    private Map<String, Object> buildOrganizationInformation(String userResideOrgId, String authorizedOrgId)
+            throws IdentityOAuth2Exception {
 
         Map<String, Object> additionalClaims = new HashMap<>();
         if (!OrganizationClaimProviderServiceComponentHolder.getInstance().isOrganizationManagementEnable()) {
             return additionalClaims;
         }
         try {
-            String organizationId = getOrganizationManager().resolveOrganizationId(tenantDomain);
-            if (StringUtils.isNotBlank(organizationId)) {
-                String organizationName = getOrganizationManager().getOrganizationNameById(organizationId);
-                additionalClaims.put(ORGANIZATION_ID_ATTRIBUTE, organizationId);
-                additionalClaims.put(ORGANIZATION_NAME_ATTRIBUTE, organizationName);
+            if (StringUtils.isNotBlank(authorizedOrgId)) {
+                String authorizedOrgName = getOrganizationManager().getOrganizationNameById(authorizedOrgId);
+                additionalClaims.put(USER_RESIDENT_ORGANIZATION_NAME_ATTRIBUTE, userResideOrgId);
+                additionalClaims.put(AUTHORIZED_ORGANIZATION_ID_ATTRIBUTE, authorizedOrgId);
+                additionalClaims.put(AUTHORIZED_ORGANIZATION_NAME_ATTRIBUTE, authorizedOrgName);
             }
+        } catch (OrganizationManagementException e) {
+            throw new IdentityOAuth2Exception("Error while resolving organization name by ID.", e);
+        }
+        return additionalClaims;
+    }
+
+    private String resolveOrganizationId(String tenantDomain) throws IdentityOAuth2Exception {
+
+        try {
+            return getOrganizationManager().resolveOrganizationId(tenantDomain);
         } catch (OrganizationManagementClientException e) {
+            // This client error handling should be removed once all the tenants have corresponding organization.
             if (ERROR_CODE_ORGANIZATION_NOT_FOUND_FOR_TENANT.getCode().equals(e.getErrorCode())) {
-                return additionalClaims;
+                return null;
             }
             throw new IdentityOAuth2Exception("Error while resolving organization id.", e);
         } catch (OrganizationManagementException e) {
             throw new IdentityOAuth2Exception("Error while resolving organization id.", e);
         }
-        return additionalClaims;
+    }
+
+    private OrganizationManager getOrganizationManager() {
+
+        return OrganizationClaimProviderServiceComponentHolder.getInstance().getOrganizationManager();
     }
 }
diff --git a/pom.xml b/pom.xml
@@ -215,6 +215,11 @@
                 <artifactId>org.wso2.carbon.identity.configuration.mgt.core</artifactId>
                 <version>${carbon.identity.framework.version}</version>
             </dependency>
+            <dependency>
+                <groupId>org.wso2.carbon.identity.framework</groupId>
+                <artifactId>org.wso2.carbon.identity.application.authentication.framework</artifactId>
+                <version>${carbon.identity.framework.version}</version>
+            </dependency>
             <dependency>
                 <groupId>org.wso2.carbon.identity.inbound.auth.oauth2</groupId>
                 <artifactId>org.wso2.carbon.identity.oauth</artifactId>
@@ -494,7 +499,7 @@
         <carbon.multitenancy.package.import.version.range>[4.7.0,5.0.0)
         </carbon.multitenancy.package.import.version.range>
 
-        <carbon.identity.framework.version>5.25.369</carbon.identity.framework.version>
+        <carbon.identity.framework.version>5.25.396</carbon.identity.framework.version>
         <carbon.identity.package.import.version.range>[5.20.0, 7.0.0)
         </carbon.identity.package.import.version.range>