fix(auth): recheck OAuth entitlement under organization lock

Author: waleedlatif1

apps/sim/lib/auth/oauth-provider-lifecycle.postgres.test.ts

@@ -35,6 +35,7 @@ async function loadRuntime() {
     { runCleanupOAuthTokens, OAUTH_TOKEN_RETENTION_DAYS },
     { isCapabilityWithheldForUser },
     { acquirePermissionGroupOrgLock },
+    { acquireOrganizationMutationLock },
   ] = await Promise.all([
     import('@sim/db'),
     import('@sim/db/schema'),
@@ -50,6 +51,7 @@ async function loadRuntime() {
     import('@/background/cleanup-oauth-tokens'),
     import('@/lib/permission-groups/user-scope.server'),
     import('@/lib/permission-groups/locks'),
+    import('@/lib/billing/organizations/membership'),
   ])
   const adapter = createSimAuthAdapter({
     plugins: [
@@ -81,6 +83,7 @@ async function loadRuntime() {
     OAUTH_TOKEN_RETENTION_DAYS,
     isCapabilityWithheldForUser,
     acquirePermissionGroupOrgLock,
+    acquireOrganizationMutationLock,
   }
 }
 
@@ -298,111 +301,157 @@ describe.skipIf(!databaseUrl)('OAuth lifecycle on the provisioned PostgreSQL sch
     }
   )
 
-  it('serializes refresh against a pending policy restriction without consuming the token', async () => {
-    const { db, schema, eq, sql, statement, acquirePermissionGroupOrgLock } = runtime
-    const groupId = await createDefaultGroup()
-    await grantConsent()
-    const family = await issueFamily()
-    const credentials = { clientId, method: 'none' as const }
-    const writerReady = Promise.withResolvers<number>()
-    const releaseWriter = Promise.withResolvers<void>()
-    const writer = db.transaction(async (tx) => {
-      await acquirePermissionGroupOrgLock(tx, organizationId!)
-      await tx
-        .update(schema.permissionGroup)
-        .set({ config: { disableOAuthAppAccess: true } })
-        .where(eq(schema.permissionGroup.id, groupId))
-      const [connection] = await tx.execute<{ pid: number }>(
-        statement`select pg_backend_pid() as pid`
-      )
-      writerReady.resolve(connection.pid)
-      await releaseWriter.promise
-    })
-    let refresh: ReturnType<typeof runtime.rotateOAuthRefreshToken> | undefined
-
-    try {
-      const writerPid = await Promise.race([
-        writerReady.promise,
-        writer.then(() => {
-          throw new Error('Policy writer finished before the concurrency check')
-        }),
-      ])
-      let refreshSettled = false
-      refresh = runtime.rotateOAuthRefreshToken({ credentials, refreshToken: family.refreshToken })
-      void refresh.then(
-        () => {
-          refreshSettled = true
-        },
-        () => {
-          refreshSettled = true
+  it.each(['policy restriction', 'enterprise activation'] as const)(
+    'serializes refresh against a pending %s without consuming the token',
+    async (change) => {
+      const {
+        db,
+        schema,
+        eq,
+        sql,
+        statement,
+        acquirePermissionGroupOrgLock,
+        acquireOrganizationMutationLock,
+      } = runtime
+      const groupId = await createDefaultGroup()
+      await grantConsent()
+      const family = await issueFamily()
+      if (change === 'enterprise activation') {
+        await db
+          .update(schema.subscription)
+          .set({
+            plan: 'team',
+            metadata: { plan: 'team', referenceId: organizationId!, seats: 5, monthlyPrice: 100 },
+          })
+          .where(eq(schema.subscription.referenceId, organizationId!))
+        await db
+          .update(schema.permissionGroup)
+          .set({ config: { disableOAuthAppAccess: true } })
+          .where(eq(schema.permissionGroup.id, groupId))
+      }
+      expect(await runtime.isCapabilityWithheldForUser(userId, 'oauth_apps.use')).toBe(false)
+      const credentials = { clientId, method: 'none' as const }
+      const writerReady = Promise.withResolvers<number>()
+      const releaseWriter = Promise.withResolvers<void>()
+      const writer = db.transaction(async (tx) => {
+        if (change === 'enterprise activation') {
+          await acquireOrganizationMutationLock(tx, organizationId!)
+          await tx
+            .update(schema.subscription)
+            .set({
+              plan: 'enterprise',
+              metadata: {
+                plan: 'enterprise',
+                referenceId: organizationId!,
+                seats: 5,
+                monthlyPrice: 100,
+              },
+            })
+            .where(eq(schema.subscription.referenceId, organizationId!))
+        } else {
+          await acquirePermissionGroupOrgLock(tx, organizationId!)
+          await tx
+            .update(schema.permissionGroup)
+            .set({ config: { disableOAuthAppAccess: true } })
+            .where(eq(schema.permissionGroup.id, groupId))
         }
-      )
-
-      let waitingForPolicy = false
-      const deadline = Date.now() + 2_000
-      while (!refreshSettled && Date.now() < deadline) {
-        const [waiter] = await sql<{ waiting: boolean }[]>`
+        const [connection] = await tx.execute<{ pid: number }>(
+          statement`select pg_backend_pid() as pid`
+        )
+        writerReady.resolve(connection.pid)
+        await releaseWriter.promise
+      })
+      let refresh: ReturnType<typeof runtime.rotateOAuthRefreshToken> | undefined
+
+      try {
+        const writerPid = await Promise.race([
+          writerReady.promise,
+          writer.then(() => {
+            throw new Error('Policy writer finished before the concurrency check')
+          }),
+        ])
+        let refreshSettled = false
+        refresh = runtime.rotateOAuthRefreshToken({
+          credentials,
+          refreshToken: family.refreshToken,
+        })
+        void refresh.then(
+          () => {
+            refreshSettled = true
+          },
+          () => {
+            refreshSettled = true
+          }
+        )
+
+        let waitingForPolicy = false
+        const deadline = Date.now() + 2_000
+        while (!refreshSettled && Date.now() < deadline) {
+          const [waiter] = await sql<{ waiting: boolean }[]>`
           SELECT EXISTS (
             SELECT 1 FROM pg_stat_activity
             WHERE ${writerPid} = ANY(pg_blocking_pids(pid))
               AND wait_event_type = 'Lock'
               AND wait_event = 'advisory'
           ) AS waiting
         `
-        if (waiter.waiting) {
-          waitingForPolicy = true
-          break
+          if (waiter.waiting) {
+            waitingForPolicy = true
+            break
+          }
+          await sleep(1)
         }
-        await sleep(1)
+        expect(waitingForPolicy, 'Refresh must wait for the organization policy writer').toBe(true)
+        expect(refreshSettled).toBe(false)
+        releaseWriter.resolve()
+        await writer
+        expect(await runtime.isCapabilityWithheldForUser(userId, 'oauth_apps.use')).toBe(true)
+        await expect(refresh).resolves.toMatchObject({ success: false, error: 'invalid_grant' })
+        expect(
+          await db
+            .select({ generation: schema.oauthTokenFamily.currentGeneration })
+            .from(schema.oauthTokenFamily)
+            .where(eq(schema.oauthTokenFamily.id, family.id))
+        ).toEqual([{ generation: 0 }])
+        expect(
+          await db
+            .select({
+              generation: schema.oauthRefreshToken.generation,
+              revoked: schema.oauthRefreshToken.revoked,
+            })
+            .from(schema.oauthRefreshToken)
+            .where(eq(schema.oauthRefreshToken.familyId, family.id))
+        ).toEqual([{ generation: 0, revoked: null }])
+        expect(
+          await db
+            .select({ id: schema.oauthAccessToken.id })
+            .from(schema.oauthAccessToken)
+            .where(eq(schema.oauthAccessToken.clientId, clientId))
+        ).toHaveLength(1)
+
+        await db.transaction(async (tx) => {
+          await acquirePermissionGroupOrgLock(tx, organizationId!)
+          await tx
+            .update(schema.permissionGroup)
+            .set({ config: {} })
+            .where(eq(schema.permissionGroup.id, groupId))
+        })
+        await expect(
+          runtime.rotateOAuthRefreshToken({ credentials, refreshToken: family.refreshToken })
+        ).resolves.toMatchObject({ success: true })
+        expect(
+          await db
+            .select({ generation: schema.oauthTokenFamily.currentGeneration })
+            .from(schema.oauthTokenFamily)
+            .where(eq(schema.oauthTokenFamily.id, family.id))
+        ).toEqual([{ generation: 1 }])
+      } finally {
+        releaseWriter.resolve()
+        await Promise.allSettled([writer, ...(refresh ? [refresh] : [])])
       }
-      expect(waitingForPolicy, 'Refresh must wait for the organization policy writer').toBe(true)
-      expect(refreshSettled).toBe(false)
-      releaseWriter.resolve()
-      await writer
-      await expect(refresh).resolves.toMatchObject({ success: false, error: 'invalid_grant' })
-      expect(
-        await db
-          .select({ generation: schema.oauthTokenFamily.currentGeneration })
-          .from(schema.oauthTokenFamily)
-          .where(eq(schema.oauthTokenFamily.id, family.id))
-      ).toEqual([{ generation: 0 }])
-      expect(
-        await db
-          .select({
-            generation: schema.oauthRefreshToken.generation,
-            revoked: schema.oauthRefreshToken.revoked,
-          })
-          .from(schema.oauthRefreshToken)
-          .where(eq(schema.oauthRefreshToken.familyId, family.id))
-      ).toEqual([{ generation: 0, revoked: null }])
-      expect(
-        await db
-          .select({ id: schema.oauthAccessToken.id })
-          .from(schema.oauthAccessToken)
-          .where(eq(schema.oauthAccessToken.clientId, clientId))
-      ).toHaveLength(1)
-
-      await db.transaction(async (tx) => {
-        await acquirePermissionGroupOrgLock(tx, organizationId!)
-        await tx
-          .update(schema.permissionGroup)
-          .set({ config: {} })
-          .where(eq(schema.permissionGroup.id, groupId))
-      })
-      await expect(
-        runtime.rotateOAuthRefreshToken({ credentials, refreshToken: family.refreshToken })
-      ).resolves.toMatchObject({ success: true })
-      expect(
-        await db
-          .select({ generation: schema.oauthTokenFamily.currentGeneration })
-          .from(schema.oauthTokenFamily)
-          .where(eq(schema.oauthTokenFamily.id, family.id))
-      ).toEqual([{ generation: 1 }])
-    } finally {
-      releaseWriter.resolve()
-      await Promise.allSettled([writer, ...(refresh ? [refresh] : [])])
-    }
-  }, 10_000)
+    },
+    10_000
+  )
 
   it('atomically converges concurrent consent submissions on one grant', async () => {
     const grants = await Promise.all([grantConsent(), grantConsent()])

apps/sim/lib/auth/oauth-token-family.ts

@@ -276,8 +276,6 @@ export async function rotateOAuthRefreshToken(
 
   const membership = await getUserOrganization(provisionalToken.userId, database)
   const organizationId = membership?.organizationId ?? null
-  const permissionRegimeActive =
-    organizationId !== null && (await isOrganizationPermissionRegimeActive(organizationId))
   const nextRefreshBody = generateSecureToken(32)
   const nextAccessBody = generateSecureToken(32)
   const nextRefreshId = generateId()
@@ -295,6 +293,8 @@ export async function rotateOAuthRefreshToken(
         'Organization membership changed. Please sign in again.'
       )
     }
+    const permissionRegimeActive =
+      organizationId !== null && (await isOrganizationPermissionRegimeActive(organizationId, tx))
 
     const [activeUser] = await tx
       .select({

apps/sim/lib/permission-groups/resolve.server.test.ts

@@ -3,6 +3,7 @@
  */
 import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing'
 import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+import type { DbOrTx } from '@/lib/db/types'
 
 const { mockIsOrganizationOnEnterprisePlan, mockGetWorkspaceWithOwner } = vi.hoisted(() => ({
   mockIsOrganizationOnEnterprisePlan: vi.fn(),
@@ -20,6 +21,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => ({
 import {
   getUserPermissionConfig,
   getUserPermissionConfigForOrganization,
+  isOrganizationPermissionRegimeActive,
   resolveVerifiedUserAccessControlContext,
 } from '@/lib/permission-groups/resolve.server'
 
@@ -101,4 +103,33 @@ describe('permission-group resolution under a failed entitlement read', () => {
     })
     await expect(getUserPermissionConfigForOrganization(ORGANIZATION_ID)).resolves.toBeNull()
   })
+
+  it('rechecks entitlement on the caller transaction after an unentitled preflight', async () => {
+    const executor = {} as DbOrTx
+    mockIsOrganizationOnEnterprisePlan.mockResolvedValueOnce(false).mockResolvedValueOnce(true)
+
+    await expect(isOrganizationPermissionRegimeActive(ORGANIZATION_ID)).resolves.toBe(false)
+    await expect(isOrganizationPermissionRegimeActive(ORGANIZATION_ID, executor)).resolves.toBe(
+      true
+    )
+    expect(mockIsOrganizationOnEnterprisePlan).toHaveBeenLastCalledWith(
+      ORGANIZATION_ID,
+      'throw',
+      executor
+    )
+  })
+
+  it('propagates a transaction entitlement read failure instead of disabling restrictions', async () => {
+    const executor = {} as DbOrTx
+    entitlementReadFails()
+
+    await expect(isOrganizationPermissionRegimeActive(ORGANIZATION_ID, executor)).rejects.toThrow(
+      'billing database unavailable'
+    )
+    expect(mockIsOrganizationOnEnterprisePlan).toHaveBeenCalledWith(
+      ORGANIZATION_ID,
+      'throw',
+      executor
+    )
+  })
 })

apps/sim/lib/permission-groups/resolve.server.ts

@@ -310,24 +310,22 @@ export async function getUserPermissionConfigForOrganization(
  * enables Access Control, and the organization holds the Enterprise entitlement
  * that turns the regime on.
  *
- * Split out of {@link getUserPermissionConfigForOrganization} so a caller that
- * must re-read the *group* under `acquirePermissionGroupOrgLock` can settle this
- * half BEFORE opening its transaction. The entitlement read cannot move into a
- * transaction: {@link isOrganizationOnEnterprisePlan} is `cache()`d on its
- * argument list, so it admits no executor, and giving it one would both miss the
- * memo on every call and — because an unentitled organization resolves to
- * `config: null`, meaning every capability ALLOWED — turn a read failure into a
- * fail-open. The lock never serialized this half either way: it guards
- * permission-group writes, not subscription changes.
+ * Callers that serialize entitlement changes with an organization mutation
+ * lock must pass their transaction after acquiring that lock. The executor is
+ * part of the entitlement cache key, so this read cannot reuse a preflight
+ * result. A permission-group lock alone only serializes group writes.
  *
  * `'throw'` for the same reason as in
  * {@link resolveUserAccessControlContextForOrganization}.
  */
 export async function isOrganizationPermissionRegimeActive(
-  organizationId: string
+  organizationId: string,
+  executor?: DbOrTx
 ): Promise<boolean> {
   if (!isHosted && !isAccessControlEnabled) return false
-  return isOrganizationOnEnterprisePlan(organizationId, 'throw')
+  return executor
+    ? isOrganizationOnEnterprisePlan(organizationId, 'throw', executor)
+    : isOrganizationOnEnterprisePlan(organizationId, 'throw')
 }
 
 /**