Improve getUserSessions conditional auth function to resolve user accessing org when getting sessions during B2B logins.

components/org.wso2.carbon.identity.conditional.auth.functions.user/pom.xml

@@ -238,7 +238,9 @@
                             net.minidev.json; version="${net.minidev.json.imp.pkg.version.range}",
                             org.wso2.carbon.identity.central.log.mgt.utils; version="${carbon.identity.package.import.version.range}",
                             org.wso2.carbon.identity.conditional.auth.functions.common.utils,
-                            org.wso2.carbon.utils*;version="${carbon.kernel.package.import.version.range}"
+                            org.wso2.carbon.utils*;version="${carbon.kernel.package.import.version.range}",
+                            org.wso2.carbon.identity.organization.management.service;version="${org.wso2.carbon.identity.organization.management.core.version.range}",
+                            org.wso2.carbon.identity.organization.management.service.exception;version="${org.wso2.carbon.identity.organization.management.core.version.range}"
                         </Import-Package>
                         <DynamicImport-Package>
                             org.graalvm.polyglot,

components/org.wso2.carbon.identity.conditional.auth.functions.user/src/main/java/org/wso2/carbon/identity/conditional/auth/functions/user/GetUserSessionsFunctionImpl.java

@@ -1,7 +1,7 @@
 /*
- * Copyright (c) 2019, WSO2 Inc. (http://www.wso2.org) All Rights Reserved.
+ * Copyright (c) 2019-2025, WSO2 LLC. (http://www.wso2.com).
  *
- * WSO2 Inc. licenses this file to you under the Apache License,
+ * 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
@@ -18,6 +18,7 @@
 
 package org.wso2.carbon.identity.conditional.auth.functions.user;
 
+import org.apache.commons.lang.StringUtils;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.graalvm.polyglot.HostAccess;
@@ -31,6 +32,7 @@
 import org.wso2.carbon.identity.conditional.auth.functions.user.internal.UserFunctionsServiceHolder;
 import org.wso2.carbon.identity.conditional.auth.functions.user.model.JsUserSession;
 import org.wso2.carbon.identity.conditional.auth.functions.user.model.JsWrapperFactoryProvider;
+import org.wso2.carbon.identity.organization.management.service.exception.OrganizationManagementException;
 import org.wso2.carbon.user.core.UserRealm;
 
 import java.util.List;
@@ -62,14 +64,14 @@ private List<UserSession> getUserSessions(AuthenticatedUser authenticatedUser)
             throws UserSessionRetrievalException {
 
         List<UserSession> userSessions = null;
-        String tenantDomain = authenticatedUser.getTenantDomain();
+        String tenantDomain = getUserTenantDomain(authenticatedUser);
 
         try {
             UserRealm userRealm = Utils.getUserRealm(tenantDomain);
             if (userRealm != null) {
                 String userId = authenticatedUser.getUserId();
                 userSessions = UserFunctionsServiceHolder.getInstance()
-                        .getUserSessionManagementService().getSessionsByUserId(userId);
+                        .getUserSessionManagementService().getSessionsByUserId(userId, tenantDomain);
             }
         } catch (SessionManagementException e) {
             throw new UserSessionRetrievalException("Error occurred while retrieving sessions: ", e);
@@ -81,4 +83,20 @@ private List<UserSession> getUserSessions(AuthenticatedUser authenticatedUser)
         return userSessions;
     }
 
+    private String getUserTenantDomain(AuthenticatedUser authenticatedUser) throws UserSessionRetrievalException {
+
+        String tenantDomain = authenticatedUser.getTenantDomain();
+        String userAccessingOrganization = authenticatedUser.getAccessingOrganization();
+        if (StringUtils.isNotBlank(userAccessingOrganization)) {
+            try {
+                tenantDomain = UserFunctionsServiceHolder.getInstance().getOrganizationManager()
+                        .resolveTenantDomain(userAccessingOrganization);
+            } catch (OrganizationManagementException e) {
+                throw new UserSessionRetrievalException(
+                        "Error occurred while resolving tenant domain of user accessing organization: " +
+                                userAccessingOrganization, e);
+            }
+        }
+        return tenantDomain;
+    }
 }

components/org.wso2.carbon.identity.conditional.auth.functions.user/src/main/java/org/wso2/carbon/identity/conditional/auth/functions/user/internal/UserFunctionsServiceComponent.java

@@ -62,6 +62,7 @@
 import org.wso2.carbon.identity.conditional.auth.functions.user.SetAccountAssociationToLocalUser;
 import org.wso2.carbon.identity.conditional.auth.functions.user.UpdateUserPasswordFunction;
 import org.wso2.carbon.identity.conditional.auth.functions.user.UpdateUserPasswordFunctionImpl;
+import org.wso2.carbon.identity.organization.management.service.OrganizationManager;
 import org.wso2.carbon.identity.role.v2.mgt.core.RoleManagementService;
 import org.wso2.carbon.idp.mgt.IdpManager;
 import org.wso2.carbon.user.core.service.RealmService;
@@ -290,4 +291,21 @@ protected void unsetRoleManagementService(RoleManagementService roleManagementSe
             LOG.debug("Role management service is unset in the conditional authentication user functions bundle");
         }
     }
+
+    @Reference(name = "identity.organization.management.component",
+            service = OrganizationManager.class,
+            cardinality = ReferenceCardinality.MANDATORY,
+            policy = ReferencePolicy.DYNAMIC,
+            unbind = "unsetOrganizationManager")
+    protected void setOrganizationManager(OrganizationManager organizationManager) {
+
+        UserFunctionsServiceHolder.getInstance().setOrganizationManager(organizationManager);
+        LOG.debug("Organization manager service is set in the conditional authentication user functions bundle.");
+    }
+
+    protected void unsetOrganizationManager(OrganizationManager organizationManager) {
+
+        UserFunctionsServiceHolder.getInstance().setOrganizationManager(null);
+        LOG.debug("Organization manager service is unset in the conditional authentication user functions bundle.");
+    }
 }

components/org.wso2.carbon.identity.conditional.auth.functions.user/src/main/java/org/wso2/carbon/identity/conditional/auth/functions/user/internal/UserFunctionsServiceHolder.java

@@ -21,6 +21,7 @@
 import org.wso2.carbon.identity.application.authentication.framework.JsFunctionRegistry;
 import org.wso2.carbon.identity.application.authentication.framework.UserSessionManagementService;
 import org.wso2.carbon.identity.application.mgt.ApplicationManagementService;
+import org.wso2.carbon.identity.organization.management.service.OrganizationManager;
 import org.wso2.carbon.identity.role.v2.mgt.core.RoleManagementService;
 import org.wso2.carbon.idp.mgt.IdpManager;
 import org.wso2.carbon.user.core.service.RealmService;
@@ -35,6 +36,7 @@ public class UserFunctionsServiceHolder {
     private ApplicationManagementService applicationManagementService;
     private RoleManagementService roleManagementService;
     private JsFunctionRegistry jsFunctionRegistry;
+    private OrganizationManager organizationManager;
 
     private UserFunctionsServiceHolder() {
 
@@ -106,4 +108,14 @@ public void setJsFunctionRegistry(JsFunctionRegistry jsFunctionRegistry) {
 
         this.jsFunctionRegistry = jsFunctionRegistry;
     }
+
+    public OrganizationManager getOrganizationManager() {
+
+        return organizationManager;
+    }
+
+    public void setOrganizationManager(OrganizationManager organizationManager) {
+
+        this.organizationManager = organizationManager;
+    }
 }

components/org.wso2.carbon.identity.conditional.auth.functions.user/src/test/java/org/wso2/carbon/identity/conditional/auth/functions/user/GetUserSessionsFunctionImplTest.java

@@ -0,0 +1,200 @@
+/*
+ * Copyright (c) 2025, 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.conditional.auth.functions.user;
+
+import org.mockito.Mock;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+import org.wso2.carbon.context.PrivilegedCarbonContext;
+import org.wso2.carbon.identity.application.authentication.framework.UserSessionManagementService;
+import org.wso2.carbon.identity.application.authentication.framework.config.model.graph.js.JsAuthenticatedUser;
+import org.wso2.carbon.identity.application.authentication.framework.config.model.graph.js.graaljs.JsGraalAuthenticatedUser;
+import org.wso2.carbon.identity.application.authentication.framework.internal.FrameworkServiceDataHolder;
+import org.wso2.carbon.identity.application.authentication.framework.model.AuthenticatedUser;
+import org.wso2.carbon.identity.application.authentication.framework.model.UserSession;
+import org.wso2.carbon.identity.common.testng.WithCarbonHome;
+import org.wso2.carbon.identity.common.testng.WithRealmService;
+import org.wso2.carbon.identity.conditional.auth.functions.user.internal.UserFunctionsServiceHolder;
+import org.wso2.carbon.identity.conditional.auth.functions.user.model.JsUserSession;
+import org.wso2.carbon.identity.core.util.IdentityTenantUtil;
+import org.wso2.carbon.identity.organization.management.service.OrganizationManager;
+import org.wso2.carbon.identity.organization.management.service.exception.OrganizationManagementException;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.mockito.MockitoAnnotations.initMocks;
+
+/**
+ * Unit tests for GetUserSessionsFunctionImpl class.
+ */
+@WithCarbonHome
+@WithRealmService(injectToSingletons = {UserFunctionsServiceHolder.class, IdentityTenantUtil.class,
+        FrameworkServiceDataHolder.class})
+public class GetUserSessionsFunctionImplTest {
+
+    private static final String TEST_USER_NAME = "testUser";
+    private static final String TENANT_DOMAIN_CARBON_SUPER = "carbon.super";
+    private static final String USER_STORE_DOMAIN_PRIMARY = "PRIMARY";
+    private static final String TEST_USER_ID_123 = "test-user-id-123";
+    private static final String TEST_USER_ID_456 = "test-user-id-456";
+    private static final String TEST_USER_ID_ORG_123 = "test-user-id-org-123";
+    private static final String TEST_USER_ID_ORG_456 = "test-user-id-org-456";
+    private static final String ORG_ID_123 = "org123";
+    private static final String ORG_TENANT_DOMAIN = "org.tenant.com";
+    private static final String INVALID_ORG = "invalidOrg";
+    private static final String ORG_NOT_FOUND_MESSAGE = "Organization not found";
+
+    @Mock
+    private UserSessionManagementService userSessionManagementService;
+    @Mock
+    private OrganizationManager organizationManager;
+    private GetUserSessionsFunctionImpl getUserSessionsFunction;
+
+    @BeforeMethod
+    public void setUp() {
+
+        initMocks(this);
+        PrivilegedCarbonContext.getThreadLocalCarbonContext().setTenantDomain(TENANT_DOMAIN_CARBON_SUPER, true);
+
+        getUserSessionsFunction = new GetUserSessionsFunctionImpl();
+        UserFunctionsServiceHolder.getInstance().setUserSessionManagementService(userSessionManagementService);
+        UserFunctionsServiceHolder.getInstance().setOrganizationManager(organizationManager);
+    }
+
+    @AfterMethod
+    public void tearDown() {
+
+        PrivilegedCarbonContext.destroyCurrentContext();
+    }
+
+    @Test
+    public void testGetUserSessionsSuccess() throws Exception {
+
+        // Create test user.
+        AuthenticatedUser authenticatedUser = new AuthenticatedUser();
+        authenticatedUser.setUserName(TEST_USER_NAME);
+        authenticatedUser.setTenantDomain(TENANT_DOMAIN_CARBON_SUPER);
+        authenticatedUser.setUserStoreDomain(USER_STORE_DOMAIN_PRIMARY);
+        authenticatedUser.setUserId(TEST_USER_ID_123);
+
+        JsAuthenticatedUser jsUser = new JsGraalAuthenticatedUser(authenticatedUser);
+
+        // Create mock user sessions.
+        List<UserSession> mockSessions = new ArrayList<>();
+        UserSession session1 = mock(UserSession.class);
+        UserSession session2 = mock(UserSession.class);
+        mockSessions.add(session1);
+        mockSessions.add(session2);
+
+        // Mock the session management service.
+        when(userSessionManagementService.getSessionsByUserId(TEST_USER_ID_123, TENANT_DOMAIN_CARBON_SUPER))
+                .thenReturn(mockSessions);
+
+        // Execute the function.
+        List<JsUserSession> result = getUserSessionsFunction.getUserSessions(jsUser);
+
+        // Verify results.
+        Assert.assertNotNull(result);
+        Assert.assertEquals(result.size(), 2);
+    }
+
+    @Test
+    public void testGetUserSessionsWithEmptyResult() throws Exception {
+
+        // Create test user.
+        AuthenticatedUser authenticatedUser = new AuthenticatedUser();
+        authenticatedUser.setUserName(TEST_USER_NAME);
+        authenticatedUser.setTenantDomain(TENANT_DOMAIN_CARBON_SUPER);
+        authenticatedUser.setUserStoreDomain(USER_STORE_DOMAIN_PRIMARY);
+        authenticatedUser.setUserId(TEST_USER_ID_456);
+
+        JsAuthenticatedUser jsUser = new JsGraalAuthenticatedUser(authenticatedUser);
+
+        // Mock empty session list.
+        when(userSessionManagementService.getSessionsByUserId(TEST_USER_ID_456, TENANT_DOMAIN_CARBON_SUPER))
+                .thenReturn(new ArrayList<>());
+
+        // Execute the function.
+        List<JsUserSession> result = getUserSessionsFunction.getUserSessions(jsUser);
+
+        // Verify results.
+        Assert.assertNotNull(result);
+        Assert.assertEquals(result.size(), 0);
+    }
+
+    @Test
+    public void testGetUserSessionsWithOrganization() throws Exception {
+
+        // Create test user with accessing organization.
+        AuthenticatedUser authenticatedUser = new AuthenticatedUser();
+        authenticatedUser.setUserName(TEST_USER_NAME);
+        authenticatedUser.setTenantDomain(TENANT_DOMAIN_CARBON_SUPER);
+        authenticatedUser.setUserStoreDomain(USER_STORE_DOMAIN_PRIMARY);
+        authenticatedUser.setUserId(TEST_USER_ID_ORG_123);
+        authenticatedUser.setAccessingOrganization(ORG_ID_123);
+
+        JsAuthenticatedUser jsUser = new JsGraalAuthenticatedUser(authenticatedUser);
+
+        // Mock organization manager.
+        when(organizationManager.resolveTenantDomain(ORG_ID_123)).thenReturn(ORG_TENANT_DOMAIN);
+
+        // Create mock user sessions
+        List<UserSession> mockSessions = new ArrayList<>();
+        UserSession session = mock(UserSession.class);
+        mockSessions.add(session);
+
+        // Mock the session management service with resolved tenant domain.
+        when(userSessionManagementService.getSessionsByUserId(TEST_USER_ID_ORG_123, ORG_TENANT_DOMAIN))
+                .thenReturn(mockSessions);
+
+        // Execute the function.
+        List<JsUserSession> result = getUserSessionsFunction.getUserSessions(jsUser);
+
+        // Verify results.
+        Assert.assertNotNull(result);
+        Assert.assertEquals(result.size(), 1);
+    }
+
+    @Test
+    public void testGetUserSessionsWithOrganizationManagementException() throws Exception {
+
+        // Create test user with accessing organization.
+        AuthenticatedUser authenticatedUser = new AuthenticatedUser();
+        authenticatedUser.setUserName(TEST_USER_NAME);
+        authenticatedUser.setTenantDomain(TENANT_DOMAIN_CARBON_SUPER);
+        authenticatedUser.setUserStoreDomain(USER_STORE_DOMAIN_PRIMARY);
+        authenticatedUser.setUserId(TEST_USER_ID_ORG_456);
+        authenticatedUser.setAccessingOrganization(INVALID_ORG);
+
+        JsAuthenticatedUser jsUser = new JsGraalAuthenticatedUser(authenticatedUser);
+
+        // Mock organization manager to throw exception.
+        when(organizationManager.resolveTenantDomain(INVALID_ORG))
+                .thenThrow(new OrganizationManagementException(ORG_NOT_FOUND_MESSAGE));
+
+        // Execute the function.
+        List<JsUserSession> result = getUserSessionsFunction.getUserSessions(jsUser);
+        Assert.assertNull(result);
+    }
+}

pom.xml

@@ -246,6 +246,11 @@
                 <artifactId>nimbus-jose-jwt</artifactId>
                 <version>${nimbusds.version}</version>
             </dependency>
+            <dependency>
+                <groupId>org.wso2.carbon.identity.organization.management.core</groupId>
+                <artifactId>org.wso2.carbon.identity.organization.management.service</artifactId>
+                <version>${identity.organization.management.core.version}</version>
+            </dependency>
 
             <dependency>
                 <groupId>org.jacoco</groupId>
@@ -321,12 +326,6 @@
                 <artifactId>org.wso2.carbon.crypto.impl</artifactId>
                 <version>${carbon.crypto.version}</version>
             </dependency>
-            <dependency>
-                <groupId>org.wso2.carbon.identity.organization.management.core</groupId>
-                <artifactId>org.wso2.carbon.identity.organization.management.service</artifactId>
-                <version>${identity.organization.management.core.version}</version>
-                <scope>test</scope>
-            </dependency>
             <dependency>
                 <groupId>org.jacoco</groupId>
                 <artifactId>jacoco-maven-plugin</artifactId>
@@ -530,6 +529,8 @@
         <carbon.user.package.import.version.range>[1.0.1, 2.0.0)</carbon.user.package.import.version.range>
         <carbon.identity.framework.version>7.7.34</carbon.identity.framework.version>
         <identity.organization.management.core.version>1.0.89</identity.organization.management.core.version>
+        <org.wso2.carbon.identity.organization.management.core.version.range>[1.0.0, 2.0.0)
+        </org.wso2.carbon.identity.organization.management.core.version.range>
         <carbon.identity.framework.testutils.version>5.20.447</carbon.identity.framework.testutils.version>
         <carbon.identity.package.import.version.range>[5.14.0, 8.0.0)</carbon.identity.package.import.version.range>
         <identity.governance.version>1.1.17</identity.governance.version>