Improve request object clean up logic to avoid unnessary delete requests

wso2-extensions · Nov 11, 2024 · edbf94e · edbf94e
1 parent bc0033a
commit edbf94e
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 18 deletions.
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
@@ -689,7 +689,7 @@ public static boolean revokeAuthzCodes(String username, UserStoreManager userSto
                                 authorizationCode.getAuthorizationCode())));
                 OAuthTokenPersistenceFactory.getInstance().getAuthorizationCodeDAO()
                         .updateAuthorizationCodeState(authorizationCode.getAuthorizationCode(),
-                                OAuthConstants.AuthorizationCodeState.REVOKED);
+                                authorizationCode.getAuthzCodeId(), OAuthConstants.AuthorizationCodeState.REVOKED);
             }
         } catch (IdentityOAuth2Exception e) {
             String errorMsg = "Error occurred while revoking authorization codes for user: " + username;

diff --git a/....identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/AccessTokenDAOImpl.java b/....identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/AccessTokenDAOImpl.java
@@ -1494,12 +1494,16 @@ public void revokeAccessTokensInBatch(String[] tokens, boolean isHashedToken) th
                 }
                 ps.executeUpdate();
 
-                // To revoke request objects which have persisted against the access token.
-                OAuth2TokenUtil.postUpdateAccessTokens(Arrays.asList(tokens), OAuthConstants.TokenStates.
-                        TOKEN_STATE_REVOKED);
+
                 if (isTokenCleanupFeatureEnabled) {
                     oldTokenCleanupObject.cleanupTokenByTokenValue(
                             getHashingPersistenceProcessor().getProcessedAccessTokenIdentifier(tokens[0]), connection);
+                    /* When token is deleted, the request objects get on delete cascade except for the SQL server.
+                        Hence, invoke the event listener to revoke the request objects.*/
+                    if (connection.getMetaData().getDriverName().contains("Microsoft")) {
+                        OAuth2TokenUtil.postUpdateAccessTokens(Arrays.asList(tokens), OAuthConstants.TokenStates.
+                                TOKEN_STATE_REVOKED);
+                    }
                 }
             } catch (SQLException e) {
                 // IdentityDatabaseUtil.rollbackTransaction(connection);

diff --git a/...dentity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/AuthorizationCodeDAO.java b/...dentity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/AuthorizationCodeDAO.java
@@ -66,6 +66,8 @@ void insertAuthorizationCode(String authzCode, String consumerKey, String appTen
     AuthorizationCodeValidationResult validateAuthorizationCode(String consumerKey, String authorizationKey)
             throws IdentityOAuth2Exception;
 
+    void updateAuthorizationCodeState(String authzCode, String codeId, String newState) throws IdentityOAuth2Exception;
+
     void updateAuthorizationCodeState(String authzCode, String newState) throws IdentityOAuth2Exception;
 
     void deactivateAuthorizationCode(AuthzCodeDO authzCodeDO) throws

diff --git a/...ity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/AuthorizationCodeDAOImpl.java b/...ity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/AuthorizationCodeDAOImpl.java
@@ -312,7 +312,8 @@ private String getTokenBindingReference(Connection connection, String tokenId, i
     }
 
     @Override
-    public void updateAuthorizationCodeState(String authzCode, String newState) throws IdentityOAuth2Exception {
+    public void updateAuthorizationCodeState(String authzCode, String codeId, String newState)
+            throws IdentityOAuth2Exception {
 
         if (log.isDebugEnabled()) {
             if (IdentityUtil.isTokenLoggable(IdentityConstants.IdentityTokens.AUTHORIZATION_CODE)) {
@@ -322,29 +323,31 @@ public void updateAuthorizationCodeState(String authzCode, String newState) thro
                 log.debug("Changing state of authorization code  to: " + newState);
             }
         }
-        boolean tokenUpdateSuccessful;
-        String authCodeStoreTable = OAuthConstants.AUTHORIZATION_CODE_STORE_TABLE;
-        Connection connection = IdentityDatabaseUtil.getDBConnection();
+        Connection connection = IdentityDatabaseUtil.getDBConnection(true);
         PreparedStatement prepStmt = null;
         try {
             prepStmt = connection.prepareStatement(SQLQueries.UPDATE_AUTHORIZATION_CODE_STATE);
             prepStmt.setString(1, newState);
             prepStmt.setString(2, getHashingPersistenceProcessor().getProcessedAuthzCode(authzCode));
             prepStmt.execute();
             IdentityDatabaseUtil.commitTransaction(connection);
-            tokenUpdateSuccessful = true;
         } catch (SQLException e) {
             IdentityDatabaseUtil.rollbackTransaction(connection);
             throw new IdentityOAuth2Exception("Error occurred while updating the state of Authorization Code : " +
-                    authzCode.toString(), e);
+                    authzCode, e);
         } finally {
             IdentityDatabaseUtil.closeAllConnections(connection, null, prepStmt);
         }
-        if (tokenUpdateSuccessful) {
-            //If the code state is updated to inactive or expired request object which is persisted against the code
-            // should be updated/removed.
-            OAuth2TokenUtil.postRevokeCode(authzCode, newState, null, null);
-        }
+        //If the code state is updated to inactive or expired request object which is persisted against the code
+        // should be updated/removed.
+        OAuth2TokenUtil.postRevokeCode(codeId, newState, null, authzCode);
+    }
+
+
+    @Override
+    public void updateAuthorizationCodeState(String authzCode, String newState) throws IdentityOAuth2Exception {
+
+        updateAuthorizationCodeState(authzCode, null, newState);
     }
 
     @Override

diff --git a/...a/org/wso2/carbon/identity/oauth2/token/handlers/grant/AuthorizationCodeGrantHandler.java b/...a/org/wso2/carbon/identity/oauth2/token/handlers/grant/AuthorizationCodeGrantHandler.java
@@ -550,7 +550,7 @@ private boolean isAuthzCodeExpired(AuthzCodeDO authzCodeBean)
     private void markAsExpired(AuthzCodeDO authzCodeBean) throws IdentityOAuth2Exception {
 
         OAuthTokenPersistenceFactory.getInstance().getAuthorizationCodeDAO()
-                .updateAuthorizationCodeState(authzCodeBean.getAuthorizationCode(),
+                .updateAuthorizationCodeState(authzCodeBean.getAuthorizationCode(), authzCodeBean.getAuthzCodeId(),
                         OAuthConstants.AuthorizationCodeState.EXPIRED);
         if (log.isDebugEnabled()) {
             log.debug("Changed state of authorization code : " + authzCodeBean.getAuthorizationCode() + " to expired");
@@ -594,8 +594,10 @@ private boolean validatePKCECode(AuthzCodeDO authzCodeBean, String verificationC
     }
 
     private void revokeAuthorizationCode(AuthzCodeDO authzCodeBean) throws IdentityOAuth2Exception {
+
         OAuthTokenPersistenceFactory.getInstance().getAuthorizationCodeDAO().updateAuthorizationCodeState(
-                authzCodeBean.getAuthorizationCode(), OAuthConstants.AuthorizationCodeState.REVOKED);
+                authzCodeBean.getAuthorizationCode(), authzCodeBean.getAuthzCodeId(),
+                OAuthConstants.AuthorizationCodeState.REVOKED);
         if (log.isDebugEnabled()) {
             log.debug("Changed state of authorization code : " + authzCodeBean.getAuthorizationCode() + " to revoked");
         }

diff --git a/...oauth/src/test/java/org/wso2/carbon/identity/oauth2/dao/AuthorizationCodeDAOImplTest.java b/...oauth/src/test/java/org/wso2/carbon/identity/oauth2/dao/AuthorizationCodeDAOImplTest.java
@@ -282,7 +282,7 @@ public void testGetActiveAuthorizationCodesByConsumerKey() throws Exception {
                     .thenAnswer(
                             (Answer<Void>) invocation -> null);
             authorizationCodeDAO.updateAuthorizationCodeState(authzCodeDO1.getAuthorizationCode(),
-                    OAuthConstants.AuthorizationCodeState.REVOKED);
+                    authzCodeDO1.getAuthzCodeId(), OAuthConstants.AuthorizationCodeState.REVOKED);
             Set<String> availableAuthzCodes = new HashSet<>();
             availableAuthzCodes.add(authzCode2);