read user_td and app_td from token.

PasanT9 · Mar 24, 2024 · fd91d75 · fd91d75
1 parent 7167866
commit fd91d75
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/...sistence/src/main/java/org/wso2/is/key/manager/tokenpersistence/PersistenceConstants.java b/...sistence/src/main/java/org/wso2/is/key/manager/tokenpersistence/PersistenceConstants.java
@@ -35,6 +35,8 @@ public static class JWTClaim {
         public static final String AUTHORIZATION_PARTY = "azp";
         public static final String IS_CONSENTED = "is_consented";
         public static final String TOKEN_TYPE_ELEM = "token_type";
+        public static final String APP_DOMAIN = "app_td";
+        public static final String USER_DOMAIN = "user_td";
     }
     public static final int SECONDS_TO_MILLISECONDS_FACTOR = 1000;
     public static final String PREV_ACCESS_TOKEN = "previousAccessToken";

diff --git a/...ersistence/src/main/java/org/wso2/is/key/manager/tokenpersistence/utils/TokenMgtUtil.java b/...ersistence/src/main/java/org/wso2/is/key/manager/tokenpersistence/utils/TokenMgtUtil.java
@@ -185,21 +185,31 @@ public static AuthenticatedUser getAuthenticatedUser(JWTClaimsSet claimsSet) thr
 
         AuthenticatedUser authenticatedUser;
         String consumerKey = null;
-        String consumerAppTenantDomain = null;
+        String appTenantDomain = null;
+        String userTenantDomain = null;
         try {
             log.debug("Getting tenant domain from OAuth app.");
-            consumerKey = (String) claimsSet.getClaim(PersistenceConstants.JWTClaim.AUTHORIZATION_PARTY);
-            if (consumerKey != null) {
-                consumerAppTenantDomain = OAuth2Util.getTenantDomainOfOauthApp(consumerKey);
+            if (claimsSet.getClaim(PersistenceConstants.JWTClaim.APP_DOMAIN) != null) {
+                appTenantDomain = (String) claimsSet.getClaim(PersistenceConstants.JWTClaim.APP_DOMAIN);
+            } else {
+                consumerKey = (String) claimsSet.getClaim(PersistenceConstants.JWTClaim.AUTHORIZATION_PARTY);
+                if (consumerKey != null) {
+                    appTenantDomain = OAuth2Util.getTenantDomainOfOauthApp(consumerKey);
+                }
             }
         } catch (InvalidOAuthClientException e) {
             throw new IdentityOAuth2Exception("Error while getting tenant domain from OAuth app with consumer key: "
                     + consumerKey);
         }
         boolean isFederated = claimsSet.getClaim(OAuth2Constants.IS_FEDERATED) != null
                 && (boolean) claimsSet.getClaim(OAuth2Constants.IS_FEDERATED);
+        if (claimsSet.getClaim(PersistenceConstants.JWTClaim.USER_DOMAIN) != null) {
+            userTenantDomain = (String) claimsSet.getClaim(PersistenceConstants.JWTClaim.USER_DOMAIN);
+        } else {
+            userTenantDomain = TokenMgtUtil.getTenantDomain();
+        }
         authenticatedUser = resolveAuthenticatedUserFromEntityId((String) claimsSet.getClaim(OAuth2Constants.ENTITY_ID),
-                consumerAppTenantDomain, isFederated, claimsSet.getSubject());
+                appTenantDomain, userTenantDomain, isFederated, claimsSet.getSubject());
         if (isFederated) {
             if (authenticatedUser == null) {
                 authenticatedUser =
@@ -220,13 +230,15 @@ public static AuthenticatedUser getAuthenticatedUser(JWTClaimsSet claimsSet) thr
      *
      * @param entityId Entity ID JWT Claim value which uniquely identifies the subject principle of the JWT. Eg: user
      * @param consumerAppTenantDomain  Tenant domain of the consumer app from the token
+     * @param userTenantDomain  Tenant domain of the user from the token
      * @param isFederated Federated user token
      * @param sub Subject claim
      * @return Username
      * @throws IdentityOAuth2Exception If an error occurs while getting the authenticated user
      */
     private static AuthenticatedUser resolveAuthenticatedUserFromEntityId(String entityId,
                                                                           String consumerAppTenantDomain,
+                                                                          String userTenantDomain,
                                                                           boolean isFederated,
                                                                           String sub)
             throws IdentityOAuth2Exception {
@@ -239,7 +251,6 @@ private static AuthenticatedUser resolveAuthenticatedUserFromEntityId(String ent
         } else {
             // Assume entity ID is userId.
             try {
-                String userTenantDomain = TokenMgtUtil.getTenantDomain();
                 String userName = getUserNameFromUserID(entityId, userTenantDomain);
                 if (StringUtils.isBlank(userName)) {
                     // if service url is not a tenant aware url, we need to get the tenant domain from the token.