fix(credentials): decide subject reuse under the lock on every keyless reconnect

A reconnect that kept its subject skipped the uniqueness check, but whether
it kept the subject was decided from the secret read before provider
verification. If another reconnect moved the credential to a new subject and a
second credential claimed the old one meanwhile, the stale reconnect wrote the
old subject back unchecked and two credentials held it.

Every keyless reconnect now locks the subject it is writing, then the
credential row (FOR UPDATE), re-reads the row's current subject, and runs the
sibling check whenever the two differ, with the decision and the write in one
transaction.

Author: waleedlatif1

apps/sim/lib/credentials/__integration__/workload-identity-subject-reuse.integration.ts

@@ -34,6 +34,8 @@ const originalFetch = globalThis.fetch
 
 /** Subjects the fixture service principal's federation policy pins. */
 const pinnedSubjects = new Set<string>()
+/** Runs once inside the next token exchange, while a save is still verifying. */
+let duringVerification: (() => Promise<void>) | undefined
 
 function grant() {
   const issued = issueWorkloadIdentitySubject('databricks', owner, userId)
@@ -102,6 +104,9 @@ describe('reusing a keyless trust setup', () => {
     vi.stubGlobal('fetch', async (input: string | URL | Request, init?: RequestInit) => {
       const target = typeof input === 'string' || input instanceof URL ? String(input) : input.url
       if (target === `https://${HOST}/oidc/v1/token`) {
+        const interleaved = duringVerification
+        duringVerification = undefined
+        await interleaved?.()
         const subjectToken = new URLSearchParams(String(init?.body)).get('subject_token') ?? ''
         const subject = subjectToken.replace(/^token-for:/, '')
         return pinnedSubjects.has(subject)
@@ -157,6 +162,22 @@ describe('reusing a keyless trust setup', () => {
     expect((await reconnect(credentialId, keyless(setup))).success).toBe(true)
   })
 
+  it('re-checks a kept subject under the lock when the credential moved off it during verification', async () => {
+    const original = grant()
+    const mover = await create('Mover', keyless(original))
+    const moverId = mover.credential?.id ?? ''
+    const moved = grant()
+    duringVerification = async () => {
+      expect((await reconnect(moverId, keyless(moved))).success).toBe(true)
+      expect((await create('Claimer', keyless(original))).success).toBe(true)
+    }
+
+    const stale = await reconnect(moverId, keyless())
+
+    expect(duringVerification).toBeUndefined()
+    expect(stale).toMatchObject({ success: false, errorCode: 'conflict' })
+  })
+
   it('lets only one of two concurrent creates of one subject through', async () => {
     const setup = grant()
 

apps/sim/lib/credentials/orchestration/index.ts

@@ -48,7 +48,7 @@ import {
 } from '@/lib/credentials/service-account-secret'
 import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors'
 import {
-  assertWorkloadIdentitySubjectUnclaimed,
+  assertReconnectedSubjectUnclaimed,
   WorkloadIdentitySubjectError,
 } from '@/lib/credentials/workload-identity/subject'
 import { invalidateEffectiveDecryptedEnvCache } from '@/lib/environment/utils'
@@ -316,7 +316,7 @@ export async function updateCredentialRecord(
 
     let rotatedSlackBotUserId: string | undefined
     let rotatedAuditMetadata: Record<string, string> | undefined
-    let rotatedWorkloadIdentitySubject: string | undefined
+    let rotatedKeylessSecret: string | undefined
     if (hasRotationSecret) {
       const providerId = params.credential.providerId ?? ''
 
@@ -458,7 +458,7 @@ export async function updateCredentialRecord(
         updates.encryptedServiceAccountKey = secret.encryptedServiceAccountKey
         rotatedSlackBotUserId = secret.botUserId
         rotatedAuditMetadata = secret.auditMetadata
-        rotatedWorkloadIdentitySubject = secret.newWorkloadIdentitySubject
+        if (needsStoredBoundSubject) rotatedKeylessSecret = secret.encryptedServiceAccountKey
 
         if (needsStoredIdentity) {
           const previousIdentity = deriveStoredDisplayName(storedBlob)
@@ -521,12 +521,12 @@ export async function updateCredentialRecord(
     }
 
     updates.updatedAt = new Date()
-    if (rotatedWorkloadIdentitySubject) {
-      const subject = rotatedWorkloadIdentitySubject
+    if (rotatedKeylessSecret) {
+      const rotatedSecret = rotatedKeylessSecret
       try {
         await db.transaction(async (tx) => {
-          await assertWorkloadIdentitySubjectUnclaimed(tx, {
-            subject,
+          await assertReconnectedSubjectUnclaimed(tx, {
+            rotatedSecret,
             scope: resourceScopeFromOwner(params.credential),
             credentialId: params.credentialId,
           })

apps/sim/lib/credentials/workload-identity/subject.ts

@@ -3,7 +3,7 @@ import { credential } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { safeCompare } from '@sim/security/compare'
 import { generateId } from '@sim/utils/id'
-import { and, inArray, isNull, ne } from 'drizzle-orm'
+import { and, eq, inArray, isNull, ne } from 'drizzle-orm'
 import { z } from 'zod'
 import { env } from '@/lib/core/config/env'
 import type { ResourceScope } from '@/lib/core/resource-scope'
@@ -183,30 +183,19 @@ async function readStoredSubject(encrypted: string): Promise<string | undefined
   }
 }
 
-/**
- * Refuses a subject that another live credential of the same owner already
- * holds. A proof stays valid for a day and is not tied to one credential, so
- * without this a member removed from a keyless credential could save a second
- * credential with the same subject and keep assuming the customer's role.
- *
- * The subject lives only inside each credential's encrypted secret, so this
- * decrypts the owner's keyless credentials (one owner holds few) rather than
- * relying on a unique index. It runs inside the transaction that writes the
- * credential: the advisory lock holds concurrent saves of one subject until
- * commit, so the later save sees the earlier row. A sibling whose secret cannot
- * be read might hold the subject, so it refuses the save rather than skip it.
- *
- * @param credentialId The credential being reconnected, which may keep its own subject.
- */
-export async function assertWorkloadIdentitySubjectUnclaimed(
-  tx: DbTransaction,
-  input: { subject: string; scope: ResourceScope; credentialId?: string }
-): Promise<void> {
-  await acquireAdvisoryXactLock(
+function lockSubject(tx: DbTransaction, subject: string): Promise<void> {
+  return acquireAdvisoryXactLock(
     tx,
     'workload_identity_subject',
-    `workload-identity-subject:${input.subject}`
+    `workload-identity-subject:${subject}`
   )
+}
+
+/** Throws unless no live keyless credential of the owner other than `credentialId` holds `subject`. */
+async function assertNoOtherCredentialHolds(
+  tx: DbTransaction,
+  input: { subject: string; scope: ResourceScope; credentialId?: string }
+): Promise<void> {
   const rows = await tx
     .select({
       id: credential.id,
@@ -240,3 +229,60 @@ export async function assertWorkloadIdentitySubjectUnclaimed(
     }
   }
 }
+
+/**
+ * Refuses a subject that another live credential of the same owner already
+ * holds. A proof stays valid for a day and is not tied to one credential, so
+ * without this a member removed from a keyless credential could save a second
+ * credential with the same subject and keep assuming the customer's role.
+ *
+ * The subject lives only inside each credential's encrypted secret, so this
+ * decrypts the owner's keyless credentials (one owner holds few) rather than
+ * relying on a unique index. It runs inside the transaction that writes the
+ * credential: the advisory lock holds concurrent saves of one subject until
+ * commit, so the later save sees the earlier row. A sibling whose secret cannot
+ * be read might hold the subject, so it refuses the save rather than skip it.
+ */
+export async function assertWorkloadIdentitySubjectUnclaimed(
+  tx: DbTransaction,
+  input: { subject: string; scope: ResourceScope }
+): Promise<void> {
+  await lockSubject(tx, input.subject)
+  await assertNoOtherCredentialHolds(tx, input)
+}
+
+/**
+ * The reconnect counterpart of {@link assertWorkloadIdentitySubjectUnclaimed},
+ * run in the transaction that writes `rotatedSecret` onto `credentialId`.
+ *
+ * Whether the reconnect keeps the credential's subject is decided here, from
+ * the row as it stands under the lock, not from the secret read before
+ * provider verification: while that verification ran, another reconnect could
+ * have moved this credential to a different subject and a second credential
+ * could have claimed the old one, so writing it back unchecked would leave two
+ * credentials holding it. Locks the subject first, then the credential row; no
+ * path takes a subject lock while holding a credential row, so the order cannot
+ * deadlock.
+ */
+export async function assertReconnectedSubjectUnclaimed(
+  tx: DbTransaction,
+  input: { rotatedSecret: string; scope: ResourceScope; credentialId: string }
+): Promise<void> {
+  const subject = await readStoredSubject(input.rotatedSecret)
+  if (!subject) return
+  await lockSubject(tx, subject)
+  const [current] = await tx
+    .select({ encryptedServiceAccountKey: credential.encryptedServiceAccountKey })
+    .from(credential)
+    .where(eq(credential.id, input.credentialId))
+    .for('update')
+  const currentSubject = current?.encryptedServiceAccountKey
+    ? await readStoredSubject(current.encryptedServiceAccountKey)
+    : undefined
+  if (currentSubject === subject) return
+  await assertNoOtherCredentialHolds(tx, {
+    subject,
+    scope: input.scope,
+    credentialId: input.credentialId,
+  })
+}