fix: possible race condition where multiple organizations could be created

controlplane/src/core/controllers/auth.ts

@@ -114,242 +114,252 @@
 
   fastify.get<{ Querystring: { code: string; code_verifier: string; redirectURL?: string; ssoSlug?: string } }>(
     '/callback',
     async (req, res) => {
       try {
         const redirectURL = req.query?.redirectURL;
         const ssoSlug = req.query?.ssoSlug;
         const { accessToken, refreshToken, idToken } = await opts.authUtils.handleAuthCallbackRequest(req);
 
         // decodeJWT will throw an error if the token is invalid or expired
         const accessTokenPayload = decodeJWT<CustomAccessTokenClaims>(accessToken);
 
         // Clear the PKCE cookie
         opts.authUtils.clearCookie(res, opts.pkce.cookieName);
         // Clear the sso cookie
         opts.authUtils.clearCookie(res, cosmoIdpHintCookieName);
 
         const sessionExpiresIn = DEFAULT_SESSION_MAX_AGE_SEC;
         const sessionExpiresDate = new Date(Date.now() + 1000 * sessionExpiresIn);
 
         const userId = accessTokenPayload.sub!;
         const userEmail = accessTokenPayload.email!;
         const firstName = accessTokenPayload.given_name || '';
         const lastName = accessTokenPayload.family_name || '';
 
         const insertedSession = await opts.db.transaction(async (tx) => {
           // Upsert the user
           await tx
             .insert(users)
             .values({
               id: userId,
               email: accessTokenPayload.email,
             })
             .onConflictDoUpdate({
               target: users.id,
               // Update the fields when the user already exists
               set: {
                 email: accessTokenPayload.email,
               },
             })
             .execute();
 
           if (accessTokenPayload.groups && accessTokenPayload.groups.length > 0) {
             const keycloakOrgs = new Set(accessTokenPayload.groups.map((grp) => grp.split('/')[1]));
             const orgRepo = new OrganizationRepository(req.log, tx, opts.defaultBillingPlanId);
             const orgGroupRepo = new OrganizationGroupRepository(tx);
 
             // delete all the org member roles
             for (const slug of keycloakOrgs) {
               const dbOrg = await orgRepo.bySlug(slug);
 
               if (!dbOrg) {
                 continue;
               }
 
               const orgMember = await orgRepo.getOrganizationMember({ organizationID: dbOrg.id, userID: userId });
               if (!orgMember) {
                 continue;
               }
 
               await tx
                 .delete(schema.organizationGroupMembers)
                 .where(eq(schema.organizationGroupMembers.organizationMemberId, orgMember.orgMemberID));
             }
 
             // upserting the members into the orgs and inserting their roles.
             for (const kcGroup of accessTokenPayload.groups) {
               const slug = kcGroup.split('/')[1];
               const dbOrg = await orgRepo.bySlug(slug);
               if (!dbOrg) {
                 continue;
               }
 
               const insertedMember = await tx
                 .insert(organizationsMembers)
                 .values({
                   userId,
                   organizationId: dbOrg.id,
                 })
                 .onConflictDoUpdate({
                   target: [organizationsMembers.userId, organizationsMembers.organizationId],
                   // Update the fields only when the org member already exists
                   set: {
                     userId,
                     organizationId: dbOrg.id,
                   },
                 })
                 .returning()
                 .execute();
 
               const groupName = kcGroup.split('/')?.[2];
               if (!groupName) {
                 continue;
               }
 
               const orgGroup = await orgGroupRepo.byName({
                 organizationId: dbOrg.id,
                 name: groupName,
               });
 
               if (!orgGroup) {
                 // The group doesn't exist for the organization, instead of failing, we'll just skip the group
                 continue;
               }
 
               await orgGroupRepo.addUserToGroup({
                 organizationMemberId: insertedMember[0].id,
                 groupId: orgGroup.groupId,
               });
             }
           }
 
           // If there is already a session for this user, update it.
           // Otherwise, insert a new session. Because we use an Idp like keycloak,
           // we can assume that the user will have only one session per client at a time.
           const insertedSessions = await tx
             .insert(sessions)
             .values({
               userId,
               idToken,
               accessToken,
               refreshToken,
               expiresAt: sessionExpiresDate,
             })
             .onConflictDoUpdate({
               target: sessions.userId,
               // Update the fields when the session already exists
               set: {
                 idToken,
                 accessToken,
                 refreshToken,
                 expiresAt: sessionExpiresDate,
                 updatedAt: new Date(),
               },
             })
             .returning({
               id: sessions.id,
               userId: sessions.userId,
             })
             .execute();
 
           return insertedSessions[0];
         });
 
-        const orgs = await opts.organizationRepository.memberships({
-          userId,
-        });
-
+        const orgs = await opts.organizationRepository.memberships({ userId });
         if (orgs.length === 0) {
-          await opts.keycloakClient.authenticateClient();
-
           const organizationSlug = uid(8);
 
+          // First, we need to create the organization and add the user as an organization member
+          const [insertedOrg, orgMember] = await opts.db.transaction(async (tx) => {
+            const orgRepo = new OrganizationRepository(req.log, tx, opts.defaultBillingPlanId);
+
+            // Create the organization...
+            const inserted = await orgRepo.createOrganization({
+              organizationName: userEmail.split('@')[0],
+              organizationSlug,
+              ownerID: userId,
+            });
+
+            // ...and add the user as an organization member.
+            const orgMember = await orgRepo.addOrganizationMember({
+              organizationID: inserted.id,
+              userID: userId,
+            });
+
+            return [inserted, orgMember];
+          });
+
+          // Finalize the organization setup by seeding the Keycloak group structure
+          await opts.keycloakClient.authenticateClient();
           const [kcRootGroupId, kcCreatedGroups] = await opts.keycloakClient.seedGroup({
             userID: userId,
             organizationSlug,
             realm: opts.keycloakRealm,
           });
 
           await opts.db.transaction(async (tx) => {
             const orgRepo = new OrganizationRepository(req.log, tx, opts.defaultBillingPlanId);
             const orgGroupRepo = new OrganizationGroupRepository(tx);
 
-            const insertedOrg = await orgRepo.createOrganization({
-              organizationName: userEmail.split('@')[0],
-              organizationSlug,
-              ownerID: userId,
+            await orgRepo.updateOrganization({
+              id: insertedOrg.id,
               kcGroupId: kcRootGroupId,
             });
 
-            const orgMember = await orgRepo.addOrganizationMember({
-              organizationID: insertedOrg.id,
-              userID: userId,
-            });
-
             await orgGroupRepo.importKeycloakGroups({
               organizationId: insertedOrg.id,
               kcGroups: kcCreatedGroups,
             });
 
             const orgAdminGroup = await orgGroupRepo.byName({
               organizationId: insertedOrg.id,
               name: 'admin',
             });
 
             if (orgAdminGroup) {
               await orgGroupRepo.addUserToGroup({
                 organizationMemberId: orgMember.id,
                 groupId: orgAdminGroup.groupId,
               });
             }
 
             const namespaceRepo = new NamespaceRepository(tx, insertedOrg.id);
             const ns = await namespaceRepo.create({
               name: DefaultNamespace,
               createdBy: userId,
             });
             if (!ns) {
               throw new Error(`Could not create ${DefaultNamespace} namespace`);
             }
           });
 
           opts.platformWebhooks.send(PlatformEventName.USER_REGISTER_SUCCESS, {
             user_id: userId,
             user_email: userEmail,
             user_first_name: firstName,
             user_last_name: lastName,
           });
         }
 
         // Create a JWT token containing the session id and user id.
         const jwt = await encrypt<UserSession>({
           maxAgeInSeconds: sessionExpiresIn,
           token: {
             iss: userId,
             sessionId: insertedSession.id,
           },
           secret: opts.jwtSecret,
         });
 
         // Set the session cookie. The cookie value is encrypted.
         opts.authUtils.createSessionCookie(res, jwt, sessionExpiresDate);
         if (ssoSlug) {
           // Set the sso cookie.
           opts.authUtils.createSsoCookie(res, ssoSlug);
         }
         if (redirectURL) {
           if (redirectURL.startsWith(opts.webBaseUrl)) {
             res.redirect(redirectURL);
           } else {
             res.redirect(opts.webBaseUrl);
           }
         } else if (orgs.length === 0) {
           res.redirect(opts.webBaseUrl + '?migrate=true');
         } else {
           res.redirect(opts.webBaseUrl);
         }
       } catch (err: any) {
         if (err instanceof AuthenticationError) {
           req.log.debug(err);
         } else {
           req.log.error(err);

controlplane/src/core/repositories/OrganizationRepository.ts

@@ -95,12 +95,13 @@ export class OrganizationRepository {
     return org;
   }
 
-  public async updateOrganization(input: { id: string; slug?: string; name?: string }) {
+  public async updateOrganization(input: { id: string; slug?: string; name?: string; kcGroupId?: string }) {
     await this.db
       .update(organizations)
       .set({
         name: input.name,
         slug: input.slug,
+        kcGroupId: input.kcGroupId,
       })
       .where(eq(organizations.id, input.id))
       .execute();

controlplane/test/utils.test.ts

@@ -1,5 +1,10 @@
 import { describe, expect, test } from 'vitest';
-import { isValidLabelMatchers, mergeUrls, normalizeLabelMatchers, isGoogleCloudStorageUrl } from '../src/core/util.js';
+import {
+  isValidLabelMatchers,
+  mergeUrls,
+  normalizeLabelMatchers,
+  isGoogleCloudStorageUrl,
+} from '../src/core/util.js';
 
 describe('Utils', () => {
   test('isValidLabelMatchers', () => {