|
| 1 | +/** |
| 2 | + * A keyless trust setup (subject plus its day-long proof) binds to one live |
| 3 | + * credential: saving it onto a second credential is refused, even concurrently, |
| 4 | + * while the credential that holds it can still be reconnected. |
| 5 | + */ |
| 6 | + |
| 7 | +import { db } from '@sim/db' |
| 8 | +import { credential, credentialMember, permissions, user, workspace } from '@sim/db/schema' |
| 9 | +import { deleteWorkspaceFixture, insertWorkspaceFixture } from '@sim/db/testing/workspace-fixtures' |
| 10 | +import { jsonResponse } from '@sim/testing/helpers/http' |
| 11 | +import { posthogServerMock } from '@sim/testing/mocks/posthog-server.mock' |
| 12 | +import { generateId } from '@sim/utils/id' |
| 13 | +import { eq } from 'drizzle-orm' |
| 14 | +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' |
| 15 | + |
| 16 | +vi.mock('@/lib/posthog/server', () => posthogServerMock) |
| 17 | +vi.mock('@/lib/credentials/workload-identity/issuer', async (importOriginal) => ({ |
| 18 | + ...(await importOriginal<typeof import('@/lib/credentials/workload-identity/issuer')>()), |
| 19 | + mintWorkloadIdentityToken: async ({ subject }: { subject: string }) => `token-for:${subject}`, |
| 20 | +})) |
| 21 | + |
| 22 | +import { updateCredentialRecord } from '@/lib/credentials/orchestration' |
| 23 | +import { createCredentialRecord } from '@/lib/credentials/orchestration/credential-create' |
| 24 | +import { issueWorkloadIdentitySubject } from '@/lib/credentials/workload-identity/subject' |
| 25 | +import { DATABRICKS_SERVICE_ACCOUNT_PROVIDER_ID } from '@/lib/oauth/types' |
| 26 | + |
| 27 | +const HOST = 'dbc-1234abcd-5678.cloud.databricks.com' |
| 28 | +const CLIENT_ID = '00000000-0000-0000-0000-000000000000' |
| 29 | +const userId = generateId() |
| 30 | +const workspaceId = generateId() |
| 31 | +const owner = { workspaceId } |
| 32 | +const originalFetch = globalThis.fetch |
| 33 | + |
| 34 | +/** Subjects the fixture service principal's federation policy pins. */ |
| 35 | +const pinnedSubjects = new Set<string>() |
| 36 | + |
| 37 | +function grant() { |
| 38 | + const issued = issueWorkloadIdentitySubject('databricks', owner, userId) |
| 39 | + pinnedSubjects.add(issued.subject) |
| 40 | + return issued |
| 41 | +} |
| 42 | + |
| 43 | +function keyless(subject?: { subject: string; subjectProof: string }) { |
| 44 | + return JSON.stringify({ |
| 45 | + authType: 'oidc_federation', |
| 46 | + host: HOST, |
| 47 | + clientId: CLIENT_ID, |
| 48 | + ...(subject ? { subject: subject.subject, subjectProof: subject.subjectProof } : {}), |
| 49 | + }) |
| 50 | +} |
| 51 | + |
| 52 | +function create(displayName: string, serviceAccountJson: string) { |
| 53 | + return createCredentialRecord( |
| 54 | + { |
| 55 | + workspaceId, |
| 56 | + type: 'service_account', |
| 57 | + userId, |
| 58 | + providerId: DATABRICKS_SERVICE_ACCOUNT_PROVIDER_ID, |
| 59 | + displayName, |
| 60 | + serviceAccountJson, |
| 61 | + }, |
| 62 | + { authorizeWorkspace: true } |
| 63 | + ) |
| 64 | +} |
| 65 | + |
| 66 | +async function reconnect(credentialId: string, serviceAccountJson: string) { |
| 67 | + const [row] = await db.select().from(credential).where(eq(credential.id, credentialId)) |
| 68 | + return updateCredentialRecord({ |
| 69 | + credentialId, |
| 70 | + credential: row, |
| 71 | + actorUserId: userId, |
| 72 | + serviceAccountJson, |
| 73 | + }) |
| 74 | +} |
| 75 | + |
| 76 | +describe('reusing a keyless trust setup', () => { |
| 77 | + beforeAll(async () => { |
| 78 | + const now = new Date() |
| 79 | + await db.insert(user).values({ |
| 80 | + id: userId, |
| 81 | + name: 'Subject reuse fixture', |
| 82 | + email: `${userId}@subject-reuse.test`, |
| 83 | + emailVerified: true, |
| 84 | + createdAt: now, |
| 85 | + updatedAt: now, |
| 86 | + }) |
| 87 | + await insertWorkspaceFixture(db, { |
| 88 | + id: workspaceId, |
| 89 | + name: 'Subject reuse fixture', |
| 90 | + ownerId: userId, |
| 91 | + billedAccountUserId: userId, |
| 92 | + }) |
| 93 | + await db.insert(permissions).values({ |
| 94 | + id: generateId(), |
| 95 | + userId, |
| 96 | + entityType: 'workspace', |
| 97 | + entityId: workspaceId, |
| 98 | + permissionType: 'admin', |
| 99 | + }) |
| 100 | + |
| 101 | + vi.stubGlobal('fetch', async (input: string | URL | Request, init?: RequestInit) => { |
| 102 | + const target = typeof input === 'string' || input instanceof URL ? String(input) : input.url |
| 103 | + if (target === `https://${HOST}/oidc/v1/token`) { |
| 104 | + const subjectToken = new URLSearchParams(String(init?.body)).get('subject_token') ?? '' |
| 105 | + const subject = subjectToken.replace(/^token-for:/, '') |
| 106 | + return pinnedSubjects.has(subject) |
| 107 | + ? jsonResponse({ access_token: `access-for:${subject}` }) |
| 108 | + : jsonResponse({ error: 'invalid_client' }, 401) |
| 109 | + } |
| 110 | + if (target === `https://${HOST}/api/2.0/preview/scim/v2/Me`) { |
| 111 | + return jsonResponse({ userName: 'service-principal@example.com' }) |
| 112 | + } |
| 113 | + return originalFetch(input, init) |
| 114 | + }) |
| 115 | + }, 60_000) |
| 116 | + |
| 117 | + afterAll(async () => { |
| 118 | + vi.unstubAllGlobals() |
| 119 | + await db.delete(credential).where(eq(credential.workspaceId, workspaceId)) |
| 120 | + await deleteWorkspaceFixture(db, eq(workspace.id, workspaceId)) |
| 121 | + await db.delete(user).where(eq(user.id, userId)) |
| 122 | + }) |
| 123 | + |
| 124 | + it('refuses a second credential with a subject another credential holds', async () => { |
| 125 | + const setup = grant() |
| 126 | + const first = await create('Lakehouse', keyless(setup)) |
| 127 | + expect(first.success).toBe(true) |
| 128 | + |
| 129 | + const second = await create('Lakehouse again', keyless(setup)) |
| 130 | + |
| 131 | + expect(second).toMatchObject({ success: false, errorCode: 'conflict' }) |
| 132 | + const rows = await db |
| 133 | + .select({ id: credential.id }) |
| 134 | + .from(credential) |
| 135 | + .where(eq(credential.workspaceId, workspaceId)) |
| 136 | + expect(rows.map((row) => row.id)).toEqual([first.credential?.id]) |
| 137 | + }) |
| 138 | + |
| 139 | + it('refuses reconnecting another credential onto a subject that is held', async () => { |
| 140 | + const held = grant() |
| 141 | + expect((await create('Held', keyless(held))).success).toBe(true) |
| 142 | + const other = await create('Other', keyless(grant())) |
| 143 | + const otherId = other.credential?.id ?? '' |
| 144 | + |
| 145 | + const result = await reconnect(otherId, keyless(held)) |
| 146 | + |
| 147 | + expect(result).toMatchObject({ success: false, errorCode: 'conflict' }) |
| 148 | + }) |
| 149 | + |
| 150 | + it('lets the holding credential reconnect with the subject it already has', async () => { |
| 151 | + const setup = grant() |
| 152 | + const created = await create('Reconnects', keyless(setup)) |
| 153 | + const credentialId = created.credential?.id ?? '' |
| 154 | + |
| 155 | + expect((await reconnect(credentialId, keyless())).success).toBe(true) |
| 156 | + expect((await reconnect(credentialId, keyless(setup))).success).toBe(true) |
| 157 | + }) |
| 158 | + |
| 159 | + it('lets only one of two concurrent creates of one subject through', async () => { |
| 160 | + const setup = grant() |
| 161 | + |
| 162 | + const results = await Promise.all([ |
| 163 | + create('Racer one', keyless(setup)), |
| 164 | + create('Racer two', keyless(setup)), |
| 165 | + ]) |
| 166 | + |
| 167 | + expect(results.filter((result) => result.success)).toHaveLength(1) |
| 168 | + expect(results.filter((result) => result.errorCode === 'conflict')).toHaveLength(1) |
| 169 | + }) |
| 170 | + |
| 171 | + it('lets only one of two concurrent reconnects onto one subject through', async () => { |
| 172 | + const first = await create('Reconnect racer one', keyless(grant())) |
| 173 | + const second = await create('Reconnect racer two', keyless(grant())) |
| 174 | + const setup = grant() |
| 175 | + |
| 176 | + const results = await Promise.all([ |
| 177 | + reconnect(first.credential?.id ?? '', keyless(setup)), |
| 178 | + reconnect(second.credential?.id ?? '', keyless(setup)), |
| 179 | + ]) |
| 180 | + |
| 181 | + expect(results.filter((result) => result.success)).toHaveLength(1) |
| 182 | + expect(results.filter((result) => result.errorCode === 'conflict')).toHaveLength(1) |
| 183 | + }) |
| 184 | + |
| 185 | + it('frees the subject once the credential holding it is deleted', async () => { |
| 186 | + const setup = grant() |
| 187 | + const created = await create('Deleted holder', keyless(setup)) |
| 188 | + const credentialId = created.credential?.id ?? '' |
| 189 | + await db.delete(credentialMember).where(eq(credentialMember.credentialId, credentialId)) |
| 190 | + await db.delete(credential).where(eq(credential.id, credentialId)) |
| 191 | + |
| 192 | + expect((await create('Successor', keyless(setup))).success).toBe(true) |
| 193 | + }) |
| 194 | +}) |
0 commit comments