Sharing the federated token in oidc authentication

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

@@ -61,6 +61,7 @@
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticationResult;
 import org.wso2.carbon.identity.application.authentication.framework.model.CommonAuthRequestWrapper;
 import org.wso2.carbon.identity.application.authentication.framework.model.CommonAuthResponseWrapper;
+import org.wso2.carbon.identity.application.authentication.framework.model.FederatedToken;
 import org.wso2.carbon.identity.application.authentication.framework.model.auth.service.AuthServiceRequest;
 import org.wso2.carbon.identity.application.authentication.framework.model.auth.service.AuthServiceResponse;
 import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants;
@@ -129,6 +130,7 @@
 import org.wso2.carbon.identity.oauth2.dto.OAuth2ClientValidationResponseDTO;
 import org.wso2.carbon.identity.oauth2.internal.OAuth2ServiceComponentHolder;
 import org.wso2.carbon.identity.oauth2.model.AccessTokenExtendedAttributes;
+import org.wso2.carbon.identity.oauth2.model.FederatedTokenDO;
 import org.wso2.carbon.identity.oauth2.model.HttpRequestHeaderHandler;
 import org.wso2.carbon.identity.oauth2.model.OAuth2Parameters;
 import org.wso2.carbon.identity.oauth2.responsemode.provider.AuthorizationResponseDTO;
@@ -173,6 +175,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 import javax.servlet.RequestDispatcher;
 import javax.servlet.ServletContext;
@@ -388,7 +391,56 @@ public Response authorize(@Context HttpServletRequest request, @Context HttpServ
         }
     }
 
+    /**
+     * Add the federated tokens comes with the authentication result to the session data cache.
+     *
+     * @param oAuthMessage         The OAuthMessage with the session data cache entry.
+     * @param authenticationResult The authentication result of authorization call.
+     */
+    private void addFederatedTokensToSessionCache(OAuthMessage oAuthMessage,
+                                                  AuthenticationResult authenticationResult) {
+
+        if (!(authenticationResult.getProperty(FrameworkConstants.FEDERATED_TOKENS) instanceof List)) {
+            return;
+        }
+        List<FederatedToken> federatedTokens =
+                (List<FederatedToken>) authenticationResult.getProperty(FrameworkConstants.FEDERATED_TOKENS);
 
+        SessionDataCacheEntry sessionDataCacheEntry = oAuthMessage.getSessionDataCacheEntry();
+        if (sessionDataCacheEntry == null || CollectionUtils.isEmpty(federatedTokens)) {
+            return;
+        }
+        sessionDataCacheEntry.setFederatedTokens(getFederatedTokenDO(federatedTokens));
+        if (log.isDebugEnabled() && authenticationResult.getSubject() != null) {
+            log.debug("Added the federated tokens 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.
+     *
+     * @param federatedTokens List of FederatedToken objects to be transformed as a list of FederatedTokenDO.
+     * @return List of FederatedTokenDO objects.
+     */
+    private List<FederatedTokenDO> getFederatedTokenDO(List<FederatedToken> federatedTokens) {
+
+        if (CollectionUtils.isEmpty(federatedTokens)) {
+            return null;
+        }
+
+        List<FederatedTokenDO>  federatedTokenDOs = federatedTokens.stream().map(federatedToken -> {
+            FederatedTokenDO federatedTokenDO =
+                    new FederatedTokenDO(federatedToken.getIdp(), federatedToken.getAccessToken());
+            federatedTokenDO.setRefreshToken(federatedToken.getRefreshToken());
+            federatedTokenDO.setScope(federatedToken.getScope());
+            federatedTokenDO.setTokenValidityPeriod(federatedToken.getTokenValidityPeriod());
+            return federatedTokenDO;
+        }).collect(Collectors.toList());
+
+        return federatedTokenDOs;
+    }
 
     private void setCommonAuthIdToRequest(HttpServletRequest request, HttpServletResponse response) {
 
@@ -1335,6 +1387,8 @@ private void addToAuthenticationResultDetailsToOAuthMessage(OAuthMessage oAuthMe
         oAuthMessage.getSessionDataCacheEntry().setAuthenticatedIdPs(authnResult.getAuthenticatedIdPs());
         oAuthMessage.getSessionDataCacheEntry().setSessionContextIdentifier((String)
                 authnResult.getProperty(FrameworkConstants.AnalyticsAttributes.SESSION_ID));
+        // Adding federated tokens come with the authentication result of the authorization call.
+        addFederatedTokensToSessionCache(oAuthMessage, authnResult);
     }
 
     private void updateAuthTimeInSessionDataCacheEntry(OAuthMessage oAuthMessage) {
@@ -2025,6 +2079,8 @@ private void addUserAttributesToOAuthMessage(OAuthMessage oAuthMessage, String c
         authorizationGrantCacheEntry.setAuthorizationCode(code);
         boolean isRequestObjectFlow = sessionDataCacheEntry.getoAuth2Parameters().isRequestObjectFlow();
         authorizationGrantCacheEntry.setRequestObjectFlow(isRequestObjectFlow);
+        authorizationGrantCacheEntry.setFederatedTokens(sessionDataCacheEntry.getFederatedTokens());
+        sessionDataCacheEntry.setFederatedTokens(null);
         oAuthMessage.setAuthorizationGrantCacheEntry(authorizationGrantCacheEntry);
     }
 

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

@@ -20,6 +20,7 @@
 
 import org.wso2.carbon.identity.application.common.model.ClaimMapping;
 import org.wso2.carbon.identity.oauth2.model.AccessTokenExtendedAttributes;
+import org.wso2.carbon.identity.oauth2.model.FederatedTokenDO;
 import org.wso2.carbon.identity.openidconnect.model.RequestObject;
 
 import java.util.ArrayList;
@@ -81,6 +82,8 @@ public class AuthorizationGrantCacheEntry extends CacheEntry {
     private AccessTokenExtendedAttributes accessTokenExtendedAttributes;
     private boolean isApiBasedAuthRequest;
 
+    private List<FederatedTokenDO> federatedTokens;
+
     public String getSubjectClaim() {
         return subjectClaim;
     }
@@ -193,6 +196,16 @@ public void setPkceCodeChallengeMethod(String pkceCodeChallengeMethod) {
         this.pkceCodeChallengeMethod = pkceCodeChallengeMethod;
     }
 
+    public List<FederatedTokenDO> getFederatedTokens() {
+
+        return federatedTokens;
+    }
+
+    public void setFederatedTokens(List<FederatedTokenDO> federatedTokens) {
+
+        this.federatedTokens = federatedTokens;
+    }
+
     /**
      * To check whether particular cache entry has non OIDC claims in it.
      *

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

@@ -20,10 +20,12 @@
 
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatedUser;
 import org.wso2.carbon.identity.oauth2.authz.OAuthAuthzReqMessageContext;
+import org.wso2.carbon.identity.oauth2.model.FederatedTokenDO;
 import org.wso2.carbon.identity.oauth2.model.OAuth2Parameters;
 
 import java.io.Serializable;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -50,6 +52,7 @@ public class SessionDataCacheEntry extends CacheEntry {
     private ConcurrentMap<String, String[]> paramMap = new ConcurrentHashMap<String, String[]>();
 
     private Map<String, Serializable> endpointParams = new HashMap<>();
+    private List<FederatedTokenDO> federatedTokens;
 
     public OAuthAuthzReqMessageContext getAuthzReqMsgCtx() {
         return authzReqMsgCtx;
@@ -159,4 +162,14 @@ public void setRemoveOnConsume(boolean removeOnConsume) {
 
         this.removeOnConsume = removeOnConsume;
     }
+
+    public List<FederatedTokenDO> getFederatedTokens() {
+
+        return federatedTokens;
+    }
+
+    public void setFederatedTokens(List<FederatedTokenDO> federatedTokens) {
+
+        this.federatedTokens = federatedTokens;
+    }
 }

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/internal/OAuth2ServiceComponent.java

@@ -86,6 +86,8 @@
 import org.wso2.carbon.identity.oauth2.token.bindings.impl.DeviceFlowTokenBinder;
 import org.wso2.carbon.identity.oauth2.token.bindings.impl.SSOSessionBasedTokenBinder;
 import org.wso2.carbon.identity.oauth2.token.handlers.claims.JWTAccessTokenClaimProvider;
+import org.wso2.carbon.identity.oauth2.token.handlers.response.AccessTokenResponseHandler;
+import org.wso2.carbon.identity.oauth2.token.handlers.response.FederatedTokenResponseHandler;
 import org.wso2.carbon.identity.oauth2.util.OAuth2Util;
 import org.wso2.carbon.identity.oauth2.validators.scope.RoleBasedScopeIssuer;
 import org.wso2.carbon.identity.oauth2.validators.scope.ScopeValidator;
@@ -382,6 +384,8 @@ protected void activate(ComponentContext context) {
             bundleContext.registerService(ScopeValidationHandler.class, new RoleBasedScopeValidationHandler(), null);
             bundleContext.registerService(ScopeValidationHandler.class, new NoPolicyScopeValidationHandler(), null);
             bundleContext.registerService(ScopeValidationHandler.class, new M2MScopeValidationHandler(), null);
+            bundleContext.registerService(AccessTokenResponseHandler.class, new FederatedTokenResponseHandler(),
+                    null);
 
             // Note : DO NOT add any activation related code below this point,
             // to make sure the server doesn't start up if any activation failures occur

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/model/FederatedTokenDO.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. 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.model;
+
+import java.io.Serializable;
+
+/**
+ * This class is model class of a federated token.
+ * A federated token is an external token obtained via an OIDC federated authenticator
+ * after a successful authentication.
+ */
+public class FederatedTokenDO implements Serializable {
+
+    private static final long serialVersionUID = 2717725650850067925L;
+    private String idp;
+    private String tokenValidityPeriod;
+    private String scope;
+    private String accessToken;
+    private String refreshToken;
+
+    // Constructor
+    public FederatedTokenDO(String idp, String accessToken) {
+
+        this.idp = idp;
+        this.accessToken = accessToken;
+    }
+
+    // Getters and setters
+    public String getIdp() {
+
+        return idp;
+    }
+
+    public void setIdp(String idp) {
+
+        this.idp = idp;
+    }
+
+    public String getTokenValidityPeriod() {
+
+        return tokenValidityPeriod;
+    }
+
+    public void setTokenValidityPeriod(String tokenValidityPeriod) {
+
+        this.tokenValidityPeriod = tokenValidityPeriod;
+    }
+
+    public String getScope() {
+
+        return scope;
+    }
+
+    public void setScope(String scope) {
+
+        this.scope = scope;
+    }
+
+    public String getAccessToken() {
+
+        return accessToken;
+    }
+
+    public void setAccessToken(String accessToken) {
+
+        this.accessToken = accessToken;
+    }
+
+    public String getRefreshToken() {
+
+        return refreshToken;
+    }
+
+    public void setRefreshToken(String refreshToken) {
+
+        this.refreshToken = refreshToken;
+    }
+}

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/token/handlers/response/FederatedTokenResponseHandler.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. 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.token.handlers.response;
+
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang.StringUtils;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkConstants;
+import org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCache;
+import org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCacheEntry;
+import org.wso2.carbon.identity.oauth.cache.AuthorizationGrantCacheKey;
+import org.wso2.carbon.identity.oauth2.model.FederatedTokenDO;
+import org.wso2.carbon.identity.oauth2.token.OAuthTokenReqMessageContext;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * This class is used to get the federated tokens for the token response.
+ * This expects the authorization code is available in the token request message context
+ * so that the federated tokens can be retrieved from the auth grant cache.
+ */
+public class FederatedTokenResponseHandler implements AccessTokenResponseHandler {
+
+    private static final Log LOG = LogFactory.getLog(FederatedTokenResponseHandler.class);
+
+    /**
+     * This method returns the federated tokens in the auth grant cache.
+     *
+     * @param tokReqMsgCtx {@link OAuthTokenReqMessageContext} Token request message context with a token request DTO.
+     * @return Map of the federated tokens.
+     */
+    @Override
+    public Map<String, Object> getAdditionalTokenResponseAttributes(OAuthTokenReqMessageContext tokReqMsgCtx) {
+
+        if (StringUtils.isBlank(tokReqMsgCtx.getOauth2AccessTokenReqDTO().getAuthorizationCode())) {
+            return null;
+        }
+        AuthorizationGrantCacheEntry cacheEntry =
+                AuthorizationGrantCache.getInstance().getValueFromCacheByCode(new AuthorizationGrantCacheKey(
+                        tokReqMsgCtx.getOauth2AccessTokenReqDTO().getAuthorizationCode()));
+
+        if (cacheEntry == null) {
+            return null;
+        }
+
+        List<FederatedTokenDO> federatedTokens = cacheEntry.getFederatedTokens();
+        if (CollectionUtils.isEmpty(federatedTokens)) {
+            return null;
+        }
+        // Removing the federated token from the session cache entry since it is no longer required.
+        cacheEntry.setFederatedTokens(null);
+        // Add federated tokens to the token response if available.
+        Map<String, Object> additionalAttributes = new HashMap<>();
+
+        additionalAttributes.putIfAbsent(FrameworkConstants.FEDERATED_TOKENS, federatedTokens);
+        if (LOG.isDebugEnabled() && tokReqMsgCtx.getAuthorizedUser() != null) {
+            LOG.debug("Federated tokens will be added to the additional attributes of the token response." +
+                    " for the user: " + tokReqMsgCtx.getAuthorizedUser().getLoggableMaskedUserId());
+        }
+
+        return additionalAttributes;
+    }
+}