From 6146c3b470dce7380da024a238eab9e6b2174296 Mon Sep 17 00:00:00 2001 From: Gustavo Aragon Date: Thu, 16 Apr 2026 15:04:36 -0300 Subject: [PATCH] postgres-engine: scope search statement_timeout to the transaction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit searchKeyword and searchVector run on a pooled postgres.js client (max: 10 by default). The original code bounded each search with await sql`SET statement_timeout = '8s'` try { await sql`` } finally { await sql`SET statement_timeout = '0'` } but every tagged template is an independent round-trip that picks an arbitrary connection from the pool. The SET, the query, and the reset could all land on DIFFERENT connections. In practice the GUC sticks to whichever connection ran the SET and then gets returned to the pool — the next unrelated caller on that connection inherits the 8s timeout (clipping legitimate long queries) or the reset-to-0 (disabling the guard for whoever expected it). A crash in the middle leaves the state set permanently. Wrap each search in sql.begin(async sql => …). postgres.js reserves a single connection for the transaction body, so the SET LOCAL, the query, and the implicit COMMIT all run on the same connection. SET LOCAL scopes the GUC to the transaction — COMMIT or ROLLBACK restores the previous value automatically, regardless of the code path out. Error paths can no longer leak the GUC. No API change. Timeout value and semantics are identical (8s cap on search queries, no effect on embed --all / bulk import which runs outside these methods). Only one transaction per search — BEGIN + COMMIT round-trips are negligible next to a ranked FTS or pgvector query. Also closes the earlier audit finding R4-F002 which reported the same pattern on searchKeyword. This PR covers both searchKeyword and searchVector so the pool-leak class is fully closed. Tests (test/postgres-engine.test.ts, new file): - No bare SET statement_timeout remains after stripping comments. - searchKeyword and searchVector each wrap their query in sql.begin. - Both use SET LOCAL. - Neither explicitly clears the timeout with SET statement_timeout=0. Source-level guardrails keep the fast unit suite DB-free. Live Postgres coverage of the search path is in test/e2e/search-quality.test.ts, which continues to exercise these methods end-to-end against pgvector when DATABASE_URL is set. --- src/core/postgres-engine.ts | 36 ++++++----- test/postgres-engine.test.ts | 112 +++++++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+), 16 deletions(-) create mode 100644 test/postgres-engine.test.ts 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}`); +}