diff --git a/src/core/postgres-engine.ts b/src/core/postgres-engine.ts index dc536c73a..ea5c4eb35 100644 --- a/src/core/postgres-engine.ts +++ b/src/core/postgres-engine.ts @@ -191,11 +191,17 @@ export class PostgresEngine implements BrainEngine { const detailLow = opts?.detail === 'low'; // Search-only timeout: prevents DoS via expensive queries without - // affecting long-running operations like embed --all or bulk import - await sql`SET statement_timeout = '8s'`; - try { + // affecting long-running operations like embed --all or bulk import. + // SET LOCAL inside sql.begin() scopes the GUC to the transaction so + // it can never leak onto a pooled connection returned to other + // callers. A bare `SET statement_timeout` goes to an arbitrary + // connection from the pool, lives past this method, and either + // clips an unrelated caller's long-running query (DoS) or — via + // `SET statement_timeout = 0` — disables the guard for them. + const rows = await sql.begin(async sql => { + await sql`SET LOCAL statement_timeout = '8s'`; // CTE: rank pages by FTS score, then pick the best chunk per page in SQL - const rows = await sql` + return await sql` WITH ranked_pages AS ( SELECT p.id, p.slug, p.title, p.type, ts_rank(p.search_vector, websearch_to_tsquery('english', ${query})) AS score @@ -221,10 +227,8 @@ export class PostgresEngine implements BrainEngine { FROM best_chunks ORDER BY score DESC `; - return rows.map(rowToSearchResult); - } finally { - await sql`SET statement_timeout = '0'`; - } + }); + return rows.map(rowToSearchResult); } async searchVector(embedding: Float32Array, opts?: SearchOpts): Promise { @@ -241,10 +245,12 @@ export class PostgresEngine implements BrainEngine { const vecStr = '[' + Array.from(embedding).join(',') + ']'; - // Search-only timeout (see searchKeyword for rationale) - await sql`SET statement_timeout = '8s'`; - try { - const rows = await sql` + // Search-only timeout (see searchKeyword for rationale). SET LOCAL + + // sql.begin ensures the GUC stays transaction-scoped on the pooled + // connection. + const rows = await sql.begin(async sql => { + await sql`SET LOCAL statement_timeout = '8s'`; + return await sql` SELECT p.slug, p.id as page_id, p.title, p.type, cc.id as chunk_id, cc.chunk_index, cc.chunk_text, cc.chunk_source, @@ -260,10 +266,8 @@ export class PostgresEngine implements BrainEngine { LIMIT ${limit} OFFSET ${offset} `; - return rows.map(rowToSearchResult); - } finally { - await sql`SET statement_timeout = '0'`; - } + }); + return rows.map(rowToSearchResult); } async getEmbeddingsByChunkIds(ids: number[]): Promise> { diff --git a/test/postgres-engine.test.ts b/test/postgres-engine.test.ts new file mode 100644 index 000000000..31e968b68 --- /dev/null +++ b/test/postgres-engine.test.ts @@ -0,0 +1,112 @@ +/** + * postgres-engine.ts source-level guardrails. + * + * Live Postgres coverage for search paths lives in test/e2e/search-quality.test.ts. + * This file stays fast and DB-free: it inspects the source of + * src/core/postgres-engine.ts to lock in decisions that protect the + * shared connection pool from per-request GUC leaks. + * + * Regression: R6-F006 / R4-F002. + * searchKeyword and searchVector used to call bare + * await sql`SET statement_timeout = '8s'` + * ...query... + * finally { await sql`SET statement_timeout = '0'` } + * against the shared pool. Each tagged template picks an arbitrary + * connection, so the SET, the query, and the reset could all land on + * DIFFERENT connections. Worst case: the 8s GUC sticks on some pooled + * connection and clips the next caller's long-running query; or the + * reset to 0 lands on a connection that other code expected to be + * protected. The fix wraps each query in sql.begin() and uses + * SET LOCAL so the GUC is transaction-scoped and auto-resets on + * COMMIT/ROLLBACK, regardless of error path. + */ + +import { describe, test, expect } from 'bun:test'; +import { readFileSync } from 'fs'; +import { join } from 'path'; + +const SRC = readFileSync( + join(import.meta.dir, '..', 'src', 'core', 'postgres-engine.ts'), + 'utf-8', +); + +describe('postgres-engine / search path timeout isolation', () => { + test('no bare `SET statement_timeout` statement survives', () => { + // Strip comments so the commentary mentioning the anti-pattern does + // not trigger a false positive. Block-comment + line-comment strip. + const stripped = SRC + .replace(/\/\*[\s\S]*?\*\//g, '') + .replace(/(^|\s)\/\/[^\n]*/g, '$1'); + + // Match a tagged-template statement of the form + // sql`SET statement_timeout = ...` + // that is NOT preceded by LOCAL. This is the exact shape that bleeds + // onto pooled connections; SET LOCAL is safe inside a transaction. + const bare = stripped.match( + /sql`\s*SET\s+(?!LOCAL\s)statement_timeout\b[^`]*`/gi, + ); + expect(bare).toBeNull(); + }); + + test('searchKeyword wraps its query in sql.begin()', () => { + const fn = extractMethod(SRC, 'searchKeyword'); + expect(fn).toMatch(/sql\.begin\s*\(\s*async\s+sql\s*=>/); + }); + + test('searchVector wraps its query in sql.begin()', () => { + const fn = extractMethod(SRC, 'searchVector'); + expect(fn).toMatch(/sql\.begin\s*\(\s*async\s+sql\s*=>/); + }); + + test('both search methods use SET LOCAL for the timeout', () => { + const keyword = extractMethod(SRC, 'searchKeyword'); + const vector = extractMethod(SRC, 'searchVector'); + expect(keyword).toMatch(/SET\s+LOCAL\s+statement_timeout/); + expect(vector).toMatch(/SET\s+LOCAL\s+statement_timeout/); + }); + + test('neither search method clears the timeout with `SET statement_timeout = 0`', () => { + // The reset-to-zero pattern was the other half of the leak: if SET + // LOCAL is in play, COMMIT handles the reset and an explicit + // `SET statement_timeout = '0'` would itself leak the GUC change + // onto the returned connection. Strip comments first so the + // commentary in the method itself (which quotes the anti-pattern + // to explain it) does not trigger a false positive. + const keyword = stripComments(extractMethod(SRC, 'searchKeyword')); + const vector = stripComments(extractMethod(SRC, 'searchVector')); + expect(keyword).not.toMatch(/SET\s+statement_timeout\s*=\s*['"]?0/); + expect(vector).not.toMatch(/SET\s+statement_timeout\s*=\s*['"]?0/); + }); +}); + +function stripComments(s: string): string { + return s + .replace(/\/\*[\s\S]*?\*\//g, '') + .replace(/(^|\s)\/\/[^\n]*/g, '$1'); +} + +// extractMethod grabs the body of a class method by brace-matching from +// its opening line. Returns the method body up to the matching closing +// brace. Good enough for the small number of methods in this file. +function extractMethod(source: string, name: string): string { + // Find "async (" at method-definition indentation (2 spaces). + const openRe = new RegExp(`^\\s+async\\s+${name}\\s*\\(`, 'm'); + const match = openRe.exec(source); + if (!match) { + throw new Error(`method ${name} not found in postgres-engine.ts`); + } + // Scan forward balancing braces. + let i = source.indexOf('{', match.index); + if (i < 0) throw new Error(`no opening brace for ${name}`); + const start = i; + let depth = 0; + for (; i < source.length; i++) { + const c = source[i]; + if (c === '{') depth++; + else if (c === '}') { + depth--; + if (depth === 0) return source.slice(start, i + 1); + } + } + throw new Error(`unbalanced braces in ${name}`); +}