Fix access token attributes for federated user

components/org.wso2.carbon.identity.oauth.endpoint/src/main/java/org/wso2/carbon/identity/oauth/endpoint/authz/OAuth2AuthzEndpoint.java

@@ -421,6 +421,36 @@ private void addFederatedTokensToSessionCache(OAuthMessage oAuthMessage,
         }
     }
 
+    /**
+     * Add mapped remote claims to session cache.
+     *
+     * @param oAuthMessage         The OAuthMessage with the session data cache entry.
+     * @param authenticationResult The authentication result of authorization call.
+     */
+    private void addMappedRemoteClaimsToSessionCache(OAuthMessage oAuthMessage,
+                                                  AuthenticationResult authenticationResult) {
+
+        Optional<Map<String, String>> mappedRemoteClaims = authenticationResult.getMappedRemoteClaims();
+        if (!mappedRemoteClaims.isPresent()) {
+            return;
+        }
+
+        SessionDataCacheEntry sessionDataCacheEntry = oAuthMessage.getSessionDataCacheEntry();
+        if (sessionDataCacheEntry == null || mappedRemoteClaims.get().isEmpty()) {
+            return;
+        }
+        Map<ClaimMapping, String> mappedRemoteClaimsMap = new HashMap<>();
+        mappedRemoteClaims.get().forEach(
+                (key, value) -> mappedRemoteClaimsMap.put(ClaimMapping.build(key, key, null,
+                        false), value));
+        sessionDataCacheEntry.setMappedRemoteClaims(mappedRemoteClaimsMap);
+        if (log.isDebugEnabled() && authenticationResult.getSubject() != null) {
+            log.debug("Added the mapped remote claims to the session data cache. " +
+                    "Session context identifier: " + sessionDataCacheEntry.getSessionContextIdentifier()
+                    + " for the user: " + authenticationResult.getSubject().getLoggableMaskedUserId());
+        }
+    }
+
     /**
      * This method creates a list of FederatedTokenDO objects from the list of FederatedToken objects.
      *
@@ -1389,6 +1419,9 @@ private void addToAuthenticationResultDetailsToOAuthMessage(OAuthMessage oAuthMe
                 authnResult.getProperty(FrameworkConstants.AnalyticsAttributes.SESSION_ID));
         // Adding federated tokens come with the authentication result of the authorization call.
         addFederatedTokensToSessionCache(oAuthMessage, authnResult);
+        // Adding mapped remoted claims come with the authentication result to resolve access token claims in
+        // federated flow.
+        addMappedRemoteClaimsToSessionCache(oAuthMessage, authnResult);
     }
 
     private void updateAuthTimeInSessionDataCacheEntry(OAuthMessage oAuthMessage) {
@@ -2143,6 +2176,10 @@ private void addUserAttributesToOAuthMessage(OAuthMessage oAuthMessage, String c
         authorizationGrantCacheEntry.setRequestObjectFlow(isRequestObjectFlow);
         authorizationGrantCacheEntry.setFederatedTokens(sessionDataCacheEntry.getFederatedTokens());
         sessionDataCacheEntry.setFederatedTokens(null);
+        Map<ClaimMapping, String> mappedRemoteClaims =  sessionDataCacheEntry.getMappedRemoteClaims();
+        if (mappedRemoteClaims != null) {
+            authorizationGrantCacheEntry.setMappedRemoteClaims(mappedRemoteClaims);
+        }
         oAuthMessage.setAuthorizationGrantCacheEntry(authorizationGrantCacheEntry);
     }
 
@@ -3785,6 +3822,7 @@ private OAuth2AuthorizeReqDTO buildAuthRequest(OAuth2Parameters oauth2Params, Se
         authzReqDTO.setState(oauth2Params.getState());
         authzReqDTO.setHttpServletRequestWrapper(new HttpServletRequestWrapper(request));
         authzReqDTO.setRequestedSubjectId(oauth2Params.getRequestedSubjectId());
+        authzReqDTO.setMappedRemoteClaims(sessionDataCacheEntry.getMappedRemoteClaims());
 
         if (sessionDataCacheEntry.getParamMap() != null && sessionDataCacheEntry.getParamMap().get(OAuthConstants
                 .AMR) != null) {
@@ -4520,6 +4558,10 @@ private void addUserAttributesToCache(SessionDataCacheEntry sessionDataCacheEntr
         DeviceAuthorizationGrantCacheKey cacheKey = new DeviceAuthorizationGrantCacheKey(deviceCode);
         DeviceAuthorizationGrantCacheEntry cacheEntry =
                 new DeviceAuthorizationGrantCacheEntry(sessionDataCacheEntry.getLoggedInUser().getUserAttributes());
+        if (sessionDataCacheEntry.getMappedRemoteClaims() != null) {
+            cacheEntry.setMappedRemoteClaims(sessionDataCacheEntry
+                    .getMappedRemoteClaims());
+        }
         DeviceAuthorizationGrantCache.getInstance().addToCache(cacheKey, cacheEntry);
     }
 

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

@@ -66,6 +66,8 @@ public class AuthorizationGrantCacheEntry extends CacheEntry {
 
     private boolean hasNonOIDCClaims;
 
+    private Map<ClaimMapping, String> mappedRemoteClaims;
+
     /*
         OIDC sub claim. This should be formatted based on the Service Provider configurations to append
         userStoreDomain and tenantDomain.
@@ -390,4 +392,15 @@ public void setPreIssueAccessTokenActionsExecuted(boolean preIssueAccessTokenAct
 
         isPreIssueAccessTokenActionsExecuted = preIssueAccessTokenActionsExecuted;
     }
+
+    public Map<ClaimMapping, String> getMappedRemoteClaims() {
+
+        return mappedRemoteClaims;
+    }
+
+    public void setMappedRemoteClaims(
+            Map<ClaimMapping, String> mappedRemoteClaims) {
+
+        this.mappedRemoteClaims = mappedRemoteClaims;
+    }
 }

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

@@ -19,6 +19,7 @@
 package org.wso2.carbon.identity.oauth.cache;
 
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatedUser;
+import org.wso2.carbon.identity.application.common.model.ClaimMapping;
 import org.wso2.carbon.identity.oauth2.authz.OAuthAuthzReqMessageContext;
 import org.wso2.carbon.identity.oauth2.model.FederatedTokenDO;
 import org.wso2.carbon.identity.oauth2.model.OAuth2Parameters;
@@ -53,6 +54,7 @@ public class SessionDataCacheEntry extends CacheEntry {
 
     private Map<String, Serializable> endpointParams = new HashMap<>();
     private List<FederatedTokenDO> federatedTokens;
+    private Map<ClaimMapping, String> mappedRemoteClaims;
 
     public OAuthAuthzReqMessageContext getAuthzReqMsgCtx() {
         return authzReqMsgCtx;
@@ -172,4 +174,14 @@ public void setFederatedTokens(List<FederatedTokenDO> federatedTokens) {
 
         this.federatedTokens = federatedTokens;
     }
+
+    public Map<ClaimMapping, String> getMappedRemoteClaims() {
+
+        return mappedRemoteClaims;
+    }
+
+    public void setMappedRemoteClaims(Map<ClaimMapping, String> mappedRemoteClaims) {
+
+        this.mappedRemoteClaims = mappedRemoteClaims;
+    }
 }

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/authz/handlers/TokenResponseTypeHandler.java

@@ -554,6 +554,11 @@ private void addUserAttributesToCache(String accessToken,
             authorizationGrantCacheEntry.setMaxAge(authorizeReqDTO.getMaxAge());
         }
 
+        if (authorizeReqDTO.getMappedRemoteClaims() != null) {
+            authorizationGrantCacheEntry.setMappedRemoteClaims(
+                    authorizeReqDTO.getMappedRemoteClaims());
+        }
+
         ClaimMapping key = new ClaimMapping();
         Claim claimOfKey = new Claim();
         claimOfKey.setClaimUri(OAuth2Util.SUB);

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/authz/handlers/util/ResponseTypeHandlerUtil.java

@@ -487,6 +487,11 @@ private static void addUserAttributesToCache(String accessToken, OAuthAuthzReqMe
             userAttributes.put(key, sub);
         }
 
+        if (authorizeReqDTO.getMappedRemoteClaims() != null) {
+            authorizationGrantCacheEntry.setMappedRemoteClaims(
+                    authorizeReqDTO.getMappedRemoteClaims());
+        }
+
         authorizationGrantCacheEntry
                 .setValidityPeriod(TimeUnit.MILLISECONDS.toNanos(accessTokenDO.getValidityPeriodInMillis()));
         AuthorizationGrantCache.getInstance().addToCacheByToken(authorizationGrantCacheKey,

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/device/cache/DeviceAuthorizationGrantCacheEntry.java

@@ -31,6 +31,7 @@ public class DeviceAuthorizationGrantCacheEntry extends CacheEntry {
     private static final long serialVersionUID = -3043225645166013281L;
 
     private Map<ClaimMapping, String> userAttributes;
+    private Map<ClaimMapping, String> mappedRemoteClaims;
 
     public DeviceAuthorizationGrantCacheEntry(Map<ClaimMapping, String> userAttributes) {
 
@@ -56,4 +57,15 @@ public void setUserAttributes(Map<ClaimMapping, String> userAttributes) {
 
         this.userAttributes = userAttributes;
     }
+
+    public Map<ClaimMapping, String> getMappedRemoteClaims() {
+
+        return mappedRemoteClaims;
+    }
+
+    public void setMappedRemoteClaims(
+            Map<ClaimMapping, String> mappedRemoteClaims) {
+
+        this.mappedRemoteClaims = mappedRemoteClaims;
+    }
 }

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dto/OAuth2AuthorizeReqDTO.java

@@ -19,10 +19,12 @@
 package org.wso2.carbon.identity.oauth2.dto;
 
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatedUser;
+import org.wso2.carbon.identity.application.common.model.ClaimMapping;
 import org.wso2.carbon.identity.oauth2.model.HttpRequestHeader;
 import org.wso2.carbon.identity.openidconnect.model.RequestObject;
 
 import java.util.LinkedHashSet;
+import java.util.Map;
 import java.util.Properties;
 
 import javax.servlet.http.Cookie;
@@ -61,6 +63,7 @@ public class OAuth2AuthorizeReqDTO {
     private boolean isRequestObjectFlow;
     private String state;
     private String requestedSubjectId;
+    private Map<ClaimMapping, String> mappedRemoteClaims;
 
     public String getRequestedSubjectId() {
 
@@ -303,4 +306,15 @@ public void setHttpServletRequestWrapper(HttpServletRequestWrapper httpServletRe
 
         this.httpServletRequestWrapper = httpServletRequestWrapper;
     }
+
+    public Map<ClaimMapping, String> getMappedRemoteClaims() {
+
+        return mappedRemoteClaims;
+    }
+
+    public void setMappedRemoteClaims(
+            Map<ClaimMapping, String> mappedRemoteClaims) {
+
+        this.mappedRemoteClaims = mappedRemoteClaims;
+    }
 }

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/AccessTokenIssuer.java

@@ -621,7 +621,13 @@ private Optional<AuthorizationGrantCacheEntry> getAuthzGrantCacheEntryFromDevice
                 DeviceAuthorizationGrantCache.getInstance().getValueFromCache(deviceCodeCacheKey);
         if (cacheEntry != null) {
             Map<ClaimMapping, String> userAttributes = cacheEntry.getUserAttributes();
-            return Optional.of(new AuthorizationGrantCacheEntry(userAttributes));
+            AuthorizationGrantCacheEntry authorizationGrantCacheEntry =
+                    new AuthorizationGrantCacheEntry(userAttributes);
+            if (cacheEntry.getMappedRemoteClaims() != null) {
+                authorizationGrantCacheEntry.setMappedRemoteClaims(cacheEntry
+                        .getMappedRemoteClaims());
+            }
+            return Optional.of(authorizationGrantCacheEntry);
         }
         return Optional.empty();
     }