Merge pull request #2673 from RusJaI/6.13.x

Fix JWT token issuing performance issue when renew_token_without_revoking_existing is enabled

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/handlers/grant/AbstractAuthorizationGrantHandler.java

@@ -70,8 +70,10 @@
 import java.util.Set;
 import java.util.UUID;
 
+import static org.wso2.carbon.identity.oauth.common.OAuthConstants.RENEW_TOKEN_WITHOUT_REVOKING_EXISTING_ENABLE_CONFIG;
 import static org.wso2.carbon.identity.oauth.common.OAuthConstants.TokenBindings.NONE;
 import static org.wso2.carbon.identity.oauth.common.OAuthConstants.TokenStates.TOKEN_STATE_ACTIVE;
+import static org.wso2.carbon.identity.oauth2.util.OAuth2Util.JWT;
 
 /**
  * Abstract authorization grant handler.
@@ -86,6 +88,9 @@ public abstract class AbstractAuthorizationGrantHandler implements Authorization
     protected static final String EXISTING_TOKEN_ISSUED = "existingTokenUsed";
     protected static final int SECONDS_TO_MILISECONDS_FACTOR = 1000;
     private boolean isHashDisabled = OAuth2Util.isHashDisabled();
+    private static final boolean renewWithoutRevokingExistingEnabled = Boolean.parseBoolean(IdentityUtil.
+            getProperty(RENEW_TOKEN_WITHOUT_REVOKING_EXISTING_ENABLE_CONFIG));
+    public static final String OAUTH_APP = "OAuthAppDO";
 
     @Override
     public void init() throws IdentityOAuth2Exception {
@@ -148,6 +153,29 @@ public OAuth2AccessTokenRespDTO issue(OAuthTokenReqMessageContext tokReqMsgCtx)
 
         synchronized ((consumerKey + ":" + authorizedUserId + ":" + scope + ":" + tokenBindingReference).intern()) {
             AccessTokenDO existingTokenBean = null;
+
+            OAuthAppDO oAuthAppDO = (OAuthAppDO) tokReqMsgCtx.getProperty(OAUTH_APP);
+            String tokenType = (oAuthAppDO != null) ? oAuthAppDO.getTokenType() : null;
+
+            /*
+            Check if the token type is JWT and renew without revoking existing tokens is enabled.
+            Additionally, ensure that the grant type used for the token request is allowed to renew without revoke,
+            based on the config.
+            */
+            if (JWT.equals(tokenType) && renewWithoutRevokingExistingEnabled &&
+                    OAuth2ServiceComponentHolder.getJwtRenewWithoutRevokeAllowedGrantTypes()
+                            .contains(tokReqMsgCtx.getOauth2AccessTokenReqDTO().getGrantType())) {
+                /*
+                If the application does not have a token binding type (i.e., no specific binding type is set),
+                binding reference will be randomly generated UUID, in that case we can generate a new access token
+                without looking up the existing tokens in the token table.
+                */
+                if (oAuthAppDO.getTokenBindingType() == null) {
+                    return generateNewAccessToken(tokReqMsgCtx, scope, consumerKey, existingTokenBean,
+                            false, oauthTokenIssuer);
+                }
+            }
+
             if (isHashDisabled) {
                 existingTokenBean = getExistingToken(tokReqMsgCtx,
                         getOAuthCacheKey(scope, consumerKey, authorizedUserId, authenticatedIDP,
@@ -1056,18 +1084,65 @@ private boolean hasValidationByApplicationScopeValidatorsFailed(OAuthTokenReqMes
      * @param tokReqMsgCtx OAuthTokenReqMessageContext.
      * @return token binding reference.
      */
+
     private String getTokenBindingReference(OAuthTokenReqMessageContext tokReqMsgCtx) {
 
-        if (tokReqMsgCtx.getTokenBinding() == null) {
-            if (log.isDebugEnabled()) {
-                log.debug("Token binding data is null.");
+        /**
+         * If OAuth.JWT.RenewTokenWithoutRevokingExisting is enabled from configurations, and current token
+         * binding is null,then we will add a new token binding (request binding) to the token binding with
+         * a value of a random UUID.
+         * The purpose of this new token binding type is to add a random value to the token binding so that
+         * "User, Application, Scope, Binding" combination will be unique for each token.
+         * Previously, if a token issue request come for the same combination of "User, Application, Scope, Binding",
+         * the existing JWT token will be revoked and issue a new token. but with this way, we can issue new tokens
+         * without revoking the old ones.
+         *
+         * Add following configuration to deployment.toml file to enable this feature.
+         *     [oauth.jwt.renew_token_without_revoking_existing]
+         *     enable = true
+         *
+         * By default, the allowed grant type for this feature is "client_credentials". If you need to enable for
+         * other grant types, add the following configuration to deployment.toml file.
+         *     [oauth.jwt.renew_token_without_revoking_existing]
+         *     enable = true
+         *     allowed_grant_types = ["client_credentials","password", ...]
+         */
+
+        OAuthAppDO oAuthAppDO = (OAuthAppDO) tokReqMsgCtx.getProperty(OAUTH_APP);
+        String tokenType = (oAuthAppDO != null) ? oAuthAppDO.getTokenType() : null;
+
+        /*
+            Check if the token type is JWT and renew without revoking existing tokens is enabled.
+            Additionally, ensure that the grant type used for the token request is allowed to renew without revoke,
+            based on the config.
+        */
+        if (JWT.equalsIgnoreCase(tokenType)) {
+
+            if (renewWithoutRevokingExistingEnabled && tokReqMsgCtx != null && (tokReqMsgCtx.getTokenBinding() == null
+                    || StringUtils.isBlank(tokReqMsgCtx.getTokenBinding().getBindingReference()))) {
+                if (OAuth2ServiceComponentHolder.getJwtRenewWithoutRevokeAllowedGrantTypes()
+                        .contains(tokReqMsgCtx.getOauth2AccessTokenReqDTO().getGrantType())) {
+                    return UUID.randomUUID().toString();
+                }
+                return NONE;
             }
-            return NONE;
         }
-        if (StringUtils.isBlank(tokReqMsgCtx.getTokenBinding().getBindingReference())) {
+        return getExistingTokenBindingReference(tokReqMsgCtx);
+    }
+
+    /**
+     * Retrieves the existing token binding reference if available, otherwise returns NONE.
+     *
+     * @param tokReqMsgCtx OAuthTokenReqMessageContext.
+     * @return token binding reference.
+     */
+    private String getExistingTokenBindingReference(OAuthTokenReqMessageContext tokReqMsgCtx) {
+
+        if (tokReqMsgCtx == null || tokReqMsgCtx.getTokenBinding() == null) {
             return NONE;
         }
-        return tokReqMsgCtx.getTokenBinding().getBindingReference();
+        String bindingReference = tokReqMsgCtx.getTokenBinding().getBindingReference();
+        return StringUtils.isBlank(bindingReference) ? NONE : bindingReference;
     }
 
     /**