Add query support to local MCP server - #328
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
konstantinoscs
left a comment
There was a problem hiding this comment.
I think we need to reconsider this. Basically this PR allows agents to run arbitary SQL that gets directly forwarded to duckDB.
DuckDB SQL can read local files, access the network, load extensions, and consume unbounded resources. DuckDB’s own security guidance says to treat untrusted SQL like Bash or Python code and recommends OS/container/WASM isolation. See:
https://duckdb.org/docs/current/operations_manual/securing_duckdb/overview
To be fair, this happens already in the human operated mode and probably this should be changed as well but an agent increases the security risk manyfold.
Btw, in the monorepo we avoid this by rusing WASM and adding read-only and single-statement "gates".
I think we need to do some sandboxing here
konstantinoscs
left a comment
There was a problem hiding this comment.
My previous security concern remains
…crate and wire the sandbox in The sandbox primitives from the previous commit had no caller: the gates, the session hardening, and the row cap all lived in an inline `#[allow(dead_code)]` module in catalogs.rs, proven only by tests. This moves Tower's DuckDB usage into a dedicated `tower-duckdb` crate and makes `tower catalogs query` its first real consumer, so the hardening runs in production rather than sitting on a shelf. The crate owns the whole DuckDB surface: opening an in-memory Session, running trusted setup, locking the session down for untrusted SQL (Session::harden), and executing a query into JSON rows with an optional row cap. The static text gates (reject write/DDL, reject multi-statement) live in its `guard` module. The value conversion, the query execution, and the full adversarial + integration test suite move with it, so the security invariants are tested next to the code they protect. tower-cmd no longer depends on duckdb directly; it goes through the crate. `catalogs query` now splits by mode. Read mode is the default and the sandboxed path: the SQL is gated before it reaches DuckDB (a smuggled second statement would otherwise run as a side effect of prepare, and a write gets a clear message instead of a raw engine error), the session is hardened after attach, and rows are capped with a truncation notice. Write mode stays the trusted power-user opt-in and runs as before. Once this lands, #328 will layer the MCP `tower_catalogs_query` tool on top of the same crate.
…tests (#331) * feat(catalogs): add DuckDB query sandbox primitives and adversarial tests Agent-issued catalog SQL is untrusted, so before an MCP query tool can run it we need the query path locked down and proof that the lockdown holds. This lands that foundation without wiring it to a caller yet. It adds a `sandbox` module with the gates a future `tower_catalogs_query` tool applies: reject write/DDL and multi-statement input (the multi-statement scanner skips separators inside strings and comments, which matters because duckdb-rs `prepare` executes every statement but the last as a side effect), and a session hardening set that disables local-filesystem access, blocks community extensions, and locks the configuration. The hardening disables only `LocalFileSystem`, so httpfs and the object-store reads an attached Iceberg catalog depends on keep working. `run_duckdb_query` gains an optional row cap that flags the result truncated, so a model cannot pull an unbounded table into memory or its context. The adversarial suite is the point: each test encodes an attack an agent query might attempt (data tampering, statement smuggling, host filesystem reads and writes, configuration escape, arbitrary extension loading, SSRF via table functions) and asserts the sandbox refuses it. A testcontainers MinIO test proves the same hardening does not break a real object-store Iceberg read while local access and config changes stay blocked, and self-skips when no Docker daemon is present so a plain `cargo test` still passes. * refactor(catalogs): extract Tower's DuckDB usage into a tower-duckdb crate and wire the sandbox in The sandbox primitives from the previous commit had no caller: the gates, the session hardening, and the row cap all lived in an inline `#[allow(dead_code)]` module in catalogs.rs, proven only by tests. This moves Tower's DuckDB usage into a dedicated `tower-duckdb` crate and makes `tower catalogs query` its first real consumer, so the hardening runs in production rather than sitting on a shelf. The crate owns the whole DuckDB surface: opening an in-memory Session, running trusted setup, locking the session down for untrusted SQL (Session::harden), and executing a query into JSON rows with an optional row cap. The static text gates (reject write/DDL, reject multi-statement) live in its `guard` module. The value conversion, the query execution, and the full adversarial + integration test suite move with it, so the security invariants are tested next to the code they protect. tower-cmd no longer depends on duckdb directly; it goes through the crate. `catalogs query` now splits by mode. Read mode is the default and the sandboxed path: the SQL is gated before it reaches DuckDB (a smuggled second statement would otherwise run as a side effect of prepare, and a write gets a clear message instead of a raw engine error), the session is hardened after attach, and rows are capped with a truncation notice. Write mode stays the trusted power-user opt-in and runs as before. Once this lands, #328 will layer the MCP `tower_catalogs_query` tool on top of the same crate. * fix(tower-duckdb): gate read-only queries with DuckDB's parser, not a keyword denylist The read-only gate checked only the first SQL keyword against a denylist. That is not a safe policy: a `--` comment ends at a carriage return as well as a newline in DuckDB, so a `-- x\rDROP …` payload looked empty to the scanner but parses as a DROP, and a statement that opens with an allowed keyword (a `WITH` CTE, for one) can still mutate. Addresses Konstantinos's review on #331. The gate now runs the SQL through DuckDB's own parser via `json_serialize_sql`, which parses without executing and serializes only SELECT statements, erroring on anything else. `classify_read_only` returns Allowed only for exactly one SELECT; everything else (writes, DDL, PRAGMA/SET, multi-statement, unparseable) is refused, fail-closed. This is an allowlist of what the executor will actually run, so the comment-terminator and CTE bypasses are caught, and the SQL is bound as a parameter rather than spliced into the parser query. The keyword denylist and the hand-rolled statement scanner are gone, along with their unit tests; new tests cover the single-SELECT allowlist, the `\r` smuggling case, the mutating-CTE case, and empty/multiple classification. * fix(tower-duckdb): bundle the json extension and skip iceberg test where it can't install The Windows test job failed because the read-only gate runs `json_serialize_sql`, which needs the `json` extension. It is not bundled, so DuckDB tried to auto-download and install it, and the install fails on the Windows runner ("Could not move file: Access is denied"). That is a real defect, not just a test problem: the gate would be broken for Windows users too. Enabling the duckdb `json` feature compiles the extension statically into the bundled build, so it is available with no autoload or network. Verified it works with `autoinstall_known_extensions` and `autoload_known_extensions` both off. The `iceberg_scan` regression test needs the `iceberg` extension, which has no such feature and must still be fetched at runtime. It now self-skips when that install cannot happen, the same way the object-store test skips without Docker, so the Windows job stops failing on an environment it can't satisfy while the test keeps running where it can. * test(tower-duckdb): co-locate read-only gate regression tests in guard.rs The gate's unit tests lived in lib.rs, away from the code they cover. This moves them next to `classify_read_only` in guard.rs and broadens them into a proper regression suite, so the parser-based read-only policy is pinned down where a reader of guard.rs will find it. The suite locks in the current DuckDB classification across the shapes that matter: the many forms of a single read (plain SELECT, CTEs, set operations, subqueries, VALUES/TABLE/FROM-first, DESCRIBE/SUMMARIZE/SHOW), that comments and whitespace and semicolons inside literals don't change the verdict, that writes and DDL and config/transaction/meta statements are rejected whatever the leading keyword, that a leading SELECT can't launder a trailing mutation (the `\r` comment-terminator and data-modifying-CTE bypasses), multi-statement and empty/comment-only classification, fail-closed rejection of unparseable input, and the shared-connection reuse path. It also documents one boundary explicitly: a SELECT that reads a local file or URL is Allowed by the gate because the gate classifies statement shape only; the session hardening is what refuses the read. The redundant gate tests are dropped from lib.rs; its adversarial suite keeps the shared `check` helper. * feat(tower-duckdb): close SELECT-shaped holes, expand the lockdown, and bound queries Parsing as a SELECT is a statement shape, not a read-only property, and the gate was treating the two as the same thing. Verified against the DuckDB we ship (v1.5.4): `SELECT nextval('s')` classifies as a plain SELECT and really does advance the sequence, and `query()`/`query_table()` classify as SELECTs while handing a string to the execution pipeline. Those are now refused by name from the parsed tree, walking `function_name` nodes rather than matching text, so comments and quoting cannot hide them. The session lockdown gains the settings it was missing: no implicit extension install or load, no unsigned extensions, secrets kept redacted, and optional memory and temp-size ceilings, with `lock_configuration` still last because it freezes every later SET. `Hardening` is now a struct so a caller that needs no object storage can set `deny_external_access` and close network egress. The default cannot: an attached Iceberg catalog is made of S3 reads, so egress stays open on the catalog path and the docs now say so plainly rather than implying the lockdown is total. A row cap turned out to bound almost nothing, since one row can carry a whole column (`string_agg`, `list`, `to_json`). Results are now bounded by rows, total bytes, and wall-clock time together, via `Limits`; the byte ceiling is what actually holds, and the timeout exists because DuckDB has no statement timeout of its own, so it is enforced by interrupting the connection. `QueryResult.truncated` carries which ceiling was hit so the CLI can say which one. The docs now state that none of this is the security boundary: the read-only credential and the READ_ONLY attach are, and everything here is defence in depth in front of them. Tests pin each judgement against the shipped DuckDB, including canaries that fail loudly on upgrade if data-modifying CTEs start parsing as SELECTs or `prepare` stops executing leading statements. * feat(catalogs): add --max-rows to override the query result ceiling Read queries are bounded by default so a runaway result cannot flood a terminal or a model's context, but there was no way to ask for more. That makes the bound a wall rather than a default: an agent (or a person) with a legitimate need for a large extract had no option short of --write, which vends read-write credentials for what is still only a read. --max-rows sets the row ceiling and lifts the size ceiling with it, because a caller who asks for a million rows should not then be cut short by a byte budget they never set. --max-rows 0 removes the ceilings entirely. The help text says plainly that a large result can exhaust memory, since that is the trade being made. The read-only gate and the session hardening are untouched: this widens how much data a read may return, not what a query is allowed to do. The decision is factored into `query_limits` so it is unit-tested directly rather than only through argument parsing. * fix(tower-duckdb): make the function gate fail-closed and the ceilings real Review found three ways the sandbox did less than it claimed, and two ways the tests hid it. The function denylist was not fail-closed. `json_execute_serialized_sql` runs whatever SQL it is handed and ships with the very `json` extension this crate enables to build the gate, so the gate supplied its own bypass: the outer statement is an ordinary SELECT and the effectful call hides in a string. `enable_logging`, `checkpoint`, `setseed` and friends were allowed for the same reason, that naming dangerous functions one at a time can never be complete. The gate now allowlists the table-function position, which is where the danger lives, so an unknown function is refused rather than waved through. Scalars are too numerous to allowlist, so it asks the engine instead: DuckDB's `has_side_effects` flags `nextval`, `setseed` and the rest, and keeps up with new versions on its own. It is NULL for every table function, which is why those are allowlisted rather than queried. An explicit list covers what neither reaches: the dynamic-SQL executors and the effectful table functions. This tightens the read path, since `read_csv`/`read_parquet` and the rest of that family are no longer permitted; a catalog query reads base tables, which are not functions. The byte ceiling measured a row after admitting it, so `SELECT repeat('x', 5e7)` returned fifty megabytes and labelled it truncated. It now measures first and withholds the row, which means a single over-budget row yields an empty truncated result. That is blunt but honest, and it is what the ceiling was for. Because that still only bounds what the caller is handed, the read path now hardens with `Hardening::agent()`, which sets an engine memory limit for what a query spends before the first row exists. The SSRF test passed for a reason production does not share: its helper never loaded `httpfs`, so there was no HTTP filesystem to block. It now loads the same extensions `attach_statements` does, and counts connections to a listener it owns rather than trusting an error string, with a control run so it cannot pass vacuously. Note the hardening blocks these today only because disabling LocalFileSystem breaks path resolution first; egress is still not closed by design, and that gap is documented rather than papered over. Also: the MinIO image is pinned to a release tag instead of `latest` and only skips when Docker is genuinely unreachable, so a registry problem fails loudly instead of silently disabling a security test; the iceberg skip is narrowed the same way; and the workspace `rust-version` is corrected to 1.88, which the toolchain and the dependency tree have required for some time. * fix(tower-duckdb): close the replacement-scan SSRF path CI found on Linux CI caught what my local testing could not. The SSRF I reported as non-reproducible does reproduce on Linux: three hardened queries reached the probe listener there. The block I measured on macOS was incidental, disabling LocalFileSystem breaks HTTP path resolution before the HTTP filesystem is consulted, and Linux does not do that. Relying on it was the mistake; the reviewer was right that the hardening is not an egress control. Two holes are closed at the gate, which is platform-independent. `SELECT * FROM 'http://…'` is a replacement scan: it parses as an ordinary base table with the URL as its name, so the table-function allowlist never saw it. Base-table references that name a file or a URL are now refused outright, with their own verdict so the error can say what is actually wrong. `iceberg_scan` is out of the allowlist. It takes a location, so allowing it left a way to reach internal services even with the rest of the family refused. Every remaining entry is a pure generator that accepts no path and no URL, which is the property to preserve: adding anything there that takes a location reopens this. The test now mirrors production by running the gate first and executing only what it allows, because asserting that hardening alone blocks the network was asserting something this crate's own documentation says is untrue. A companion test records that hardening does not close egress, so the next person does not rediscover it in CI. Separately, the MinIO test failed on Windows because MinIO publishes Linux images only, and my previous change had turned that into a hard failure. It now skips when no Linux container runtime is available, decided from the platform rather than inferred from a pull error, and still fails loudly for anything else.
Adds a tower_catalogs_query MCP tool that runs one read-only SQL statement against a Tower-managed storage catalog, and extends tower_catalogs_show to list the catalog's namespaces and tables. Agent queries go through the tower-duckdb guard (single read-only statement, denied functions and file/URL table references rejected), run hardened (Hardening::agent()), and are capped by Limits::agent() (1000 rows / 1 MiB / 60s). Rows come back as positional arrays so duplicate column names survive.
Behave scenarios for the MCP catalog tools: show lists tables, a read-only query succeeds, writes and multi-statement SQL are rejected. Skipped unless TOWER_TEST_CATALOG (and a real TOWER_URL) is set, since they exercise the real attach -> Iceberg -> query path.
af073f3 to
0e2b86b
Compare
There was a problem hiding this comment.
My main concerns about security are solved, thanks!
There are a few nits regarding resource consumption constraints that I think remain. I tested them locally with the help of Sol. You can have your agent take a stab at them but feel free to merge afterwards:
-
Nested JSON bypasses the 1 MiB response cap.
The limiter undercounts arrays, strings requiring JSON escaping, separators, and pretty-printing. The entire nested DuckDB value is also materialized before the estimate is checked.
A local probe with:
SELECT list_transform(range(10000000), x -> '') AS xs
was allowed and returned one row with truncated:false, but the MCP-equivalent serialized response was 120,000,115 bytes. -
DuckDB errors bypass the result ceiling entirely.
query_catalog_for_agentforwardserr.to_string()directly through MCP.
This allowed query:
SELECT CAST(repeat('x', 2000000) AS INTEGER)
generated a 2,000,126-byte error containing the offending value. The same pattern can use an aggregate over a catalog column, bypassing the intended output cap through the error channel.
Return a bounded, stable error to MCP and retain detailed errors only in protected logs. Token redaction is present and correct, but it does not bound the message. -
Parallel calls multiply the per-query resource limits.
Every request starts a new blocking DuckDB session atcatalogs.rsline 690. MCP requests are dispatched concurrently, and there is no service-wide semaphore in TowerService.
Each session may consume 1 GiB RAM and 2 GiB temporary storage, so a handful of concurrent agent calls can exhaust the machine. Cancelling the MCP future also does not terminate an already-running spawn_blocking task.
This differs from the storage web product, which explicitly serializes every engine operation through a single promise chain incatalogService.ts(line 171). Add a process-wide shared concurrency budget and cancellation-aware execution. -
Catalog discovery remains unbounded.
list_catalog_tablesrunsSHOW ALL TABLESwithLimits::none(), after whichtower_catalogs_showserializes every table into model context.
Give the MCP discovery path pagination or row/byte limits and return tables_truncated or a cursor. -
My understanding of reading the new tests you've added is that they won't run in the CI? Is that intentional?
Title says it all. Should make our output a little nicer/easier to use in some cases!