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 token grant, we map the cache against the access token.
+     *
+     * @param key Authorization code or access token.
+     * @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();
+                }
             }
         }
 

components/org.wso2.carbon.identity.oauth/src/test/java/org/wso2/carbon/identity/oauth2/dao/AccessTokenDAOImplTest.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.org).
+ *
+ * WSO2 Inc. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.carbon.identity.oauth2.dao;
+
+import org.apache.commons.dbcp.BasicDataSource;
+import org.mockito.MockedStatic;
+import org.mockito.testng.MockitoTestNGListener;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Listeners;
+import org.testng.annotations.Test;
+import org.wso2.carbon.identity.common.testng.WithCarbonHome;
+import org.wso2.carbon.identity.core.util.IdentityDatabaseUtil;
+import org.wso2.carbon.identity.oauth2.IdentityOAuth2Exception;
+import org.wso2.carbon.identity.oauth2.dao.util.DAOUtils;
+
+import java.sql.Connection;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.mockito.Mockito.mockStatic;
+import static org.testng.Assert.assertEquals;
+
+@WithCarbonHome
+@Listeners(MockitoTestNGListener.class)
+public class AccessTokenDAOImplTest {
+
+    private AccessTokenDAOImpl accessTokenDAO;
+    private static Map<String, BasicDataSource> dataSourceMap = new HashMap<>();
+
+    public static final String H2_SCRIPT_NAME = "identity.sql";
+    public static final String H2_SCRIPT2_NAME = "insert_token_binding.sql";
+    public static final String DB_NAME = "AccessTokenDB";
+    Connection connection = null;
+    private MockedStatic<IdentityDatabaseUtil> identityDatabaseUtil;
+
+    @BeforeClass
+    public void initTest() throws Exception {
+
+        try {
+            DAOUtils.initializeBatchDataSource(DB_NAME, H2_SCRIPT_NAME, H2_SCRIPT2_NAME);
+        } catch (Exception e) {
+            throw new IdentityOAuth2Exception("Error while initializing the data source", e);
+        }
+        accessTokenDAO = new AccessTokenDAOImpl();
+    }
+
+    @BeforeMethod
+    public void setUp() throws Exception {
+
+        connection = DAOUtils.getConnection(DB_NAME);
+    }
+
+    @AfterMethod
+    public void closeup() throws Exception {
+
+        if (connection != null) {
+            connection.close();
+        }
+        identityDatabaseUtil.close();
+    }
+
+    @AfterClass
+    public void tearDown() throws Exception {
+
+        closeH2Base(DB_NAME);
+    }
+
+    @Test
+    public void getSessionIdentifierByTokenId() throws Exception {
+
+        connection = DAOUtils.getConnection(DB_NAME);
+        identityDatabaseUtil = mockStatic(IdentityDatabaseUtil.class);
+
+        identityDatabaseUtil.when(() -> IdentityDatabaseUtil.getDBConnection(false))
+                .thenReturn(connection);
+        assertEquals(accessTokenDAO.getSessionIdentifierByTokenId("2sa9a678f890877856y66e75f605d456"),
+                    "4503eb1561bfd6bf237b7e05c15afaff21f511d81135423015a747ee7e3f0bc0");
+    }
+
+    private static void closeH2Base(String databaseName) throws Exception {
+
+        BasicDataSource dataSource = dataSourceMap.get(databaseName);
+        if (dataSource != null) {
+            dataSource.close();
+        }
+    }
+}

components/org.wso2.carbon.identity.oauth/src/test/resources/dbScripts/insert_token_binding.sql

@@ -0,0 +1,22 @@
+INSERT INTO IDN_OAUTH_CONSUMER_APPS (CONSUMER_KEY, CONSUMER_SECRET, USERNAME, TENANT_ID, USER_DOMAIN, APP_NAME,
+            OAUTH_VERSION, CALLBACK_URL, GRANT_TYPES, APP_STATE) VALUES
+            ('ca19a540f544777860e44e75f605d924', '87n9a540f544777860e44e75f605d425', 'user1', 1234, 'PRIMARY',
+            'myApp', 'OAuth-2.0', 'http://localhost:8080/redirect',
+            'refresh_token implicit password iwa:ntlm client_credentials authorization_code', 'ACTIVE');
+
+INSERT INTO IDN_OAUTH2_ACCESS_TOKEN (TOKEN_ID, ACCESS_TOKEN, REFRESH_TOKEN, CONSUMER_KEY_ID, AUTHZ_USER, TENANT_ID,
+            USER_DOMAIN, USER_TYPE, GRANT_TYPE, TIME_CREATED, REFRESH_TOKEN_TIME_CREATED, VALIDITY_PERIOD,
+            REFRESH_TOKEN_VALIDITY_PERIOD, TOKEN_SCOPE_HASH, TOKEN_STATE, TOKEN_STATE_ID, SUBJECT_IDENTIFIER,
+            ACCESS_TOKEN_HASH, REFRESH_TOKEN_HASH, IDP_ID, AUTHORIZED_ORGANIZATION) VALUES
+            ('2sa9a678f890877856y66e75f605d456', 'd43e8da324a33bdc941b9b95cad6a6a2',
+            '2881c5a375d03dc0ba12787386451b29', 1, 'user1', 1234, 'PRIMARY', 'APPLICATION_USER', 'password', NOW(),
+            NOW(), 3600, 14400, '369db21a386ae433e65c0ff34d35708d', 'ACTIVE', 'NONE', 'user1', NULL, NULL, 1, 'NONE');
+
+INSERT INTO IDN_OAUTH2_ACCESS_TOKEN_SCOPE (TOKEN_ID, TOKEN_SCOPE, TENANT_ID) VALUES
+    ('2sa9a678f890877856y66e75f605d456', 'default', 1234);
+
+INSERT INTO IDP (TENANT_ID, NAME, UUID) VALUES (1234, 'LOCAL', 5678);
+
+INSERT INTO IDN_OAUTH2_TOKEN_BINDING (TOKEN_ID, TOKEN_BINDING_TYPE, TOKEN_BINDING_REF, TOKEN_BINDING_VALUE,
+            TENANT_ID) VALUES('2sa9a678f890877856y66e75f605d456', 'DEFAULT', 'bb9e4ed2cd4342ac0d6563e5fa6d2300',
+            '4503eb1561bfd6bf237b7e05c15afaff21f511d81135423015a747ee7e3f0bc0', '-1234');

components/org.wso2.carbon.identity.oauth/src/test/resources/testng.xml

@@ -94,6 +94,7 @@
             <class name="org.wso2.carbon.identity.oauth2.cache.JWKSCacheKeyTest"/>
             <class name="org.wso2.carbon.identity.oauth2.cache.JWKSCacheTest"/>
             <class name="org.wso2.carbon.identity.oauth2.dao.AccessContextTokenDOTest"/>
+            <class name="org.wso2.carbon.identity.oauth2.dao.AccessTokenDAOImplTest"/>
             <class name="org.wso2.carbon.identity.oauth2.dao.AuthContextTokenDOTest"/>
             <class name="org.wso2.carbon.identity.oauth2.dao.ScopeMgtDAOTest"/>
             <class name="org.wso2.carbon.identity.oauth2.dao.AuthorizationCodeDAOImplTest"/>