Support SSO for shared users across sub-organizations.

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/cache/SessionContextCache.java

@@ -75,13 +75,32 @@
 
         if (log.isDebugEnabled()) {
             log.debug("Adding session context corresponding to the key : " + key.getContextId() +
                 " with accessed time " + entry.getAccessedTime() + " and validity time " + entry.getValidityPeriod());
         }
         entry.setAccessedTime();
         super.addToCache(key, entry, resolveLoginTenantDomain(loginTenantDomain));
         optimizeAndStoreSessionData(key, entry);
     }
 
+    /**
+     * Add the given session context entry to the cache only, without persisting it to the session data store.
+     *
+     * @param key               Key which the cache entry is indexed by.
+     * @param entry             Value to be stored in the cache.
+     * @param loginTenantDomain Login tenant domain under which the cache entry is stored.
+     */
+    public void addToCacheWithoutPersisting(SessionContextCacheKey key, SessionContextCacheEntry entry,
+                                            String loginTenantDomain) {
+
+        if (log.isDebugEnabled()) {
+            log.debug("Adding session context to cache without updating the session data store corresponding " +
+                    "to the key : " + key.getContextId() + " with accessed time " + entry.getAccessedTime() +
+                    " and validity time " + entry.getValidityPeriod());
+        }
+        entry.setAccessedTime();
+        super.addToCache(key, entry, resolveLoginTenantDomain(loginTenantDomain));
+    }
+
     /**
      * Add a cache entry during a READ operation.
      * <p>

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/model/PrimaryAppData.java

@@ -28,6 +28,7 @@ public class PrimaryAppData implements Serializable {
     private static final long serialVersionUID = -3925617004643909936L;
 
     private int id;
+    private String tenantDomain;
 
     public int getId() {
 
@@ -38,4 +39,14 @@ public void setId(int id) {
 
         this.id = id;
     }
+
+    public String getTenantDomain() {
+
+        return tenantDomain;
+    }
+
+    public void setTenantDomain(String tenantDomain) {
+
+        this.tenantDomain = tenantDomain;
+    }
 }

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/main/java/org/wso2/carbon/identity/application/authentication/framework/util/FrameworkUtils.java

@@ -1324,6 +1324,46 @@ public static void addSessionContextToCache(String key, SessionContext sessionCo
                                                 String loginTenantDomain, String orgId) {
 
         SessionContextCacheKey cacheKey = new SessionContextCacheKey(key);
+        SessionContextCacheEntry cacheEntry = buildSessionContextCacheEntry(key, sessionContext, tenantDomain, orgId);
+        SessionContextCache.getInstance().addToCache(cacheKey, cacheEntry, loginTenantDomain);
+    }
+
+    /**
+     * Adds the given session context to the session context cache only, without persisting it to the session data
+     * store.
+     *
+     * @param key               Session context cache key.
+     * @param sessionContext    Session context to be cached.
+     * @param tenantDomain      Application tenant domain used to resolve session timeout configurations.
+     * @param loginTenantDomain Login tenant domain under which the cache entry is stored.
+     * @param orgId             Organization id used to look up authenticated sequences from
+     *                          {@link SessionContext#getAuthenticatedOrgData()}.
+     */
+    public static void addSessionContextToCacheWithoutPersisting(String key, SessionContext sessionContext,
+                                                                 String tenantDomain, String loginTenantDomain,
+                                                                 String orgId) {
+
+        SessionContextCacheKey cacheKey = new SessionContextCacheKey(key);
+        SessionContextCacheEntry cacheEntry = buildSessionContextCacheEntry(key, sessionContext, tenantDomain, orgId);
+        SessionContextCache.getInstance().addToCacheWithoutPersisting(cacheKey, cacheEntry, loginTenantDomain);
+    }
+
+    /**
+     * Builds a {@link SessionContextCacheEntry} for the given session context. When an {@code orgId} is provided and
+     * the session context holds an {@link AuthenticatedOrgData} entry for that organization, the authenticated
+     * sequences of that organization are used for cleanup (clearing user attributes and the authentication graph).
+     * Otherwise, the top-level authenticated sequences of the session context are used.
+     *
+     * @param key            Session context cache key.
+     * @param sessionContext Session context to be cached.
+     * @param tenantDomain   Application tenant domain used to resolve session timeout configurations.
+     * @param orgId          Organization id used to look up authenticated sequences from
+     *                       {@link SessionContext#getAuthenticatedOrgData()}.
+     * @return The built session context cache entry.
+     */
+    private static SessionContextCacheEntry buildSessionContextCacheEntry(String key, SessionContext sessionContext,
+                                                                          String tenantDomain, String orgId) {
+
         SessionContextCacheEntry cacheEntry = new SessionContextCacheEntry();
         cacheEntry.setContextIdentifier(key);
 
@@ -1364,7 +1404,7 @@ public static void addSessionContextToCache(String key, SessionContext sessionCo
 
         cacheEntry.setContext(sessionContext);
         cacheEntry.setValidityPeriod(timeoutPeriod);
-        SessionContextCache.getInstance().addToCache(cacheKey, cacheEntry, loginTenantDomain);
+        return cacheEntry;
     }
 
     /**

components/authentication-framework/org.wso2.carbon.identity.application.authentication.framework/src/test/java/org/wso2/carbon/identity/application/authentication/framework/handler/request/impl/DefaultAuthenticationRequestHandlerTest.java

@@ -28,12 +28,14 @@
 import org.testng.annotations.Test;
 import org.wso2.carbon.identity.application.authentication.framework.cache.AuthenticationResultCacheEntry;
 import org.wso2.carbon.identity.application.authentication.framework.config.model.ApplicationConfig;
+import org.wso2.carbon.identity.application.authentication.framework.config.model.AuthenticatorConfig;
 import org.wso2.carbon.identity.application.authentication.framework.config.model.SequenceConfig;
 import org.wso2.carbon.identity.application.authentication.framework.context.AuthenticationContext;
 import org.wso2.carbon.identity.application.authentication.framework.exception.FrameworkException;
 import org.wso2.carbon.identity.application.authentication.framework.exception.PostAuthenticationFailedException;
 import org.wso2.carbon.identity.application.authentication.framework.handler.sequence.impl.DefaultStepBasedSequenceHandler;
 import org.wso2.carbon.identity.application.authentication.framework.internal.FrameworkServiceDataHolder;
+import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatedIdPData;
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatedUser;
 import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticationResult;
 import org.wso2.carbon.identity.application.authentication.framework.model.CommonAuthResponseWrapper;
@@ -51,7 +53,9 @@
 import org.wso2.carbon.identity.common.testng.WithCarbonHome;
 
 import java.io.IOException;
+import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -75,6 +79,7 @@
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNotSame;
 import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
 import static org.wso2.carbon.base.MultitenantConstants.SUPER_TENANT_DOMAIN_NAME;
@@ -633,5 +638,90 @@ public void testConcludeFlowAllowsTenantMismatchForSharedUser() throws Exception
                     + "when user is a shared user");
         }
     }
+
+    @SuppressWarnings("unchecked")
+    private Map<String, AuthenticatedIdPData> invokeMergeAuthenticatedIdPs(
+            Map<String, AuthenticatedIdPData> previous,
+            Map<String, AuthenticatedIdPData> current) throws Exception {
+
+        Method method = DefaultAuthenticationRequestHandler.class.getDeclaredMethod(
+                "mergeAuthenticatedIdPs", Map.class, Map.class);
+        method.setAccessible(true);
+        return (Map<String, AuthenticatedIdPData>) method.invoke(
+                new DefaultAuthenticationRequestHandler(), previous, current);
+    }
+
+    private AuthenticatedIdPData buildAuthenticatedIdPData(String idpName, String... authenticatorNames) {
+
+        AuthenticatedIdPData authenticatedIdPData = new AuthenticatedIdPData();
+        authenticatedIdPData.setIdpName(idpName);
+        authenticatedIdPData.setUser(new AuthenticatedUser());
+        List<AuthenticatorConfig> authenticators = new ArrayList<>();
+        for (String authenticatorName : authenticatorNames) {
+            authenticators.add(new AuthenticatorConfig(authenticatorName, true, null));
+        }
+        authenticatedIdPData.setAuthenticators(authenticators);
+        return authenticatedIdPData;
+    }
+
+    private List<String> authenticatorNames(AuthenticatedIdPData authenticatedIdPData) {
+
+        List<String> names = new ArrayList<>();
+        for (AuthenticatorConfig authenticatorConfig : authenticatedIdPData.getAuthenticators()) {
+            names.add(authenticatorConfig.getName());
+        }
+        return names;
+    }
+
+    @Test(description = "Merging disjoint authenticated IdP maps keeps all IdPs from both maps.")
+    public void testMergeAuthenticatedIdPsWithDisjointIdPs() throws Exception {
+
+        Map<String, AuthenticatedIdPData> previous = new HashMap<>();
+        previous.put("LOCAL", buildAuthenticatedIdPData("LOCAL", "BasicAuthenticator"));
+        Map<String, AuthenticatedIdPData> current = new HashMap<>();
+        current.put("FederatedIdP", buildAuthenticatedIdPData("FederatedIdP", "OpenIDConnectAuthenticator"));
+
+        Map<String, AuthenticatedIdPData> merged = invokeMergeAuthenticatedIdPs(previous, current);
+
+        assertEquals(merged.size(), 2);
+        assertTrue(merged.containsKey("LOCAL"));
+        assertTrue(merged.containsKey("FederatedIdP"));
+        assertNotNull(merged.get("LOCAL"));
+        assertNotSame(merged.get("LOCAL"), previous.get("LOCAL"));
+    }
+
+    @Test(description = "Merging IdP maps that share an IdP appends only the authenticators not already present.")
+    public void testMergeAuthenticatedIdPsAppendsMissingAuthenticators() throws Exception {
+
+        Map<String, AuthenticatedIdPData> previous = new HashMap<>();
+        previous.put("LOCAL", buildAuthenticatedIdPData("LOCAL", "BasicAuthenticator"));
+        Map<String, AuthenticatedIdPData> current = new HashMap<>();
+        // Same IdP with a duplicate authenticator and a new authenticator.
+        current.put("LOCAL", buildAuthenticatedIdPData("LOCAL", "BasicAuthenticator", "TOTPAuthenticator"));
+
+        Map<String, AuthenticatedIdPData> merged = invokeMergeAuthenticatedIdPs(previous, current);
+
+        assertEquals(merged.size(), 1);
+        List<String> mergedAuthenticators = authenticatorNames(merged.get("LOCAL"));
+        assertEquals(mergedAuthenticators.size(), 2, "Duplicate authenticator should not be added twice.");
+        assertTrue(mergedAuthenticators.contains("BasicAuthenticator"));
+        assertTrue(mergedAuthenticators.contains("TOTPAuthenticator"));
+    }
+
+    @Test(description = "Merging handles null and empty authenticated IdP maps gracefully.")
+    public void testMergeAuthenticatedIdPsWithNullAndEmptyMaps() throws Exception {
+
+        Map<String, AuthenticatedIdPData> current = new HashMap<>();
+        current.put("LOCAL", buildAuthenticatedIdPData("LOCAL", "BasicAuthenticator"));
+
+        // Previous map is null.
+        Map<String, AuthenticatedIdPData> merged = invokeMergeAuthenticatedIdPs(null, current);
+        assertEquals(merged.size(), 1);
+        assertTrue(merged.containsKey("LOCAL"));
+
+        // Both maps empty/null.
+        Map<String, AuthenticatedIdPData> emptyMerge = invokeMergeAuthenticatedIdPs(null, new HashMap<>());
+        assertTrue(emptyMerge.isEmpty());
+    }
 }