feat(db): add manual Trigger.dev task for the account token sealing backfill (#9393)

* feat(db): add manual Trigger.dev task for the account token sealing backfill

* fix(db): reject seal-account-tokens payloads without a boolean dryRun

Author: waleedlatif1

apps/sim/background/seal-account-tokens.test.ts

@@ -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 })
+  })
+})

apps/sim/background/seal-account-tokens.ts

@@ -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)}`)
+  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),
+})

packages/db/sealed-text.ts

@@ -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