feat(core): organization jit sso

packages/core/src/libraries/user.test.ts

@@ -33,7 +33,8 @@ mockEsm('#src/utils/password.js', () => ({
 }));
 
 const { MockQueries } = await import('#src/test-utils/tenant.js');
-const { encryptUserPassword, createUserLibrary } = await import('./user.js');
+const { createUserLibrary } = await import('./user.js');
+const { encryptUserPassword } = await import('./user.utils.js');
 
 const hasUserWithId = jest.fn();
 const updateUserById = jest.fn();

packages/core/src/libraries/user.ts

@@ -1,7 +1,7 @@
-import type { BindMfa, CreateUser, MfaVerification, Scope, User } from '@logto/schemas';
-import { MfaFactor, RoleType, Users, UsersPasswordEncryptionMethod } from '@logto/schemas';
+import type { BindMfa, CreateUser, Scope, User } from '@logto/schemas';
+import { RoleType, Users, UsersPasswordEncryptionMethod } from '@logto/schemas';
 import { generateStandardId, generateStandardShortId } from '@logto/shared';
-import { deduplicateByKey, type Nullable } from '@silverhand/essentials';
+import { condArray, deduplicateByKey, type Nullable } from '@silverhand/essentials';
 import { argon2Verify, bcryptVerify, md5, sha1, sha256 } from 'hash-wasm';
 import pRetry from 'p-retry';
 
@@ -9,75 +9,14 @@ import { buildInsertIntoWithPool } from '#src/database/insert-into.js';
 import { EnvSet } from '#src/env-set/index.js';
 import RequestError from '#src/errors/RequestError/index.js';
 import { type JitOrganization } from '#src/queries/organization/email-domains.js';
-import OrganizationQueries from '#src/queries/organization/index.js';
 import { createUsersRolesQueries } from '#src/queries/users-roles.js';
 import type Queries from '#src/tenants/Queries.js';
 import assertThat from '#src/utils/assert-that.js';
-import { encryptPassword } from '#src/utils/password.js';
 import type { OmitAutoSetFields } from '#src/utils/sql.js';
 
-export const encryptUserPassword = async (
-  password: string
-): Promise<{
-  passwordEncrypted: string;
-  passwordEncryptionMethod: UsersPasswordEncryptionMethod;
-}> => {
-  const passwordEncryptionMethod = UsersPasswordEncryptionMethod.Argon2i;
-  const passwordEncrypted = await encryptPassword(password, passwordEncryptionMethod);
+import { convertBindMfaToMfaVerification, encryptUserPassword } from './user.utils.js';
 
-  return { passwordEncrypted, passwordEncryptionMethod };
-};
-
-/**
- * Convert bindMfa to mfaVerification, add common fields like "id" and "createdAt"
- * and transpile formats like "codes" to "code" for backup code
- */
-const converBindMfaToMfaVerification = (bindMfa: BindMfa): MfaVerification => {
-  const { type } = bindMfa;
-  const base = {
-    id: generateStandardId(),
-    createdAt: new Date().toISOString(),
-  };
-
-  if (type === MfaFactor.BackupCode) {
-    const { codes } = bindMfa;
-
-    return {
-      ...base,
-      type,
-      codes: codes.map((code) => ({ code })),
-    };
-  }
-
-  if (type === MfaFactor.TOTP) {
-    const { secret } = bindMfa;
-
-    return {
-      ...base,
-      type,
-      key: secret,
-    };
-  }
-
-  const { credentialId, counter, publicKey, transports, agent } = bindMfa;
-  return {
-    ...base,
-    type,
-    credentialId,
-    counter,
-    publicKey,
-    transports,
-    agent,
-  };
-};
-
-export type InsertUserResult = [
-  User,
-  {
-    /** The organizations and organization roles that the user has been provisioned into. */
-    organizations: readonly JitOrganization[];
-  },
-];
+export type InsertUserResult = [User];
 
 export type UserLibrary = ReturnType<typeof createUserLibrary>;
 
@@ -143,43 +82,7 @@ export const createUserLibrary = (queries: Queries) => {
         );
       }
 
-      // TODO: If the user's email is not verified, we should not provision the user into any organization.
-      const provisionOrganizations = async (): Promise<readonly JitOrganization[]> => {
-        // Just-in-time organization provisioning
-        const userEmailDomain = data.primaryEmail?.split('@')[1];
-        if (userEmailDomain) {
-          const organizationQueries = new OrganizationQueries(connection);
-          const organizations = await organizationQueries.jit.emailDomains.getJitOrganizations(
-            userEmailDomain
-          );
-
-          if (organizations.length > 0) {
-            await organizationQueries.relations.users.insert(
-              ...organizations.map(({ organizationId }) => ({
-                organizationId,
-                userId: user.id,
-              }))
-            );
-
-            const data = organizations.flatMap(({ organizationId, organizationRoleIds }) =>
-              organizationRoleIds.map((organizationRoleId) => ({
-                organizationId,
-                organizationRoleId,
-                userId: user.id,
-              }))
-            );
-            if (data.length > 0) {
-              await organizationQueries.relations.rolesUsers.insert(...data);
-            }
-
-            return organizations;
-          }
-        }
-
-        return [];
-      };
-
-      return [user, { organizations: await provisionOrganizations() }];
+      return [user];
     });
   };
 
@@ -261,7 +164,7 @@ export const createUserLibrary = (queries: Queries) => {
     // TODO @sijie use jsonb array append
     const { mfaVerifications } = await findUserById(userId);
     await updateUserById(userId, {
-      mfaVerifications: [...mfaVerifications, converBindMfaToMfaVerification(payload)],
+      mfaVerifications: [...mfaVerifications, convertBindMfaToMfaVerification(payload)],
     });
   };
 
@@ -338,6 +241,66 @@ export const createUserLibrary = (queries: Queries) => {
   const findUserSsoIdentities = async (userId: string) =>
     userSsoIdentities.findUserSsoIdentitiesByUserId(userId);
 
+  type ProvisionOrganizationsParams =
+    | {
+        /** The user ID to provision organizations for. */
+        userId: string;
+        /** The user's email to determine JIT organizations. */
+        email: string;
+        /** The SSO connector ID to determine JIT organizations. */
+        ssoConnectorId?: undefined;
+      }
+    | {
+        /** The user ID to provision organizations for. */
+        userId: string;
+        /** The user's email to determine JIT organizations. */
+        email?: undefined;
+        /** The SSO connector ID to determine JIT organizations. */
+        ssoConnectorId: string;
+      };
+
+  // TODO: If the user's email is not verified, we should not provision the user into any organization.
+  /**
+   * Provision the user with JIT organizations and roles based on the user's email domain and the
+   * enterprise SSO connector.
+   */
+  const provisionOrganizations = async ({
+    userId,
+    email,
+    ssoConnectorId,
+  }: ProvisionOrganizationsParams): Promise<readonly JitOrganization[]> => {
+    const userEmailDomain = email?.split('@')[1];
+    const jitOrganizations = condArray(
+      userEmailDomain &&
+        (await organizations.jit.emailDomains.getJitOrganizations(userEmailDomain)),
+      ssoConnectorId && (await organizations.jit.ssoConnectors.getJitOrganizations(ssoConnectorId))
+    );
+
+    if (jitOrganizations.length === 0) {
+      return [];
+    }
+
+    await organizations.relations.users.insert(
+      ...jitOrganizations.map(({ organizationId }) => ({
+        organizationId,
+        userId,
+      }))
+    );
+
+    const data = jitOrganizations.flatMap(({ organizationId, organizationRoleIds }) =>
+      organizationRoleIds.map((organizationRoleId) => ({
+        organizationId,
+        organizationRoleId,
+        userId,
+      }))
+    );
+    if (data.length > 0) {
+      await organizations.relations.rolesUsers.insert(...data);
+    }
+
+    return jitOrganizations;
+  };
+
   return {
     generateUserId,
     insertUser,
@@ -349,5 +312,6 @@ export const createUserLibrary = (queries: Queries) => {
     verifyUserPassword,
     signOutUser,
     findUserSsoIdentities,
+    provisionOrganizations,
   };
 };

packages/core/src/libraries/user.utils.ts

@@ -0,0 +1,60 @@
+import type { BindMfa, MfaVerification } from '@logto/schemas';
+import { MfaFactor, UsersPasswordEncryptionMethod } from '@logto/schemas';
+import { generateStandardId } from '@logto/shared';
+
+import { encryptPassword } from '#src/utils/password.js';
+
+export const encryptUserPassword = async (
+  password: string
+): Promise<{
+  passwordEncrypted: string;
+  passwordEncryptionMethod: UsersPasswordEncryptionMethod;
+}> => {
+  const passwordEncryptionMethod = UsersPasswordEncryptionMethod.Argon2i;
+  const passwordEncrypted = await encryptPassword(password, passwordEncryptionMethod);
+
+  return { passwordEncrypted, passwordEncryptionMethod };
+};
+
+/**
+ * Convert bindMfa to mfaVerification, add common fields like "id" and "createdAt"
+ * and transpile formats like "codes" to "code" for backup code
+ */
+export const convertBindMfaToMfaVerification = (bindMfa: BindMfa): MfaVerification => {
+  const { type } = bindMfa;
+  const base = {
+    id: generateStandardId(),
+    createdAt: new Date().toISOString(),
+  };
+
+  if (type === MfaFactor.BackupCode) {
+    const { codes } = bindMfa;
+
+    return {
+      ...base,
+      type,
+      codes: codes.map((code) => ({ code })),
+    };
+  }
+
+  if (type === MfaFactor.TOTP) {
+    const { secret } = bindMfa;
+
+    return {
+      ...base,
+      type,
+      key: secret,
+    };
+  }
+
+  const { credentialId, counter, publicKey, transports, agent } = bindMfa;
+  return {
+    ...base,
+    type,
+    credentialId,
+    counter,
+    publicKey,
+    transports,
+    agent,
+  };
+};

packages/core/src/queries/organization/email-domains.ts

@@ -54,7 +54,11 @@ export class EmailDomainQueries {
    * Given an email domain, return the organizations and organization roles that need to be
    * provisioned.
    */
-  async getJitOrganizations(emailDomain: string): Promise<readonly JitOrganization[]> {
+  async getJitOrganizations(emailDomain?: string): Promise<readonly JitOrganization[]> {
+    if (!emailDomain) {
+      return [];
+    }
+
     const { fields } = convertToIdentifiers(OrganizationJitEmailDomains, true);
     const organizationJitRoles = convertToIdentifiers(OrganizationJitRoles, true);
     return this.pool.any<JitOrganization>(sql`

packages/core/src/queries/organization/index.ts

@@ -24,8 +24,6 @@ import {
   OrganizationJitRoles,
   OrganizationApplicationRelations,
   Applications,
-  OrganizationJitSsoConnectors,
-  SsoConnectors,
 } from '@logto/schemas';
 import { sql, type CommonQueryMethods } from '@silverhand/slonik';
 
@@ -36,6 +34,7 @@ import { conditionalSql, convertToIdentifiers } from '#src/utils/sql.js';
 
 import { EmailDomainQueries } from './email-domains.js';
 import { RoleUserRelationQueries } from './role-user-relations.js';
+import { SsoConnectorQueries } from './sso-connectors.js';
 import { UserRelationQueries } from './user-relations.js';
 
 /**
@@ -311,12 +310,7 @@ export default class OrganizationQueries extends SchemaQueries<
       Organizations,
       OrganizationRoles
     ),
-    ssoConnectors: new TwoRelationsQueries(
-      this.pool,
-      OrganizationJitSsoConnectors.table,
-      Organizations,
-      SsoConnectors
-    ),
+    ssoConnectors: new SsoConnectorQueries(this.pool),
   };
 
   constructor(pool: CommonQueryMethods) {

packages/core/src/queries/organization/sso-connectors.ts

@@ -0,0 +1,45 @@
+import {
+  OrganizationJitRoles,
+  OrganizationJitSsoConnectors,
+  Organizations,
+  SsoConnectors,
+} from '@logto/schemas';
+import { type CommonQueryMethods, sql } from '@silverhand/slonik';
+
+import { TwoRelationsQueries } from '#src/utils/RelationQueries.js';
+import { convertToIdentifiers } from '#src/utils/sql.js';
+
+import { type JitOrganization } from './email-domains.js';
+
+const { table, fields } = convertToIdentifiers(OrganizationJitSsoConnectors);
+
+export class SsoConnectorQueries extends TwoRelationsQueries<
+  typeof Organizations,
+  typeof SsoConnectors
+> {
+  constructor(pool: CommonQueryMethods) {
+    super(pool, OrganizationJitSsoConnectors.table, Organizations, SsoConnectors);
+  }
+
+  async getJitOrganizations(ssoConnectorId?: string): Promise<readonly JitOrganization[]> {
+    if (!ssoConnectorId) {
+      return [];
+    }
+
+    const { fields } = convertToIdentifiers(OrganizationJitSsoConnectors, true);
+    const organizationJitRoles = convertToIdentifiers(OrganizationJitRoles, true);
+    return this.pool.any<JitOrganization>(sql`
+      select
+        ${fields.organizationId},
+        array_remove(
+          array_agg(${organizationJitRoles.fields.organizationRoleId}),
+          null
+        ) as "organizationRoleIds"
+      from ${table}
+      left join ${organizationJitRoles.table}
+        on ${fields.organizationId} = ${organizationJitRoles.fields.organizationId}
+      where ${fields.ssoConnectorId} = ${ssoConnectorId}
+      group by ${fields.organizationId}
+    `);
+  }
+}

packages/core/src/routes-me/user.ts

@@ -4,7 +4,7 @@ import { conditional, pick } from '@silverhand/essentials';
 import { literal, object, string } from 'zod';
 
 import RequestError from '#src/errors/RequestError/index.js';
-import { encryptUserPassword } from '#src/libraries/user.js';
+import { encryptUserPassword } from '#src/libraries/user.utils.js';
 import koaGuard from '#src/middleware/koa-guard.js';
 import assertThat from '#src/utils/assert-that.js';
 

packages/core/src/routes/admin-user/basics.test.ts

@@ -70,7 +70,7 @@ const mockHasUserWithPhone = jest.fn(async () => false);
 const { hasUser, findUserById, updateUserById, deleteUserIdentity, deleteUserById } =
   mockedQueries.users;
 
-const { encryptUserPassword } = await mockEsmWithActual('#src/libraries/user.js', () => ({
+const { encryptUserPassword } = await mockEsmWithActual('#src/libraries/user.utils.js', () => ({
   encryptUserPassword: jest.fn(() => ({
     passwordEncrypted: 'password',
     passwordEncryptionMethod: 'Argon2i',
@@ -87,7 +87,6 @@ const usersLibraries = {
         ...mockUser,
         ...removeUndefinedKeys(user), // No undefined values will be returned from database
       },
-      { organizations: [] },
     ]
   ),
   verifyUserPassword,

packages/core/src/routes/admin-user/basics.ts

@@ -12,7 +12,7 @@ import { boolean, literal, nativeEnum, object, string } from 'zod';
 
 import RequestError from '#src/errors/RequestError/index.js';
 import { buildManagementApiContext } from '#src/libraries/hook/utils.js';
-import { encryptUserPassword } from '#src/libraries/user.js';
+import { encryptUserPassword } from '#src/libraries/user.utils.js';
 import koaGuard from '#src/middleware/koa-guard.js';
 import assertThat from '#src/utils/assert-that.js';
 
@@ -200,7 +200,7 @@ export default function adminUserBasicsRoutes<T extends ManagementApiRouter>(
 
       const id = await generateUserId();
 
-      const [user, { organizations }] = await insertUser(
+      const [user] = await insertUser(
         {
           id,
           primaryEmail,
@@ -221,13 +221,6 @@ export default function adminUserBasicsRoutes<T extends ManagementApiRouter>(
         []
       );
 
-      for (const { organizationId } of organizations) {
-        ctx.appendDataHookContext('Organization.Membership.Updated', {
-          ...buildManagementApiContext(ctx),
-          organizationId,
-        });
-      }
-
       ctx.body = pick(user, ...userInfoSelectFields);
       return next();
     }

packages/core/src/routes/admin-user/mfa-verifications.test.ts

@@ -55,7 +55,6 @@ const usersLibraries = {
         ...mockUser,
         ...removeUndefinedKeys(user), // No undefined values will be returned from database
       },
-      { organizations: [] },
     ]
   ),
 } satisfies Partial<Libraries['users']>;

packages/core/src/routes/admin-user/search.test.ts

@@ -58,7 +58,6 @@ const usersLibraries = {
         ...mockUser,
         ...removeUndefinedKeys(user), // No undefined values will be returned from database
       },
-      { organizations: [] },
     ]
   ),
 } satisfies Partial<Libraries['users']>;

packages/core/src/routes/admin-user/social.test.ts

@@ -54,7 +54,6 @@ const usersLibraries = {
         ...mockUser,
         ...removeUndefinedKeys(user), // No undefined values will be returned from database
       },
-      { organizations: [] },
     ]
   ),
 } satisfies Partial<Libraries['users']>;

packages/core/src/routes/interaction/actions/helpers.ts

@@ -6,7 +6,7 @@ import { type IRouterContext } from 'koa-router';
 import { EnvSet } from '#src/env-set/index.js';
 import { type CloudConnectionLibrary } from '#src/libraries/cloud-connection.js';
 import { type ConnectorLibrary } from '#src/libraries/connector.js';
-import { encryptUserPassword } from '#src/libraries/user.js';
+import { encryptUserPassword } from '#src/libraries/user.utils.js';
 import type Queries from '#src/tenants/Queries.js';
 import type TenantContext from '#src/tenants/TenantContext.js';
 import { getConsoleLogFromContext } from '#src/utils/console.js';

packages/core/src/routes/interaction/actions/submit-interaction.mfa.test.ts

@@ -24,7 +24,7 @@ const { assignInteractionResults } = mockEsm('#src/libraries/session.js', () =>
   assignInteractionResults: jest.fn(),
 }));
 
-mockEsm('#src/libraries/user.js', () => ({
+mockEsm('#src/libraries/user.utils.js', () => ({
   encryptUserPassword: jest.fn().mockResolvedValue({
     passwordEncrypted: 'passwordEncrypted',
     passwordEncryptionMethod: 'plain',

packages/core/src/routes/interaction/actions/submit-interaction.test.ts

@@ -1,4 +1,3 @@
-/* eslint-disable max-lines */
 import {
   InteractionEvent,
   adminConsoleApplicationId,
@@ -33,7 +32,7 @@ const { assignInteractionResults } = mockEsm('#src/libraries/session.js', () =>
   assignInteractionResults: jest.fn(),
 }));
 
-const { encryptUserPassword } = mockEsm('#src/libraries/user.js', () => ({
+const { encryptUserPassword } = mockEsm('#src/libraries/user.utils.js', () => ({
   encryptUserPassword: jest.fn().mockResolvedValue({
     passwordEncrypted: 'passwordEncrypted',
     passwordEncryptionMethod: 'plain',
@@ -63,9 +62,7 @@ const { hasActiveUsers, updateUserById, hasUserWithEmail, hasUserWithPhone } = u
 
 const userLibraries = {
   generateUserId: jest.fn().mockResolvedValue('uid'),
-  insertUser: jest.fn(
-    async (user: CreateUser): Promise<InsertUserResult> => [user as User, { organizations: [] }]
-  ),
+  insertUser: jest.fn(async (user: CreateUser): Promise<InsertUserResult> => [user as User]),
 };
 const { generateUserId, insertUser } = userLibraries;
 
@@ -465,4 +462,3 @@ describe('submit action', () => {
     });
   });
 });
-/* eslint-enable max-lines */

packages/core/src/routes/interaction/actions/submit-interaction.ts

@@ -21,7 +21,7 @@ import { conditional, conditionalArray, trySafe } from '@silverhand/essentials';
 
 import { EnvSet } from '#src/env-set/index.js';
 import { assignInteractionResults } from '#src/libraries/session.js';
-import { encryptUserPassword } from '#src/libraries/user.js';
+import { encryptUserPassword } from '#src/libraries/user.utils.js';
 import type { LogEntry, WithLogContext } from '#src/middleware/koa-audit-log.js';
 import type TenantContext from '#src/tenants/TenantContext.js';
 import { getConsoleLogFromContext } from '#src/utils/console.js';
@@ -135,7 +135,7 @@ async function handleSubmitRegister(
     (invitation) => invitation.status === OrganizationInvitationStatus.Pending
   );
 
-  const [user, { organizations: provisionedOrganizations }] = await insertUser(
+  const [user] = await insertUser(
     {
       id,
       ...userProfile,
@@ -190,10 +190,18 @@ async function handleSubmitRegister(
   ctx.assignInteractionHookResult({ userId: id });
   ctx.appendDataHookContext('User.Created', { user });
 
-  for (const { organizationId } of provisionedOrganizations) {
-    ctx.appendDataHookContext('Organization.Membership.Updated', {
-      organizationId,
+  // JIT provisioning for email domain
+  if (user.primaryEmail) {
+    const provisionedOrganizations = await libraries.users.provisionOrganizations({
+      userId: id,
+      email: user.primaryEmail,
     });
+
+    for (const { organizationId } of provisionedOrganizations) {
+      ctx.appendDataHookContext('Organization.Membership.Updated', {
+        organizationId,
+      });
+    }
   }
 
   log?.append({ userId: id });

packages/core/src/routes/interaction/utils/single-sign-on.ts

@@ -35,7 +35,7 @@ type AuthorizationUrlPayload = z.infer<typeof authorizationUrlPayloadGuard>;
 
 // Get the authorization url for the SSO provider
 export const getSsoAuthorizationUrl = async (
-  ctx: WithLogContext & WithInteractionDetailsContext,
+  ctx: WithInteractionDetailsContext<WithLogContext>,
   { provider, id: tenantId }: TenantContext,
   connectorData: SupportedSsoConnector,
   payload: AuthorizationUrlPayload
@@ -84,7 +84,7 @@ type SsoAuthenticationResult = {
 
 // Get the user authentication result from the SSO provider
 export const getSsoAuthentication = async (
-  ctx: WithLogContext,
+  ctx: WithInteractionHooksContext<WithLogContext>,
   { provider, id: tenantId }: TenantContext,
   connectorData: SupportedSsoConnector,
   data: Record<string, unknown>
@@ -133,7 +133,7 @@ export const getSsoAuthentication = async (
 
 // Handle the SSO authentication result and return the user id
 export const handleSsoAuthentication = async (
-  ctx: WithLogContext,
+  ctx: WithInteractionHooksContext<WithLogContext>,
   tenant: TenantContext,
   connectorData: SupportedSsoConnector,
   ssoAuthentication: SsoAuthenticationResult
@@ -161,7 +161,7 @@ export const handleSsoAuthentication = async (
 
   // SignIn and link with existing user account with a same email
   if (user) {
-    return signInAndLinkWithSsoAuthentication(ctx, queries, {
+    return signInAndLinkWithSsoAuthentication(ctx, tenant, {
       connectorData,
       user,
       ssoAuthentication,
@@ -178,7 +178,7 @@ export const handleSsoAuthentication = async (
 };
 
 const signInWithSsoAuthentication = async (
-  ctx: WithLogContext,
+  ctx: WithInteractionHooksContext<WithLogContext>,
   { userSsoIdentities: userSsoIdentitiesQueries, users: usersQueries }: Queries,
   {
     connectorData: { id: connectorId, syncProfile },
@@ -232,8 +232,11 @@ const signInWithSsoAuthentication = async (
 };
 
 const signInAndLinkWithSsoAuthentication = async (
-  ctx: WithLogContext,
-  { userSsoIdentities: userSsoIdentitiesQueries, users: usersQueries }: Queries,
+  ctx: WithInteractionHooksContext<WithLogContext>,
+  {
+    queries: { userSsoIdentities: userSsoIdentitiesQueries, users: usersQueries },
+    libraries: { users: usersLibrary },
+  }: TenantContext,
   {
     connectorData: { id: connectorId, syncProfile },
     user: { id: userId },
@@ -273,6 +276,18 @@ const signInAndLinkWithSsoAuthentication = async (
     lastSignInAt: Date.now(),
   });
 
+  // JIT provision for existing users signing in with SSO for the first time
+  const provisionedOrganizations = await usersLibrary.provisionOrganizations({
+    userId,
+    ssoConnectorId: connectorId,
+  });
+
+  for (const { organizationId } of provisionedOrganizations) {
+    ctx.appendDataHookContext('Organization.Membership.Updated', {
+      organizationId,
+    });
+  }
+
   log.append({
     userId,
     interaction: {
@@ -309,7 +324,7 @@ export const registerWithSsoAuthentication = async (
   };
 
   // Insert new user
-  const [user, { organizations }] = await usersLibrary.insertUser(
+  const [user] = await usersLibrary.insertUser(
     {
       id: await usersLibrary.generateUserId(),
       ...syncingProfile,
@@ -317,11 +332,6 @@ export const registerWithSsoAuthentication = async (
     },
     []
   );
-  for (const { organizationId } of organizations) {
-    ctx.appendDataHookContext('Organization.Membership.Updated', {
-      organizationId,
-    });
-  }
 
   const { id: userId } = user;
 
@@ -335,6 +345,18 @@ export const registerWithSsoAuthentication = async (
     detail: userInfo,
   });
 
+  // JIT provision for new users signing up with SSO
+  const provisionedOrganizations = await usersLibrary.provisionOrganizations({
+    userId: user.id,
+    ssoConnectorId: connectorId,
+  });
+
+  for (const { organizationId } of provisionedOrganizations) {
+    ctx.appendDataHookContext('Organization.Membership.Updated', {
+      organizationId,
+    });
+  }
+
   log.append({
     userId,
     interaction: {

packages/integration-tests/src/helpers/interactions.ts

@@ -77,6 +77,28 @@ export const registerWithEmail = async (email: string) => {
   return { client, id };
 };
 
+export const signInWithEmail = async (email: string) => {
+  const client = await initClient();
+
+  await client.successSend(putInteraction, {
+    event: InteractionEvent.SignIn,
+  });
+  await client.successSend(sendVerificationCode, {
+    email,
+  });
+
+  const { code } = await readConnectorMessage('Email');
+
+  await client.successSend(patchInteractionIdentifiers, {
+    email,
+    verificationCode: code,
+  });
+
+  const { redirectTo } = await client.submitInteraction();
+  const id = await processSession(client, redirectTo);
+  return { client, id };
+};
+
 export const signInWithPassword = async (
   payload: UsernamePasswordPayload | EmailPasswordPayload | PhonePasswordPayload
 ) => {

packages/integration-tests/src/tests/api/admin-user.organization-jit.test.ts

@@ -11,7 +11,7 @@ describe('organization just-in-time provisioning', () => {
     await Promise.all([organizationApi.cleanUp(), userApi.cleanUp()]);
   });
 
-  it('should automatically provision a user to the organizations with roles', async () => {
+  it('should not automatically provision a user to the organizations when email domain matches', async () => {
     const organizations = await Promise.all([
       organizationApi.create({ name: 'foo' }),
       organizationApi.create({ name: 'bar' }),
@@ -36,60 +36,6 @@ describe('organization just-in-time provisioning', () => {
     const { id } = await userApi.create({ primaryEmail: email });
 
     const userOrganizations = await getUserOrganizations(id);
-    expect(userOrganizations).toEqual(
-      expect.arrayContaining([
-        expect.objectContaining({
-          id: organizations[0].id,
-          // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-          organizationRoles: expect.arrayContaining([
-            expect.objectContaining({ id: roles[0].id }),
-            expect.objectContaining({ id: roles[1].id }),
-          ]),
-        }),
-        expect.objectContaining({
-          id: organizations[1].id,
-          organizationRoles: [expect.objectContaining({ id: roles[0].id })],
-        }),
-        expect.objectContaining({
-          id: organizations[2].id,
-          organizationRoles: [],
-        }),
-      ])
-    );
-  });
-
-  it('should automatically provision a user to the organizations without roles', async () => {
-    const organizations = await Promise.all([
-      organizationApi.create({ name: 'foo' }),
-      organizationApi.create({ name: 'bar' }),
-      organizationApi.create({ name: 'baz' }),
-    ]);
-    const emailDomain = 'foo.com';
-    await Promise.all(
-      organizations.map(async (organization) =>
-        organizationApi.jit.addEmailDomain(organization.id, emailDomain)
-      )
-    );
-
-    const email = randomString() + '@' + emailDomain;
-    const { id } = await userApi.create({ primaryEmail: email });
-
-    const userOrganizations = await getUserOrganizations(id);
-    expect(userOrganizations).toEqual(
-      expect.arrayContaining([
-        expect.objectContaining({
-          id: organizations[0].id,
-          organizationRoles: [],
-        }),
-        expect.objectContaining({
-          id: organizations[1].id,
-          organizationRoles: [],
-        }),
-        expect.objectContaining({
-          id: organizations[2].id,
-          organizationRoles: [],
-        }),
-      ])
-    );
+    expect(userOrganizations).toEqual([]);
   });
 });

packages/integration-tests/src/tests/api/hook/hook.trigger.custom.data.test.ts

@@ -129,15 +129,6 @@ describe('manual data hook tests', () => {
     });
 
   describe('organization membership update by just-in-time organization provisioning', () => {
-    it('should trigger `Organization.Membership.Updated` event when user is provisioned by Management API', async () => {
-      const organization = await organizationApi.create({ name: 'foo' });
-      const domain = 'example.com';
-      await organizationApi.jit.addEmailDomain(organization.id, domain);
-
-      await userApi.create({ primaryEmail: `${randomString()}@${domain}` });
-      await assertOrganizationMembershipUpdated(organization.id);
-    });
-
     // TODO: Add user deletion test case
 
     it('should trigger `Organization.Membership.Updated` event when user is provisioned by experience', async () => {
@@ -160,9 +151,9 @@ describe('manual data hook tests', () => {
     it('should trigger `Organization.Membership.Updated` event when user is provisioned by SSO', async () => {
       const organization = await organizationApi.create({ name: 'bar' });
       const domain = 'sso_example.com';
-      await organizationApi.jit.addEmailDomain(organization.id, domain);
-
       const connector = await ssoConnectorApi.createMockOidcConnector([domain]);
+      await organizationApi.jit.ssoConnectors.add(organization.id, [connector.id]);
+
       await updateSignInExperience({
         singleSignOnEnabled: true,
       });

packages/integration-tests/src/tests/api/interaction/organization-jit.test.ts

@@ -14,13 +14,13 @@ import {
   setEmailConnector,
   setSmsConnector,
 } from '#src/helpers/connector.js';
-import { registerWithEmail } from '#src/helpers/interactions.js';
+import { registerWithEmail, signInWithEmail } from '#src/helpers/interactions.js';
 import { OrganizationApiTest } from '#src/helpers/organization.js';
 import {
   enableAllVerificationCodeSignInMethods,
   resetPasswordPolicy,
 } from '#src/helpers/sign-in-experience.js';
-import { registerNewUserWithSso } from '#src/helpers/single-sign-on.js';
+import { registerNewUserWithSso, signInWithSso } from '#src/helpers/single-sign-on.js';
 import { generateEmail, randomString } from '#src/utils.js';
 
 describe('organization just-in-time provisioning', () => {
@@ -45,7 +45,7 @@ describe('organization just-in-time provisioning', () => {
     });
   });
 
-  it('should automatically provision a user to the organization with roles', async () => {
+  const prepare = async (emailDomain = `foo-${randomString()}.com`, ssoConnectorId?: string) => {
     const organizations = await Promise.all([
       organizationApi.create({ name: 'foo' }),
       organizationApi.create({ name: 'bar' }),
@@ -55,47 +55,83 @@ describe('organization just-in-time provisioning', () => {
       organizationApi.roleApi.create({ name: randomString() }),
       organizationApi.roleApi.create({ name: randomString() }),
     ]);
-    const emailDomain = 'foo.com';
     await Promise.all(
-      organizations.map(async (organization) =>
-        organizationApi.jit.addEmailDomain(organization.id, emailDomain)
-      )
+      organizations.map(async (organization) => {
+        if (emailDomain) {
+          await organizationApi.jit.addEmailDomain(organization.id, emailDomain);
+        }
+
+        if (ssoConnectorId) {
+          await organizationApi.jit.ssoConnectors.add(organization.id, [ssoConnectorId]);
+        }
+      })
     );
     await Promise.all([
       organizationApi.jit.roles.add(organizations[0].id, [roles[0].id, roles[1].id]),
       organizationApi.jit.roles.add(organizations[1].id, [roles[0].id]),
     ]);
+    return {
+      organizations,
+      roles,
+      emailDomain,
+      expectOrganizations: () =>
+        expect.arrayContaining([
+          expect.objectContaining({
+            id: organizations[0].id,
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+            organizationRoles: expect.arrayContaining([
+              expect.objectContaining({ id: roles[0].id }),
+              expect.objectContaining({ id: roles[1].id }),
+            ]),
+          }),
+          expect.objectContaining({
+            id: organizations[1].id,
+            organizationRoles: [expect.objectContaining({ id: roles[0].id })],
+          }),
+          expect.objectContaining({
+            id: organizations[2].id,
+            organizationRoles: [],
+          }),
+        ]),
+    };
+  };
 
+  beforeEach(async () => {
+    await updateSignInExperience({
+      singleSignOnEnabled: false,
+    });
+  });
+
+  afterEach(async () => {
+    await Promise.all([organizationApi.cleanUp(), ssoConnectorApi.cleanUp()]);
+  });
+
+  beforeAll(async () => {
+    await clearConnectorsByTypes([ConnectorType.Email, ConnectorType.Sms]);
+    await Promise.all([setEmailConnector(), setSmsConnector()]);
+
+    await resetPasswordPolicy();
+    // Run it sequentially to avoid race condition
+    await enableAllVerificationCodeSignInMethods({
+      identifiers: [SignInIdentifier.Email],
+      password: false,
+      verify: true,
+    });
+  });
+
+  it('should automatically provision a user with roles', async () => {
+    const { emailDomain, expectOrganizations } = await prepare();
     const email = randomString() + '@' + emailDomain;
     const { client, id } = await registerWithEmail(email);
 
     const userOrganizations = await getUserOrganizations(id);
-    expect(userOrganizations).toEqual(
-      expect.arrayContaining([
-        expect.objectContaining({
-          id: organizations[0].id,
-          // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-          organizationRoles: expect.arrayContaining([
-            expect.objectContaining({ id: roles[0].id }),
-            expect.objectContaining({ id: roles[1].id }),
-          ]),
-        }),
-        expect.objectContaining({
-          id: organizations[1].id,
-          organizationRoles: [expect.objectContaining({ id: roles[0].id })],
-        }),
-        expect.objectContaining({
-          id: organizations[2].id,
-          organizationRoles: [],
-        }),
-      ])
-    );
+    expect(userOrganizations).toEqual(expectOrganizations());
 
     await logoutClient(client);
     await deleteUser(id);
   });
 
-  it('should automatically provision a user with the matched email to the organization from a SSO identity', async () => {
+  it('should not automatically provision a user with the matched email from a SSO identity', async () => {
     const organization = await organizationApi.create({ name: 'sso_foo' });
     const domain = 'sso_example.com';
     await organizationApi.jit.addEmailDomain(organization.id, domain);
@@ -113,6 +149,28 @@ describe('organization just-in-time provisioning', () => {
       },
     });
 
+    const userOrganizations = await getUserOrganizations(userId);
+    expect(userOrganizations).toEqual([]);
+    await deleteUser(userId);
+  });
+
+  it('should automatically provision a user with the matched SSO connector', async () => {
+    const organization = await organizationApi.create({ name: 'sso_foo' });
+    const domain = 'sso_example.com';
+    const connector = await ssoConnectorApi.createMockOidcConnector([domain]);
+    await organizationApi.jit.ssoConnectors.add(organization.id, [connector.id]);
+    await updateSignInExperience({
+      singleSignOnEnabled: true,
+    });
+
+    const userId = await registerNewUserWithSso(connector.id, {
+      authData: {
+        sub: randomString(),
+        email: generateEmail(domain),
+        email_verified: true,
+      },
+    });
+
     const userOrganizations = await getUserOrganizations(userId);
 
     expect(userOrganizations).toEqual(
@@ -121,4 +179,48 @@ describe('organization just-in-time provisioning', () => {
 
     await deleteUser(userId);
   });
+
+  it('should not automatically provision an existing user when the user is an existing user', async () => {
+    const emailDomain = `foo-${randomString()}.com`;
+    const email = randomString() + '@' + emailDomain;
+    const { client } = await registerWithEmail(email);
+    await client.signOut();
+
+    await prepare(emailDomain);
+    const { id } = await signInWithEmail(email);
+
+    const userOrganizations = await getUserOrganizations(id);
+    expect(userOrganizations).toEqual([]);
+    await deleteUser(id);
+  });
+
+  it('should automatically provision a user with new SSO identity', async () => {
+    const domain = `sso-${randomString()}.com`;
+    const email = randomString() + '@' + domain;
+
+    // No organization should be provisioned at this point
+    const { client, id: userId } = await registerWithEmail(email);
+    expect(await getUserOrganizations(userId)).toEqual([]);
+    await client.signOut();
+
+    // Configure the SSO connector
+    const connector = await ssoConnectorApi.createMockOidcConnector([domain]);
+    const { expectOrganizations } = await prepare(undefined, connector.id);
+    await updateSignInExperience({
+      singleSignOnEnabled: true,
+    });
+
+    await signInWithSso(connector.id, {
+      authData: {
+        sub: randomString(),
+        email,
+        email_verified: true,
+      },
+    });
+
+    const userOrganizations = await getUserOrganizations(userId);
+    expect(userOrganizations).toEqual(expectOrganizations());
+
+    await deleteUser(userId);
+  });
 });