fix(sso): regenerate the primary provider migration after staging, and harden provider switching

Author: waleedlatif1

apps/docs/content/docs/platform/enterprise/sso.mdx

@@ -101,7 +101,7 @@ Click **Save**. To test, sign out and use the **Sign in with SSO** button on the
 
 ## Editing and advanced configuration
 
-For a saved connection, open **Sign-in**, select the provider, and select **Edit**. The Provider ID remains fixed. **Delete** removes that sign-in path only: accounts and memberships it admitted stay. If the domain has another provider, it becomes primary; otherwise people at the domain sign in another way until a provider serves it again. A saved OIDC client secret appears as a mask with a suffix when available; **Replace** lets you enter a new secret, and **Keep saved** cancels that replacement. Select **Update** to save the provider, or **Discard** to abandon changes.
+For a saved connection, open **Sign-in**, select the provider, and select **Edit**. The Provider ID remains fixed. **Delete** removes that sign-in path only: accounts and memberships it admitted stay. If you delete the primary provider and the domain has another verified provider, that one becomes primary; otherwise people at the domain sign in another way until a provider serves it again. A saved OIDC client secret appears as a mask with a suffix when available; **Replace** lets you enter a new secret, and **Keep saved** cancels that replacement. Select **Update** to save the provider, or **Discard** to abandon changes.
 
 **Advanced options** contains OIDC scopes and optional authorization, token, and JWKS endpoint overrides. For SAML, it contains Audience, Callback URL override, signed-assertion requirements, NameID format, and optional IdP metadata XML. **Attribute mapping** lets either protocol override the email, name, and stable user-ID claim names. Leave a mapping blank to use the protocol default.
 

apps/sim/app/api/auth/sso/providers/[providerId]/route.test.ts

@@ -10,13 +10,23 @@ import {
   schemaMock,
 } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { OrchestrationError } from '@/lib/core/orchestration/types'
 
-const { mockGetSession } = vi.hoisted(() => ({
+const { mockGetSession, mockSetPrimary } = vi.hoisted(() => ({
   mockGetSession: vi.fn(),
+  mockSetPrimary: vi.fn(),
 }))
 
 vi.mock('@sim/db', () => ({ ...dbChainMock, ...schemaMock }))
 vi.mock('@/lib/auth', () => ({ getSession: mockGetSession }))
+/** Authorization and the primary switch are the use case's; its own tests and the PostgreSQL suite cover them. */
+vi.mock('@/lib/auth/sso/application/set-primary-provider', () => ({
+  setPrimarySsoProviderOperation: { id: 'organization.sso.set_primary_provider' },
+  setPrimarySsoProvider: {
+    operation: { id: 'organization.sso.set_primary_provider' },
+    execute: mockSetPrimary,
+  },
+}))
 
 import { DELETE, PATCH } from '@/app/api/auth/sso/providers/[providerId]/route'
 
@@ -136,59 +146,45 @@ describe('PATCH /api/auth/sso/providers/[providerId]', () => {
 
   beforeEach(() => {
     vi.clearAllMocks()
-    resetDbChainMock()
-    mockGetSession.mockResolvedValue({ user: { id: 'u1' } })
-    dbChainMockFns.returning.mockResolvedValue([{ id: 'domain-1' }])
+    mockGetSession.mockResolvedValue({ user: { id: 'u1' }, session: { id: 's1' } })
+    mockSetPrimary.mockResolvedValue({
+      providerId: 'acme-okta',
+      organizationId: 'org1',
+      domain: 'acme.com',
+    })
   })
 
-  function queueOrgProvider(role: string) {
-    queueTableRows(schemaMock.ssoProvider, [
-      { id: 'row-1', organizationId: 'org1', userId: 'u-other', domain: 'acme.com' },
-    ])
-    queueTableRows(schemaMock.member, [{ role }])
-  }
-
-  it('requires a session', async () => {
+  it('requires a session before the use case runs', async () => {
     mockGetSession.mockResolvedValue(null)
     expect((await patch()).status).toBe(401)
-    expect(dbChainMockFns.update).not.toHaveBeenCalled()
+    expect(mockSetPrimary).not.toHaveBeenCalled()
   })
 
   it('only accepts making a provider primary', async () => {
     expect((await patch({ isPrimary: false })).status).toBe(400)
-    expect(dbChainMockFns.update).not.toHaveBeenCalled()
-  })
-
-  it("refuses a member who is not the organization's owner or admin", async () => {
-    queueOrgProvider('member')
-    expect((await patch()).status).toBe(403)
-    expect(dbChainMockFns.update).not.toHaveBeenCalled()
+    expect(mockSetPrimary).not.toHaveBeenCalled()
   })
 
-  it('names the provider on its verified domain', async () => {
-    queueOrgProvider('admin')
+  it('passes the routed provider to the use case and presents its result', async () => {
     const res = await patch()
     expect(res.status).toBe(200)
     await expect(res.json()).resolves.toEqual({ success: true, providerId: 'acme-okta' })
-    expect(dbChainMockFns.update).toHaveBeenCalledWith(schemaMock.ssoDomain)
-    expect(dbChainMockFns.set).toHaveBeenCalledWith(
-      expect.objectContaining({ primaryProviderId: 'acme-okta' })
+    expect(mockSetPrimary).toHaveBeenCalledWith(
+      expect.objectContaining({
+        principal: expect.objectContaining({ kind: 'session', userId: 'u1' }),
+        input: { providerId: 'acme-okta' },
+      })
     )
   })
 
-  it('refuses a provider whose domain is not verified', async () => {
-    queueOrgProvider('owner')
-    dbChainMockFns.returning.mockResolvedValue([])
+  it.each([
+    ['conflict', 409, 'Verify acme.com before making this provider primary.'],
+    ['forbidden', 403, 'Organization administrator access is required'],
+    ['not_found', 404, 'Provider not found'],
+  ] as const)('projects a %s refusal with its message', async (code, status, message) => {
+    mockSetPrimary.mockRejectedValue(new OrchestrationError(code, message))
     const res = await patch()
-    expect(res.status).toBe(409)
-    await expect(res.json()).resolves.toMatchObject({ error: expect.stringContaining('acme.com') })
-  })
-
-  it('refuses a personal provider, which has no primary', async () => {
-    queueTableRows(schemaMock.ssoProvider, [
-      { id: 'row-1', organizationId: null, userId: 'u1', domain: 'acme.com' },
-    ])
-    expect((await patch()).status).toBe(400)
-    expect(dbChainMockFns.update).not.toHaveBeenCalled()
+    expect(res.status).toBe(status)
+    await expect(res.json()).resolves.toEqual({ error: message })
   })
 })

apps/sim/app/api/auth/sso/providers/[providerId]/route.ts

@@ -1,11 +1,21 @@
-import { db, ssoDomain, ssoProvider } from '@sim/db'
-import { forgetPrimaryProviders, verifiedDomainOfProvider } from '@sim/db/sso-primary-provider'
+import { db, ssoProvider } from '@sim/db'
+import { forgetPrimaryProvider } from '@sim/db/sso-primary-provider'
 import { createLogger } from '@sim/logger'
-import { and, eq, exists, isNull, sql } from 'drizzle-orm'
+import { and, eq, isNull } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import { deleteSsoProviderContract, setPrimarySsoProviderContract } from '@/lib/api/contracts/auth'
 import { parseRequest } from '@/lib/api/server'
+import {
+  defineInternalJsonRoute,
+  internalOrchestrationErrorPolicy,
+  internalRateLimits,
+  internalSessionAuth,
+} from '@/lib/api/server/routes'
 import { getSession } from '@/lib/auth'
+import {
+  setPrimarySsoProvider,
+  setPrimarySsoProviderOperation,
+} from '@/lib/auth/sso/application/set-primary-provider'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { isOrganizationAdminOrOwner } from '@/lib/workspaces/permissions/utils'
 
@@ -15,10 +25,10 @@ type RouteContext = { params: Promise<{ providerId: string }> }
 
 /**
  * Loads a provider the caller may manage: an organization provider for its
- * owners and admins, a personal provider for its creator. Sim owns these
- * mutations rather than exposing the SSO plugin's own, which
- * `/api/auth/[...all]` blocks by design: the plugin gates only on the row's
- * creator, while an organization's providers belong to the organization.
+ * owners and admins, a personal provider for its creator. Sim owns deleting
+ * rather than exposing the SSO plugin's own delete, which `/api/auth/[...all]`
+ * blocks by design: the plugin gates only on the row's creator, while an
+ * organization's providers belong to the organization.
  */
 async function loadManagedProvider(userId: string, providerId: string) {
   const [provider] = await db
@@ -41,69 +51,16 @@ async function loadManagedProvider(userId: string, providerId: string) {
   return provider
 }
 
-/**
- * Makes this provider the one its domain signs in through. The previous primary
- * stays configured and reachable by test link, so the switch can be reversed
- * the same way.
- */
-export const PATCH = withRouteHandler(async (request: NextRequest, context: RouteContext) => {
-  const session = await getSession()
-  if (!session?.user?.id) {
-    return NextResponse.json({ error: 'Authentication required' }, { status: 401 })
-  }
-
-  const parsed = await parseRequest(setPrimarySsoProviderContract, request, context)
-  if (!parsed.success) return parsed.response
-  const { providerId } = parsed.data.params
-
-  const provider = await loadManagedProvider(session.user.id, providerId)
-  if (provider instanceof NextResponse) return provider
-  if (!provider.organizationId) {
-    return NextResponse.json(
-      { error: 'Only organization identity providers have a primary provider' },
-      { status: 400 }
-    )
-  }
-
-  /**
-   * Names the provider on the one domain record sign-in joins it to, so a
-   * provider this succeeds for is exactly the one sign-in then uses.
-   */
-  const named = await db
-    .update(ssoDomain)
-    .set({ primaryProviderId: providerId, updatedAt: new Date() })
-    .where(
-      and(
-        eq(ssoDomain.organizationId, provider.organizationId),
-        exists(
-          db
-            .select({ found: sql`1` })
-            .from(ssoProvider)
-            .where(
-              and(
-                eq(ssoProvider.id, provider.id),
-                eq(ssoProvider.domainVerified, true),
-                verifiedDomainOfProvider
-              )
-            )
-        )
-      )
-    )
-    .returning({ id: ssoDomain.id })
-  if (named.length === 0) {
-    return NextResponse.json(
-      { error: `Verify ${provider.domain} before making this provider primary.` },
-      { status: 409 }
-    )
-  }
-
-  logger.info('Set primary SSO provider', {
-    providerId,
-    organizationId: provider.organizationId,
-    domain: provider.domain,
-    userId: session.user.id,
-  })
-  return NextResponse.json({ success: true, providerId })
+/** Makes this provider the one its domain signs in through. */
+export const PATCH = defineInternalJsonRoute({
+  contract: setPrimarySsoProviderContract,
+  auth: internalSessionAuth,
+  operation: setPrimarySsoProviderOperation,
+  rateLimit: internalRateLimits.user({ bucketName: 'sso-set-primary-provider' }),
+  errorPolicy: internalOrchestrationErrorPolicy,
+  mapInput: ({ params }) => ({ providerId: params.providerId }),
+  useCase: setPrimarySsoProvider,
+  present: ({ providerId }) => ({ success: true as const, providerId }),
 })
 
 /**
@@ -140,7 +97,7 @@ export const DELETE = withRouteHandler(async (request: NextRequest, context: Rou
       .where(and(eq(ssoProvider.id, provider.id), ownerClause))
       .returning({ id: ssoProvider.id })
     if (deleted.length > 0 && organizationId) {
-      await forgetPrimaryProviders(tx, organizationId, [providerId])
+      await forgetPrimaryProvider(tx, organizationId, providerId)
     }
     return deleted
   })

apps/sim/app/api/auth/sso/register/route.test.ts

@@ -514,12 +514,25 @@ describe('POST /api/auth/sso/register', () => {
 
   it("does not report the caller's own org-less provider as another tenant's claim", async () => {
     queueMembers([{ organizationId: 'org1', role: 'owner' }])
-    queueProviders([{ domain: 'acme.com', userId: 'u1', organizationId: null }])
+    queueProviders([
+      { domain: 'acme.com', userId: 'u1', organizationId: null, providerId: 'acme-oidc' },
+    ])
     const res = await POST(request(OIDC_BODY))
     expect(res.status).toBe(200)
     expect(mockRegisterSSOProvider).toHaveBeenCalledTimes(1)
   })
 
+  it("refuses a second provider on a domain the caller's own org-less provider signs in", async () => {
+    queueMembers([{ organizationId: 'org1', role: 'owner' }])
+    queueProviders([
+      { domain: 'acme.com', userId: 'u1', organizationId: null, providerId: 'acme-personal' },
+    ])
+    const res = await POST(request(OIDC_BODY))
+    expect(res.status).toBe(409)
+    await expect(res.json()).resolves.toMatchObject({ code: 'SSO_DOMAIN_ALREADY_ROUTED' })
+    expect(mockRegisterSSOProvider).not.toHaveBeenCalled()
+  })
+
   it("still blocks an org admin from claiming another user's user-scoped domain", async () => {
     queueMembers([{ organizationId: 'org1', role: 'owner' }])
     queueProviders([{ domain: 'acme.com', userId: 'someone-else', organizationId: null }])

apps/sim/app/api/auth/sso/register/route.ts

@@ -1,5 +1,5 @@
 import { db, member, ssoDomain, ssoProvider } from '@sim/db'
-import { nameIncumbentSignInProvider, ssoProviderDomainKey } from '@sim/db/sso-primary-provider'
+import { keepDomainSignInProvider, ssoProviderDomainKey } from '@sim/db/sso-primary-provider'
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import { normalizeSSODomain } from '@sim/utils/sso-domain'
@@ -191,16 +191,17 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
     )
 
     /**
-     * Refuses the domain when another tenant has claimed it. The caller's own
-     * organization may add a provider to a domain it already signs in through:
-     * the new provider waits, reachable by test link, until an admin makes it
-     * the domain's primary.
+     * Refuses the domain when another tenant has claimed it, or when the caller's
+     * own personal provider signs it in. The caller's organization may add a
+     * provider to a domain it already signs in through: the new provider waits,
+     * reachable by test link, until an admin makes it the domain's primary.
      */
     const findDomainRefusal = async (): Promise<NextResponse | null> => {
       const claims = await db
         .select({
           userId: ssoProvider.userId,
           organizationId: ssoProvider.organizationId,
+          providerId: ssoProvider.providerId,
         })
         .from(ssoProvider)
         .where(sql`${ssoProviderDomainKey} = ${domain}`)
@@ -218,6 +219,21 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
           { status: 409 }
         )
       }
+      const personal = claims.find(
+        (provider) =>
+          !provider.organizationId &&
+          typeof provider.providerId === 'string' &&
+          provider.providerId !== providerId
+      )
+      if (personal) {
+        return NextResponse.json(
+          {
+            error: `${domain} already signs in through the provider "${personal.providerId}". Edit that provider, or give this one a different verified domain.`,
+            code: 'SSO_DOMAIN_ALREADY_ROUTED',
+          },
+          { status: 409 }
+        )
+      }
       return null
     }
 
@@ -625,6 +641,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         id: ssoProvider.id,
         issuer: ssoProvider.issuer,
         domain: ssoProvider.domain,
+        domainVerified: ssoProvider.domainVerified,
         oidcConfig: ssoProvider.oidcConfig,
         samlConfig: ssoProvider.samlConfig,
         jitProvisioningEnabled: ssoProvider.jitProvisioningEnabled,
@@ -643,14 +660,15 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
      * first the SELECT finds nothing.
      *
      * A provider joining a domain another provider already signs in does not
-     * take over by sorting first: when the domain names no primary, the provider
-     * signing it in until now is named, in the same transaction. The lock is
-     * `FOR UPDATE` so two providers joining at once name it one after the other.
+     * take over by sorting first: unless the domain's named primary still signs
+     * it in, the provider signing it in until now is named, in the same
+     * transaction. The lock is `FOR UPDATE` so two providers joining at once
+     * settle it one after the other.
      */
     const grantProviderDomainTrust = (joinsDomain: boolean): Promise<boolean> =>
       db.transaction(async (tx) => {
         const [proof] = await tx
-          .select({ id: ssoDomain.id, primaryProviderId: ssoDomain.primaryProviderId })
+          .select({ id: ssoDomain.id })
           .from(ssoDomain)
           .where(verifiedDomainClause)
           .limit(1)
@@ -664,8 +682,8 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
           .returning({ id: ssoProvider.id })
         if (granted.length === 0) return false
 
-        if (joinsDomain && proof.primaryProviderId === null) {
-          await nameIncumbentSignInProvider(tx, {
+        if (joinsDomain) {
+          await keepDomainSignInProvider(tx, {
             domainRecordId: proof.id,
             organizationId: orgId,
             domain,
@@ -703,8 +721,10 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
 
       let domainTrustGranted: boolean
       try {
+        /** An owned provider joins the domain when it moves to it or is not yet trusted on it. */
         domainTrustGranted = await grantProviderDomainTrust(
-          normalizeSSODomain(existingOwnedProvider.domain) !== domain
+          !existingOwnedProvider.domainVerified ||
+            normalizeSSODomain(existingOwnedProvider.domain) !== domain
         )
       } catch (error) {
         try {

apps/sim/ee/sso/components/sso-form.test.tsx

@@ -183,6 +183,8 @@ describe('SSOForm sign-in errors', () => {
     expect(mockSsoSignIn).toHaveBeenCalledWith(
       expect.objectContaining({ email: 'user@example.com', providerId: 'example-okta' })
     )
+    const [signIn] = mockSsoSignIn.mock.calls[0]
+    expect(signIn.errorCallbackURL).not.toContain('provider=')
   })
 
   it('signs in through the provider a test link names, and returns to that link on failure', async () => {
@@ -199,6 +201,9 @@ describe('SSOForm sign-in errors', () => {
     const [signIn] = mockSsoSignIn.mock.calls[0]
     expect(signIn.providerId).toBe('example-okta')
     expect(signIn.errorCallbackURL).toContain('provider=example-okta')
+    /** The SSO plugin appends `?error=…` to the URL, so the test link must survive that suffix. */
+    const retry = new URL(`${signIn.errorCallbackURL}?error=invalid_provider`, 'https://sim.test')
+    expect(retry.searchParams.get('provider')).toBe('example-okta')
   })
 
   it('explains a test link that does not match the email domain', async () => {