Repository navigation
feat(db): add manual Trigger.dev task for the account token sealing backfill #9393
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,67 @@ | ||
| /** | ||
| * @vitest-environment node | ||
| */ | ||
|
|
||
| import { sealingKeyId } from '@sim/security/encryption' | ||
| import { databaseMockFns } from '@sim/testing/mocks/database.mock' | ||
| import { beforeEach, describe, expect, it, vi } from 'vitest' | ||
|
|
||
| const { mockPostgres } = vi.hoisted(() => ({ mockPostgres: vi.fn() })) | ||
|
|
||
| vi.mock('postgres', () => ({ default: mockPostgres })) | ||
| /** The pass derives each column's sealing context from the real column definitions. */ | ||
| vi.unmock('@sim/db/schema') | ||
|
|
||
| import { | ||
| runSealAccountTokens, | ||
| type SealAccountTokensPayload, | ||
| } from '@/background/seal-account-tokens' | ||
|
|
||
| const ENCRYPTION_KEY = 'ab'.repeat(32) | ||
| const WORKER_KEY_ID = sealingKeyId(Buffer.from(ENCRYPTION_KEY, 'hex')) | ||
|
|
||
| /** A postgres.js client over an empty `account` table. */ | ||
| function emptyAccountTable() { | ||
| return Object.assign(async () => [], { end: vi.fn(async () => {}) }) | ||
| } | ||
|
|
||
| describe('runSealAccountTokens', () => { | ||
| beforeEach(() => { | ||
| vi.stubEnv('ENCRYPTION_KEY', ENCRYPTION_KEY) | ||
| vi.stubEnv('ENCRYPT_ACCOUNT_TOKENS', 'true') | ||
| databaseMockFns.mockResolveDbUrl.mockReturnValue('postgresql://fixture/sim') | ||
| mockPostgres.mockReset() | ||
| mockPostgres.mockImplementation(emptyAccountTable) | ||
| }) | ||
|
|
||
| it('refuses a real seal under a key id other than the worker’s, before connecting', async () => { | ||
| await expect( | ||
| runSealAccountTokens({ mode: 'seal', dryRun: false, expectedKeyId: 'deadbeef' }) | ||
| ).rejects.toThrow(`ENCRYPTION_KEY has key id ${WORKER_KEY_ID}, not the expected deadbeef`) | ||
| expect(mockPostgres).not.toHaveBeenCalled() | ||
| }) | ||
|
|
||
| it('refuses a payload without dryRun instead of treating it as a real run', async () => { | ||
| const payload = { mode: 'seal', expectedKeyId: WORKER_KEY_ID } as SealAccountTokensPayload | ||
|
|
||
| await expect(runSealAccountTokens(payload)).rejects.toThrow('dryRun must be a boolean') | ||
| expect(mockPostgres).not.toHaveBeenCalled() | ||
| }) | ||
|
|
||
| it('refuses to seal while this worker writes plaintext tokens, before connecting', async () => { | ||
| vi.stubEnv('ENCRYPT_ACCOUNT_TOKENS', 'false') | ||
|
|
||
| await expect( | ||
| runSealAccountTokens({ mode: 'seal', dryRun: false, expectedKeyId: WORKER_KEY_ID }) | ||
| ).rejects.toThrow('ENCRYPT_ACCOUNT_TOKENS is not enabled') | ||
| expect(mockPostgres).not.toHaveBeenCalled() | ||
| }) | ||
|
|
||
| it('still allows unsealing while writers store plaintext, which is the rollback', async () => { | ||
| vi.stubEnv('ENCRYPT_ACCOUNT_TOKENS', 'false') | ||
|
|
||
| await expect( | ||
| runSealAccountTokens({ mode: 'unseal', dryRun: false, expectedKeyId: WORKER_KEY_ID }) | ||
| ).resolves.toEqual({ keyId: WORKER_KEY_ID, scanned: 0, pending: 0, updated: 0, foreignKey: 0 }) | ||
| }) | ||
| }) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,98 @@ | ||
| import { resolveDbUrl } from '@sim/db' | ||
| import { sealAccountTokens } from '@sim/db/script-migrations/0037_seal_account_tokens' | ||
| import { isAccountTokenWriteSealingEnabled, readEncryptionKey } from '@sim/db/sealed-text' | ||
| import { createLogger } from '@sim/logger' | ||
| import { sealingKeyId } from '@sim/security/encryption' | ||
| import { task } from '@trigger.dev/sdk' | ||
| import postgres from 'postgres' | ||
|
|
||
| const logger = createLogger('SealAccountTokensTask') | ||
|
|
||
| const MAX_BATCH_SIZE = 1000 | ||
| const MAX_PAUSE_MS = 10_000 | ||
|
|
||
| export interface SealAccountTokensPayload { | ||
| mode: 'seal' | 'unseal' | ||
| dryRun: boolean | ||
| /** The key id the fleet logs at startup (`Account token sealing`); must match this worker's. */ | ||
| expectedKeyId: string | ||
| batchSize?: number | ||
| pauseMs?: number | ||
| } | ||
|
|
||
| function assertIntegerInRange(name: string, value: number | undefined, min: number, max: number) { | ||
| if (value === undefined) return | ||
| if (!Number.isInteger(value) || value < min || value > max) { | ||
| throw new Error(`${name} must be an integer between ${min} and ${max}`) | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Runs the account-token sealing backfill (`0037_seal_account_tokens`) from a worker, which already | ||
| * holds the database and the fleet's `ENCRYPTION_KEY`. Refuses before connecting when the payload | ||
| * names a key id other than this worker's, so a pass can never seal rows under a key the fleet | ||
| * cannot open, and refuses `seal` while this worker's writers still store plaintext. | ||
| * | ||
| * Opens its own single connection rather than reusing the `@sim/db` pool: the pass binds JS arrays | ||
| * through the postgres.js tag, which needs the array type map the shared pool skips | ||
| * (`fetch_types: false`). It resolves the same URL as the worker's `db` client and sets no startup | ||
| * `search_path`, which PgBouncer rejects. | ||
| */ | ||
| export async function runSealAccountTokens(payload: SealAccountTokensPayload) { | ||
| const { mode, dryRun, expectedKeyId, batchSize, pauseMs } = payload | ||
| if (mode !== 'seal' && mode !== 'unseal') throw new Error(`Unknown mode ${String(mode)}`) | ||
|
waleedlatif1 marked this conversation as resolved.
|
||
| if (typeof dryRun !== 'boolean') throw new Error('dryRun must be a boolean') | ||
| if (typeof expectedKeyId !== 'string') throw new Error('expectedKeyId must be a string') | ||
| assertIntegerInRange('batchSize', batchSize, 1, MAX_BATCH_SIZE) | ||
| assertIntegerInRange('pauseMs', pauseMs, 0, MAX_PAUSE_MS) | ||
|
|
||
| const keyId = sealingKeyId(readEncryptionKey()) | ||
| if (expectedKeyId !== keyId) { | ||
| throw new Error(`ENCRYPTION_KEY has key id ${keyId}, not the expected ${expectedKeyId}`) | ||
| } | ||
| if (mode === 'seal' && !isAccountTokenWriteSealingEnabled()) { | ||
| throw new Error( | ||
| 'Refusing to seal: ENCRYPT_ACCOUNT_TOKENS is not enabled on this worker, so writers still store plaintext' | ||
| ) | ||
| } | ||
|
|
||
| const role = process.env.SIM_DB_ROLE?.trim() || 'web' | ||
| const url = resolveDbUrl('DATABASE_URL', role) | ||
| if (!url) throw new Error('Missing DATABASE_URL environment variable') | ||
|
|
||
| logger.info('Account token sealing started', { mode, dryRun, keyId, batchSize, pauseMs }) | ||
| const sql = postgres(url, { | ||
| max: 1, | ||
| prepare: false, | ||
| connect_timeout: 30, | ||
| idle_timeout: 20, | ||
| onnotice: () => {}, | ||
| connection: { application_name: 'sim-account-token-sealing' }, | ||
| }) | ||
| try { | ||
| const result = await sealAccountTokens(sql, { | ||
| mode, | ||
| expectedKeyId, | ||
| dryRun, | ||
| batchSize, | ||
| pauseMs, | ||
| }) | ||
| logger.info('Account token sealing finished', { mode, dryRun, ...result }) | ||
| return result | ||
| } finally { | ||
| await sql.end({ timeout: 5 }) | ||
| } | ||
| } | ||
|
|
||
| /** Manual-only: never scheduled or dispatched by app code; trigger it from the dashboard or API. */ | ||
| export const sealAccountTokensTask = task({ | ||
| id: 'seal-account-tokens', | ||
| machine: 'small-1x', | ||
| maxDuration: 10_800, | ||
| retry: { maxAttempts: 1 }, | ||
| queue: { | ||
| name: 'seal-account-tokens', | ||
| concurrencyLimit: 1, | ||
| }, | ||
| run: (payload: SealAccountTokensPayload) => runSealAccountTokens(payload), | ||
| }) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.