feat(core): jit organization roles (#6049)

packages/core/src/libraries/user.ts

@@ -8,6 +8,7 @@ import pRetry from 'p-retry';
 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';
@@ -73,8 +74,8 @@ const converBindMfaToMfaVerification = (bindMfa: BindMfa): MfaVerification => {
 export type InsertUserResult = [
   User,
   {
-    /** The organization IDs that the user has been provisioned into. */
-    organizationIds: readonly string[];
+    /** The organizations and organization roles that the user has been provisioned into. */
+    organizations: readonly JitOrganization[];
   },
 ];
 
@@ -143,29 +144,42 @@ 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 string[]> => {
+      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 organizationIds =
-            await organizationQueries.jit.emailDomains.getOrganizationIdsByDomain(userEmailDomain);
+          const organizations = await organizationQueries.jit.emailDomains.getJitOrganizations(
+            userEmailDomain
+          );
 
-          if (organizationIds.length > 0) {
+          if (organizations.length > 0) {
             await organizationQueries.relations.users.insert(
-              ...organizationIds.map<[string, string]>((organizationId) => [
+              ...organizations.map<[string, string]>(({ organizationId }) => [
                 organizationId,
                 user.id,
               ])
             );
-            return organizationIds;
+
+            const data = organizations.flatMap(({ organizationId, organizationRoleIds }) =>
+              organizationRoleIds.map<[string, string, string]>((organizationRoleId) => [
+                organizationId,
+                organizationRoleId,
+                user.id,
+              ])
+            );
+            if (data.length > 0) {
+              await organizationQueries.relations.rolesUsers.insert(...data);
+            }
+
+            return organizations;
           }
         }
 
         return [];
       };
 
-      return [user, { organizationIds: await provisionOrganizations() }];
+      return [user, { organizations: await provisionOrganizations() }];
     });
   };
 

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

@@ -2,6 +2,7 @@ import {
   type OrganizationJitEmailDomain,
   OrganizationJitEmailDomains,
   type CreateOrganizationJitEmailDomain,
+  OrganizationJitRoles,
 } from '@logto/schemas';
 import { type CommonQueryMethods, sql } from '@silverhand/slonik';
 
@@ -49,13 +50,26 @@ export class EmailDomainQueries {
     return [Number(count), rows];
   }
 
-  async getOrganizationIdsByDomain(emailDomain: string): Promise<readonly string[]> {
-    const rows = await this.pool.any<Pick<OrganizationJitEmailDomain, 'organizationId'>>(sql`
-      select ${fields.organizationId}
+  /**
+   * Given an email domain, return the organizations and organization roles that need to be
+   * provisioned.
+   */
+  async getJitOrganizations(emailDomain: string): Promise<readonly JitOrganization[]> {
+    const { fields } = convertToIdentifiers(OrganizationJitEmailDomains, 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.emailDomain} = ${emailDomain}
+      group by ${fields.organizationId}
     `);
-    return rows.map((row) => row.organizationId);
   }
 
   async insert(organizationId: string, emailDomain: string): Promise<OrganizationJitEmailDomain> {
@@ -111,3 +125,8 @@ export class EmailDomainQueries {
     });
   }
 }
+
+export type JitOrganization = {
+  organizationId: string;
+  organizationRoleIds: string[];
+};

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

@@ -87,7 +87,7 @@ const usersLibraries = {
         ...mockUser,
         ...removeUndefinedKeys(user), // No undefined values will be returned from database
       },
-      { organizationIds: [] },
+      { organizations: [] },
     ]
   ),
   verifyUserPassword,

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

@@ -200,7 +200,7 @@ export default function adminUserBasicsRoutes<T extends ManagementApiRouter>(
 
       const id = await generateUserId();
 
-      const [user, { organizationIds }] = await insertUser(
+      const [user, { organizations }] = await insertUser(
         {
           id,
           primaryEmail,
@@ -221,7 +221,7 @@ export default function adminUserBasicsRoutes<T extends ManagementApiRouter>(
         []
       );
 
-      for (const organizationId of organizationIds) {
+      for (const { organizationId } of organizations) {
         ctx.appendDataHookContext('Organization.Membership.Updated', {
           ...buildManagementApiContext(ctx),
           organizationId,

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

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

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

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

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

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

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

@@ -54,7 +54,7 @@ const { hasActiveUsers, updateUserById } = userQueries;
 
 const userLibraries = {
   generateUserId: jest.fn().mockResolvedValue('uid'),
-  insertUser: jest.fn().mockResolvedValue([{}, { organizationIds: [] }]),
+  insertUser: jest.fn().mockResolvedValue([{}, { organizations: [] }]),
 };
 const { generateUserId, insertUser } = userLibraries;
 

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

@@ -64,7 +64,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, { organizationIds: [] }]
+    async (user: CreateUser): Promise<InsertUserResult> => [user as User, { organizations: [] }]
   ),
 };
 const { generateUserId, insertUser } = userLibraries;

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

@@ -135,7 +135,7 @@ async function handleSubmitRegister(
     (invitation) => invitation.status === OrganizationInvitationStatus.Pending
   );
 
-  const [user, { organizationIds }] = await insertUser(
+  const [user, { organizations: provisionedOrganizations }] = await insertUser(
     {
       id,
       ...userProfile,
@@ -190,7 +190,7 @@ async function handleSubmitRegister(
   ctx.assignInteractionHookResult({ userId: id });
   ctx.appendDataHookContext('User.Created', { user });
 
-  for (const organizationId of organizationIds) {
+  for (const { organizationId } of provisionedOrganizations) {
     ctx.appendDataHookContext('Organization.Membership.Updated', {
       organizationId,
     });

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

@@ -26,7 +26,7 @@ const updateUserSsoIdentityMock = jest.fn();
 const insertUserSsoIdentityMock = jest.fn();
 const updateUserMock = jest.fn();
 const findUserByEmailMock = jest.fn();
-const insertUserMock = jest.fn().mockResolvedValue([{ id: 'foo' }, { organizationIds: [] }]);
+const insertUserMock = jest.fn().mockResolvedValue([{ id: 'foo' }, { organizations: [] }]);
 const generateUserIdMock = jest.fn().mockResolvedValue('foo');
 const getAvailableSsoConnectorsMock = jest.fn();
 
@@ -291,7 +291,7 @@ describe('Single sign on util methods tests', () => {
 
   describe('registerWithSsoAuthentication tests', () => {
     it('should register if no related user account found', async () => {
-      insertUserMock.mockResolvedValueOnce([{ id: 'foo' }, { organizationIds: [] }]);
+      insertUserMock.mockResolvedValueOnce([{ id: 'foo' }, { organizations: [] }]);
 
       const { id } = await registerWithSsoAuthentication(mockContext, tenant, {
         connectorId: wellConfiguredSsoConnector.id,

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

@@ -309,7 +309,7 @@ export const registerWithSsoAuthentication = async (
   };
 
   // Insert new user
-  const [user, { organizationIds }] = await usersLibrary.insertUser(
+  const [user, { organizations }] = await usersLibrary.insertUser(
     {
       id: await usersLibrary.generateUserId(),
       ...syncingProfile,
@@ -317,7 +317,7 @@ export const registerWithSsoAuthentication = async (
     },
     []
   );
-  for (const organizationId of organizationIds) {
+  for (const { organizationId } of organizations) {
     ctx.appendDataHookContext('Organization.Membership.Updated', {
       organizationId,
     });

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

@@ -11,10 +11,58 @@ describe('organization just-in-time provisioning', () => {
     await Promise.all([organizationApi.cleanUp(), userApi.cleanUp()]);
   });
 
-  it('should automatically provision a user to the organization with the matched email domain', async () => {
+  it('should automatically provision a user to the organizations with roles', async () => {
     const organizations = await Promise.all([
       organizationApi.create({ name: 'foo' }),
       organizationApi.create({ name: 'bar' }),
+      organizationApi.create({ name: 'baz' }),
+    ]);
+    const roles = await Promise.all([
+      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)
+      )
+    );
+    await Promise.all([
+      organizationApi.jit.addRole(organizations[0].id, [roles[0].id, roles[1].id]),
+      organizationApi.jit.addRole(organizations[1].id, [roles[0].id]),
+    ]);
+
+    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,
+          // 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(
@@ -28,7 +76,20 @@ describe('organization just-in-time provisioning', () => {
 
     const userOrganizations = await getUserOrganizations(id);
     expect(userOrganizations).toEqual(
-      expect.arrayContaining(organizations.map((item) => expect.objectContaining({ id: item.id })))
+      expect.arrayContaining([
+        expect.objectContaining({
+          id: organizations[0].id,
+          organizationRoles: [],
+        }),
+        expect.objectContaining({
+          id: organizations[1].id,
+          organizationRoles: [],
+        }),
+        expect.objectContaining({
+          id: organizations[2].id,
+          organizationRoles: [],
+        }),
+      ])
     );
   });
 });

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

@@ -45,10 +45,15 @@ describe('organization just-in-time provisioning', () => {
     });
   });
 
-  it('should automatically provision a user to the organization with the matched email domain', async () => {
+  it('should automatically provision a user to the organization with roles', async () => {
     const organizations = await Promise.all([
       organizationApi.create({ name: 'foo' }),
       organizationApi.create({ name: 'bar' }),
+      organizationApi.create({ name: 'baz' }),
+    ]);
+    const roles = await Promise.all([
+      organizationApi.roleApi.create({ name: randomString() }),
+      organizationApi.roleApi.create({ name: randomString() }),
     ]);
     const emailDomain = 'foo.com';
     await Promise.all(
@@ -56,20 +61,41 @@ describe('organization just-in-time provisioning', () => {
         organizationApi.jit.addEmailDomain(organization.id, emailDomain)
       )
     );
+    await Promise.all([
+      organizationApi.jit.addRole(organizations[0].id, [roles[0].id, roles[1].id]),
+      organizationApi.jit.addRole(organizations[1].id, [roles[0].id]),
+    ]);
 
     const email = randomString() + '@' + emailDomain;
     const { client, id } = await registerWithEmail(email);
 
     const userOrganizations = await getUserOrganizations(id);
     expect(userOrganizations).toEqual(
-      expect.arrayContaining(organizations.map((item) => expect.objectContaining({ id: item.id })))
+      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: [],
+        }),
+      ])
     );
 
     await logoutClient(client);
     await deleteUser(id);
   });
 
-  it('should automatically provision a user to the organization with the matched email from a SSO identity', async () => {
+  it('should automatically provision a user with the matched email to the organization 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);