Fix missing claims for federated users for jwt token issuers

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth/cache/AuthorizationGrantCache.java

@@ -18,6 +18,8 @@
 
 package org.wso2.carbon.identity.oauth.cache;
 
+import com.nimbusds.jwt.JWT;
+import com.nimbusds.jwt.JWTParser;
 import org.apache.commons.codec.digest.DigestUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.logging.Log;
@@ -29,10 +31,10 @@
 import org.wso2.carbon.identity.oauth.config.OAuthServerConfiguration;
 import org.wso2.carbon.identity.oauth2.IdentityOAuth2Exception;
 import org.wso2.carbon.identity.oauth2.dao.OAuthTokenPersistenceFactory;
-import org.wso2.carbon.identity.oauth2.model.AccessTokenDO;
 import org.wso2.carbon.identity.oauth2.util.OAuth2Util;
 import org.wso2.carbon.utils.CarbonUtils;
 
+import java.text.ParseException;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -237,11 +239,20 @@ private String replaceFromCodeId(String authzCode) {
      * @return TOKEN_ID from the database
      */
     private String replaceFromTokenId(String keyValue) {
-        try {
-            AccessTokenDO accessTokenDO = OAuth2Util.findAccessToken(keyValue, true);
-            if (accessTokenDO != null) {
-                return accessTokenDO.getTokenId();
+
+        if (OAuth2Util.isJWT(keyValue)) {
+            try {
+                JWT parsedJwtToken = JWTParser.parse(keyValue);
+                keyValue = parsedJwtToken.getJWTClaimsSet().getJWTID();
+            } catch (ParseException e) {
+                if (log.isDebugEnabled() && IdentityUtil.isTokenLoggable(
+                        IdentityConstants.IdentityTokens.ACCESS_TOKEN)) {
+                    log.debug("Error while getting JWTID from token: " + keyValue, e);
+                }
             }
+        }
+        try {
+            return OAuthTokenPersistenceFactory.getInstance().getAccessTokenDAO().getTokenIdByAccessToken(keyValue);
         } catch (IdentityOAuth2Exception e) {
             log.error("Failed to retrieve token id by token from store for - ." + keyValue, e);
         }

components/org.wso2.carbon.identity.oauth/src/test/java/org/wso2/carbon/identity/oauth/cache/AuthorizationGrantCacheTest.java

@@ -0,0 +1,112 @@
+package org.wso2.carbon.identity.oauth.cache;
+
+import com.nimbusds.jwt.JWT;
+import com.nimbusds.jwt.JWTClaimsSet;
+import com.nimbusds.jwt.JWTParser;
+import org.mockito.Mock;
+import org.mockito.MockedStatic;
+import org.mockito.MockitoAnnotations;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+import org.wso2.carbon.identity.application.authentication.framework.store.SessionDataStore;
+import org.wso2.carbon.identity.oauth2.IdentityOAuth2Exception;
+import org.wso2.carbon.identity.oauth2.dao.AccessTokenDAO;
+import org.wso2.carbon.identity.oauth2.dao.OAuthTokenPersistenceFactory;
+
+import java.text.ParseException;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class AuthorizationGrantCacheTest {
+
+    @Mock
+    private AccessTokenDAO accessTokenDAO;
+
+    private AuthorizationGrantCache cache;
+
+    @Mock
+    private OAuthTokenPersistenceFactory mockedOAuthTokenPersistenceFactory;
+
+    @Mock
+    private SessionDataStore sessionDataStore;
+
+    private static final String AUTHORIZATION_GRANT_CACHE_NAME = "AuthorizationGrantCache";
+
+    @BeforeMethod
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        cache = AuthorizationGrantCache.getInstance();
+    }
+
+    @Test(dataProvider = "replaceFromTokenIdDataProvider")
+    public void testReplaceFromTokenId(String accessToken, String jwtId, String tokenId, boolean isJwtToken,
+                                       boolean isInvalidJWTToken, boolean isFailedTokenRetrieval) throws Exception {
+
+        try (MockedStatic<OAuthTokenPersistenceFactory> mockedFactory = mockStatic(OAuthTokenPersistenceFactory.class);
+             MockedStatic<JWTParser> mockedJwtParser = mockStatic(JWTParser.class);
+             MockedStatic<SessionDataStore> mockedSessionDataStore = mockStatic(SessionDataStore.class)) {
+
+            mockedFactory.when(OAuthTokenPersistenceFactory::getInstance).thenReturn(
+                    mockedOAuthTokenPersistenceFactory);
+
+            when(mockedOAuthTokenPersistenceFactory.getAccessTokenDAO()).thenReturn(accessTokenDAO);
+
+            if (isJwtToken) {
+                JWT jwtMock = mock(JWT.class);
+                JWTClaimsSet claimsSetMock = mock(JWTClaimsSet.class);
+
+                if (isInvalidJWTToken) {
+                    when(JWTParser.parse(accessToken)).thenThrow(new ParseException("Invalid JWT", 0));
+                } else {
+                    mockedJwtParser.when(() -> JWTParser.parse(accessToken)).thenReturn(jwtMock);
+                    when(jwtMock.getJWTClaimsSet()).thenReturn(claimsSetMock);
+                    when(claimsSetMock.getJWTID()).thenReturn(jwtId);
+                }
+            }
+
+            if (isFailedTokenRetrieval) {
+                when(accessTokenDAO.getTokenIdByAccessToken(jwtId)).thenThrow(
+                        new IdentityOAuth2Exception("Failed to retrieve token id by token from store"));
+            } else {
+                when(accessTokenDAO.getTokenIdByAccessToken(jwtId != null ? jwtId : accessToken)).thenReturn(tokenId);
+            }
+
+            // Mock SessionDataStore static instance and return a mock session data store
+            mockedSessionDataStore.when(SessionDataStore::getInstance).thenReturn(sessionDataStore);
+
+            // Mock SessionDataStore return for session data (from getFromSessionStore)
+            AuthorizationGrantCacheEntry mockCacheEntry = new AuthorizationGrantCacheEntry();
+            mockCacheEntry.setTokenId(tokenId);
+
+            when(sessionDataStore.getSessionData(tokenId, AUTHORIZATION_GRANT_CACHE_NAME)).thenReturn(mockCacheEntry);
+
+            AuthorizationGrantCacheKey key = new AuthorizationGrantCacheKey(accessToken);
+            AuthorizationGrantCacheEntry result = cache.getValueFromCacheByToken(key);
+
+            // Verify the token ID returned from the DAO is as expected
+            assertEquals(tokenId, result.getTokenId());
+
+            // Verify that the JWT token was parsed and the correct claim was retrieved if it was a JWT
+            if (isJwtToken && !isInvalidJWTToken) {
+                verify(accessTokenDAO).getTokenIdByAccessToken(jwtId);
+            } else {
+                verify(accessTokenDAO).getTokenIdByAccessToken(accessToken);
+            }
+        }
+    }
+
+    @DataProvider(name = "replaceFromTokenIdDataProvider")
+    public Object[][] getReplaceFromTokenIdData() {
+        return new Object[][]{
+                {"jwt.Access.Token", "jwtId", "jwtTokenId", true, false, false},
+                {"nonJWTAccessToken", null, "nonJWTTokenId", false, false, false},
+                {"invalid.JWT.Token", null, "invalid.JWT.Token", true, true, false},
+                {"fail.Store.TokenId", "jwtId", "jwtId", true, false, true}
+        };
+    }
+}

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

@@ -123,6 +123,7 @@
             <class name="org.wso2.carbon.identity.oauth2.device.codegenerator.GenerateKeysTest"/>
             <class name="org.wso2.carbon.identity.oauth2.responsemode.provider.ResponseModeProviderTest"/>
             <class name="org.wso2.carbon.identity.oauth2.token.handlers.claims.ImpersonatedAccessTokenClaimProviderTest"/>
+            <class name="org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCacheTest"/>
         </classes>
     </test>
 
@@ -197,6 +198,7 @@
             <class name="org.wso2.carbon.identity.openidconnect.JWTAccessTokenOIDCClaimsHandler"/>
             <class name="org.wso2.carbon.identity.openidconnect.OpenIDConnectSystemClaimImplTest"/>
             <class name="org.wso2.carbon.identity.openidconnect.OpenIDConnectClaimFilterImplTest"/>
+            <class name="org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCacheTest"/>
         </classes>
     </test>
 </suite>