diff --git a/.github/extensions/database-explorer/extension.mjs b/.github/extensions/database-explorer/extension.mjs index 6a6bb54..e5c4886 100644 --- a/.github/extensions/database-explorer/extension.mjs +++ b/.github/extensions/database-explorer/extension.mjs @@ -3,6 +3,7 @@ import { fileURLToPath } from "node:url"; import { dirname, join, resolve } from "node:path"; import { createClient } from "@libsql/client"; import { CanvasError, createCanvas, joinSession } from "@github/copilot-sdk/extension"; +import { executeReadOnlyQuery } from "./query-policy.mjs"; const servers = new Map(); const MAX_DISPLAY_ROWS = 250; @@ -28,86 +29,6 @@ function readRequestBody(request) { }); } -function stripCommentsAndValidateSingleStatement(sql) { - let output = ""; - let quote = ""; - let statementEnded = false; - - for (let index = 0; index < sql.length; index += 1) { - const character = sql[index]; - const next = sql[index + 1]; - - if (quote) { - output += character; - if (character === quote) { - if (next === quote) { - output += next; - index += 1; - } else { - quote = ""; - } - } - continue; - } - - if (character === "'" || character === '"' || character === "`") { - quote = character; - output += character; - continue; - } - - if (character === "-" && next === "-") { - index = sql.indexOf("\n", index + 2); - if (index === -1) { - break; - } - output += " "; - continue; - } - - if (character === "/" && next === "*") { - const end = sql.indexOf("*/", index + 2); - if (end === -1) { - throw new Error("The SQL comment is not closed."); - } - index = end + 1; - output += " "; - continue; - } - - if (character === ";") { - statementEnded = true; - continue; - } - - if (statementEnded && !/\s/.test(character)) { - throw new Error("Run one SQL statement at a time."); - } - - output += character; - } - - if (quote) { - throw new Error("The SQL string is not closed."); - } - - return output.trim(); -} - -function validateReadOnlyQuery(query) { - if (typeof query !== "string" || !query.trim()) { - throw new Error("Enter a SQL query."); - } - - const statement = stripCommentsAndValidateSingleStatement(query); - const keyword = statement.match(/^([a-zA-Z]+)/)?.[1]?.toUpperCase(); - if (keyword !== "SELECT" && keyword !== "WITH") { - throw new Error("Only read-only SELECT and WITH queries are allowed."); - } - - return statement; -} - function normalizeValue(value) { return typeof value === "bigint" ? value.toString() : value; } @@ -294,7 +215,7 @@ async function startServer(url) { if (request.url === "/api/query" && request.method === "POST") { const body = JSON.parse(await readRequestBody(request)); - const data = await executeQuery(url, validateReadOnlyQuery(body.query)); + const data = await executeReadOnlyQuery(body.query, (query) => executeQuery(url, query)); response.writeHead(200, { "Content-Type": "application/json; charset=utf-8" }); response.end(JSON.stringify(data)); return; @@ -336,7 +257,7 @@ await joinSession({ }, handler: async (ctx) => { try { - return executeQuery(databaseUrl(), validateReadOnlyQuery(ctx.input?.query)); + return executeReadOnlyQuery(ctx.input?.query, (query) => executeQuery(databaseUrl(), query)); } catch (error) { throw new CanvasError("database_query_invalid", error instanceof Error ? error.message : "The query is invalid."); } diff --git a/.github/extensions/database-explorer/query-policy.mjs b/.github/extensions/database-explorer/query-policy.mjs new file mode 100644 index 0000000..a8eaf59 --- /dev/null +++ b/.github/extensions/database-explorer/query-policy.mjs @@ -0,0 +1,83 @@ +function stripCommentsAndValidateSingleStatement(sql) { + let output = ""; + let quote = ""; + let statementEnded = false; + + for (let index = 0; index < sql.length; index += 1) { + const character = sql[index]; + const next = sql[index + 1]; + + if (quote) { + output += character; + if (character === quote) { + if (next === quote) { + output += next; + index += 1; + } else { + quote = ""; + } + } + continue; + } + + if (character === "'" || character === '"' || character === "`") { + quote = character; + output += character; + continue; + } + + if (character === "-" && next === "-") { + index = sql.indexOf("\n", index + 2); + if (index === -1) { + break; + } + output += " "; + continue; + } + + if (character === "/" && next === "*") { + const end = sql.indexOf("*/", index + 2); + if (end === -1) { + throw new Error("The SQL comment is not closed."); + } + index = end + 1; + output += " "; + continue; + } + + if (character === ";") { + statementEnded = true; + continue; + } + + if (statementEnded && !/\s/.test(character)) { + throw new Error("Run one SQL statement at a time."); + } + + output += character; + } + + if (quote) { + throw new Error("The SQL string is not closed."); + } + + return output.trim(); +} + +export function validateReadOnlyQuery(query) { + if (typeof query !== "string" || !query.trim()) { + throw new Error("Enter a SQL query."); + } + + const statement = stripCommentsAndValidateSingleStatement(query); + const keyword = statement.match(/^([a-zA-Z]+)/)?.[1]?.toUpperCase(); + if (keyword !== "SELECT" && keyword !== "WITH") { + throw new Error("Only read-only SELECT and WITH queries are allowed."); + } + + return statement; +} + +export function executeReadOnlyQuery(query, execute) { + return execute(validateReadOnlyQuery(query)); +} diff --git a/.github/extensions/database-explorer/query-policy.test.mjs b/.github/extensions/database-explorer/query-policy.test.mjs new file mode 100644 index 0000000..b20e487 --- /dev/null +++ b/.github/extensions/database-explorer/query-policy.test.mjs @@ -0,0 +1,47 @@ +import { describe, expect, it, vi } from "vitest"; +import { executeReadOnlyQuery, validateReadOnlyQuery } from "./query-policy.mjs"; + +describe("Database Explorer query policy", () => { + it.each([ + ["plain SELECT", "SELECT id FROM games", "SELECT id FROM games"], + ["quoted semicolon", "SELECT ';' AS value;", "SELECT ';' AS value"], + ["comments", "/* note */ SELECT id FROM games -- trailing comment", "SELECT id FROM games"], + [ + "read-only CTE", + "WITH selected AS (SELECT id FROM games) SELECT id FROM selected", + "WITH selected AS (SELECT id FROM games) SELECT id FROM selected", + ], + ])("normalizes a supported %s query", (_name, input, expected) => { + expect(validateReadOnlyQuery(input)).toBe(expected); + }); + + it.each([ + ["non-string input", null], + ["empty input", ""], + ["whitespace input", " \t "], + ["direct mutation", "DELETE FROM games"], + ["multiple statements", "SELECT id FROM games; DELETE FROM games"], + ["unterminated string", "SELECT 'unfinished"], + ["unterminated comment", "SELECT 1 /* unfinished"], + ])("rejects %s", (_name, input) => { + expect(() => validateReadOnlyQuery(input)).toThrow(); + }); + + it("passes a normalized supported query to the executor", () => { + const execute = vi.fn((query) => query); + const rows = executeReadOnlyQuery( + "SELECT id FROM games;", + execute, + ); + + expect(rows).toBe("SELECT id FROM games"); + expect(execute).toHaveBeenCalledExactlyOnceWith("SELECT id FROM games"); + }); + + it("does not invoke the executor for a rejected query", () => { + const execute = vi.fn(); + + expect(() => executeReadOnlyQuery("DELETE FROM games", execute)).toThrow(); + expect(execute).not.toHaveBeenCalled(); + }); +}); diff --git a/vitest.config.ts b/vitest.config.ts index 585d8e2..384e7ff 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -2,7 +2,11 @@ import { defineConfig } from 'vitest/config'; export default defineConfig({ test: { - include: ['db/**/*.test.ts', 'src/**/*.test.ts'], + include: [ + 'db/**/*.test.ts', + 'src/**/*.test.ts', + '.github/extensions/database-explorer/**/*.test.mjs', + ], environment: 'node', }, });