|
| 1 | +import { sealAccountTokens } from '@sim/db/script-migrations/0037_seal_account_tokens' |
| 2 | +import { openText } from '@sim/db/sealed-text' |
| 3 | +import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure' |
| 4 | +import { generateId } from '@sim/utils/id' |
| 5 | +import postgres, { type Sql } from 'postgres' |
| 6 | +import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest' |
| 7 | + |
| 8 | +const databaseUrl = readTestDatabaseUrl() |
| 9 | + |
| 10 | +interface StoredRow { |
| 11 | + id: string |
| 12 | + access_token: string | null |
| 13 | + refresh_token: string | null |
| 14 | + updated_at: string |
| 15 | +} |
| 16 | + |
| 17 | +/** Copies the account table into an isolated disposable schema; no public rows are touched. */ |
| 18 | +describe('sealing stored account tokens in PostgreSQL', () => { |
| 19 | + let sql: Sql |
| 20 | + let admin: Sql |
| 21 | + const schemaName = `account_sealing_${generateId().replaceAll('-', '')}` |
| 22 | + const previousKey = process.env.ENCRYPTION_KEY |
| 23 | + |
| 24 | + beforeAll(async () => { |
| 25 | + process.env.ENCRYPTION_KEY = 'a'.repeat(64) |
| 26 | + admin = postgres(databaseUrl, { max: 1, onnotice: () => undefined }) |
| 27 | + await admin.unsafe(`CREATE SCHEMA "${schemaName}"`) |
| 28 | + await admin.unsafe( |
| 29 | + `CREATE TABLE "${schemaName}"."account" (LIKE public."account" INCLUDING DEFAULTS)` |
| 30 | + ) |
| 31 | + sql = postgres(databaseUrl, { |
| 32 | + max: 2, |
| 33 | + connection: { search_path: schemaName }, |
| 34 | + onnotice: () => undefined, |
| 35 | + }) |
| 36 | + }) |
| 37 | + |
| 38 | + afterAll(async () => { |
| 39 | + process.env.ENCRYPTION_KEY = previousKey |
| 40 | + await sql?.end() |
| 41 | + if (admin) { |
| 42 | + await admin.unsafe(`DROP SCHEMA IF EXISTS "${schemaName}" CASCADE`) |
| 43 | + await admin.end() |
| 44 | + } |
| 45 | + }) |
| 46 | + |
| 47 | + beforeEach(async () => { |
| 48 | + await sql`TRUNCATE account` |
| 49 | + }) |
| 50 | + |
| 51 | + async function insertAccount(accessToken: string | null, refreshToken: string | null) { |
| 52 | + const id = generateId() |
| 53 | + await sql`INSERT INTO account (id, account_id, provider_id, user_id, access_token, refresh_token, created_at, updated_at) |
| 54 | + VALUES (${id}, ${`subject-${id}`}, 'google-drive', ${`user-${id}`}, ${accessToken}, ${refreshToken}, |
| 55 | + '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z')` |
| 56 | + return id |
| 57 | + } |
| 58 | + |
| 59 | + async function storedRows() { |
| 60 | + return sql<StoredRow[]>` |
| 61 | + SELECT id, access_token, refresh_token, updated_at::text FROM account ORDER BY id` |
| 62 | + } |
| 63 | + |
| 64 | + it('seals every stored token across batches without moving the chain version', async () => { |
| 65 | + await Promise.all( |
| 66 | + Array.from({ length: 5 }, (_, index) => insertAccount(`access-${index}`, `refresh-${index}`)) |
| 67 | + ) |
| 68 | + const accessOnly = await insertAccount('access-only', null) |
| 69 | + const before = await storedRows() |
| 70 | + |
| 71 | + expect(await sealAccountTokens(sql, { mode: 'seal', batchSize: 2, pauseMs: 0 })).toEqual({ |
| 72 | + scanned: 6, |
| 73 | + pending: 6, |
| 74 | + updated: 6, |
| 75 | + }) |
| 76 | + |
| 77 | + const after = await storedRows() |
| 78 | + for (const row of after) { |
| 79 | + const original = before.find((candidate) => candidate.id === row.id) |
| 80 | + expect(row.updated_at).toBe(original?.updated_at) |
| 81 | + expect(row.access_token).toMatch(/^enc:v1:/) |
| 82 | + expect(openText(row.access_token ?? '')).toBe(original?.access_token) |
| 83 | + if (row.id === accessOnly) { |
| 84 | + expect(row.refresh_token).toBeNull() |
| 85 | + } else { |
| 86 | + expect(row.refresh_token).toMatch(/^enc:v1:/) |
| 87 | + expect(openText(row.refresh_token ?? '')).toBe(original?.refresh_token) |
| 88 | + } |
| 89 | + } |
| 90 | + }) |
| 91 | + |
| 92 | + it('finds nothing left to do on a rerun or a dry run after a completed pass', async () => { |
| 93 | + await insertAccount('access', 'refresh') |
| 94 | + await sealAccountTokens(sql, { mode: 'seal', pauseMs: 0 }) |
| 95 | + const sealed = await storedRows() |
| 96 | + |
| 97 | + expect(await sealAccountTokens(sql, { mode: 'seal', pauseMs: 0 })).toMatchObject({ |
| 98 | + pending: 0, |
| 99 | + updated: 0, |
| 100 | + }) |
| 101 | + expect(await sealAccountTokens(sql, { mode: 'seal', dryRun: true, pauseMs: 0 })).toMatchObject({ |
| 102 | + pending: 0, |
| 103 | + }) |
| 104 | + expect(await storedRows()).toEqual(sealed) |
| 105 | + }) |
| 106 | + |
| 107 | + it('reports pending rows on a dry run without writing them', async () => { |
| 108 | + await insertAccount('access', 'refresh') |
| 109 | + const before = await storedRows() |
| 110 | + |
| 111 | + expect(await sealAccountTokens(sql, { mode: 'seal', dryRun: true, pauseMs: 0 })).toEqual({ |
| 112 | + scanned: 1, |
| 113 | + pending: 1, |
| 114 | + updated: 0, |
| 115 | + }) |
| 116 | + expect(await storedRows()).toEqual(before) |
| 117 | + }) |
| 118 | + |
| 119 | + it('restores the exact plaintext when unsealing for a rollback', async () => { |
| 120 | + await insertAccount('ya29.access:with:colons', 'refresh') |
| 121 | + await insertAccount(null, 'refresh-only') |
| 122 | + const plaintext = await storedRows() |
| 123 | + await sealAccountTokens(sql, { mode: 'seal', pauseMs: 0 }) |
| 124 | + |
| 125 | + expect(await sealAccountTokens(sql, { mode: 'unseal', pauseMs: 0 })).toMatchObject({ |
| 126 | + updated: 2, |
| 127 | + }) |
| 128 | + expect(await storedRows()).toEqual(plaintext) |
| 129 | + }) |
| 130 | +}) |
0 commit comments