diff --git a/package-lock.json b/package-lock.json index 7f85727..d05f78a 100644 --- a/package-lock.json +++ b/package-lock.json @@ -12,6 +12,7 @@ "@anthropic-ai/mcpb": "^2.1.2", "@modelcontextprotocol/sdk": "^1.29.0", "cdb-converter": "^0.1.2", + "sql-query-identifier": "^3.1.1", "sql.js": "^1.14.1", "zod": "^3.25.76" }, @@ -4746,6 +4747,15 @@ "node": ">=0.10.0" } }, + "node_modules/sql-query-identifier": { + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/sql-query-identifier/-/sql-query-identifier-3.1.1.tgz", + "integrity": "sha512-yKLx1hMFM5RKyNkFRCJrOFCPX/ZLhVMl45PWdOH1+UuEX3khELbCdgZuBnSaU9hpVvdYGMPqWc5JAAE27aybRA==", + "license": "MIT", + "engines": { + "node": ">= 10.13" + } + }, "node_modules/sql.js": { "version": "1.14.1", "resolved": "https://registry.npmjs.org/sql.js/-/sql.js-1.14.1.tgz", diff --git a/package.json b/package.json index facf5da..ecd5de0 100644 --- a/package.json +++ b/package.json @@ -51,6 +51,7 @@ "@anthropic-ai/mcpb": "^2.1.2", "@modelcontextprotocol/sdk": "^1.29.0", "cdb-converter": "^0.1.2", + "sql-query-identifier": "^3.1.1", "sql.js": "^1.14.1", "zod": "^3.25.76" } diff --git a/src/save-db.ts b/src/save-db.ts index 3c1a5ee..cbdae83 100644 --- a/src/save-db.ts +++ b/src/save-db.ts @@ -66,6 +66,8 @@ export async function withSaveDb>( const cdbBuffer = await readFile(save.path); db = cdbToSql(cdbBuffer, SQL); + db.run("PRAGMA query_only = ON;"); + const output = await fn(db, save); return validResponse(output); diff --git a/src/tools/query-save.ts b/src/tools/query-save.ts index ea0a138..5f0fa86 100644 --- a/src/tools/query-save.ts +++ b/src/tools/query-save.ts @@ -1,4 +1,5 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { identify } from "sql-query-identifier"; import { z } from "zod"; import { withSaveDb } from "../save-db"; @@ -109,42 +110,51 @@ function explainQueryError(error: unknown): Error { ); } + // Raised by `PRAGMA query_only = ON` when a statement tries to write. + if (/readonly database|not authorized/i.test(message)) { + return new Error( + "This tool is read-only — the query attempted to modify the save, which is not allowed.", + ); + } + return error instanceof Error ? error : new Error(message); } /** - * Reject anything that isn't a single read-only `SELECT`/`WITH` statement. - * The save is loaded into an in-memory sql.js database (changes are never - * written back to disk), but we still enforce read-only intent defensively. + * Enforce that a query is a single read-only statement. + * + * Parsing is delegated to `sql-query-identifier`, which tokenizes SQL properly: + * a `;` inside a string literal, comment or quoted identifier is not mistaken + * for a statement separator, and CTEs are classified by their leaf operation — + * `WITH … SELECT` reads (`LISTING`) while `WITH … DELETE` writes + * (`MODIFICATION`). Anything that isn't exactly one `LISTING` statement is + * rejected here; `PRAGMA query_only = ON` (see {@link withSaveDb}) stays as the + * engine-level backstop. */ export function assertReadOnlyQuery(rawQuery: string): string { - // Strip a single trailing semicolon, then reject any further statement - // separators to prevent stacked statements. + // Strip a single trailing semicolon so a normal `SELECT …;` is accepted; + // the returned query is what gets prepared. const query = rawQuery.trim().replace(/;\s*$/, ""); if (query.length === 0) { throw new Error("Query is empty."); } - if (query.includes(";")) { - throw new Error( - "Only a single statement is allowed — remove extra semicolons.", - ); + const statements = identify(query, { strict: false, dialect: "sqlite" }); + + if (statements.length === 0) { + throw new Error("Query is empty."); } - if (!/^(select|with)\b/i.test(query)) { + if (statements.length > 1) { throw new Error( - "Only read-only SELECT (or WITH … SELECT) queries are allowed.", + "Only a single statement is allowed — remove extra semicolons.", ); } - // Defense in depth: reject statements that could mutate or attach data. - const forbidden = - /\b(insert|update|delete|drop|create|alter|replace|attach|detach|reindex|vacuum|pragma|truncate)\b/i; - const match = forbidden.exec(query); - if (match) { + if (statements[0].executionType !== "LISTING") { throw new Error( - `Write/DDL keyword "${match[0].toUpperCase()}" is not allowed — this tool is read-only.`, + "Only read-only SELECT (or WITH … SELECT) queries are allowed.", ); } diff --git a/test/query-save.test.ts b/test/query-save.test.ts deleted file mode 100644 index 5ae6a9f..0000000 --- a/test/query-save.test.ts +++ /dev/null @@ -1,190 +0,0 @@ -import { describe, expect, it } from "vitest"; -import { assertReadOnlyQuery } from "../src/tools/query-save"; - -describe("assertReadOnlyQuery", () => { - describe("allowed queries", () => { - it("accepts a simple SELECT", () => { - const q = "SELECT * FROM DYN_cyclist"; - expect(assertReadOnlyQuery(q)).toBe(q); - }); - - it("accepts SELECT with a trailing semicolon (strips it)", () => { - expect(assertReadOnlyQuery("SELECT 1;")).toBe("SELECT 1"); - }); - - it("accepts SELECT with trailing semicolon and whitespace", () => { - // trim() runs before the semicolon strip, so a space before ';' is preserved - expect(assertReadOnlyQuery("SELECT 1 ; ")).toBe("SELECT 1 "); - }); - - it("accepts SELECT with leading/trailing whitespace", () => { - expect(assertReadOnlyQuery(" SELECT id FROM foo ")).toBe( - "SELECT id FROM foo", - ); - }); - - it("accepts a WITH … SELECT (CTE)", () => { - const q = - "WITH cte AS (SELECT id FROM foo) SELECT * FROM cte WHERE id > 5"; - expect(assertReadOnlyQuery(q)).toBe(q); - }); - - it("accepts lowercase select", () => { - expect(assertReadOnlyQuery("select * from bar")).toBe( - "select * from bar", - ); - }); - - it("accepts mixed-case SELECT", () => { - expect(assertReadOnlyQuery("Select id From foo")).toBe( - "Select id From foo", - ); - }); - - it("accepts SELECT containing a column named 'update_date' (keyword inside identifier)", () => { - const q = "SELECT update_date FROM DYN_cyclist"; - expect(assertReadOnlyQuery(q)).toBe(q); - }); - - it("rejects SELECT with a forbidden keyword inside a LIKE string (known false-positive: \\b matches across %)", () => { - // '\bcreate\b' matches 'create' in '%create%' because '%' is a non-word char. - // This is a known limitation of the regex-based guard. - expect(() => - assertReadOnlyQuery("SELECT * FROM foo WHERE name LIKE '%create%'"), - ).toThrowError('Write/DDL keyword "CREATE" is not allowed'); - }); - - it("accepts SELECT with ORDER BY, LIMIT, GROUP BY", () => { - const q = - "SELECT name, COUNT(*) AS n FROM foo GROUP BY name ORDER BY n DESC LIMIT 10"; - expect(assertReadOnlyQuery(q)).toBe(q); - }); - - it("accepts SELECT with a JOIN", () => { - const q = "SELECT a.id, b.name FROM a INNER JOIN b ON a.id = b.a_id"; - expect(assertReadOnlyQuery(q)).toBe(q); - }); - }); - - describe("empty / blank queries", () => { - it("rejects an empty string", () => { - expect(() => assertReadOnlyQuery("")).toThrowError("Query is empty."); - }); - - it("rejects a string that is only whitespace", () => { - expect(() => assertReadOnlyQuery(" ")).toThrowError("Query is empty."); - }); - - it("rejects a bare semicolon", () => { - expect(() => assertReadOnlyQuery(";")).toThrowError("Query is empty."); - }); - - it("rejects whitespace + semicolon", () => { - expect(() => assertReadOnlyQuery(" ; ")).toThrowError( - "Query is empty.", - ); - }); - }); - - describe("multiple statements", () => { - it("rejects two SELECT statements separated by a semicolon", () => { - expect(() => assertReadOnlyQuery("SELECT 1; SELECT 2")).toThrowError( - "Only a single statement is allowed", - ); - }); - - it("rejects SELECT followed by a write statement", () => { - expect(() => - assertReadOnlyQuery("SELECT 1; DROP TABLE foo"), - ).toThrowError("Only a single statement is allowed"); - }); - }); - - describe("non-SELECT openers", () => { - it("rejects a bare INSERT", () => { - expect(() => - assertReadOnlyQuery("INSERT INTO foo VALUES (1)"), - ).toThrowError("Only read-only SELECT"); - }); - - it("rejects UPDATE", () => { - expect(() => assertReadOnlyQuery("UPDATE foo SET bar = 1")).toThrowError( - "Only read-only SELECT", - ); - }); - - it("rejects DELETE", () => { - expect(() => assertReadOnlyQuery("DELETE FROM foo")).toThrowError( - "Only read-only SELECT", - ); - }); - - it("rejects DROP TABLE", () => { - expect(() => assertReadOnlyQuery("DROP TABLE foo")).toThrowError( - "Only read-only SELECT", - ); - }); - - it("rejects CREATE TABLE", () => { - expect(() => - assertReadOnlyQuery("CREATE TABLE foo (id INTEGER)"), - ).toThrowError("Only read-only SELECT"); - }); - - it("rejects PRAGMA", () => { - expect(() => assertReadOnlyQuery("PRAGMA table_info(foo)")).toThrowError( - "Only read-only SELECT", - ); - }); - }); - - describe("forbidden keywords inside a SELECT", () => { - it("rejects SELECT with inline INSERT via INSERT INTO (subquery trick)", () => { - expect(() => - assertReadOnlyQuery("SELECT * FROM (INSERT INTO foo VALUES (1)) AS t"), - ).toThrowError('Write/DDL keyword "INSERT" is not allowed'); - }); - - it("rejects SELECT containing DROP keyword", () => { - expect(() => assertReadOnlyQuery("SELECT DROP FROM foo")).toThrowError( - 'Write/DDL keyword "DROP" is not allowed', - ); - }); - - it("rejects SELECT containing ATTACH", () => { - expect(() => - assertReadOnlyQuery("SELECT * FROM foo; ATTACH DATABASE 'x' AS y"), - ).toThrowError("Only a single statement is allowed"); - }); - - it("rejects ATTACH as a standalone statement opener", () => { - expect(() => - assertReadOnlyQuery("ATTACH DATABASE 'evil.db' AS e"), - ).toThrowError("Only read-only SELECT"); - }); - - it("rejects VACUUM", () => { - expect(() => assertReadOnlyQuery("SELECT 1 WHERE VACUUM")).toThrowError( - 'Write/DDL keyword "VACUUM" is not allowed', - ); - }); - - it("rejects ALTER", () => { - expect(() => - assertReadOnlyQuery("SELECT * FROM foo ALTER TABLE bar"), - ).toThrowError('Write/DDL keyword "ALTER" is not allowed'); - }); - - it("rejects REPLACE", () => { - expect(() => - assertReadOnlyQuery("SELECT REPLACE(name,'a','b') FROM foo"), - ).toThrowError('Write/DDL keyword "REPLACE" is not allowed'); - }); - - it("rejects TRUNCATE", () => { - expect(() => - assertReadOnlyQuery("SELECT * FROM foo TRUNCATE"), - ).toThrowError('Write/DDL keyword "TRUNCATE" is not allowed'); - }); - }); -}); diff --git a/test/save-db.test.ts b/test/save-db.test.ts index fe152f4..ddfdfdd 100644 --- a/test/save-db.test.ts +++ b/test/save-db.test.ts @@ -22,8 +22,8 @@ const cdbToSqlMock = cdbToSql as Mock; let dir: string; let savePath: string; -/** A fake sql.js database; we only care that it gets closed. */ -const fakeDb = { close: vi.fn() }; +/** A fake sql.js database; we only care that it gets closed and configured. */ +const fakeDb = { close: vi.fn(), run: vi.fn() }; beforeEach(async () => { dir = await mkdtemp(join(tmpdir(), "pcm-save-db-")); @@ -33,6 +33,7 @@ beforeEach(async () => { cdbToSqlMock.mockReset(); cdbToSqlMock.mockReturnValue(fakeDb); fakeDb.close.mockReset(); + fakeDb.run.mockReset(); }); afterEach(async () => { @@ -60,6 +61,19 @@ describe("withSaveDb", () => { expect(save.path).toBe(savePath); }); + it("puts the database in read-only mode before running the callback", async () => { + const runOrder: string[] = []; + fakeDb.run.mockImplementation((sql: string) => runOrder.push(sql)); + + await withSaveDb(savePath, () => { + runOrder.push("callback"); + return {}; + }); + + expect(fakeDb.run).toHaveBeenCalledWith("PRAGMA query_only = ON;"); + expect(runOrder).toEqual(["PRAGMA query_only = ON;", "callback"]); + }); + it("supports async callbacks", async () => { const result = await withSaveDb(savePath, async () => ({ async: true })); diff --git a/test/tools/query-save.test.ts b/test/tools/query-save.test.ts new file mode 100644 index 0000000..4605c63 --- /dev/null +++ b/test/tools/query-save.test.ts @@ -0,0 +1,260 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + assertReadOnlyQuery, + registerQuerySave, +} from "../../src/tools/query-save"; +import { saveFixtures } from "../fixtures/save.fixture"; +import { createMockMcpServer } from "../mocks/mock-mcp-server"; +import type { MockMcpServer } from "../mocks/mock-mcp-server"; + +describe("querySave", () => { + let mcp: MockMcpServer; + + beforeEach(() => { + mcp = createMockMcpServer(); + registerQuerySave(mcp.server); + }); + + it("registers the pcm_query_save tool", () => { + expect(mcp.getTool("pcm_query_save")).toBeDefined(); + expect(mcp.registerTool).toHaveBeenCalledOnce(); + }); + + it.each( + saveFixtures, + )("runs a read-only SELECT against %s", async (_name, path) => { + const result = await mcp.callTool("pcm_query_save", { + savePath: path, + query: "SELECT COUNT(*) AS n FROM STA_race", + }); + + expect(result.isError).toBeUndefined(); + expect(result.structuredContent).toMatchObject({ rowCount: 1 }); + }); + + it.each( + saveFixtures, + )("rejects a WITH … DELETE CTE for %s", async (_name, path) => { + const result = await mcp.callTool("pcm_query_save", { + savePath: path, + query: "WITH x AS (SELECT 1) DELETE FROM STA_race", + }); + + expect(result.isError).toBe(true); + expect(result.content).toEqual([ + { + type: "text", + text: "Only read-only SELECT (or WITH … SELECT) queries are allowed.", + }, + ]); + }); + + describe("assertReadOnlyQuery", () => { + describe("allowed queries", () => { + it("accepts a simple SELECT", () => { + const q = "SELECT * FROM DYN_cyclist"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts SELECT with a trailing semicolon (strips it)", () => { + expect(assertReadOnlyQuery("SELECT 1;")).toBe("SELECT 1"); + }); + + it("accepts SELECT with trailing semicolon and whitespace", () => { + // trim() runs before the semicolon strip, so a space before ';' is preserved + expect(assertReadOnlyQuery("SELECT 1 ; ")).toBe("SELECT 1 "); + }); + + it("accepts SELECT with leading/trailing whitespace", () => { + expect(assertReadOnlyQuery(" SELECT id FROM foo ")).toBe( + "SELECT id FROM foo", + ); + }); + + it("accepts a WITH … SELECT (CTE)", () => { + const q = + "WITH cte AS (SELECT id FROM foo) SELECT * FROM cte WHERE id > 5"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts lowercase select", () => { + expect(assertReadOnlyQuery("select * from bar")).toBe( + "select * from bar", + ); + }); + + it("accepts mixed-case SELECT", () => { + expect(assertReadOnlyQuery("Select id From foo")).toBe( + "Select id From foo", + ); + }); + + it("accepts SELECT containing a column named 'update_date' (keyword inside identifier)", () => { + const q = "SELECT update_date FROM DYN_cyclist"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts a SELECT whose string literal contains a keyword like 'create'", () => { + const q = "SELECT * FROM foo WHERE name LIKE '%create%'"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts a SELECT with a semicolon inside a string literal", () => { + const q = "SELECT ';' AS semi"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts a SELECT whose WHERE literal contains a semicolon", () => { + const q = "SELECT * FROM foo WHERE note = 'a;b'"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts a SELECT with a semicolon inside a quoted identifier", () => { + const q = 'SELECT "weird;col" FROM foo'; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts a SELECT with a semicolon inside a comment", () => { + const q = "SELECT 1 -- ; not a statement"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("keeps the trailing-semicolon strip working alongside a literal semicolon", () => { + expect(assertReadOnlyQuery("SELECT ';' AS semi;")).toBe( + "SELECT ';' AS semi", + ); + }); + + it("accepts a SELECT using the read-only REPLACE() function", () => { + const q = "SELECT REPLACE(name,'a','b') FROM foo"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts SELECT with ORDER BY, LIMIT, GROUP BY", () => { + const q = + "SELECT name, COUNT(*) AS n FROM foo GROUP BY name ORDER BY n DESC LIMIT 10"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + + it("accepts SELECT with a JOIN", () => { + const q = "SELECT a.id, b.name FROM a INNER JOIN b ON a.id = b.a_id"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + }); + + describe("empty / blank queries", () => { + it("rejects an empty string", () => { + expect(() => assertReadOnlyQuery("")).toThrowError("Query is empty."); + }); + + it("rejects a string that is only whitespace", () => { + expect(() => assertReadOnlyQuery(" ")).toThrowError( + "Query is empty.", + ); + }); + + it("rejects a bare semicolon", () => { + expect(() => assertReadOnlyQuery(";")).toThrowError("Query is empty."); + }); + + it("rejects whitespace + semicolon", () => { + expect(() => assertReadOnlyQuery(" ; ")).toThrowError( + "Query is empty.", + ); + }); + }); + + describe("multiple statements", () => { + it("rejects two SELECT statements separated by a semicolon", () => { + expect(() => assertReadOnlyQuery("SELECT 1; SELECT 2")).toThrowError( + "Only a single statement is allowed", + ); + }); + + it("rejects SELECT followed by a write statement", () => { + expect(() => + assertReadOnlyQuery("SELECT 1; DROP TABLE foo"), + ).toThrowError("Only a single statement is allowed"); + }); + }); + + describe("non-SELECT openers", () => { + it("rejects a bare INSERT", () => { + expect(() => + assertReadOnlyQuery("INSERT INTO foo VALUES (1)"), + ).toThrowError("Only read-only SELECT"); + }); + + it("rejects UPDATE", () => { + expect(() => + assertReadOnlyQuery("UPDATE foo SET bar = 1"), + ).toThrowError("Only read-only SELECT"); + }); + + it("rejects DELETE", () => { + expect(() => assertReadOnlyQuery("DELETE FROM foo")).toThrowError( + "Only read-only SELECT", + ); + }); + + it("rejects DROP TABLE", () => { + expect(() => assertReadOnlyQuery("DROP TABLE foo")).toThrowError( + "Only read-only SELECT", + ); + }); + + it("rejects CREATE TABLE", () => { + expect(() => + assertReadOnlyQuery("CREATE TABLE foo (id INTEGER)"), + ).toThrowError("Only read-only SELECT"); + }); + + it("rejects PRAGMA", () => { + expect(() => + assertReadOnlyQuery("PRAGMA table_info(foo)"), + ).toThrowError("Only read-only SELECT"); + }); + }); + + describe("ATTACH / DETACH", () => { + it("rejects ATTACH stacked after a SELECT", () => { + expect(() => + assertReadOnlyQuery("SELECT * FROM foo; ATTACH DATABASE 'x' AS y"), + ).toThrowError("Only a single statement is allowed"); + }); + + it("rejects ATTACH as a standalone statement opener", () => { + expect(() => + assertReadOnlyQuery("ATTACH DATABASE 'evil.db' AS e"), + ).toThrowError("Only read-only SELECT"); + }); + + it("rejects DETACH as a standalone statement opener", () => { + expect(() => assertReadOnlyQuery("DETACH DATABASE e")).toThrowError( + "Only read-only SELECT", + ); + }); + }); + + describe("write detection via the statement identifier", () => { + it("rejects a WITH … DELETE CTE (write behind a read opener)", () => { + expect(() => + assertReadOnlyQuery("WITH x AS (SELECT 1) DELETE FROM foo"), + ).toThrowError("Only read-only SELECT"); + }); + + it("rejects a WITH … INSERT CTE", () => { + expect(() => + assertReadOnlyQuery( + "WITH x AS (SELECT 1) INSERT INTO foo SELECT * FROM x", + ), + ).toThrowError("Only read-only SELECT"); + }); + + it("accepts a SELECT merely mentioning a write keyword in a literal", () => { + const q = "SELECT * FROM foo WHERE note = 'please update'"; + expect(assertReadOnlyQuery(q)).toBe(q); + }); + }); + }); +});