import { describe, expect, test } from 'bun:test'
import { Buffer } from 'node:buffer'
import {
decryptWebhookSecret,
isWebhookSecretEnvelopeCandidate,
parseWebhookSecretEncryptionKeyring,
} from '@wispplace/atproto-utils'
import type { SQL } from 'bun'
import {
type LegacyWebhookSecretRow,
migrateLegacyWebhookSecretRows,
migrateWebhookSecretEnvelopes,
type WebhookSecretEnvelopeMigrationStore,
} from './webhook-secret-encryption'
const keyring = parseWebhookSecretEncryptionKeyring({
WEBHOOK_SECRET_ENCRYPTION_KEY: Buffer.alloc(32, 91).toString('base64url'),
})
interface StoredRow extends LegacyWebhookSecretRow {
token: string
}
class MemorySecretStore implements WebhookSecretEnvelopeMigrationStore {
rows: StoredRow[]
beforeCompareAndSwap?: (row: LegacyWebhookSecretRow) => void
constructor(rows: StoredRow[]) {
this.rows = rows.map((row) => ({ ...row }))
}
async transaction<T>(operation: () => Promise<T>): Promise<T> {
const snapshot = this.rows.map((row) => ({ ...row }))
try {
return await operation()
} catch (error) {
this.rows = snapshot
throw error
}
}
async selectLegacyBatch(limit: number): Promise<ReadonlyArray<LegacyWebhookSecretRow>> {
return this.rows
.filter((row) => !isWebhookSecretEnvelopeCandidate(row.token))
.slice(0, limit)
.map((row) => ({ ...row }))
}
async compareAndSwap(row: LegacyWebhookSecretRow, envelope: string): Promise<boolean> {
this.beforeCompareAndSwap?.(row)
const stored = this.rows.find((candidate) => candidate.did === row.did && candidate.name === row.name)
if (!stored || stored.token !== row.token) return false
stored.token = envelope
return true
}
async countLegacy(): Promise<number> {
return this.rows.filter((row) => !isWebhookSecretEnvelopeCandidate(row.token)).length
}
}
const legacyRows = (): StoredRow[] => [
{ did: 'did:plc:alice', name: 'first', token: 'wsk_plaintext-first' },
{ did: 'did:plc:alice', name: 'second', token: 'wsk_plaintext-second' },
]
describe('webhook secret envelope migration', () => {
test('does not silently preserve legacy plaintext when no key is available', async () => {
const store = new MemorySecretStore(legacyRows())
const result = await migrateLegacyWebhookSecretRows(store, null)
expect(result).toEqual({ encrypted: 0, legacyRemaining: 2, encryptionAvailable: false })
expect(store.rows.map((row) => row.token)).toEqual(['wsk_plaintext-first', 'wsk_plaintext-second'])
})
test('stores only envelopes after a successful migration, never the raw database token', async () => {
const store = new MemorySecretStore(legacyRows())
const rawTokens = store.rows.map((row) => row.token)
const result = await migrateLegacyWebhookSecretRows(store, keyring, { batchSize: 1 })
expect(result).toEqual({ encrypted: 2, legacyRemaining: 0, encryptionAvailable: true })
for (const [index, row] of store.rows.entries()) {
expect(isWebhookSecretEnvelopeCandidate(row.token)).toBe(true)
expect(row.token).not.toContain(rawTokens[index]!)
expect(decryptWebhookSecret(row.token, keyring)).toBe(rawTokens[index])
}
})
test('resumes after a crash following a committed batch', async () => {
const store = new MemorySecretStore(legacyRows())
let crashOnce = true
await expect(
migrateLegacyWebhookSecretRows(store, keyring, {
batchSize: 1,
afterBatchCommitted: () => {
if (crashOnce) {
crashOnce = false
throw new Error('simulated process stop')
}
},
}),
).rejects.toThrow('simulated process stop')
// The first transaction was committed before the process stopped.
expect(isWebhookSecretEnvelopeCandidate(store.rows[0]!.token)).toBe(true)
expect(store.rows[1]!.token).toBe('wsk_plaintext-second')
const resumed = await migrateLegacyWebhookSecretRows(store, keyring, { batchSize: 1 })
expect(resumed).toEqual({ encrypted: 1, legacyRemaining: 0, encryptionAvailable: true })
expect(store.rows.map((row) => decryptWebhookSecret(row.token, keyring))).toEqual([
'wsk_plaintext-first',
'wsk_plaintext-second',
])
})
test('leaves a changed row for the next idempotent compare-and-swap pass', async () => {
const store = new MemorySecretStore([{ did: 'did:plc:alice', name: 'first', token: 'wsk_original' }])
store.beforeCompareAndSwap = () => {
store.beforeCompareAndSwap = undefined
store.rows[0]!.token = 'wsk_changed-by-old-instance'
}
const first = await migrateLegacyWebhookSecretRows(store, keyring)
expect(first).toEqual({ encrypted: 0, legacyRemaining: 1, encryptionAvailable: true })
expect(store.rows[0]!.token).toBe('wsk_changed-by-old-instance')
const resumed = await migrateLegacyWebhookSecretRows(store, keyring)
expect(resumed).toEqual({ encrypted: 1, legacyRemaining: 0, encryptionAvailable: true })
expect(decryptWebhookSecret(store.rows[0]!.token, keyring)).toBe('wsk_changed-by-old-instance')
})
test('uses a reserved transaction and a separate advisory lock around production storage migration', async () => {
const store = new MemorySecretStore([{ did: 'did:plc:alice', name: 'first', token: 'wsk_raw-db-token' }])
const statements: string[] = []
let released = false
let closed = false
let reserved: unknown
const tag = async (strings: TemplateStringsArray, ...values: unknown[]) => {
const statement = strings.join('?').replace(/\s+/g, ' ').trim()
statements.push(statement)
if (statement.includes('SELECT did, name, token')) return await store.selectLegacyBatch(values[1] as number)
if (statement.includes('UPDATE webhook_secrets')) {
return (await store.compareAndSwap(
{ did: values[1] as string, name: values[2] as string, token: values[3] as string },
values[0] as string,
))
? [{ did: values[1] as string }]
: []
}
if (statement.includes('SELECT COUNT(*)')) return [{ count: await store.countLegacy() }]
if (statement.includes('pg_advisory_unlock')) return [{ unlocked: true }]
return []
}
reserved = Object.assign(tag, {
begin: async <T>(operation: () => Promise<T>): Promise<T> => await store.transaction(operation),
release: () => {
released = true
},
close: async () => {
closed = true
},
})
const database = { reserve: async () => reserved } as unknown as SQL
const result = await migrateWebhookSecretEnvelopes(database, keyring)
expect(result).toEqual({ encrypted: 1, legacyRemaining: 0, encryptionAvailable: true })
expect(store.rows[0]!.token).not.toContain('wsk_raw-db-token')
expect(decryptWebhookSecret(store.rows[0]!.token, keyring)).toBe('wsk_raw-db-token')
expect(statements.some((statement) => statement === 'BEGIN')).toBe(true)
expect(statements.some((statement) => statement === 'COMMIT')).toBe(true)
expect(statements.some((statement) => statement.includes('pg_advisory_lock'))).toBe(true)
expect(statements.some((statement) => statement.includes('pg_advisory_unlock'))).toBe(true)
expect(released).toBe(true)
expect(closed).toBe(false)
})
})