Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 67 additions & 0 deletions apps/sim/background/seal-account-tokens.test.ts
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 })
})
})
98 changes: 98 additions & 0 deletions apps/sim/background/seal-account-tokens.ts
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
Comment thread
waleedlatif1 marked this conversation as resolved.
if (mode !== 'seal' && mode !== 'unseal') throw new Error(`Unknown mode ${String(mode)}`)
Comment thread
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),
})
5 changes: 5 additions & 0 deletions packages/db/sealed-text.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,11 @@ function writeSealingKey(): Buffer | null {
return key
}

/** Whether this process seals account tokens it writes: the flag is on and the key is usable. */
export function isAccountTokenWriteSealingEnabled(): boolean {
return writeSealingKey() !== null
}

/** Logs once per process whether writes seal account tokens, and the key id (never the key). */
export function reportAccountTokenSealing(): void {
if (reportedSealing) return
Expand Down
Loading