back port wso2-extensions#2317

dushaniw · Jan 30, 2024 · a217f7b · a217f7b
1 parent 05002e7
commit a217f7b
Show file tree

Hide file tree

Showing 8 changed files with 102 additions and 43 deletions.
diff --git a/...java/org/wso2/carbon/identity/oauth/tokenprocessor/DefaultRefreshTokenGrantProcessor.java b/...java/org/wso2/carbon/identity/oauth/tokenprocessor/DefaultRefreshTokenGrantProcessor.java
@@ -19,10 +19,14 @@
 package org.wso2.carbon.identity.oauth.tokenprocessor;
 
 import org.apache.commons.codec.digest.DigestUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.wso2.carbon.identity.base.IdentityConstants;
 import org.wso2.carbon.identity.core.util.IdentityUtil;
+import org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCache;
+import org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCacheEntry;
+import org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCacheKey;
 import org.wso2.carbon.identity.oauth.common.OAuthConstants;
 import org.wso2.carbon.identity.oauth2.IdentityOAuth2Exception;
 import org.wso2.carbon.identity.oauth2.dao.OAuthTokenPersistenceFactory;
@@ -38,6 +42,7 @@
 import java.util.Date;
 import java.util.List;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Default implementation of @RefreshTokenProcessor responsible for handling refresh token persistence logic.
@@ -184,4 +189,44 @@ private List<AccessTokenDO> getAccessTokenBeans(OAuth2AccessTokenReqDTO tokenReq
         }
         return accessTokenBeans;
     }
+
+    public void addUserAttributesToCache(AccessTokenDO accessTokenBean, OAuthTokenReqMessageContext msgCtx) {
+
+        RefreshTokenValidationDataDO oldAccessToken =
+                (RefreshTokenValidationDataDO) msgCtx.getProperty(PREV_ACCESS_TOKEN);
+        if (oldAccessToken.getAccessToken() == null) {
+            return;
+        }
+        AuthorizationGrantCacheKey oldAuthorizationGrantCacheKey = new AuthorizationGrantCacheKey(oldAccessToken
+                .getAccessToken());
+        if (log.isDebugEnabled()) {
+            log.debug("Getting AuthorizationGrantCacheEntry using access token id: " + accessTokenBean.getTokenId());
+        }
+        AuthorizationGrantCacheEntry grantCacheEntry =
+                AuthorizationGrantCache.getInstance().getValueFromCacheByTokenId(oldAuthorizationGrantCacheKey,
+                        oldAccessToken.getTokenId());
+
+        if (grantCacheEntry != null) {
+            if (log.isDebugEnabled()) {
+                log.debug("Getting user attributes cached against the previous access token with access token id: " +
+                        oldAccessToken.getTokenId());
+            }
+            AuthorizationGrantCacheKey authorizationGrantCacheKey = new AuthorizationGrantCacheKey(accessTokenBean
+                    .getAccessToken());
+
+            if (StringUtils.isNotBlank(accessTokenBean.getTokenId())) {
+                grantCacheEntry.setTokenId(accessTokenBean.getTokenId());
+            } else {
+                grantCacheEntry.setTokenId(null);
+            }
+
+            grantCacheEntry.setValidityPeriod(
+                    TimeUnit.MILLISECONDS.toNanos(accessTokenBean.getValidityPeriodInMillis()));
+
+            // This new method has introduced in order to resolve a regression occurred : wso2/product-is#4366.
+            AuthorizationGrantCache.getInstance().clearCacheEntryByTokenId(oldAuthorizationGrantCacheKey,
+                    oldAccessToken.getTokenId());
+            AuthorizationGrantCache.getInstance().addToCacheByToken(authorizationGrantCacheKey, grantCacheEntry);
+        }
+    }
 }
diff --git a/...c/main/java/org/wso2/carbon/identity/oauth/tokenprocessor/RefreshTokenGrantProcessor.java b/...c/main/java/org/wso2/carbon/identity/oauth/tokenprocessor/RefreshTokenGrantProcessor.java
@@ -76,4 +76,15 @@ AccessTokenDO createAccessTokenBean(OAuthTokenReqMessageContext tokReqMsgCtx, OA
      */
     boolean isLatestRefreshToken(OAuth2AccessTokenReqDTO tokenReq, RefreshTokenValidationDataDO validationBean,
                                  String userStoreDomain) throws IdentityOAuth2Exception;
+
+    /**
+     * After new refresh token and access token generated, this method will be called to add user attributes to the
+     * grant cache.
+     *
+     * @param accessTokenBean Access token data object.
+     * @param msgCtx Token request message context.
+     * @throws IdentityOAuth2Exception If an error occurred while adding user attributes to the grant cache.
+     */
+    public void addUserAttributesToCache(AccessTokenDO accessTokenBean, OAuthTokenReqMessageContext msgCtx)
+            throws IdentityOAuth2Exception;
 }
diff --git a/....carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/OAuth2Constants.java b/....carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/OAuth2Constants.java
@@ -36,6 +36,7 @@ public static class TokenBinderType {
     public static final String ENTITY_ID = "entity_id";
     public static final String IS_CONSENTED = "is_consented";
     public static final String IS_FEDERATED = "is_federated";
+    public static final String USER_SESSION_ID = "usid";
     public static final boolean DEFAULT_PERSIST_ENABLED = true;
     public static final String OAUTH_TOKEN_PERSISTENCE_ENABLE = "OAuth.TokenPersistence.Enable";
     public static final String OAUTH_CODE_PERSISTENCE_ENABLE = "OAuth.EnableAuthCodePersistence";

diff --git a/...ain/java/org/wso2/carbon/identity/oauth2/authz/handlers/util/ResponseTypeHandlerUtil.java b/...ain/java/org/wso2/carbon/identity/oauth2/authz/handlers/util/ResponseTypeHandlerUtil.java
@@ -45,6 +45,7 @@
 import org.wso2.carbon.identity.oauth.event.OAuthEventInterceptor;
 import org.wso2.carbon.identity.oauth.internal.OAuthComponentServiceHolder;
 import org.wso2.carbon.identity.oauth2.IdentityOAuth2Exception;
+import org.wso2.carbon.identity.oauth2.OAuth2Constants;
 import org.wso2.carbon.identity.oauth2.OAuth2Service;
 import org.wso2.carbon.identity.oauth2.authz.OAuthAuthzReqMessageContext;
 import org.wso2.carbon.identity.oauth2.bean.OAuthClientAuthnContext;
@@ -607,6 +608,9 @@ private static AccessTokenDO createNewTokenBean(OAuthAuthzReqMessageContext oaut
         newTokenBean.setScope(oauthAuthzMsgCtx.getApprovedScope());
         newTokenBean.setTokenId(UUID.randomUUID().toString());
         newTokenBean.setTokenType(OAuthConstants.UserType.APPLICATION_USER);
+        String tokenId = UUID.randomUUID().toString();
+        newTokenBean.setTokenId(tokenId);
+        oauthAuthzMsgCtx.addProperty(OAuth2Constants.USER_SESSION_ID, tokenId);
         newTokenBean.setIssuedTime(timestamp);
         newTokenBean.setValidityPeriodInMillis(validityPeriodInMillis);
         newTokenBean.setValidityPeriod(validityPeriodInMillis / SECOND_TO_MILLISECONDS_FACTOR);

diff --git a/...on.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/JWTTokenIssuer.java b/...on.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/JWTTokenIssuer.java
@@ -806,6 +806,13 @@ protected void setClaimsForNonPersistence(JWTClaimsSet.Builder jwtClaimsSetBuild
                 // when no persistence of tokens, there is no existing token to check the consented value for.
                 jwtClaimsSetBuilder.claim(OAuth2Constants.IS_CONSENTED, isConsented);
                 jwtClaimsSetBuilder.claim(OAuth2Constants.IS_FEDERATED, authenticatedUser.isFederatedUser());
+                if (tokenReqMessageContext != null) {
+                    jwtClaimsSetBuilder.claim(OAuth2Constants.USER_SESSION_ID,
+                            tokenReqMessageContext.getProperty(OAuth2Constants.USER_SESSION_ID));
+                } else {
+                    jwtClaimsSetBuilder.claim(OAuth2Constants.USER_SESSION_ID,
+                            authAuthzReqMessageContext.getProperty(OAuth2Constants.USER_SESSION_ID));
+                }
             }
         }
     }

diff --git a/...g/wso2/carbon/identity/oauth2/token/handlers/grant/AbstractAuthorizationGrantHandler.java b/...g/wso2/carbon/identity/oauth2/token/handlers/grant/AbstractAuthorizationGrantHandler.java
@@ -44,6 +44,7 @@
 import org.wso2.carbon.identity.oauth.dao.OAuthAppDO;
 import org.wso2.carbon.identity.oauth.internal.OAuthComponentServiceHolder;
 import org.wso2.carbon.identity.oauth2.IdentityOAuth2Exception;
+import org.wso2.carbon.identity.oauth2.OAuth2Constants;
 import org.wso2.carbon.identity.oauth2.OAuth2Service;
 import org.wso2.carbon.identity.oauth2.dao.OAuthTokenPersistenceFactory;
 import org.wso2.carbon.identity.oauth2.dto.OAuth2AccessTokenReqDTO;
@@ -463,6 +464,9 @@ private AccessTokenDO createNewTokenBean(OAuthTokenReqMessageContext tokReqMsgCt
         newTokenBean.setTenantID(OAuth2Util.getTenantId(tenantDomain));
         newTokenBean.setTokenId(UUID.randomUUID().toString());
         newTokenBean.setGrantType(tokenReq.getGrantType());
+        String tokenId = UUID.randomUUID().toString();
+        tokReqMsgCtx.addProperty(OAuth2Constants.USER_SESSION_ID, tokenId);
+        newTokenBean.setTokenId(tokenId);
         /* If the existing token is available, the consented token flag will be extracted from that. Otherwise,
         from the current grant. */
         if (OAuth2ServiceComponentHolder.isConsentedTokenColumnEnabled()) {

diff --git a/...c/main/java/org/wso2/carbon/identity/oauth2/token/handlers/grant/RefreshGrantHandler.java b/...c/main/java/org/wso2/carbon/identity/oauth2/token/handlers/grant/RefreshGrantHandler.java
@@ -127,7 +127,7 @@ public OAuth2AccessTokenRespDTO issue(OAuthTokenReqMessageContext tokReqMsgCtx)
         }
 
         setTokenDataToMessageContext(tokReqMsgCtx, accessTokenBean);
-        addUserAttributesToCache(accessTokenBean, tokReqMsgCtx);
+        getRefreshTokenGrantProcessor().addUserAttributesToCache(accessTokenBean, tokReqMsgCtx);
         return buildTokenResponse(tokReqMsgCtx, accessTokenBean);
     }
 
@@ -573,47 +573,6 @@ private long getRefreshTokenValidityPeriod(long refreshTokenValidityPeriod, OAut
         return refreshTokenValidityPeriod;
     }
 
-    private static void addUserAttributesToCache(AccessTokenDO accessTokenBean,
-                                                 OAuthTokenReqMessageContext msgCtx) {
-
-        RefreshTokenValidationDataDO oldAccessToken =
-                (RefreshTokenValidationDataDO) msgCtx.getProperty(PREV_ACCESS_TOKEN);
-        if (oldAccessToken.getAccessToken() == null) {
-            return;
-        }
-        AuthorizationGrantCacheKey oldAuthorizationGrantCacheKey = new AuthorizationGrantCacheKey(oldAccessToken
-                .getAccessToken());
-        if (log.isDebugEnabled()) {
-            log.debug("Getting AuthorizationGrantCacheEntry using access token id: " + accessTokenBean.getTokenId());
-        }
-        AuthorizationGrantCacheEntry grantCacheEntry =
-                AuthorizationGrantCache.getInstance().getValueFromCacheByTokenId(oldAuthorizationGrantCacheKey,
-                        oldAccessToken.getTokenId());
-
-        if (grantCacheEntry != null) {
-            if (log.isDebugEnabled()) {
-                log.debug("Getting user attributes cached against the previous access token with access token id: " +
-                        oldAccessToken.getTokenId());
-            }
-            AuthorizationGrantCacheKey authorizationGrantCacheKey = new AuthorizationGrantCacheKey(accessTokenBean
-                    .getAccessToken());
-
-            if (StringUtils.isNotBlank(accessTokenBean.getTokenId())) {
-                grantCacheEntry.setTokenId(accessTokenBean.getTokenId());
-            } else {
-                grantCacheEntry.setTokenId(null);
-            }
-
-            grantCacheEntry.setValidityPeriod(
-                    TimeUnit.MILLISECONDS.toNanos(accessTokenBean.getValidityPeriodInMillis()));
-
-            // This new method has introduced in order to resolve a regression occurred : wso2/product-is#4366.
-            AuthorizationGrantCache.getInstance().clearCacheEntryByTokenId(oldAuthorizationGrantCacheKey,
-                    oldAccessToken.getTokenId());
-            AuthorizationGrantCache.getInstance().addToCacheByToken(authorizationGrantCacheKey, grantCacheEntry);
-        }
-    }
-
     private boolean isRenewRefreshToken(String renewRefreshToken) {
 
         if (StringUtils.isNotBlank(renewRefreshToken)) {

diff --git a/...rc/main/java/org/wso2/carbon/identity/openidconnect/DefaultOIDCClaimsCallbackHandler.java b/...rc/main/java/org/wso2/carbon/identity/openidconnect/DefaultOIDCClaimsCallbackHandler.java
@@ -311,14 +311,42 @@ private Map<ClaimMapping, String> getCachedUserAttributes(OAuthTokenReqMessageCo
                 }
                 RefreshTokenValidationDataDO refreshTokenValidationDataDO =
                         (RefreshTokenValidationDataDO) previousAccessTokenObject;
-                userAttributes = getUserAttributesCachedAgainstToken(refreshTokenValidationDataDO.getAccessToken());
+                if (refreshTokenValidationDataDO.getAccessToken() != null) {
+                    userAttributes = getUserAttributesCachedAgainstToken(refreshTokenValidationDataDO.getAccessToken());
+                } else if (refreshTokenValidationDataDO.getTokenId() != null) {
+                    userAttributes = getUserAttributesCachedAgainstTokenId(refreshTokenValidationDataDO.getTokenId());
+                } else {
+                    userAttributes = new HashMap<>();
+                }
                 requestMsgCtx.addProperty(OIDCConstants.HAS_NON_OIDC_CLAIMS,
                         isTokenHasCustomUserClaims(refreshTokenValidationDataDO));
             }
         }
         return userAttributes;
     }
 
+    /**
+     * Get the user attributes cached against the token ID when access token is not persisted.
+     *
+     * @param tokenId Token Id
+     * @return Map of user attributes
+     */
+    private Map<ClaimMapping, String> getUserAttributesCachedAgainstTokenId(String tokenId) {
+        AuthorizationGrantCacheEntry cacheEntry;
+        if (log.isDebugEnabled()) {
+            if (IdentityUtil.isTokenLoggable(IdentityConstants.IdentityTokens.ACCESS_TOKEN)) {
+                log.debug("Retrieving user attributes cached against access token Id : "
+                        + tokenId);
+            } else {
+                log.debug("Retrieving user attributes cached against access token Id.");
+            }
+        }
+        // Since access token is not persisted, we are passing null as the grant cache key. User attributes will
+        // be cached against the token id in session store.
+        cacheEntry = AuthorizationGrantCache.getInstance().getValueFromCacheByTokenId(null, tokenId);
+        return cacheEntry == null ? new HashMap<>() : cacheEntry.getUserAttributes();
+    }
+
     private Map<String, Object> retrieveClaimsForLocalUser(OAuthTokenReqMessageContext requestMsgCtx)
             throws IdentityOAuth2Exception {