From 03d9dbe790d7296ab271e86e008a92ed4c41c199 Mon Sep 17 00:00:00 2001
From: Sougata Mandal <iaamsougata@gmail.com>
Date: Tue, 6 Oct 2026 14:43:00 +0530
Subject: [PATCH] fix(viewer): enforce readonly queries in data_query MCP tool

The data_query tool claimed to be readonly but executed arbitrary SQL. Validate the statement client-side: allow only SELECT/WITH/VALUES/DESCRIBE/SHOW/EXPLAIN, reject stacked statements and DDL/DML keywords, ignoring strings and comments.
---
 .../viewer/src/model_context/model_context.ts |  9 +++-
 .../src/model_context/readonly_query.ts       | 47 +++++++++++++++++++
 packages/viewer/test/model_context.test.ts    | 34 ++++++++++++++
 3 files changed, 88 insertions(+), 2 deletions(-)
 create mode 100644 packages/viewer/src/model_context/readonly_query.ts
 create mode 100644 packages/viewer/test/model_context.test.ts

diff --git a/packages/viewer/src/model_context/model_context.ts b/packages/viewer/src/model_context/model_context.ts
index 4303690e..a265475f 100644
--- a/packages/viewer/src/model_context/model_context.ts
+++ b/packages/viewer/src/model_context/model_context.ts
@@ -15,6 +15,7 @@ import {
 } from "../schemas.js";
 import type { EmbeddingAtlasStore } from "../stores/embedding_atlas_store.js";
 import { screenshot, type ScreenshotOptions } from "../utils/screenshot.js";
+import { isReadonlyQuery } from "./readonly_query.js";
 
 export interface ModelContextDelegate {
   container: HTMLDivElement;
@@ -71,11 +72,13 @@ export class EmbeddingAtlasControl {
 
     this.register("data_query", {
       args: {
-        query: z.string().describe("The SQL query to run. Must keep this readonly - the server does not enforce it."),
+        query: z.string().describe("The readonly SQL query to run (SELECT/WITH/VALUES/DESCRIBE/SHOW/EXPLAIN only)."),
       },
       description: "Run a readonly SQL query in DuckDB",
       handler: async ({ query }) => {
-        // TODO: enforce readonly query.
+        if (!isReadonlyQuery(query)) {
+          return { error: "only readonly queries (SELECT/WITH/VALUES/DESCRIBE/SHOW/EXPLAIN) are allowed" };
+        }
         let result = await store.coordinator.query(query);
         return result.toArray();
       },
@@ -429,3 +432,5 @@ function parseImageDataUrl(dataUrl: string): { mimeType: string; data: string }
 
   return { mimeType, data: base64Content };
 }
+
+export { isReadonlyQuery } from "./readonly_query.js";
diff --git a/packages/viewer/src/model_context/readonly_query.ts b/packages/viewer/src/model_context/readonly_query.ts
new file mode 100644
index 00000000..e6a5202e
--- /dev/null
+++ b/packages/viewer/src/model_context/readonly_query.ts
@@ -0,0 +1,47 @@
+// Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+
+const READONLY_FIRST_WORD = new Set([
+  "SELECT",
+  "WITH",
+  "VALUES",
+  "TABLE",
+  "DESCRIBE",
+  "DESC",
+  "SHOW",
+  "EXPLAIN",
+  "SUMMARIZE",
+  "PIVOT",
+]);
+
+const WRITE_KEYWORDS =
+  /\b(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|COPY|ATTACH|DETACH|INSTALL|LOAD|CALL|PRAGMA|SET|VACUUM|CHECKPOINT|TRUNCATE|GRANT|REVOKE|USE)\b/i;
+
+function stripStringsAndComments(sql: string): string {
+  // Remove '...', "...", `...`, /* ... */, -- ... to avoid false positives.
+  return sql
+    .replace(/'(?:[^']|'')*'/g, " ")
+    .replace(/"(?:[^"\\]|\\.)*"/g, " ")
+    .replace(/`(?:[^`\\]|\\.)*`/g, " ")
+    .replace(/\/\*[\s\S]*?\*\//g, " ")
+    .replace(/--[^\n]*/g, " ");
+}
+
+export function isReadonlyQuery(query: string): boolean {
+  let cleaned = stripStringsAndComments(query).trim();
+  // Allow a single trailing semicolon; reject stacked statements.
+  cleaned = cleaned.replace(/;+\s*$/, "").trim();
+  if (cleaned === "" || cleaned.includes(";")) {
+    return false;
+  }
+  let firstWord = cleaned
+    .split(/\s+/, 1)[0]
+    ?.replace(/[^A-Za-z]/g, "")
+    .toUpperCase();
+  if (firstWord == null || !READONLY_FIRST_WORD.has(firstWord)) {
+    return false;
+  }
+  if (WRITE_KEYWORDS.test(cleaned)) {
+    return false;
+  }
+  return true;
+}
diff --git a/packages/viewer/test/model_context.test.ts b/packages/viewer/test/model_context.test.ts
new file mode 100644
index 00000000..02c6887e
--- /dev/null
+++ b/packages/viewer/test/model_context.test.ts
@@ -0,0 +1,34 @@
+// Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+
+import { describe, expect, it } from "vitest";
+
+import { isReadonlyQuery } from "../src/model_context/readonly_query.js";
+
+describe("isReadonlyQuery", () => {
+  it("allows plain reads", () => {
+    expect(isReadonlyQuery("SELECT * FROM dataset")).toBe(true);
+    expect(isReadonlyQuery("  with t AS (SELECT 1) SELECT * FROM t")).toBe(true);
+    expect(isReadonlyQuery("VALUES (1), (2);")).toBe(true);
+    expect(isReadonlyQuery("DESCRIBE dataset")).toBe(true);
+    expect(isReadonlyQuery("SHOW TABLES")).toBe(true);
+    expect(isReadonlyQuery("EXPLAIN SELECT 1")).toBe(true);
+    expect(isReadonlyQuery("-- comment\nSELECT version() AS version")).toBe(true);
+  });
+
+  it("rejects writes and stacked statements", () => {
+    expect(isReadonlyQuery("DROP TABLE dataset")).toBe(false);
+    expect(isReadonlyQuery("DELETE FROM dataset")).toBe(false);
+    expect(isReadonlyQuery("INSERT INTO dataset VALUES (1)")).toBe(false);
+    expect(isReadonlyQuery("UPDATE dataset SET x = 1")).toBe(false);
+    expect(isReadonlyQuery("CREATE TABLE t (x INT)")).toBe(false);
+    expect(isReadonlyQuery("COPY dataset TO 'out.parquet'")).toBe(false);
+    expect(isReadonlyQuery("ATTACH 'other.db'")).toBe(false);
+    expect(isReadonlyQuery("SELECT * FROM dataset; DROP TABLE dataset")).toBe(false);
+    expect(isReadonlyQuery("")).toBe(false);
+  });
+
+  it("ignores keywords inside strings and comments", () => {
+    expect(isReadonlyQuery("SELECT 'DROP TABLE not_a_table' AS x")).toBe(true);
+    expect(isReadonlyQuery("SELECT * FROM t -- DELETE")).toBe(true);
+  });
+});