Fix isk claim removal from refreshed id token when the user attributes are updated

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/AccessTokenDAO.java

@@ -82,6 +82,18 @@ default void storeTokenToSessionMapping(String sessionIdentifier, String tokenId
 
     }
 
+    /**
+     * Get session identifier by token identifier.
+     *
+     * @param tokenId Token identifier.
+     * @return Session identifier.
+     * @throws IdentityOAuth2Exception
+     */
+    default String getSessionIdentifierByTokenId(String tokenId) throws IdentityOAuth2Exception {
+
+        return null;
+    }
+
     Set<AccessTokenDO> getAccessTokens(String consumerKey, AuthenticatedUser userName,
                                        String userStoreDomain, boolean includeExpired) throws IdentityOAuth2Exception;
 

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/AccessTokenDAOImpl.java

@@ -1265,6 +1265,37 @@ public Set<String> getTokenIdBySessionIdentifier(String sessionId) throws Identi
         return tokenIds;
     }
 
+    /**
+     * Get the session identifier of the given token ID.
+     *
+     * @param tokenId Token ID
+     * @throws IdentityOAuth2Exception
+     * @return
+     */
+    public String getSessionIdentifierByTokenId(String tokenId) throws IdentityOAuth2Exception {
+
+        String sql = SQLQueries.RETRIEVE_SESSION_ID_BY_TOKEN_ID;
+        Connection connection = IdentityDatabaseUtil.getDBConnection(false);
+        PreparedStatement prepStmt = null;
+        String sessionId = null;
+        try {
+            prepStmt = connection.prepareStatement(sql);
+            prepStmt.setString(1, tokenId);
+            try (ResultSet resultSet = prepStmt.executeQuery()) {
+                while (resultSet.next()) {
+                    sessionId = resultSet.getString("TOKEN_BINDING_VALUE");
+                }
+            }
+        } catch (SQLException e) {
+            String errorMsg = "Error occurred while retrieving 'session id' for " +
+                    "token id : " + tokenId;
+            throw new IdentityOAuth2Exception(errorMsg, e);
+        } finally {
+            IdentityDatabaseUtil.closeAllConnections(connection, null, prepStmt);
+        }
+        return sessionId;
+    }
+
     public void updateAccessTokenState(String tokenId, String tokenState) throws IdentityOAuth2Exception {
         updateAccessTokenState(tokenId, tokenState, null);
     }

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/SQLQueries.java

@@ -1625,6 +1625,9 @@ public class SQLQueries {
     public static final String RETRIEVE_TOKENS_MAPPED_FOR_TOKEN_BINDING_VALUE = "SELECT TOKEN_ID FROM " +
             "IDN_OAUTH2_TOKEN_BINDING WHERE TOKEN_BINDING_VALUE = ? AND TOKEN_BINDING_TYPE = ?";
 
+    public static final String RETRIEVE_SESSION_ID_BY_TOKEN_ID = "SELECT TOKEN_BINDING_VALUE FROM " +
+            "IDN_OAUTH2_TOKEN_BINDING WHERE  TOKEN_BINDING_TYPE = 'DEFAULT' AND TOKEN_ID = ?";
+
     public static final String RETRIEVE_TOKEN_BINDING_REF_EXISTS = "SELECT COUNT(*) AS TOTAL FROM "
             + "IDN_OAUTH2_TOKEN_BINDING WHERE TOKEN_BINDING_REF = ?";
 

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

@@ -96,6 +96,7 @@
 public class RefreshGrantHandler extends AbstractAuthorizationGrantHandler {
 
     public static final String PREV_ACCESS_TOKEN = "previousAccessToken";
+    public static final String SESSION_IDENTIFIER = "sessionIdentifier";
     public static final int LAST_ACCESS_TOKEN_RETRIEVAL_LIMIT = 10;
     public static final int ALLOWED_MINIMUM_VALIDITY_PERIOD = 1000;
     public static final String DEACTIVATED_ACCESS_TOKEN = "DeactivatedAccessToken";
@@ -251,6 +252,47 @@ private void setPropertiesForTokenGeneration(OAuthTokenReqMessageContext tokReqM
         // Store the old access token as a OAuthTokenReqMessageContext property, this is already
         // a preprocessed token.
         tokReqMsgCtx.addProperty(PREV_ACCESS_TOKEN, validationBean);
+
+        /*
+        Add the session id from the last access token to OAuthTokenReqMessageContext. First check whether the
+        session Id can be resolved from the authorization grant cache. If not resolve the session id from the token
+        id session id mapping in the token binding table. Here we are assigning the session id of the refreshed
+        token as same as the previously issued access token.
+        */
+        String sessionId = getSessionContextIdentifier(validationBean.getAccessToken());
+        if (sessionId == null) {
+            String oldTokenId = validationBean.getTokenId();
+            sessionId = OAuthTokenPersistenceFactory.getInstance().getAccessTokenDAO()
+                    .getSessionIdentifierByTokenId(oldTokenId);
+        }
+        if (sessionId != null) {
+            tokReqMsgCtx.addProperty(SESSION_IDENTIFIER, sessionId);
+        }
+    }
+
+    /**
+     * Return session context identifier from authorization grant cache. For authorization code flow, we mapped it
+     * against auth_code. For refresh key grant, we map the cache against the accesstoken.
+     *
+     * @param key Authorization code or accesstoken.
+     * @return SessionContextIdentifier.
+     */
+    private static String getSessionContextIdentifier(String key) {
+
+        String sessionContextIdentifier = null;
+        if (StringUtils.isNotBlank(key)) {
+            AuthorizationGrantCacheKey cacheKey = new AuthorizationGrantCacheKey(key);
+            AuthorizationGrantCacheEntry cacheEntry =
+                    AuthorizationGrantCache.getInstance().getValueFromCacheByToken(cacheKey);
+            if (cacheEntry != null) {
+                sessionContextIdentifier = cacheEntry.getSessionContextIdentifier();
+                if (log.isDebugEnabled()) {
+                    log.debug(String.format("Found session context identifier: %s for the obtained authorization code",
+                            sessionContextIdentifier));
+                }
+            }
+        }
+        return sessionContextIdentifier;
     }
 
     private boolean validateRefreshTokenInRequest(OAuth2AccessTokenReqDTO tokenReq,

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/openidconnect/DefaultIDTokenBuilder.java

@@ -67,6 +67,7 @@
 import static org.wso2.carbon.identity.oauth.common.OAuthConstants.OIDCClaims.AZP;
 import static org.wso2.carbon.identity.oauth.common.OAuthConstants.OIDCClaims.IDP_SESSION_KEY;
 import static org.wso2.carbon.identity.oauth.common.OAuthConstants.OIDCClaims.NONCE;
+import static org.wso2.carbon.identity.oauth2.token.handlers.grant.RefreshGrantHandler.SESSION_IDENTIFIER;
 
 /**
  * Default IDToken generator for the OpenID Connect Implementation.
@@ -173,6 +174,9 @@ public String buildIDToken(OAuthTokenReqMessageContext tokenReqMsgCtxt,
             if (!OAuthConstants.GrantTypes.PASSWORD.equalsIgnoreCase(
                     tokenReqMsgCtxt.getOauth2AccessTokenReqDTO().getGrantType())) {
                 idpSessionKey = getIdpSessionKey(accessToken);
+                if (idpSessionKey == null && tokenReqMsgCtxt.getProperty(SESSION_IDENTIFIER) != null) {
+                    idpSessionKey = tokenReqMsgCtxt.getProperty(SESSION_IDENTIFIER).toString();
+                }
             }
         }