Skip to content

Commit d344d56

Browse files
waleedlatif1waleedlatif1
committed
improvement(auth): read the sealed context from the column definition
1 parent 699b984 commit d344d56

4 files changed

Lines changed: 160 additions & 69 deletions

File tree

‎packages/db/sealed-text.test.ts‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import { account } from '@sim/db/schema'
2+
import { sealedTextEquals } from '@sim/db/sealed-text'
3+
import { alias, PgDialect } from 'drizzle-orm/pg-core'
4+
import { beforeEach, describe, expect, it, vi } from 'vitest'
5+
6+
describe('sealedTextEquals', () => {
7+
beforeEach(() => {
8+
vi.stubEnv('ENCRYPTION_KEY', 'c'.repeat(64))
9+
vi.stubEnv('ENCRYPT_ACCOUNT_TOKENS', 'true')
10+
})
11+
12+
it('matches the sealed form a write stores, through an alias of the table', () => {
13+
const stored = account.refreshToken.mapToDriverValue('refresh-token')
14+
const sibling = alias(account, 'sibling')
15+
16+
const { params } = new PgDialect().sqlToQuery(
17+
sealedTextEquals(sibling.refreshToken, 'refresh-token')
18+
)
19+
expect(params).toEqual(['refresh-token', stored])
20+
})
21+
})

‎packages/db/sealed-text.ts‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { createLogger } from '@sim/logger'
22
import { isSealed, seal, sealingKeyId, unseal } from '@sim/security/encryption'
3-
import { getTableName, type SQL, sql } from 'drizzle-orm'
3+
import { toRecord } from '@sim/utils/object'
4+
import { type SQL, sql } from 'drizzle-orm'
45
import { type AnyPgColumn, customType } from 'drizzle-orm/pg-core'
56

67
const logger = createLogger('SealedText')
@@ -56,9 +57,17 @@ export function reportAccountTokenSealing(): void {
5657
})
5758
}
5859

59-
/** The context a sealed-text column authenticates its values with, e.g. `account.access_token`. */
60+
/**
61+
* The context a sealed-text column authenticates its values with (`account.access_token`), read
62+
* from the column's own definition rather than rebuilt from its table, so an `alias()` of the
63+
* table or a later rename never changes it. Drizzle keeps the field config on the column's
64+
* `config`, which aliased columns share.
65+
*/
6066
export function sealedTextContext(column: AnyPgColumn): string {
61-
return `${getTableName(column.table)}.${column.name}`
67+
const context = toRecord(toRecord(Reflect.get(column, 'config')).fieldConfig).context
68+
if (typeof context !== 'string')
69+
throw new Error(`Column ${column.name} is not a sealed-text column`)
70+
return context
6271
}
6372

6473
/** Seals `plaintext` for `context` with `ENCRYPTION_KEY`; an already sealed value is returned as is. */
@@ -80,7 +89,7 @@ export function openText(stored: string, context: string): string {
8089
*/
8190
export function sealedText(table: string, column: string) {
8291
const context = `${table}.${column}`
83-
return customType<{ data: string; driverData: string }>({
92+
return customType<{ data: string; driverData: string; config: { context: string } }>({
8493
dataType() {
8594
return 'text'
8695
},
@@ -91,7 +100,7 @@ export function sealedText(table: string, column: string) {
91100
fromDriver(value) {
92101
return openText(value, context)
93102
},
94-
})(column)
103+
})(column, { context })
95104
}
96105

97106
/**

‎scripts/check-sealed-columns.test.ts‎

Lines changed: 47 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,53 +1,74 @@
11
import { describe, expect, it } from 'vitest'
22
import { findViolations, readSealedColumns } from './check-sealed-columns'
33

4-
const SCHEMA = `
4+
const COLUMNS = readSealedColumns(`
55
export const account = pgTable(
66
'account',
77
{
88
accessToken: sealedText('account', 'access_token'),
9-
refreshToken: sealedText('acount', 'refresh_token'),
9+
refreshToken: sealedText('account', 'refresh_token'),
1010
}
1111
)
12-
`
12+
`)
1313

14-
describe('check-sealed-columns', () => {
15-
it('rejects a sealed column whose context names another table', () => {
16-
const { columns, violations } = readSealedColumns(SCHEMA)
17-
expect(columns.map((column) => column.reference)).toEqual([
18-
'account.accessToken',
19-
'account.refreshToken',
20-
])
21-
expect(violations).toMatchObject([{ rule: 'context-mismatch', line: 6 }])
22-
})
14+
function found(source: string, file = 'lib/file.ts') {
15+
return findViolations(file, source, COLUMNS)
16+
.map(({ rule, line }) => [rule, line])
17+
.sort(([, a], [, b]) => Number(a) - Number(b))
18+
}
2319

24-
it('follows a renamed import of the table', () => {
25-
const { columns } = readSealedColumns(SCHEMA)
26-
const source = [
27-
"import { account as accountTable, user } from '@sim/db/schema'",
28-
'where(eq(accountTable.refreshToken, token))',
29-
].join('\n')
30-
expect(findViolations('file.ts', source, columns)).toMatchObject([
31-
{ rule: 'comparison', line: 2 },
20+
describe('check-sealed-columns', () => {
21+
it('reads every sealed column with its physical name', () => {
22+
expect(COLUMNS).toMatchObject([
23+
{ reference: 'account.accessToken', columnName: 'access_token' },
24+
{ reference: 'account.refreshToken', columnName: 'refresh_token' },
3225
])
3326
})
3427

35-
it('flags comparisons and raw interpolations but not the sealed helper', () => {
36-
const { columns } = readSealedColumns(SCHEMA)
28+
it('flags comparison operators and raw interpolations but not the sealed helper', () => {
3729
const source = [
3830
'where(eq(account.refreshToken, token))',
3931
'where(inArray(schema.account.accessToken, tokens))',
4032
`sql\`\${account.accessToken} IS NOT NULL\``,
33+
'where(like(account.accessToken, prefix))',
34+
'where(gte(account.refreshToken, token))',
4135
'where(sealedTextEquals(account.refreshToken, token))',
36+
'where(isNotNull(account.refreshToken))',
4237
'select({ accessToken: account.accessToken })',
4338
].join('\n')
44-
const found = findViolations('file.ts', source, columns)
45-
expect(
46-
found.map(({ rule, line }) => [rule, line]).sort(([, a], [, b]) => Number(a) - Number(b))
47-
).toEqual([
39+
expect(found(source)).toEqual([
4840
['comparison', 1],
4941
['comparison', 2],
5042
['raw-interpolation', 3],
43+
['comparison', 4],
44+
['comparison', 5],
45+
])
46+
})
47+
48+
it('follows renamed imports and alias() bindings of the table', () => {
49+
const source = [
50+
"import { account as accountTable, user } from '@sim/db/schema'",
51+
'where(eq(accountTable.refreshToken, token))',
52+
"const sibling = alias(account, 'sibling')",
53+
'where(ne(sibling.accessToken, token))',
54+
].join('\n')
55+
expect(found(source)).toEqual([
56+
['comparison', 2],
57+
['comparison', 4],
58+
])
59+
})
60+
61+
it('flags physical column names in raw SQL outside tests and the backfill', () => {
62+
const source = [
63+
`await db.execute(sql\`SELECT access_token FROM account WHERE id = \${id}\`)`,
64+
"sql.raw('refresh_token IS NULL')",
65+
`sql\`SELECT id FROM oauth_refresh_token WHERE grant = \${'refresh_token'}\``,
66+
].join('\n')
67+
expect(found(source)).toEqual([
68+
['raw-column-name', 1],
69+
['raw-column-name', 2],
5170
])
71+
expect(found(source, 'lib/file.integration.ts')).toEqual([])
72+
expect(found(source, 'packages/db/script-migrations/0037_seal_account_tokens.ts')).toEqual([])
5273
})
5374
})

‎scripts/check-sealed-columns.ts‎

Lines changed: 78 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,20 @@
11
#!/usr/bin/env bun
22
/**
33
* Fails when a sealed-text column (`sealedText(...)` in `packages/db/schema.ts`) is compared or
4-
* interpolated in a way that bypasses its sealing.
4+
* read in a way that bypasses its sealing.
55
*
6-
* - `comparison`: `eq`/`ne`/`inArray`/`notInArray` against the column. The bound parameter goes
6+
* - `comparison`: a Drizzle comparison operator against the column. Its bound parameter goes
77
* through the column's mapper, so a process that seals writes compares the sealed form and
8-
* misses rows stored before sealing. `sealedTextEquals` from `@sim/db/sealed-text` matches both.
8+
* misses rows stored before sealing (ordering and pattern operators are meaningless on
9+
* ciphertext). `sealedTextEquals` from `@sim/db/sealed-text` matches both forms; `isNull` and
10+
* `isNotNull` are unaffected and allowed.
911
* - `raw-interpolation`: `${table.column}` inside a raw `sql` template. A raw expression is read
1012
* or compared as stored, so it sees ciphertext; select the column itself, or use the helper.
11-
* - `context-mismatch`: a `sealedText(table, column)` whose table is not the one it is declared in.
12-
* The pair is the authenticated context `sealedTextEquals` rebuilds from the column, so a
13-
* mismatch makes every equality miss.
13+
* - `raw-column-name`: the physical column name inside a `sql` template or `sql.raw(...)` body,
14+
* which reads the stored form the same way. Tests and the sealing backfill read it on purpose.
15+
*
16+
* Table aliases count as the table: renamed imports from `@sim/db` / `@sim/db/schema` and
17+
* `alias(table, ...)` bindings.
1418
*
1519
* Swept clean when it landed, so there is no baseline: every hit fails.
1620
*
@@ -23,15 +27,42 @@ import { escapeRegExp } from '@sim/utils/string'
2327

2428
const ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '..')
2529
const SCHEMA = 'packages/db/schema.ts'
26-
const HELPER = 'packages/db/sealed-text.ts'
30+
const SKIPPED_FILES = new Set([
31+
'packages/db/sealed-text.ts',
32+
'scripts/check-sealed-columns.test.ts',
33+
])
34+
/** Files that read the stored form on purpose, so only `raw-column-name` skips them. */
35+
const STORED_FORM_READERS =
36+
/(?:\.test\.tsx?|\.integration\.ts|\/scripts\/test-[\w-]+-e2e\.ts)$|^packages\/db\/script-migrations\/0037_seal_account_tokens\.ts$/
2737
const SCAN_DIRS = ['apps', 'packages', 'scripts']
2838
const SKIP_DIRS = new Set(['node_modules', '.next', '.turbo', 'coverage', 'dist', 'build', 'out'])
2939
const SOURCE_FILE = /\.(?:ts|tsx|mts|cts)$/
3040

3141
const TABLE_DECLARATION = /export const (\w+) = pgTable\(\s*'(\w+)'/g
3242
const SEALED_COLUMN = /(\w+):\s*sealedText\(\s*'(\w+)'\s*,\s*'(\w+)'\s*\)/g
43+
const COMPARISON_OPERATORS = [
44+
'eq',
45+
'ne',
46+
'gt',
47+
'gte',
48+
'lt',
49+
'lte',
50+
'between',
51+
'notBetween',
52+
'like',
53+
'notLike',
54+
'ilike',
55+
'notIlike',
56+
'inArray',
57+
'notInArray',
58+
'arrayContains',
59+
'arrayContained',
60+
'arrayOverlaps',
61+
]
62+
const SQL_TEMPLATE = /\bsql(?:<[^>\n]*>)?`((?:[^`\\]|\\.)*)`/g
63+
const SQL_RAW = /\bsql\.raw\(\s*(['"`])((?:(?!\1)[\s\S])*)\1/g
3364

34-
type Rule = 'comparison' | 'raw-interpolation' | 'context-mismatch'
65+
type Rule = 'comparison' | 'raw-interpolation' | 'raw-column-name'
3566

3667
interface Violation {
3768
file: string
@@ -45,6 +76,8 @@ interface SealedColumn {
4576
table: string
4677
/** The column's property (`accessToken`). */
4778
property: string
79+
/** The physical column name (`access_token`). */
80+
columnName: string
4881
/** How code references it, for reports (`account.accessToken`). */
4982
reference: string
5083
}
@@ -53,63 +86,69 @@ function lineOf(source: string, index: number): number {
5386
return source.slice(0, index).split('\n').length
5487
}
5588

56-
/** Reads every sealed column from the schema and checks each names the table it is declared in. */
57-
export function readSealedColumns(schemaSource: string): {
58-
columns: SealedColumn[]
59-
violations: Violation[]
60-
} {
89+
/** Reads every sealed column from the schema. */
90+
export function readSealedColumns(schemaSource: string): SealedColumn[] {
6191
const tables = [...schemaSource.matchAll(TABLE_DECLARATION)].map((match) => ({
6292
exportName: match[1],
63-
tableName: match[2],
6493
index: match.index ?? 0,
6594
}))
6695
const columns: SealedColumn[] = []
67-
const violations: Violation[] = []
6896
for (const match of schemaSource.matchAll(SEALED_COLUMN)) {
69-
const index = match.index ?? 0
70-
const [text, property, tableArg] = match
71-
const table = tables.filter((candidate) => candidate.index < index).at(-1)
97+
const [, property, , columnName] = match
98+
const table = tables.filter((candidate) => candidate.index < (match.index ?? 0)).at(-1)
7299
if (!table) continue
73100
columns.push({
74101
table: table.exportName,
75102
property,
103+
columnName,
76104
reference: `${table.exportName}.${property}`,
77105
})
78-
if (tableArg !== table.tableName) {
79-
violations.push({
80-
file: SCHEMA,
81-
line: lineOf(schemaSource, index),
82-
rule: 'context-mismatch',
83-
text,
84-
})
85-
}
86106
}
87-
return { columns, violations }
107+
return columns
88108
}
89109

90-
/** The local names a file binds `table` to: itself, plus any `import { table as alias }` rename. */
110+
/**
111+
* The local names a file binds `table` to: itself, any `import { table as x }` rename, and any
112+
* `const x = alias(table, ...)`.
113+
*/
91114
function tableNames(source: string, table: string): string[] {
92-
const rename = new RegExp(`\\b${escapeRegExp(table)}\\s+as\\s+(\\w+)`, 'g')
115+
const escaped = escapeRegExp(table)
116+
const rename = new RegExp(`\\b${escaped}\\s+as\\s+(\\w+)`, 'g')
93117
const imports =
94118
source.match(/import\s+(?:type\s+)?\{[^}]*\}\s*from\s*'@sim\/db(?:\/schema)?'/g) ?? []
95-
const aliases = imports.flatMap((statement) => [...statement.matchAll(rename)].map((m) => m[1]))
96-
return [table, ...aliases]
119+
const renames = imports.flatMap((statement) => [...statement.matchAll(rename)].map((m) => m[1]))
120+
const aliasBinding = new RegExp(`\\b(\\w+)\\s*=\\s*alias\\(\\s*(?:\\w+\\.)?${escaped}\\b`, 'g')
121+
const aliases = [...source.matchAll(aliasBinding)].map((m) => m[1])
122+
return [table, ...renames, ...aliases]
97123
}
98124

99-
/** Finds comparisons and raw interpolations of `columns` in one source file. */
125+
/** Finds comparisons and raw reads of `columns` in one source file. */
100126
export function findViolations(file: string, source: string, columns: SealedColumn[]): Violation[] {
101127
const violations: Violation[] = []
128+
const report = (rule: Rule, index: number, text: string) =>
129+
violations.push({ file, line: lineOf(source, index), rule, text })
130+
const operators = COMPARISON_OPERATORS.join('|')
102131
for (const { table, property } of columns) {
103132
const names = tableNames(source, table).map(escapeRegExp).join('|')
104133
const column = `(?:\\w+\\.)?(?:${names})\\.${escapeRegExp(property)}\\b`
105134
const patterns: [Rule, RegExp][] = [
106-
['comparison', new RegExp(`\\b(?:eq|ne|inArray|notInArray)\\(\\s*${column}`, 'g')],
135+
['comparison', new RegExp(`\\b(?:${operators})\\(\\s*${column}`, 'g')],
107136
['raw-interpolation', new RegExp(`\\$\\{\\s*${column}\\s*\\}`, 'g')],
108137
]
109138
for (const [rule, pattern] of patterns) {
110-
for (const match of source.matchAll(pattern)) {
111-
violations.push({ file, line: lineOf(source, match.index ?? 0), rule, text: match[0] })
112-
}
139+
for (const match of source.matchAll(pattern)) report(rule, match.index ?? 0, match[0])
140+
}
141+
}
142+
if (STORED_FORM_READERS.test(file)) return violations
143+
const names = new RegExp(`\\b(?:${columns.map((c) => escapeRegExp(c.columnName)).join('|')})\\b`)
144+
for (const pattern of [SQL_TEMPLATE, SQL_RAW]) {
145+
for (const match of source.matchAll(pattern)) {
146+
const body = match[match.length - 1].replace(/\$\{[^}]*\}/g, (value) =>
147+
' '.repeat(value.length)
148+
)
149+
const hit = names.exec(body)
150+
if (hit)
151+
report('raw-column-name', (match.index ?? 0) + match[0].indexOf(body) + hit.index, hit[0])
113152
}
114153
}
115154
return violations
@@ -126,15 +165,16 @@ function* sourceFiles(dir: string): Generator<string> {
126165
}
127166

128167
if (import.meta.main) {
129-
const { columns, violations } = readSealedColumns(readFileSync(path.join(ROOT, SCHEMA), 'utf8'))
168+
const columns = readSealedColumns(readFileSync(path.join(ROOT, SCHEMA), 'utf8'))
169+
const violations: Violation[] = []
130170
if (columns.length === 0) {
131171
console.error(`No sealedText columns found in ${SCHEMA}; update this audit with the schema.`)
132172
process.exit(1)
133173
}
134174
for (const dir of SCAN_DIRS) {
135175
for (const absolute of sourceFiles(path.join(ROOT, dir))) {
136176
const file = path.relative(ROOT, absolute)
137-
if (file === HELPER || file === 'scripts/check-sealed-columns.test.ts') continue
177+
if (SKIPPED_FILES.has(file)) continue
138178
violations.push(...findViolations(file, readFileSync(absolute, 'utf8'), columns))
139179
}
140180
}

0 commit comments

Comments
 (0)