From 0821cf7870c901441e865424e04262bb7bc8cba0 Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Fri, 24 Jul 2026 09:23:33 +0100 Subject: [PATCH 1/4] 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. --- Cargo.lock | 302 +++++++++++++++-- crates/tower-cmd/Cargo.toml | 1 + crates/tower-cmd/src/catalogs.rs | 557 ++++++++++++++++++++++++++++++- 3 files changed, 834 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b445be24..41b1409a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -301,7 +301,7 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", ] [[package]] @@ -501,9 +501,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.9.1" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" [[package]] name = "bitvec" @@ -526,6 +526,56 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bollard" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97ccca1260af6a459d75994ad5acc1651bcabcbdbc41467cc9786519ab854c30" +dependencies = [ + "base64", + "bollard-stubs", + "bytes", + "futures-core", + "futures-util", + "hex", + "home", + "http", + "http-body-util", + "hyper", + "hyper-named-pipe", + "hyper-rustls", + "hyper-util", + "hyperlocal", + "log", + "pin-project-lite", + "rustls", + "rustls-native-certs", + "rustls-pemfile", + "rustls-pki-types", + "serde", + "serde_derive", + "serde_json", + "serde_repr", + "serde_urlencoded", + "thiserror 2.0.12", + "tokio", + "tokio-util", + "tower-service", + "url", + "winapi", +] + +[[package]] +name = "bollard-stubs" +version = "1.47.1-rc.27.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f179cfbddb6e77a5472703d4b30436bff32929c0aa8a9008ecf23d1d3cdd0da" +dependencies = [ + "serde", + "serde_repr", + "serde_with", +] + [[package]] name = "borsh" version = "1.7.0" @@ -664,7 +714,7 @@ dependencies = [ "num-traits", "serde", "wasm-bindgen", - "windows-link", + "windows-link 0.1.3", ] [[package]] @@ -941,7 +991,7 @@ version = "0.28.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "crossterm_winapi", "parking_lot", "rustix 0.38.44", @@ -954,7 +1004,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "crossterm_winapi", "derive_more 2.1.1", "document-features", @@ -1310,6 +1360,17 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10" +[[package]] +name = "docker_credential" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29547a1dc60885a552306986316bc9701ba120c1a8db6769fa68691529ad373d" +dependencies = [ + "base64", + "serde", + "serde_json", +] + [[package]] name = "document-features" version = "0.2.12" @@ -1393,6 +1454,17 @@ dependencies = [ "str-buf", ] +[[package]] +name = "etcetera" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26c7b13d0780cb82722fd59f6f57f925e143427e4a75313a6c77243bf5326ae6" +dependencies = [ + "cfg-if", + "home", + "windows-sys 0.59.0", +] + [[package]] name = "eventsource-stream" version = "0.2.3" @@ -1717,7 +1789,7 @@ version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf760ebf69878d9fd8f110c89703d90ce35095324d1f1edcb595c63945ee757" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "ignore", "walkdir", ] @@ -1785,6 +1857,15 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "http" version = "1.3.1" @@ -1857,6 +1938,20 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-named-pipe" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fab3637d6b04a8037af8a266fdf6cf92ea957e8c53981a2bf6136572531025bf" +dependencies = [ + "hex", + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-rustls" version = "0.27.7" @@ -1898,6 +1993,21 @@ dependencies = [ "tracing", ] +[[package]] +name = "hyperlocal" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "986c5ce3b994526b3cd75578e62554abd09f0899d6206de48b3e96ab34ccc8c7" +dependencies = [ + "hex", + "http-body-util", + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "iana-time-zone" version = "0.1.63" @@ -2107,7 +2217,7 @@ version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6654738b8024300cf062d04a1c13c10c8e2cea598ec1c47dc9b6641159429756" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "crossterm 0.29.0", "dyn-clone", "fuzzy-matcher", @@ -2130,7 +2240,7 @@ version = "0.7.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d93587f37623a1a17d94ef2bc9ada592f5465fe7732084ab7beefabe5c77c0c4" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "cfg-if", "libc", ] @@ -2338,7 +2448,7 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4488594b9328dee448adb906d8b126d9b7deb7cf5c22161ee591610bb1be83c0" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "libc", "redox_syscall 0.5.15", ] @@ -2517,7 +2627,7 @@ version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "cfg-if", "cfg_aliases", "libc", @@ -2653,7 +2763,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "900831247d2fe1a09a683278e5384cfb8c80c79fe6b166f9d14bfdde0ea1b03c" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "objc2", ] @@ -2684,6 +2794,12 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + [[package]] name = "option-ext" version = "0.2.0" @@ -2725,6 +2841,31 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "parse-display" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "914a1c2265c98e2446911282c6ac86d8524f495792c38c5bd884f80499c7538a" +dependencies = [ + "parse-display-derive", + "regex", + "regex-syntax 0.8.5", +] + +[[package]] +name = "parse-display-derive" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ae7800a4c974efd12df917266338e79a7a74415173caf7e70aa0a0707345281" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "regex-syntax 0.8.5", + "structmeta", + "syn 2.0.104", +] + [[package]] name = "paste" version = "1.0.15" @@ -3182,7 +3323,7 @@ version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e8af0dde094006011e6a740d4879319439489813bd0bcdc7d821beaeeff48ec" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", ] [[package]] @@ -3506,7 +3647,7 @@ version = "0.38.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys 0.4.15", @@ -3519,7 +3660,7 @@ version = "1.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "11181fbabf243db407ef8df94a6ce0b2f9a733bd8be4ad02b4eda9602296cac8" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys 0.9.4", @@ -3540,6 +3681,27 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pemfile" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "rustls-pki-types" version = "1.12.0" @@ -3606,6 +3768,15 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "schemars" version = "0.9.0" @@ -3668,6 +3839,29 @@ dependencies = [ "syn 2.0.104", ] +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags 2.13.1", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "semver" version = "1.0.28" @@ -4029,6 +4223,29 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "structmeta" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e1575d8d40908d70f6fd05537266b90ae71b15dbbe7a8b7dffa2b759306d329" +dependencies = [ + "proc-macro2", + "quote", + "structmeta-derive", + "syn 2.0.104", +] + +[[package]] +name = "structmeta-derive" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "152a0b65a590ff6c3da95cabe2353ee04e6167c896b28e3b14478c2636c922fc" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + [[package]] name = "strum" version = "0.24.1" @@ -4208,6 +4425,35 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "testcontainers" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23bb7577dca13ad86a78e8271ef5d322f37229ec83b8d98da6d996c588a1ddb1" +dependencies = [ + "async-trait", + "bollard", + "bollard-stubs", + "bytes", + "docker_credential", + "either", + "etcetera", + "futures", + "log", + "memchr", + "parse-display", + "pin-project-lite", + "serde", + "serde_json", + "serde_with", + "thiserror 2.0.12", + "tokio", + "tokio-stream", + "tokio-tar", + "tokio-util", + "url", +] + [[package]] name = "testutils" version = "0.3.70-rc.1" @@ -4590,6 +4836,7 @@ dependencies = [ "snafu", "spinners", "tempfile", + "testcontainers", "testutils", "tokio", "tokio-test", @@ -4612,7 +4859,7 @@ version = "0.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adc82fd73de2a9722ac5da747f12383d2bfdb93591ee6c58486e0097890f05f2" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "bytes", "futures-util", "http", @@ -5147,7 +5394,7 @@ checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" dependencies = [ "windows-implement", "windows-interface", - "windows-link", + "windows-link 0.1.3", "windows-result", "windows-strings", ] @@ -5180,13 +5427,19 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + [[package]] name = "windows-result" version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" dependencies = [ - "windows-link", + "windows-link 0.1.3", ] [[package]] @@ -5195,7 +5448,7 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" dependencies = [ - "windows-link", + "windows-link 0.1.3", ] [[package]] @@ -5243,6 +5496,15 @@ dependencies = [ "windows-targets 0.53.2", ] +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link 0.2.1", +] + [[package]] name = "windows-targets" version = "0.42.2" @@ -5506,7 +5768,7 @@ version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", ] [[package]] diff --git a/crates/tower-cmd/Cargo.toml b/crates/tower-cmd/Cargo.toml index 6d5bcf5e..db90795f 100644 --- a/crates/tower-cmd/Cargo.toml +++ b/crates/tower-cmd/Cargo.toml @@ -49,3 +49,4 @@ uuid = { version = "1.0", features = ["v4"] } futures = { workspace = true } cucumber = { version = "0.21", features = ["macros"] } tokio-test = "0.4" +testcontainers = { version = "0.24", features = ["blocking"] } diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index 4dcf6fd7..44592667 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -13,6 +13,163 @@ use crate::{api, beta, output, util::cmd}; const STORAGE_CATALOG_TYPE: &str = "tower-catalog"; +/// Query-sandbox primitives for untrusted (agent) callers. These are the gates +/// and the DuckDB session lockdown a future MCP `tower_catalogs_query` tool will +/// apply to agent-issued SQL: reject write/DDL and multi-statement input, cap +/// the row count, and lock the session down so a query cannot read the local +/// filesystem, load community extensions, or unwind the settings. Nothing wires +/// these into the query path yet, so `dead_code` is allowed here; the +/// adversarial test suite exercises them directly. +#[allow(dead_code)] +mod sandbox { + /// Row cap for agent-issued queries. Rows past this are dropped and the + /// caller is told the set was truncated, so a model cannot pull an unbounded + /// table into memory or its context. + pub(super) const AGENT_MAX_ROWS: usize = 1_000; + + /// Leading keywords that write data or schema, or repoint the session. An + /// agent query starting with one of these is refused, as defence in depth on + /// top of the read-only credentials it is given. + const WRITE_LEADING_KEYWORDS: &[&str] = &[ + "insert", "update", "delete", "merge", "create", "drop", "alter", "truncate", "replace", + "copy", "attach", "detach", + ]; + + /// Statements that lock an attached DuckDB session down before an untrusted + /// query runs: no local-filesystem access (so `read_csv('/etc/passwd')` and + /// friends fail), no community extensions, and a configuration lock so the + /// query cannot unwind any of it. These run after the catalog is attached, + /// which is what installs extensions and reaches the network. Network access + /// stays on because httpfs is how the attached Iceberg catalog reads its + /// data, so this narrows but does not eliminate the query surface. + pub(super) fn agent_hardening_statements() -> Vec { + vec![ + "SET disabled_filesystems = 'LocalFileSystem'".to_string(), + "SET allow_community_extensions = false".to_string(), + "SET lock_configuration = true".to_string(), + ] + } + + /// The leading SQL keyword, lowercased, after skipping leading whitespace and + /// `--` / `/* */` comments. + pub(super) fn first_sql_keyword(sql: &str) -> String { + let mut s = sql.trim_start(); + loop { + if let Some(rest) = s.strip_prefix("--") { + match rest.find('\n') { + Some(nl) => s = rest[nl + 1..].trim_start(), + None => return String::new(), + } + } else if let Some(rest) = s.strip_prefix("/*") { + match rest.find("*/") { + Some(end) => s = rest[end + 2..].trim_start(), + None => return String::new(), + } + } else { + break; + } + } + s.chars() + .take_while(|c| c.is_ascii_alphabetic()) + .collect::() + .to_lowercase() + } + + /// True when `sql` starts with a write/DDL keyword. + pub(super) fn is_write_statement(sql: &str) -> bool { + WRITE_LEADING_KEYWORDS.contains(&first_sql_keyword(sql).as_str()) + } + + /// True when `sql` holds more than one statement. A `;` inside a string + /// literal or comment is data, not a separator, so those spans are skipped. + /// This gate matters because duckdb-rs `prepare` runs every statement but the + /// last as a side effect, so unguarded multi-statement SQL would execute its + /// leading statements even though only the final one is returned. + pub(super) fn contains_multiple_statements(sql: &str) -> bool { + #[derive(PartialEq)] + enum State { + Normal, + Single, + Double, + Line, + Block, + } + + let mut state = State::Normal; + let mut statements = 0usize; + let mut current_has_content = false; + let mut chars = sql.chars().peekable(); + + while let Some(c) = chars.next() { + match state { + State::Normal => match c { + '\'' => { + state = State::Single; + current_has_content = true; + } + '"' => { + state = State::Double; + current_has_content = true; + } + '-' if chars.peek() == Some(&'-') => { + chars.next(); + state = State::Line; + } + '/' if chars.peek() == Some(&'*') => { + chars.next(); + state = State::Block; + } + ';' => { + if current_has_content { + statements += 1; + if statements > 1 { + return true; + } + } + current_has_content = false; + } + c if c.is_whitespace() => {} + _ => current_has_content = true, + }, + State::Single => { + if c == '\'' { + if chars.peek() == Some(&'\'') { + chars.next(); + } else { + state = State::Normal; + } + } + } + State::Double => { + if c == '"' { + if chars.peek() == Some(&'"') { + chars.next(); + } else { + state = State::Normal; + } + } + } + State::Line => { + if c == '\n' { + state = State::Normal; + } + } + State::Block => { + if c == '*' && chars.peek() == Some(&'/') { + chars.next(); + state = State::Normal; + } + } + } + } + + if current_has_content { + statements += 1; + } + statements > 1 + } +} + pub fn catalogs_cmd() -> Command { Command::new("catalogs") .about(format!( @@ -406,6 +563,7 @@ async fn fetch_catalog_tables( &setup, "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", duckdb::params![db_name], + None, ) }) .await @@ -497,7 +655,7 @@ async fn execute_catalog_query( let token = response.credentials.oauth_token.clone(); let setup = attach_statements(name, &response.credentials, mode); - let result = tokio::task::spawn_blocking(move || run_duckdb_query(&setup, &sql, [])).await; + let result = tokio::task::spawn_blocking(move || run_duckdb_query(&setup, &sql, [], None)).await; match result { Ok(Ok(query_result)) => { @@ -533,9 +691,15 @@ fn read_sql_from_stdin(out: &output::Out) -> String { sql } +#[derive(Debug)] struct QueryResult { columns: Vec, rows: Vec>, + /// Rows were dropped to honour a caller-supplied row cap. Populated by + /// `run_duckdb_query`; the reader lands with the agent query path, so it is + /// exercised only by the sandbox tests today. + #[allow(dead_code)] + truncated: bool, } /// Statements that install the Iceberg support and attach the catalog under @@ -573,11 +737,14 @@ fn attach_statements( /// Runs `setup` statements one at a time, then `query` as a prepared statement /// with `params` bound. Values that fit a bind position should go through -/// `params` rather than into the query text. +/// `params` rather than into the query text. When `max_rows` is set, rows past +/// it are dropped and the result is flagged truncated, so an untrusted caller +/// cannot pull an unbounded table into memory. fn run_duckdb_query( setup: &[String], query: &str, params: P, + max_rows: Option, ) -> Result { let conn = duckdb::Connection::open_in_memory()?; // Setup statements embed the vended OAuth token, so time them without logging @@ -596,6 +763,7 @@ fn run_duckdb_query( let mut stmt = conn.prepare(query)?; let mut columns: Vec = Vec::new(); let mut rows = Vec::new(); + let mut truncated = false; { let mut result_rows = stmt.query(params)?; @@ -603,6 +771,10 @@ fn run_duckdb_query( if columns.is_empty() { columns = row.as_ref().column_names(); } + if max_rows.is_some_and(|max| rows.len() >= max) { + truncated = true; + break; + } let mut record = Vec::with_capacity(columns.len()); for idx in 0..columns.len() { let value: duckdb::types::Value = row.get(idx)?; @@ -624,7 +796,11 @@ fn run_duckdb_query( query ); - Ok(QueryResult { columns, rows }) + Ok(QueryResult { + columns, + rows, + truncated, + }) } /// How many `loadTable` requests `--full` runs against the Iceberg REST catalog @@ -733,6 +909,7 @@ async fn fetch_catalog_columns_via_rest( "column_types".to_string(), ], rows, + truncated: false, }) } @@ -1374,6 +1551,10 @@ fn snippets( #[cfg(test)] mod tests { + use super::sandbox::{ + agent_hardening_statements, contains_multiple_statements, first_sql_keyword, + is_write_statement, + }; use super::{ attach_statements, catalogs_cmd, duckdb_value_to_json, is_storage_catalog_type, parse_mode, run_duckdb_query, snippets, token_export_command, @@ -1594,6 +1775,7 @@ mod tests { &setup, "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", duckdb::params!["memory"], + None, ) .expect("query should succeed"); @@ -1742,7 +1924,7 @@ mod tests { "CREATE TABLE t (id INTEGER, name VARCHAR); INSERT INTO t VALUES (1, 'a'), (2, NULL);" .to_string(), ]; - let result = run_duckdb_query(&setup, "SELECT id, name FROM t ORDER BY id", []) + let result = run_duckdb_query(&setup, "SELECT id, name FROM t ORDER BY id", [], None) .expect("query should succeed"); assert_eq!(result.columns, vec!["id", "name"]); @@ -1763,6 +1945,7 @@ mod tests { &[], "SELECT [1, 2] AS l, {'a': 1, 'b': 'x'} AS s, MAP {'k': 2} AS m", [], + None, ) .expect("query should succeed"); @@ -1779,13 +1962,375 @@ mod tests { #[test] fn run_duckdb_query_reports_columns_for_empty_results() { - let result = - run_duckdb_query(&[], "SELECT 1 AS x WHERE 1 = 0", []).expect("query should succeed"); + let result = run_duckdb_query(&[], "SELECT 1 AS x WHERE 1 = 0", [], None) + .expect("query should succeed"); assert_eq!(result.columns, vec!["x"]); assert!(result.rows.is_empty()); } + #[test] + fn run_duckdb_query_caps_rows_and_flags_truncation() { + let capped = run_duckdb_query(&[], "SELECT * FROM range(5) AS t(i)", [], Some(3)) + .expect("query should succeed"); + assert_eq!(capped.rows.len(), 3); + assert!(capped.truncated); + + let exact = run_duckdb_query(&[], "SELECT * FROM range(3) AS t(i)", [], Some(3)) + .expect("query should succeed"); + assert_eq!(exact.rows.len(), 3); + assert!(!exact.truncated); + } + + // --- Sandbox primitive gates ----------------------------------------- + + #[test] + fn contains_multiple_statements_ignores_separators_in_strings_and_comments() { + assert!(!contains_multiple_statements("SELECT 1")); + assert!(!contains_multiple_statements("SELECT 1;")); + assert!(!contains_multiple_statements(" SELECT 1 ; ")); + assert!(!contains_multiple_statements("SELECT 'a;b'")); + assert!(!contains_multiple_statements("SELECT 1 -- ; not a statement")); + assert!(!contains_multiple_statements("SELECT 1; -- trailing comment")); + assert!(!contains_multiple_statements("SELECT /* ; */ 1")); + + assert!(contains_multiple_statements("SELECT 1; SELECT 2")); + assert!(contains_multiple_statements("SELECT 1; DROP TABLE t")); + assert!(contains_multiple_statements("SELECT 'a;b'; SELECT 2")); + } + + #[test] + fn write_statements_are_detected_through_case_and_comments() { + assert_eq!(first_sql_keyword(" SELECT 1"), "select"); + assert_eq!(first_sql_keyword("/* c */ INSERT INTO t VALUES (1)"), "insert"); + assert_eq!(first_sql_keyword("-- lead\nDELETE FROM t"), "delete"); + + assert!(!is_write_statement("SELECT * FROM t")); + assert!(!is_write_statement("WITH x AS (SELECT 1) SELECT * FROM x")); + assert!(is_write_statement("insert into t values (1)")); + assert!(is_write_statement(" DROP TABLE t")); + assert!(is_write_statement("COPY t TO 'out.csv'")); + assert!(is_write_statement("ATTACH 'x' AS y")); + } + + #[test] + fn agent_hardening_allows_selects_but_blocks_local_files_and_config_changes() { + let setup = agent_hardening_statements(); + + let ok = run_duckdb_query(&setup, "SELECT 1 AS x", [], None) + .expect("a plain select should still run under the hardened session"); + assert_eq!(ok.columns, vec!["x"]); + assert_eq!(ok.rows, vec![vec![serde_json::json!(1)]]); + + let fs_err = run_duckdb_query(&setup, "SELECT * FROM read_csv('Cargo.toml')", [], None) + .expect_err("local filesystem access should be blocked"); + assert!( + fs_err.to_string().to_lowercase().contains("disabled"), + "unexpected error: {fs_err}" + ); + + let cfg_err = run_duckdb_query(&setup, "SET memory_limit = '1GB'", [], None) + .expect_err("configuration should be locked"); + let cfg_msg = cfg_err.to_string().to_lowercase(); + assert!( + cfg_msg.contains("lock") || cfg_msg.contains("configuration"), + "unexpected error: {cfg_err}" + ); + } + + /// Regression check against real Iceberg data. The other `run_duckdb_query` + /// tests use plain in-memory tables; this one writes a small Iceberg table + /// with DuckDB's `COPY … (FORMAT iceberg)` (real metadata + manifests + + /// parquet) and reads it back, so the iceberg reader and our column/row + /// extraction are exercised end to end, NULLs and the row cap included. + /// Needs the `iceberg` extension, fetched on first run and then cached. + #[test] + fn iceberg_scan_reads_written_table_through_run_duckdb_query() { + let dir = std::env::temp_dir().join(format!("tower_iceberg_test_{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).expect("create temp dir"); + let table = dir.join("events").to_string_lossy().replace('\'', "''"); + + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + conn.execute_batch("INSTALL iceberg").expect("install iceberg"); + conn.execute_batch("LOAD iceberg").expect("load iceberg"); + conn.execute_batch(&format!( + "COPY (SELECT * FROM (VALUES (1, 'a'), (2, 'b'), (3, NULL)) AS t(id, name)) \ + TO '{table}' (FORMAT iceberg)" + )) + .expect("write iceberg table"); + + let setup = vec!["INSTALL iceberg".to_string(), "LOAD iceberg".to_string()]; + + let result = run_duckdb_query( + &setup, + &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), + [], + None, + ) + .expect("iceberg_scan should read the table"); + assert_eq!(result.columns, vec!["id", "name"]); + assert_eq!( + result.rows, + vec![ + vec![serde_json::json!(1), serde_json::json!("a")], + vec![serde_json::json!(2), serde_json::json!("b")], + vec![serde_json::json!(3), serde_json::Value::Null], + ] + ); + assert!(!result.truncated); + + let capped = run_duckdb_query( + &setup, + &format!("SELECT id FROM iceberg_scan('{table}') ORDER BY id"), + [], + Some(2), + ) + .expect("iceberg_scan should read the table"); + assert_eq!(capped.rows.len(), 2); + assert!(capped.truncated); + + let _ = std::fs::remove_dir_all(&dir); + } + + // --- Adversarial regression suite ------------------------------------ + // + // Each test encodes an attack an agent-issued query might attempt and + // asserts the sandbox refuses it. These are the security invariants: if a + // future change weakens the gates or the hardening, one of these fails. + + /// A statement run under the full agent hardening. Extensions are not + /// loaded, so the attacks below must be blocked by the session lockdown + /// alone. + fn run_hardened(sql: &str) -> Result { + run_duckdb_query(&agent_hardening_statements(), sql, [], None) + } + + #[test] + fn sandbox_gate_rejects_data_tampering() { + for sql in [ + "DROP TABLE runs", + "DELETE FROM runs", + "UPDATE runs SET id = 0", + "INSERT INTO runs VALUES (1)", + "CREATE TABLE evil AS SELECT 1", + "ALTER TABLE runs ADD COLUMN x INTEGER", + "TRUNCATE runs", + "MERGE INTO runs USING x ON true WHEN MATCHED THEN DELETE", + "COPY runs TO '/tmp/exfil.csv'", + "ATTACH '/tmp/evil.db' AS e", + "DETACH runs", + " drop TABLE runs", + "/* sneaky */ DELETE FROM runs", + ] { + assert!(is_write_statement(sql), "not rejected as write/DDL: {sql}"); + } + for sql in [ + "SELECT * FROM runs", + "WITH t AS (SELECT 1) SELECT * FROM t", + "SELECT count(*) FROM runs", + ] { + assert!(!is_write_statement(sql), "legit read wrongly flagged: {sql}"); + } + } + + #[test] + fn sandbox_gate_rejects_statement_smuggling() { + for sql in [ + "SELECT 1; DROP TABLE runs", + "SELECT 1; DELETE FROM runs", + "SELECT 'a;b'; DROP TABLE runs", + "SELECT 1;\n-- c\nUPDATE runs SET id = 0", + ] { + assert!(contains_multiple_statements(sql), "smuggled statement not caught: {sql}"); + } + assert!(!contains_multiple_statements( + "SELECT * FROM runs -- a trailing ; comment" + )); + } + + #[test] + fn sandbox_blocks_host_filesystem_reads() { + for sql in [ + "SELECT * FROM read_csv('/etc/passwd')", + "SELECT * FROM read_text('/etc/hostname')", + "SELECT * FROM read_json('/etc/passwd')", + "SELECT * FROM read_parquet('/tmp/x.parquet')", + "SELECT * FROM read_csv('/etc/*')", + "SELECT * FROM '/etc/passwd'", + ] { + assert!(run_hardened(sql).is_err(), "host file read NOT blocked: {sql}"); + } + let err = run_hardened("SELECT * FROM read_csv('/etc/passwd')").unwrap_err(); + assert!(err.to_string().to_lowercase().contains("disabled"), "{err}"); + } + + #[test] + fn sandbox_blocks_host_filesystem_writes() { + for sql in [ + "COPY (SELECT 1) TO '/tmp/pwned.csv'", + "COPY (SELECT 1) TO '/tmp/pwned.parquet' (FORMAT parquet)", + ] { + assert!(run_hardened(sql).is_err(), "host file write NOT blocked: {sql}"); + } + } + + #[test] + fn sandbox_blocks_configuration_escape() { + for sql in [ + "SET disabled_filesystems = ''", + "RESET disabled_filesystems", + "SET enable_external_access = true", + "SET allow_community_extensions = true", + "SET lock_configuration = false", + "PRAGMA disabled_filesystems=''", + ] { + assert!(run_hardened(sql).is_err(), "configuration escape NOT blocked: {sql}"); + } + } + + #[test] + fn sandbox_blocks_arbitrary_extension_loading() { + for sql in [ + "LOAD '/tmp/evil.duckdb_extension'", + "INSTALL some_untrusted_extension_xyz", + ] { + assert!(run_hardened(sql).is_err(), "extension load NOT blocked: {sql}"); + } + } + + #[test] + fn sandbox_blocks_network_ssrf_via_table_functions() { + for sql in [ + "SELECT * FROM read_csv('http://169.254.169.254/latest/meta-data/')", + "SELECT * FROM read_parquet('https://attacker.example/x.parquet')", + "SELECT * FROM read_csv('http://localhost:8080/internal')", + ] { + assert!(run_hardened(sql).is_err(), "SSRF NOT blocked: {sql}"); + } + } + + /// Polls MinIO's health endpoint over a raw socket until it answers 200. + fn wait_for_minio(port: u16) -> bool { + use std::io::{Read, Write}; + for _ in 0..60 { + if let Ok(mut stream) = std::net::TcpStream::connect(("127.0.0.1", port)) { + let _ = stream.set_read_timeout(Some(std::time::Duration::from_secs(2))); + let request = + "GET /minio/health/live HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"; + let mut response = String::new(); + if stream.write_all(request.as_bytes()).is_ok() + && stream.read_to_string(&mut response).is_ok() + && response.contains(" 200 ") + { + return true; + } + } + std::thread::sleep(std::time::Duration::from_millis(500)); + } + false + } + + /// The coverage the local-only tests can't give: it proves the hardening + /// does NOT break a real object-store Iceberg read (the reader uses + /// S3FileSystem, which the hardening leaves enabled) while local-filesystem + /// access and configuration changes stay blocked in the same session. Starts + /// a MinIO container via testcontainers and self-skips when no Docker daemon + /// is available, so a plain `cargo test` still passes without Docker. + #[test] + fn sandbox_holds_over_object_store_iceberg() { + use testcontainers::core::{IntoContainerPort, WaitFor}; + use testcontainers::runners::SyncRunner; + use testcontainers::{GenericImage, ImageExt}; + + let bucket = "warehouse"; + let image = GenericImage::new("minio/minio", "latest") + .with_wait_for(WaitFor::seconds(1)) + .with_exposed_port(9000.tcp()) + .with_entrypoint("sh") + .with_env_var("MINIO_ROOT_USER", "minioadmin") + .with_env_var("MINIO_ROOT_PASSWORD", "minioadmin") + .with_cmd([ + "-c".to_string(), + format!("mkdir -p /data/{bucket} && exec minio server /data"), + ]); + + let container = match image.start() { + Ok(container) => container, + Err(err) => { + eprintln!( + "skipping sandbox_holds_over_object_store_iceberg (no Docker daemon?): {err}" + ); + return; + } + }; + let port = container + .get_host_port_ipv4(9000.tcp()) + .expect("mapped MinIO port"); + assert!(wait_for_minio(port), "MinIO did not become healthy"); + + let secret = format!( + "CREATE SECRET s3sec (TYPE s3, KEY_ID 'minioadmin', SECRET 'minioadmin', \ + ENDPOINT '127.0.0.1:{port}', URL_STYLE 'path', USE_SSL false, REGION 'us-east-1')" + ); + let table = format!("s3://{bucket}/tbl"); + let extensions = || { + vec![ + "INSTALL httpfs".to_string(), + "LOAD httpfs".to_string(), + "INSTALL iceberg".to_string(), + "LOAD iceberg".to_string(), + ] + }; + + let seed = duckdb::Connection::open_in_memory().expect("open duckdb"); + for stmt in extensions() { + seed.execute_batch(&stmt).expect("load extension"); + } + seed.execute_batch(&secret).expect("create s3 secret"); + seed.execute_batch(&format!( + "COPY (SELECT * FROM (VALUES (1,'a'),(2,'b'),(3,NULL)) t(id,name)) \ + TO '{table}' (FORMAT iceberg)" + )) + .expect("seed iceberg table on object storage"); + + let mut setup = extensions(); + setup.push(secret); + setup.extend(agent_hardening_statements()); + + let read = run_duckdb_query( + &setup, + &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), + [], + None, + ) + .expect("hardening must not break object-store Iceberg reads"); + assert_eq!(read.columns, vec!["id", "name"]); + assert_eq!( + read.rows, + vec![ + vec![serde_json::json!(1), serde_json::json!("a")], + vec![serde_json::json!(2), serde_json::json!("b")], + vec![serde_json::json!(3), serde_json::Value::Null], + ] + ); + + let local = run_duckdb_query(&setup, "SELECT * FROM read_csv('/etc/hostname')", [], None) + .expect_err("local filesystem reads must stay blocked"); + assert!( + local.to_string().to_lowercase().contains("disabled"), + "unexpected error: {local}" + ); + + let cfg = run_duckdb_query(&setup, "SET memory_limit = '1GB'", [], None) + .expect_err("configuration must stay locked"); + let cfg = cfg.to_string().to_lowercase(); + assert!( + cfg.contains("lock") || cfg.contains("configuration"), + "unexpected error: {cfg}" + ); + } + #[test] fn token_export_command_fetches_token_without_printing_it() { let credentials = CatalogCredentials::new( From f35c4e7110f11ea89d1a1dd17026eb0cb3437ede Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Fri, 24 Jul 2026 10:33:11 +0100 Subject: [PATCH 2/4] 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. --- Cargo.lock | 14 +- Cargo.toml | 1 + crates/tower-cmd/Cargo.toml | 3 +- crates/tower-cmd/src/catalogs.rs | 815 ++----------------------------- crates/tower-duckdb/Cargo.toml | 17 + crates/tower-duckdb/src/guard.rs | 135 +++++ crates/tower-duckdb/src/lib.rs | 700 ++++++++++++++++++++++++++ 7 files changed, 912 insertions(+), 773 deletions(-) create mode 100644 crates/tower-duckdb/Cargo.toml create mode 100644 crates/tower-duckdb/src/guard.rs create mode 100644 crates/tower-duckdb/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 41b1409a..f728341f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4817,7 +4817,6 @@ dependencies = [ "crypto", "ctrlc", "cucumber", - "duckdb", "futures", "futures-util", "http", @@ -4836,7 +4835,6 @@ dependencies = [ "snafu", "spinners", "tempfile", - "testcontainers", "testutils", "tokio", "tokio-test", @@ -4844,6 +4842,7 @@ dependencies = [ "toml", "toml_edit 0.22.27", "tower-api", + "tower-duckdb", "tower-package", "tower-runtime", "tower-telemetry", @@ -4853,6 +4852,17 @@ dependencies = [ "webbrowser", ] +[[package]] +name = "tower-duckdb" +version = "0.3.70-rc.1" +dependencies = [ + "chrono", + "duckdb", + "serde_json", + "testcontainers", + "tower-telemetry", +] + [[package]] name = "tower-http" version = "0.6.6" diff --git a/Cargo.toml b/Cargo.toml index a005bb5a..b737e493 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -68,6 +68,7 @@ toml = "0.8" toml_edit = "0.22" tower-api = { path = "crates/tower-api" } tower-cmd = { path = "crates/tower-cmd" } +tower-duckdb = { path = "crates/tower-duckdb" } tower-package = { path = "crates/tower-package" } tower-runtime = { path = "crates/tower-runtime" } tower-telemetry = { path = "crates/tower-telemetry" } diff --git a/crates/tower-cmd/Cargo.toml b/crates/tower-cmd/Cargo.toml index db90795f..93c5fe01 100644 --- a/crates/tower-cmd/Cargo.toml +++ b/crates/tower-cmd/Cargo.toml @@ -11,7 +11,6 @@ clap = { workspace = true } cli-table = { workspace = true } colored = { workspace = true } config = { workspace = true } -duckdb = { workspace = true } crypto = { workspace = true } ctrlc = { workspace = true } futures-util = { workspace = true } @@ -30,6 +29,7 @@ spinners = { workspace = true } tokio = { workspace = true } tokio-util = { workspace = true } tower-api = { workspace = true } +tower-duckdb = { workspace = true } tower-package = { workspace = true } tower-runtime = { workspace = true } tower-telemetry = { workspace = true } @@ -49,4 +49,3 @@ uuid = { version = "1.0", features = ["v4"] } futures = { workspace = true } cucumber = { version = "0.21", features = ["macros"] } tokio-test = "0.4" -testcontainers = { version = "0.24", features = ["blocking"] } diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index 44592667..0ab260b8 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -7,169 +7,13 @@ use std::time::{Duration, Instant}; use tower_api::models::{ vend_catalog_credentials_body, CatalogCredentials, DescribeCatalogResponse, }; +use tower_duckdb::{guard, params, run_query, QueryResult, Session}; use tower_telemetry::debug; use crate::{api, beta, output, util::cmd}; const STORAGE_CATALOG_TYPE: &str = "tower-catalog"; -/// Query-sandbox primitives for untrusted (agent) callers. These are the gates -/// and the DuckDB session lockdown a future MCP `tower_catalogs_query` tool will -/// apply to agent-issued SQL: reject write/DDL and multi-statement input, cap -/// the row count, and lock the session down so a query cannot read the local -/// filesystem, load community extensions, or unwind the settings. Nothing wires -/// these into the query path yet, so `dead_code` is allowed here; the -/// adversarial test suite exercises them directly. -#[allow(dead_code)] -mod sandbox { - /// Row cap for agent-issued queries. Rows past this are dropped and the - /// caller is told the set was truncated, so a model cannot pull an unbounded - /// table into memory or its context. - pub(super) const AGENT_MAX_ROWS: usize = 1_000; - - /// Leading keywords that write data or schema, or repoint the session. An - /// agent query starting with one of these is refused, as defence in depth on - /// top of the read-only credentials it is given. - const WRITE_LEADING_KEYWORDS: &[&str] = &[ - "insert", "update", "delete", "merge", "create", "drop", "alter", "truncate", "replace", - "copy", "attach", "detach", - ]; - - /// Statements that lock an attached DuckDB session down before an untrusted - /// query runs: no local-filesystem access (so `read_csv('/etc/passwd')` and - /// friends fail), no community extensions, and a configuration lock so the - /// query cannot unwind any of it. These run after the catalog is attached, - /// which is what installs extensions and reaches the network. Network access - /// stays on because httpfs is how the attached Iceberg catalog reads its - /// data, so this narrows but does not eliminate the query surface. - pub(super) fn agent_hardening_statements() -> Vec { - vec![ - "SET disabled_filesystems = 'LocalFileSystem'".to_string(), - "SET allow_community_extensions = false".to_string(), - "SET lock_configuration = true".to_string(), - ] - } - - /// The leading SQL keyword, lowercased, after skipping leading whitespace and - /// `--` / `/* */` comments. - pub(super) fn first_sql_keyword(sql: &str) -> String { - let mut s = sql.trim_start(); - loop { - if let Some(rest) = s.strip_prefix("--") { - match rest.find('\n') { - Some(nl) => s = rest[nl + 1..].trim_start(), - None => return String::new(), - } - } else if let Some(rest) = s.strip_prefix("/*") { - match rest.find("*/") { - Some(end) => s = rest[end + 2..].trim_start(), - None => return String::new(), - } - } else { - break; - } - } - s.chars() - .take_while(|c| c.is_ascii_alphabetic()) - .collect::() - .to_lowercase() - } - - /// True when `sql` starts with a write/DDL keyword. - pub(super) fn is_write_statement(sql: &str) -> bool { - WRITE_LEADING_KEYWORDS.contains(&first_sql_keyword(sql).as_str()) - } - - /// True when `sql` holds more than one statement. A `;` inside a string - /// literal or comment is data, not a separator, so those spans are skipped. - /// This gate matters because duckdb-rs `prepare` runs every statement but the - /// last as a side effect, so unguarded multi-statement SQL would execute its - /// leading statements even though only the final one is returned. - pub(super) fn contains_multiple_statements(sql: &str) -> bool { - #[derive(PartialEq)] - enum State { - Normal, - Single, - Double, - Line, - Block, - } - - let mut state = State::Normal; - let mut statements = 0usize; - let mut current_has_content = false; - let mut chars = sql.chars().peekable(); - - while let Some(c) = chars.next() { - match state { - State::Normal => match c { - '\'' => { - state = State::Single; - current_has_content = true; - } - '"' => { - state = State::Double; - current_has_content = true; - } - '-' if chars.peek() == Some(&'-') => { - chars.next(); - state = State::Line; - } - '/' if chars.peek() == Some(&'*') => { - chars.next(); - state = State::Block; - } - ';' => { - if current_has_content { - statements += 1; - if statements > 1 { - return true; - } - } - current_has_content = false; - } - c if c.is_whitespace() => {} - _ => current_has_content = true, - }, - State::Single => { - if c == '\'' { - if chars.peek() == Some(&'\'') { - chars.next(); - } else { - state = State::Normal; - } - } - } - State::Double => { - if c == '"' { - if chars.peek() == Some(&'"') { - chars.next(); - } else { - state = State::Normal; - } - } - } - State::Line => { - if c == '\n' { - state = State::Normal; - } - } - State::Block => { - if c == '*' && chars.peek() == Some(&'/') { - chars.next(); - state = State::Normal; - } - } - } - } - - if current_has_content { - statements += 1; - } - statements > 1 - } -} - pub fn catalogs_cmd() -> Command { Command::new("catalogs") .about(format!( @@ -559,10 +403,10 @@ async fn fetch_catalog_tables( ); let db_name = name.to_string(); tokio::task::spawn_blocking(move || { - run_duckdb_query( + run_query( &setup, "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", - duckdb::params![db_name], + params![db_name], None, ) }) @@ -621,13 +465,27 @@ pub async fn do_query(out: &output::Out, config: Config, args: &ArgMatches) { } let write = cmd::get_bool_flag(args, "write"); + // Read mode runs untrusted SQL, so gate it before it reaches DuckDB: a + // smuggled second statement would otherwise execute as a side effect of + // `prepare`, and a write in read mode should fail with a clear message rather + // than a raw engine error. Write mode is the trusted power-user path. + if !write { + if guard::is_multi_statement(&sql) { + out.die("Only a single SQL statement can be run at a time. Remove the extra statement(s)."); + } + if guard::is_write_statement(&sql) { + out.die("This command runs read-only queries. The statement looks like it writes or changes data; re-run with --write to modify the catalog."); + } + } let query_result = execute_catalog_query(out, &config, name, &env, sql, write).await; output_query_result(out, &query_result); } /// Vends credentials for the catalog, attaches it in an in-memory DuckDB, and -/// runs `sql` against it. Read-only unless `write` is set, in which case -/// read-write credentials are vended and the attach allows writes. Dies with a +/// runs `sql` against it. In read mode (the default) the session is hardened +/// after attach and the result row count is capped, so an untrusted query cannot +/// read the host or pull an unbounded table back. `write` vends read-write +/// credentials, lets the attach write, and runs the query trusted. Dies with a /// user-facing error on failure. async fn execute_catalog_query( out: &output::Out, @@ -655,7 +513,20 @@ async fn execute_catalog_query( let token = response.credentials.oauth_token.clone(); let setup = attach_statements(name, &response.credentials, mode); - let result = tokio::task::spawn_blocking(move || run_duckdb_query(&setup, &sql, [], None)).await; + // Read mode is the sandboxed path: lock the session down after attach and cap + // the rows, so the query cannot read the host, escape the config, or pull an + // unbounded table back. Write mode is trusted and runs the setup as-is. + let harden = !write; + let max_rows = (!write).then_some(guard::AGENT_MAX_ROWS); + let result = tokio::task::spawn_blocking(move || -> Result { + let session = Session::open()?; + session.run_setup(&setup)?; + if harden { + session.harden()?; + } + session.query(&sql, [], max_rows) + }) + .await; match result { Ok(Ok(query_result)) => { @@ -691,17 +562,6 @@ fn read_sql_from_stdin(out: &output::Out) -> String { sql } -#[derive(Debug)] -struct QueryResult { - columns: Vec, - rows: Vec>, - /// Rows were dropped to honour a caller-supplied row cap. Populated by - /// `run_duckdb_query`; the reader lands with the agent query path, so it is - /// exercised only by the sandbox tests today. - #[allow(dead_code)] - truncated: bool, -} - /// Statements that install the Iceberg support and attach the catalog under /// its Tower name — mirrors `templates/duckdb.sql.tmpl`. The attach is /// READ_ONLY unless read-write credentials were vended. No `USE`: DuckDB's @@ -735,74 +595,6 @@ fn attach_statements( ] } -/// Runs `setup` statements one at a time, then `query` as a prepared statement -/// with `params` bound. Values that fit a bind position should go through -/// `params` rather than into the query text. When `max_rows` is set, rows past -/// it are dropped and the result is flagged truncated, so an untrusted caller -/// cannot pull an unbounded table into memory. -fn run_duckdb_query( - setup: &[String], - query: &str, - params: P, - max_rows: Option, -) -> Result { - let conn = duckdb::Connection::open_in_memory()?; - // Setup statements embed the vended OAuth token, so time them without logging - // their text. - let setup_start = Instant::now(); - for statement in setup { - conn.execute_batch(statement)?; - } - debug!( - "duckdb: setup ({} statements) took {:?}", - setup.len(), - setup_start.elapsed() - ); - - let query_start = Instant::now(); - let mut stmt = conn.prepare(query)?; - let mut columns: Vec = Vec::new(); - let mut rows = Vec::new(); - let mut truncated = false; - - { - let mut result_rows = stmt.query(params)?; - while let Some(row) = result_rows.next()? { - if columns.is_empty() { - columns = row.as_ref().column_names(); - } - if max_rows.is_some_and(|max| rows.len() >= max) { - truncated = true; - break; - } - let mut record = Vec::with_capacity(columns.len()); - for idx in 0..columns.len() { - let value: duckdb::types::Value = row.get(idx)?; - record.push(duckdb_value_to_json(value)); - } - rows.push(record); - } - } - - // A query with no result rows never populates columns above. - if columns.is_empty() { - columns = stmt.column_names(); - } - - debug!( - "duckdb: query took {:?} ({} rows): {}", - query_start.elapsed(), - rows.len(), - query - ); - - Ok(QueryResult { - columns, - rows, - truncated, - }) -} - /// How many `loadTable` requests `--full` runs against the Iceberg REST catalog /// at once. Each is a single metadata fetch; kept modest because Polaris rate /// limits (HTTP 429) aggressive fan-out. Throttled requests are retried, so this @@ -1176,89 +968,6 @@ fn iceberg_primitive_to_display(name: &str) -> String { } } -fn duckdb_value_to_json(value: duckdb::types::Value) -> serde_json::Value { - use duckdb::types::{TimeUnit, Value}; - use serde_json::json; - - match value { - Value::Null => serde_json::Value::Null, - Value::Boolean(v) => json!(v), - Value::TinyInt(v) => json!(v), - Value::SmallInt(v) => json!(v), - Value::Int(v) => json!(v), - Value::BigInt(v) => json!(v), - Value::HugeInt(v) => json!(v.to_string()), - Value::UTinyInt(v) => json!(v), - Value::USmallInt(v) => json!(v), - Value::UInt(v) => json!(v), - Value::UBigInt(v) => json!(v), - Value::Float(v) => json!(v), - Value::Double(v) => json!(v), - Value::Decimal(v) => json!(v.to_string()), - Value::Text(v) => json!(v), - Value::Timestamp(unit, v) => { - let micros = match unit { - TimeUnit::Second => v.checked_mul(1_000_000), - TimeUnit::Millisecond => v.checked_mul(1_000), - TimeUnit::Microsecond => Some(v), - TimeUnit::Nanosecond => Some(v / 1_000), - }; - match micros.and_then(chrono::DateTime::from_timestamp_micros) { - Some(ts) => json!(ts.naive_utc().to_string()), - None => json!(format!("{:?}", Value::Timestamp(unit, v))), - } - } - Value::Date32(days) => { - let date = chrono::DateTime::from_timestamp(i64::from(days) * 86_400, 0); - match date { - Some(d) => json!(d.date_naive().to_string()), - None => json!(format!("{:?}", Value::Date32(days))), - } - } - Value::Time64(unit, v) => { - let micros = match unit { - TimeUnit::Second => v.checked_mul(1_000_000), - TimeUnit::Millisecond => v.checked_mul(1_000), - TimeUnit::Microsecond => Some(v), - TimeUnit::Nanosecond => Some(v / 1_000), - }; - let time = micros.and_then(|m| { - chrono::NaiveTime::from_num_seconds_from_midnight_opt( - (m / 1_000_000) as u32, - ((m % 1_000_000) * 1_000) as u32, - ) - }); - match time { - Some(t) => json!(t.to_string()), - None => json!(format!("{:?}", Value::Time64(unit, v))), - } - } - Value::Enum(v) => json!(v), - Value::List(items) | Value::Array(items) => { - serde_json::Value::Array(items.into_iter().map(duckdb_value_to_json).collect()) - } - Value::Struct(fields) => serde_json::Value::Object( - fields - .iter() - .map(|(name, value)| (name.clone(), duckdb_value_to_json(value.clone()))) - .collect(), - ), - Value::Map(entries) => serde_json::Value::Object( - entries - .iter() - .map(|(key, value)| { - ( - json_value_to_cell(&duckdb_value_to_json(key.clone())), - duckdb_value_to_json(value.clone()), - ) - }) - .collect(), - ), - Value::Union(inner) => duckdb_value_to_json(*inner), - other => json!(format!("{:?}", other)), - } -} - fn output_query_result(out: &output::Out, result: &QueryResult) { let json_rows: Vec> = result .rows @@ -1280,7 +989,14 @@ fn output_query_result(out: &output::Out, result: &QueryResult) { .collect(); out.table(result.columns.clone(), data, Some(&json_rows)); - out.note(&format!("\n{} row(s)\n", result.rows.len())); + if result.truncated { + out.note(&format!( + "\nShowing the first {} row(s); result truncated. Add a LIMIT or filter to narrow it.\n", + result.rows.len() + )); + } else { + out.note(&format!("\n{} row(s)\n", result.rows.len())); + } } fn json_value_to_cell(value: &serde_json::Value) -> String { @@ -1551,15 +1267,12 @@ fn snippets( #[cfg(test)] mod tests { - use super::sandbox::{ - agent_hardening_statements, contains_multiple_statements, first_sql_keyword, - is_write_statement, - }; use super::{ - attach_statements, catalogs_cmd, duckdb_value_to_json, is_storage_catalog_type, parse_mode, - run_duckdb_query, snippets, token_export_command, + attach_statements, catalogs_cmd, is_storage_catalog_type, parse_mode, snippets, + token_export_command, }; use tower_api::models::{vend_catalog_credentials_body, CatalogCredentials}; + use tower_duckdb::{params, run_query}; #[test] fn list_defaults_to_default_environment() { @@ -1771,10 +1484,10 @@ mod tests { "CREATE SCHEMA s; CREATE TABLE s.t1 (i INTEGER); CREATE TABLE s.t2 (i INTEGER);" .to_string(), ]; - let result = run_duckdb_query( + let result = run_query( &setup, "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", - duckdb::params!["memory"], + params!["memory"], None, ) .expect("query should succeed"); @@ -1895,442 +1608,6 @@ mod tests { assert_eq!(query_args.get_one::("write").copied(), Some(true)); } - #[test] - fn duckdb_values_convert_to_json() { - use duckdb::types::{TimeUnit, Value}; - - assert_eq!(duckdb_value_to_json(Value::Null), serde_json::Value::Null); - assert_eq!( - duckdb_value_to_json(Value::BigInt(42)), - serde_json::json!(42) - ); - assert_eq!( - duckdb_value_to_json(Value::Text("hi".to_string())), - serde_json::json!("hi") - ); - assert_eq!( - duckdb_value_to_json(Value::Timestamp(TimeUnit::Microsecond, 0)), - serde_json::json!("1970-01-01 00:00:00") - ); - assert_eq!( - duckdb_value_to_json(Value::Date32(1)), - serde_json::json!("1970-01-02") - ); - } - - #[test] - fn run_duckdb_query_returns_columns_and_rows() { - let setup = vec![ - "CREATE TABLE t (id INTEGER, name VARCHAR); INSERT INTO t VALUES (1, 'a'), (2, NULL);" - .to_string(), - ]; - let result = run_duckdb_query(&setup, "SELECT id, name FROM t ORDER BY id", [], None) - .expect("query should succeed"); - - assert_eq!(result.columns, vec!["id", "name"]); - assert_eq!(result.rows.len(), 2); - assert_eq!( - result.rows[0], - vec![serde_json::json!(1), serde_json::json!("a")] - ); - assert_eq!( - result.rows[1], - vec![serde_json::json!(2), serde_json::Value::Null] - ); - } - - #[test] - fn nested_duckdb_values_convert_to_json_structures() { - let result = run_duckdb_query( - &[], - "SELECT [1, 2] AS l, {'a': 1, 'b': 'x'} AS s, MAP {'k': 2} AS m", - [], - None, - ) - .expect("query should succeed"); - - assert_eq!(result.columns, vec!["l", "s", "m"]); - assert_eq!( - result.rows[0], - vec![ - serde_json::json!([1, 2]), - serde_json::json!({"a": 1, "b": "x"}), - serde_json::json!({"k": 2}), - ] - ); - } - - #[test] - fn run_duckdb_query_reports_columns_for_empty_results() { - let result = run_duckdb_query(&[], "SELECT 1 AS x WHERE 1 = 0", [], None) - .expect("query should succeed"); - - assert_eq!(result.columns, vec!["x"]); - assert!(result.rows.is_empty()); - } - - #[test] - fn run_duckdb_query_caps_rows_and_flags_truncation() { - let capped = run_duckdb_query(&[], "SELECT * FROM range(5) AS t(i)", [], Some(3)) - .expect("query should succeed"); - assert_eq!(capped.rows.len(), 3); - assert!(capped.truncated); - - let exact = run_duckdb_query(&[], "SELECT * FROM range(3) AS t(i)", [], Some(3)) - .expect("query should succeed"); - assert_eq!(exact.rows.len(), 3); - assert!(!exact.truncated); - } - - // --- Sandbox primitive gates ----------------------------------------- - - #[test] - fn contains_multiple_statements_ignores_separators_in_strings_and_comments() { - assert!(!contains_multiple_statements("SELECT 1")); - assert!(!contains_multiple_statements("SELECT 1;")); - assert!(!contains_multiple_statements(" SELECT 1 ; ")); - assert!(!contains_multiple_statements("SELECT 'a;b'")); - assert!(!contains_multiple_statements("SELECT 1 -- ; not a statement")); - assert!(!contains_multiple_statements("SELECT 1; -- trailing comment")); - assert!(!contains_multiple_statements("SELECT /* ; */ 1")); - - assert!(contains_multiple_statements("SELECT 1; SELECT 2")); - assert!(contains_multiple_statements("SELECT 1; DROP TABLE t")); - assert!(contains_multiple_statements("SELECT 'a;b'; SELECT 2")); - } - - #[test] - fn write_statements_are_detected_through_case_and_comments() { - assert_eq!(first_sql_keyword(" SELECT 1"), "select"); - assert_eq!(first_sql_keyword("/* c */ INSERT INTO t VALUES (1)"), "insert"); - assert_eq!(first_sql_keyword("-- lead\nDELETE FROM t"), "delete"); - - assert!(!is_write_statement("SELECT * FROM t")); - assert!(!is_write_statement("WITH x AS (SELECT 1) SELECT * FROM x")); - assert!(is_write_statement("insert into t values (1)")); - assert!(is_write_statement(" DROP TABLE t")); - assert!(is_write_statement("COPY t TO 'out.csv'")); - assert!(is_write_statement("ATTACH 'x' AS y")); - } - - #[test] - fn agent_hardening_allows_selects_but_blocks_local_files_and_config_changes() { - let setup = agent_hardening_statements(); - - let ok = run_duckdb_query(&setup, "SELECT 1 AS x", [], None) - .expect("a plain select should still run under the hardened session"); - assert_eq!(ok.columns, vec!["x"]); - assert_eq!(ok.rows, vec![vec![serde_json::json!(1)]]); - - let fs_err = run_duckdb_query(&setup, "SELECT * FROM read_csv('Cargo.toml')", [], None) - .expect_err("local filesystem access should be blocked"); - assert!( - fs_err.to_string().to_lowercase().contains("disabled"), - "unexpected error: {fs_err}" - ); - - let cfg_err = run_duckdb_query(&setup, "SET memory_limit = '1GB'", [], None) - .expect_err("configuration should be locked"); - let cfg_msg = cfg_err.to_string().to_lowercase(); - assert!( - cfg_msg.contains("lock") || cfg_msg.contains("configuration"), - "unexpected error: {cfg_err}" - ); - } - - /// Regression check against real Iceberg data. The other `run_duckdb_query` - /// tests use plain in-memory tables; this one writes a small Iceberg table - /// with DuckDB's `COPY … (FORMAT iceberg)` (real metadata + manifests + - /// parquet) and reads it back, so the iceberg reader and our column/row - /// extraction are exercised end to end, NULLs and the row cap included. - /// Needs the `iceberg` extension, fetched on first run and then cached. - #[test] - fn iceberg_scan_reads_written_table_through_run_duckdb_query() { - let dir = std::env::temp_dir().join(format!("tower_iceberg_test_{}", std::process::id())); - let _ = std::fs::remove_dir_all(&dir); - std::fs::create_dir_all(&dir).expect("create temp dir"); - let table = dir.join("events").to_string_lossy().replace('\'', "''"); - - let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); - conn.execute_batch("INSTALL iceberg").expect("install iceberg"); - conn.execute_batch("LOAD iceberg").expect("load iceberg"); - conn.execute_batch(&format!( - "COPY (SELECT * FROM (VALUES (1, 'a'), (2, 'b'), (3, NULL)) AS t(id, name)) \ - TO '{table}' (FORMAT iceberg)" - )) - .expect("write iceberg table"); - - let setup = vec!["INSTALL iceberg".to_string(), "LOAD iceberg".to_string()]; - - let result = run_duckdb_query( - &setup, - &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), - [], - None, - ) - .expect("iceberg_scan should read the table"); - assert_eq!(result.columns, vec!["id", "name"]); - assert_eq!( - result.rows, - vec![ - vec![serde_json::json!(1), serde_json::json!("a")], - vec![serde_json::json!(2), serde_json::json!("b")], - vec![serde_json::json!(3), serde_json::Value::Null], - ] - ); - assert!(!result.truncated); - - let capped = run_duckdb_query( - &setup, - &format!("SELECT id FROM iceberg_scan('{table}') ORDER BY id"), - [], - Some(2), - ) - .expect("iceberg_scan should read the table"); - assert_eq!(capped.rows.len(), 2); - assert!(capped.truncated); - - let _ = std::fs::remove_dir_all(&dir); - } - - // --- Adversarial regression suite ------------------------------------ - // - // Each test encodes an attack an agent-issued query might attempt and - // asserts the sandbox refuses it. These are the security invariants: if a - // future change weakens the gates or the hardening, one of these fails. - - /// A statement run under the full agent hardening. Extensions are not - /// loaded, so the attacks below must be blocked by the session lockdown - /// alone. - fn run_hardened(sql: &str) -> Result { - run_duckdb_query(&agent_hardening_statements(), sql, [], None) - } - - #[test] - fn sandbox_gate_rejects_data_tampering() { - for sql in [ - "DROP TABLE runs", - "DELETE FROM runs", - "UPDATE runs SET id = 0", - "INSERT INTO runs VALUES (1)", - "CREATE TABLE evil AS SELECT 1", - "ALTER TABLE runs ADD COLUMN x INTEGER", - "TRUNCATE runs", - "MERGE INTO runs USING x ON true WHEN MATCHED THEN DELETE", - "COPY runs TO '/tmp/exfil.csv'", - "ATTACH '/tmp/evil.db' AS e", - "DETACH runs", - " drop TABLE runs", - "/* sneaky */ DELETE FROM runs", - ] { - assert!(is_write_statement(sql), "not rejected as write/DDL: {sql}"); - } - for sql in [ - "SELECT * FROM runs", - "WITH t AS (SELECT 1) SELECT * FROM t", - "SELECT count(*) FROM runs", - ] { - assert!(!is_write_statement(sql), "legit read wrongly flagged: {sql}"); - } - } - - #[test] - fn sandbox_gate_rejects_statement_smuggling() { - for sql in [ - "SELECT 1; DROP TABLE runs", - "SELECT 1; DELETE FROM runs", - "SELECT 'a;b'; DROP TABLE runs", - "SELECT 1;\n-- c\nUPDATE runs SET id = 0", - ] { - assert!(contains_multiple_statements(sql), "smuggled statement not caught: {sql}"); - } - assert!(!contains_multiple_statements( - "SELECT * FROM runs -- a trailing ; comment" - )); - } - - #[test] - fn sandbox_blocks_host_filesystem_reads() { - for sql in [ - "SELECT * FROM read_csv('/etc/passwd')", - "SELECT * FROM read_text('/etc/hostname')", - "SELECT * FROM read_json('/etc/passwd')", - "SELECT * FROM read_parquet('/tmp/x.parquet')", - "SELECT * FROM read_csv('/etc/*')", - "SELECT * FROM '/etc/passwd'", - ] { - assert!(run_hardened(sql).is_err(), "host file read NOT blocked: {sql}"); - } - let err = run_hardened("SELECT * FROM read_csv('/etc/passwd')").unwrap_err(); - assert!(err.to_string().to_lowercase().contains("disabled"), "{err}"); - } - - #[test] - fn sandbox_blocks_host_filesystem_writes() { - for sql in [ - "COPY (SELECT 1) TO '/tmp/pwned.csv'", - "COPY (SELECT 1) TO '/tmp/pwned.parquet' (FORMAT parquet)", - ] { - assert!(run_hardened(sql).is_err(), "host file write NOT blocked: {sql}"); - } - } - - #[test] - fn sandbox_blocks_configuration_escape() { - for sql in [ - "SET disabled_filesystems = ''", - "RESET disabled_filesystems", - "SET enable_external_access = true", - "SET allow_community_extensions = true", - "SET lock_configuration = false", - "PRAGMA disabled_filesystems=''", - ] { - assert!(run_hardened(sql).is_err(), "configuration escape NOT blocked: {sql}"); - } - } - - #[test] - fn sandbox_blocks_arbitrary_extension_loading() { - for sql in [ - "LOAD '/tmp/evil.duckdb_extension'", - "INSTALL some_untrusted_extension_xyz", - ] { - assert!(run_hardened(sql).is_err(), "extension load NOT blocked: {sql}"); - } - } - - #[test] - fn sandbox_blocks_network_ssrf_via_table_functions() { - for sql in [ - "SELECT * FROM read_csv('http://169.254.169.254/latest/meta-data/')", - "SELECT * FROM read_parquet('https://attacker.example/x.parquet')", - "SELECT * FROM read_csv('http://localhost:8080/internal')", - ] { - assert!(run_hardened(sql).is_err(), "SSRF NOT blocked: {sql}"); - } - } - - /// Polls MinIO's health endpoint over a raw socket until it answers 200. - fn wait_for_minio(port: u16) -> bool { - use std::io::{Read, Write}; - for _ in 0..60 { - if let Ok(mut stream) = std::net::TcpStream::connect(("127.0.0.1", port)) { - let _ = stream.set_read_timeout(Some(std::time::Duration::from_secs(2))); - let request = - "GET /minio/health/live HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"; - let mut response = String::new(); - if stream.write_all(request.as_bytes()).is_ok() - && stream.read_to_string(&mut response).is_ok() - && response.contains(" 200 ") - { - return true; - } - } - std::thread::sleep(std::time::Duration::from_millis(500)); - } - false - } - - /// The coverage the local-only tests can't give: it proves the hardening - /// does NOT break a real object-store Iceberg read (the reader uses - /// S3FileSystem, which the hardening leaves enabled) while local-filesystem - /// access and configuration changes stay blocked in the same session. Starts - /// a MinIO container via testcontainers and self-skips when no Docker daemon - /// is available, so a plain `cargo test` still passes without Docker. - #[test] - fn sandbox_holds_over_object_store_iceberg() { - use testcontainers::core::{IntoContainerPort, WaitFor}; - use testcontainers::runners::SyncRunner; - use testcontainers::{GenericImage, ImageExt}; - - let bucket = "warehouse"; - let image = GenericImage::new("minio/minio", "latest") - .with_wait_for(WaitFor::seconds(1)) - .with_exposed_port(9000.tcp()) - .with_entrypoint("sh") - .with_env_var("MINIO_ROOT_USER", "minioadmin") - .with_env_var("MINIO_ROOT_PASSWORD", "minioadmin") - .with_cmd([ - "-c".to_string(), - format!("mkdir -p /data/{bucket} && exec minio server /data"), - ]); - - let container = match image.start() { - Ok(container) => container, - Err(err) => { - eprintln!( - "skipping sandbox_holds_over_object_store_iceberg (no Docker daemon?): {err}" - ); - return; - } - }; - let port = container - .get_host_port_ipv4(9000.tcp()) - .expect("mapped MinIO port"); - assert!(wait_for_minio(port), "MinIO did not become healthy"); - - let secret = format!( - "CREATE SECRET s3sec (TYPE s3, KEY_ID 'minioadmin', SECRET 'minioadmin', \ - ENDPOINT '127.0.0.1:{port}', URL_STYLE 'path', USE_SSL false, REGION 'us-east-1')" - ); - let table = format!("s3://{bucket}/tbl"); - let extensions = || { - vec![ - "INSTALL httpfs".to_string(), - "LOAD httpfs".to_string(), - "INSTALL iceberg".to_string(), - "LOAD iceberg".to_string(), - ] - }; - - let seed = duckdb::Connection::open_in_memory().expect("open duckdb"); - for stmt in extensions() { - seed.execute_batch(&stmt).expect("load extension"); - } - seed.execute_batch(&secret).expect("create s3 secret"); - seed.execute_batch(&format!( - "COPY (SELECT * FROM (VALUES (1,'a'),(2,'b'),(3,NULL)) t(id,name)) \ - TO '{table}' (FORMAT iceberg)" - )) - .expect("seed iceberg table on object storage"); - - let mut setup = extensions(); - setup.push(secret); - setup.extend(agent_hardening_statements()); - - let read = run_duckdb_query( - &setup, - &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), - [], - None, - ) - .expect("hardening must not break object-store Iceberg reads"); - assert_eq!(read.columns, vec!["id", "name"]); - assert_eq!( - read.rows, - vec![ - vec![serde_json::json!(1), serde_json::json!("a")], - vec![serde_json::json!(2), serde_json::json!("b")], - vec![serde_json::json!(3), serde_json::Value::Null], - ] - ); - - let local = run_duckdb_query(&setup, "SELECT * FROM read_csv('/etc/hostname')", [], None) - .expect_err("local filesystem reads must stay blocked"); - assert!( - local.to_string().to_lowercase().contains("disabled"), - "unexpected error: {local}" - ); - - let cfg = run_duckdb_query(&setup, "SET memory_limit = '1GB'", [], None) - .expect_err("configuration must stay locked"); - let cfg = cfg.to_string().to_lowercase(); - assert!( - cfg.contains("lock") || cfg.contains("configuration"), - "unexpected error: {cfg}" - ); - } - #[test] fn token_export_command_fetches_token_without_printing_it() { let credentials = CatalogCredentials::new( diff --git a/crates/tower-duckdb/Cargo.toml b/crates/tower-duckdb/Cargo.toml new file mode 100644 index 00000000..7ba0dcea --- /dev/null +++ b/crates/tower-duckdb/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "tower-duckdb" +version = { workspace = true } +authors = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } +license = { workspace = true } +repository = { workspace = true } + +[dependencies] +chrono = { workspace = true } +duckdb = { workspace = true } +serde_json = { workspace = true } +tower-telemetry = { workspace = true } + +[dev-dependencies] +testcontainers = { version = "0.24", features = ["blocking"] } diff --git a/crates/tower-duckdb/src/guard.rs b/crates/tower-duckdb/src/guard.rs new file mode 100644 index 00000000..8e5b2f3a --- /dev/null +++ b/crates/tower-duckdb/src/guard.rs @@ -0,0 +1,135 @@ +//! Static gates on untrusted SQL text, applied before a statement runs. They are +//! defence in depth on top of read-only credentials and the session hardening: +//! they reject write/DDL input and multi-statement input, and cap how many rows +//! an agent query may pull back. + +/// Row cap for agent-issued queries. Rows past this are dropped and the result +/// is flagged truncated, so a model cannot pull an unbounded table into memory +/// or its context. +pub const AGENT_MAX_ROWS: usize = 1_000; + +/// Leading keywords that write data or schema, or repoint the session. A query +/// starting with one of these is refused before it runs. +const WRITE_LEADING_KEYWORDS: &[&str] = &[ + "insert", "update", "delete", "merge", "create", "drop", "alter", "truncate", "replace", + "copy", "attach", "detach", +]; + +/// The leading SQL keyword, lowercased, after skipping leading whitespace and +/// `--` / `/* */` comments. +pub fn first_keyword(sql: &str) -> String { + let mut s = sql.trim_start(); + loop { + if let Some(rest) = s.strip_prefix("--") { + match rest.find('\n') { + Some(nl) => s = rest[nl + 1..].trim_start(), + None => return String::new(), + } + } else if let Some(rest) = s.strip_prefix("/*") { + match rest.find("*/") { + Some(end) => s = rest[end + 2..].trim_start(), + None => return String::new(), + } + } else { + break; + } + } + s.chars() + .take_while(|c| c.is_ascii_alphabetic()) + .collect::() + .to_lowercase() +} + +/// True when `sql` starts with a write/DDL keyword. +pub fn is_write_statement(sql: &str) -> bool { + WRITE_LEADING_KEYWORDS.contains(&first_keyword(sql).as_str()) +} + +/// True when `sql` holds more than one statement. A `;` inside a string literal +/// or comment is data, not a separator, so those spans are skipped. This gate +/// matters because duckdb-rs `prepare` runs every statement but the last as a +/// side effect, so unguarded multi-statement SQL would execute its leading +/// statements even though only the final one is returned. +pub fn is_multi_statement(sql: &str) -> bool { + #[derive(PartialEq)] + enum State { + Normal, + Single, + Double, + Line, + Block, + } + + let mut state = State::Normal; + let mut statements = 0usize; + let mut current_has_content = false; + let mut chars = sql.chars().peekable(); + + while let Some(c) = chars.next() { + match state { + State::Normal => match c { + '\'' => { + state = State::Single; + current_has_content = true; + } + '"' => { + state = State::Double; + current_has_content = true; + } + '-' if chars.peek() == Some(&'-') => { + chars.next(); + state = State::Line; + } + '/' if chars.peek() == Some(&'*') => { + chars.next(); + state = State::Block; + } + ';' => { + if current_has_content { + statements += 1; + if statements > 1 { + return true; + } + } + current_has_content = false; + } + c if c.is_whitespace() => {} + _ => current_has_content = true, + }, + State::Single => { + if c == '\'' { + if chars.peek() == Some(&'\'') { + chars.next(); + } else { + state = State::Normal; + } + } + } + State::Double => { + if c == '"' { + if chars.peek() == Some(&'"') { + chars.next(); + } else { + state = State::Normal; + } + } + } + State::Line => { + if c == '\n' { + state = State::Normal; + } + } + State::Block => { + if c == '*' && chars.peek() == Some(&'/') { + chars.next(); + state = State::Normal; + } + } + } + } + + if current_has_content { + statements += 1; + } + statements > 1 +} diff --git a/crates/tower-duckdb/src/lib.rs b/crates/tower-duckdb/src/lib.rs new file mode 100644 index 00000000..5dffaa0f --- /dev/null +++ b/crates/tower-duckdb/src/lib.rs @@ -0,0 +1,700 @@ +//! Tower's usage of DuckDB in one place: opening a session, running trusted +//! setup, locking the session down for untrusted SQL, and executing a query +//! into JSON rows with an optional row cap. +//! +//! The point of the crate is the hardened query path. Agent-issued SQL runs on +//! a customer's machine with their catalog credentials, so before it runs we +//! reject write/DDL and multi-statement input (the [`guard`] module), lock the +//! session down so a query cannot read the local filesystem, load community +//! extensions, or unwind the settings ([`Session::harden`]), and cap the rows a +//! result can carry back ([`Session::query`]). The adversarial tests exercise +//! each of these invariants directly. + +use std::time::Instant; + +use tower_telemetry::debug; + +pub use duckdb::{params, Error, Params}; + +pub mod guard; + +/// A tabular query result: column names, rows as positional JSON values, and a +/// flag set when rows were dropped to honour a caller-supplied cap. +#[derive(Debug, Clone)] +pub struct QueryResult { + pub columns: Vec, + pub rows: Vec>, + /// Rows were dropped to honour the `max_rows` passed to [`Session::query`]. + pub truncated: bool, +} + +/// The statements that lock a session down before untrusted SQL runs: no +/// local-filesystem access (so `read_csv('/etc/passwd')` and friends fail), no +/// community extensions, and a configuration lock so the query cannot unwind any +/// of it. These run after setup, because attaching a catalog is what installs +/// extensions and reaches the network. Only `LocalFileSystem` is disabled, so +/// httpfs and the object-store reads an attached Iceberg catalog depends on keep +/// working; this narrows the query surface without breaking those reads. +pub fn hardening_statements() -> Vec { + vec![ + "SET disabled_filesystems = 'LocalFileSystem'".to_string(), + "SET allow_community_extensions = false".to_string(), + "SET lock_configuration = true".to_string(), + ] +} + +/// An in-memory DuckDB connection Tower runs queries through. +/// +/// The lifecycle is setup, then optionally harden, then query: [`run_setup`] +/// installs extensions and attaches catalogs (the access [`harden`] removes), +/// [`harden`] locks the session down, and [`query`] runs a single statement. +/// +/// [`run_setup`]: Session::run_setup +/// [`harden`]: Session::harden +/// [`query`]: Session::query +pub struct Session { + conn: duckdb::Connection, +} + +impl Session { + /// Open a fresh in-memory session. + pub fn open() -> Result { + Ok(Self { + conn: duckdb::Connection::open_in_memory()?, + }) + } + + /// Run trusted setup statements one at a time. These may embed secrets + /// (vended tokens, credentials), so their text is timed but never logged. + pub fn run_setup(&self, statements: &[String]) -> Result<(), Error> { + let start = Instant::now(); + for statement in statements { + self.conn.execute_batch(statement)?; + } + debug!( + "tower-duckdb: setup ({} statements) took {:?}", + statements.len(), + start.elapsed() + ); + Ok(()) + } + + /// Lock the session down for untrusted SQL. Apply after [`run_setup`], since + /// attaching a catalog needs the access this removes. + /// + /// [`run_setup`]: Session::run_setup + pub fn harden(&self) -> Result<(), Error> { + for statement in hardening_statements() { + self.conn.execute_batch(&statement)?; + } + Ok(()) + } + + /// Execute a single query as a prepared statement with `params` bound. Values + /// that fit a bind position should go through `params` rather than the query + /// text. When `max_rows` is set, rows past it are dropped and the result is + /// flagged truncated, so an untrusted caller cannot pull an unbounded table + /// into memory or a model's context. + pub fn query( + &self, + sql: &str, + params: P, + max_rows: Option, + ) -> Result { + let query_start = Instant::now(); + let mut stmt = self.conn.prepare(sql)?; + let mut columns: Vec = Vec::new(); + let mut rows = Vec::new(); + let mut truncated = false; + + { + let mut result_rows = stmt.query(params)?; + while let Some(row) = result_rows.next()? { + if columns.is_empty() { + columns = row.as_ref().column_names(); + } + if max_rows.is_some_and(|max| rows.len() >= max) { + truncated = true; + break; + } + let mut record = Vec::with_capacity(columns.len()); + for idx in 0..columns.len() { + let value: duckdb::types::Value = row.get(idx)?; + record.push(value_to_json(value)); + } + rows.push(record); + } + } + + // A query with no result rows never populates columns above. + if columns.is_empty() { + columns = stmt.column_names(); + } + + debug!( + "tower-duckdb: query took {:?} ({} rows): {}", + query_start.elapsed(), + rows.len(), + sql + ); + + Ok(QueryResult { + columns, + rows, + truncated, + }) + } +} + +/// Open a session, run `setup`, and execute `query`. Convenience for one-shot +/// callers that do not need to hold the session. Callers running untrusted SQL +/// should build a [`Session`] and call [`Session::harden`] between setup and +/// query instead. +pub fn run_query( + setup: &[String], + query: &str, + params: P, + max_rows: Option, +) -> Result { + let session = Session::open()?; + session.run_setup(setup)?; + session.query(query, params, max_rows) +} + +/// Converts a DuckDB value into a `serde_json::Value`. Integers that overflow an +/// f64 (HugeInt, Decimal) and temporal types are rendered as strings so no +/// precision is lost through JSON's number type. +pub fn value_to_json(value: duckdb::types::Value) -> serde_json::Value { + use duckdb::types::{TimeUnit, Value}; + use serde_json::json; + + match value { + Value::Null => serde_json::Value::Null, + Value::Boolean(v) => json!(v), + Value::TinyInt(v) => json!(v), + Value::SmallInt(v) => json!(v), + Value::Int(v) => json!(v), + Value::BigInt(v) => json!(v), + Value::HugeInt(v) => json!(v.to_string()), + Value::UTinyInt(v) => json!(v), + Value::USmallInt(v) => json!(v), + Value::UInt(v) => json!(v), + Value::UBigInt(v) => json!(v), + Value::Float(v) => json!(v), + Value::Double(v) => json!(v), + Value::Decimal(v) => json!(v.to_string()), + Value::Text(v) => json!(v), + Value::Timestamp(unit, v) => { + let micros = match unit { + TimeUnit::Second => v.checked_mul(1_000_000), + TimeUnit::Millisecond => v.checked_mul(1_000), + TimeUnit::Microsecond => Some(v), + TimeUnit::Nanosecond => Some(v / 1_000), + }; + match micros.and_then(chrono::DateTime::from_timestamp_micros) { + Some(ts) => json!(ts.naive_utc().to_string()), + None => json!(format!("{:?}", Value::Timestamp(unit, v))), + } + } + Value::Date32(days) => { + let date = chrono::DateTime::from_timestamp(i64::from(days) * 86_400, 0); + match date { + Some(d) => json!(d.date_naive().to_string()), + None => json!(format!("{:?}", Value::Date32(days))), + } + } + Value::Time64(unit, v) => { + let micros = match unit { + TimeUnit::Second => v.checked_mul(1_000_000), + TimeUnit::Millisecond => v.checked_mul(1_000), + TimeUnit::Microsecond => Some(v), + TimeUnit::Nanosecond => Some(v / 1_000), + }; + let time = micros.and_then(|m| { + chrono::NaiveTime::from_num_seconds_from_midnight_opt( + (m / 1_000_000) as u32, + ((m % 1_000_000) * 1_000) as u32, + ) + }); + match time { + Some(t) => json!(t.to_string()), + None => json!(format!("{:?}", Value::Time64(unit, v))), + } + } + Value::Enum(v) => json!(v), + Value::List(items) | Value::Array(items) => { + serde_json::Value::Array(items.into_iter().map(value_to_json).collect()) + } + Value::Struct(fields) => serde_json::Value::Object( + fields + .iter() + .map(|(name, value)| (name.clone(), value_to_json(value.clone()))) + .collect(), + ), + Value::Map(entries) => serde_json::Value::Object( + entries + .iter() + .map(|(key, value)| { + ( + stringify_key(&value_to_json(key.clone())), + value_to_json(value.clone()), + ) + }) + .collect(), + ), + Value::Union(inner) => value_to_json(*inner), + other => json!(format!("{:?}", other)), + } +} + +/// Renders a JSON value as a plain string for use as a map key: strings pass +/// through unquoted, everything else uses its JSON form. +fn stringify_key(value: &serde_json::Value) -> String { + match value { + serde_json::Value::Null => String::new(), + serde_json::Value::String(s) => s.clone(), + other => other.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::guard::{first_keyword, is_multi_statement, is_write_statement}; + use super::{hardening_statements, run_query, value_to_json, QueryResult}; + + #[test] + fn duckdb_values_convert_to_json() { + use duckdb::types::{TimeUnit, Value}; + + assert_eq!(value_to_json(Value::Null), serde_json::Value::Null); + assert_eq!(value_to_json(Value::Boolean(true)), serde_json::json!(true)); + assert_eq!(value_to_json(Value::Int(7)), serde_json::json!(7)); + assert_eq!( + value_to_json(Value::HugeInt(170141183460469231731687303715884105727)), + serde_json::json!("170141183460469231731687303715884105727") + ); + assert_eq!( + value_to_json(Value::Text("hello".to_string())), + serde_json::json!("hello") + ); + assert_eq!( + value_to_json(Value::Timestamp(TimeUnit::Microsecond, 0)), + serde_json::json!("1970-01-01 00:00:00") + ); + assert_eq!( + value_to_json(Value::Date32(0)), + serde_json::json!("1970-01-01") + ); + } + + #[test] + fn nested_values_convert_to_json_structures() { + // Round-trip through DuckDB so the list/struct/map values are built the + // way the engine builds them, rather than by hand. + let result = run_query( + &[], + "SELECT [1, 2] AS l, {'a': 1, 'b': 'x'} AS s, MAP {'k': 2} AS m", + [], + None, + ) + .expect("query should succeed"); + + assert_eq!(result.columns, vec!["l", "s", "m"]); + assert_eq!( + result.rows[0], + vec![ + serde_json::json!([1, 2]), + serde_json::json!({"a": 1, "b": "x"}), + serde_json::json!({"k": 2}), + ] + ); + } + + #[test] + fn run_query_returns_columns_and_rows() { + let setup = vec![ + "CREATE TABLE t (id INTEGER, name TEXT)".to_string(), + "INSERT INTO t VALUES (1, 'a'), (2, 'b')".to_string(), + ]; + let result = run_query(&setup, "SELECT id, name FROM t ORDER BY id", [], None) + .expect("query should succeed"); + + assert_eq!(result.columns, vec!["id", "name"]); + assert_eq!( + result.rows, + vec![ + vec![serde_json::json!(1), serde_json::json!("a")], + vec![serde_json::json!(2), serde_json::json!("b")], + ] + ); + assert!(!result.truncated); + } + + #[test] + fn run_query_reports_columns_for_empty_results() { + let result = + run_query(&[], "SELECT 1 AS x WHERE 1 = 0", [], None).expect("query should succeed"); + + assert_eq!(result.columns, vec!["x"]); + assert!(result.rows.is_empty()); + } + + #[test] + fn run_query_caps_rows_and_flags_truncation() { + let capped = run_query(&[], "SELECT * FROM range(5) AS t(i)", [], Some(3)) + .expect("query should succeed"); + assert_eq!(capped.rows.len(), 3); + assert!(capped.truncated); + + let exact = run_query(&[], "SELECT * FROM range(3) AS t(i)", [], Some(3)) + .expect("query should succeed"); + assert_eq!(exact.rows.len(), 3); + assert!(!exact.truncated); + } + + // --- Guard gates ----------------------------------------------------- + + #[test] + fn multi_statement_ignores_separators_in_strings_and_comments() { + assert!(!is_multi_statement("SELECT 1")); + assert!(!is_multi_statement("SELECT 1;")); + assert!(!is_multi_statement(" SELECT 1 ; ")); + assert!(!is_multi_statement("SELECT 'a;b'")); + assert!(!is_multi_statement("SELECT 1 -- ; not a statement")); + assert!(!is_multi_statement("SELECT 1; -- trailing comment")); + assert!(!is_multi_statement("SELECT /* ; */ 1")); + + assert!(is_multi_statement("SELECT 1; SELECT 2")); + assert!(is_multi_statement("SELECT 1; DROP TABLE t")); + assert!(is_multi_statement("SELECT 'a;b'; SELECT 2")); + } + + #[test] + fn write_statements_are_detected_through_case_and_comments() { + assert_eq!(first_keyword(" SELECT 1"), "select"); + assert_eq!(first_keyword("/* c */ INSERT INTO t VALUES (1)"), "insert"); + assert_eq!(first_keyword("-- lead\nDELETE FROM t"), "delete"); + + assert!(!is_write_statement("SELECT * FROM t")); + assert!(!is_write_statement("WITH x AS (SELECT 1) SELECT * FROM x")); + assert!(is_write_statement("insert into t values (1)")); + assert!(is_write_statement(" DROP TABLE t")); + assert!(is_write_statement("COPY t TO 'out.csv'")); + assert!(is_write_statement("ATTACH 'x' AS y")); + } + + #[test] + fn harden_allows_selects_but_blocks_local_files_and_config_changes() { + let setup = hardening_statements(); + + let ok = run_query(&setup, "SELECT 1 AS x", [], None) + .expect("a plain select should still run under the hardened session"); + assert_eq!(ok.columns, vec!["x"]); + assert_eq!(ok.rows, vec![vec![serde_json::json!(1)]]); + + let fs_err = run_query(&setup, "SELECT * FROM read_csv('Cargo.toml')", [], None) + .expect_err("local filesystem access should be blocked"); + assert!( + fs_err.to_string().to_lowercase().contains("disabled"), + "unexpected error: {fs_err}" + ); + + let cfg_err = run_query(&setup, "SET memory_limit = '1GB'", [], None) + .expect_err("configuration should be locked"); + let cfg_msg = cfg_err.to_string().to_lowercase(); + assert!( + cfg_msg.contains("lock") || cfg_msg.contains("configuration"), + "unexpected error: {cfg_err}" + ); + } + + /// Regression check against real Iceberg data. The other query tests use + /// plain in-memory tables; this one writes a small Iceberg table with + /// DuckDB's `COPY … (FORMAT iceberg)` (real metadata + manifests + parquet) + /// and reads it back, so the iceberg reader and our column/row extraction are + /// exercised end to end, NULLs and the row cap included. Needs the `iceberg` + /// extension, fetched on first run and then cached. + #[test] + fn iceberg_scan_reads_written_table_through_run_query() { + let dir = std::env::temp_dir().join(format!("tower_iceberg_test_{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).expect("create temp dir"); + let table = dir.join("events").to_string_lossy().replace('\'', "''"); + + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + conn.execute_batch("INSTALL iceberg").expect("install iceberg"); + conn.execute_batch("LOAD iceberg").expect("load iceberg"); + conn.execute_batch(&format!( + "COPY (SELECT * FROM (VALUES (1, 'a'), (2, 'b'), (3, NULL)) AS t(id, name)) \ + TO '{table}' (FORMAT iceberg)" + )) + .expect("write iceberg table"); + + let setup = vec!["INSTALL iceberg".to_string(), "LOAD iceberg".to_string()]; + + let result = run_query( + &setup, + &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), + [], + None, + ) + .expect("iceberg_scan should read the table"); + assert_eq!(result.columns, vec!["id", "name"]); + assert_eq!( + result.rows, + vec![ + vec![serde_json::json!(1), serde_json::json!("a")], + vec![serde_json::json!(2), serde_json::json!("b")], + vec![serde_json::json!(3), serde_json::Value::Null], + ] + ); + assert!(!result.truncated); + + let capped = run_query( + &setup, + &format!("SELECT id FROM iceberg_scan('{table}') ORDER BY id"), + [], + Some(2), + ) + .expect("iceberg_scan should read the table"); + assert_eq!(capped.rows.len(), 2); + assert!(capped.truncated); + + let _ = std::fs::remove_dir_all(&dir); + } + + // --- Adversarial regression suite ------------------------------------ + // + // Each test encodes an attack an agent-issued query might attempt and + // asserts the sandbox refuses it. These are the security invariants: if a + // future change weakens the gates or the hardening, one of these fails. + + /// A statement run under the session hardening. Extensions are not loaded, so + /// the attacks below must be blocked by the session lockdown alone. + fn run_hardened(sql: &str) -> Result { + run_query(&hardening_statements(), sql, [], None) + } + + #[test] + fn sandbox_gate_rejects_data_tampering() { + for sql in [ + "DROP TABLE runs", + "DELETE FROM runs", + "UPDATE runs SET id = 0", + "INSERT INTO runs VALUES (1)", + "CREATE TABLE evil AS SELECT 1", + "ALTER TABLE runs ADD COLUMN x INTEGER", + "TRUNCATE runs", + "MERGE INTO runs USING x ON true WHEN MATCHED THEN DELETE", + "COPY runs TO '/tmp/exfil.csv'", + "ATTACH '/tmp/evil.db' AS e", + "DETACH runs", + " drop TABLE runs", + "/* sneaky */ DELETE FROM runs", + ] { + assert!(is_write_statement(sql), "not rejected as write/DDL: {sql}"); + } + for sql in [ + "SELECT * FROM runs", + "WITH t AS (SELECT 1) SELECT * FROM t", + "SELECT count(*) FROM runs", + ] { + assert!(!is_write_statement(sql), "legit read wrongly flagged: {sql}"); + } + } + + #[test] + fn sandbox_gate_rejects_statement_smuggling() { + for sql in [ + "SELECT 1; DROP TABLE runs", + "SELECT 1; DELETE FROM runs", + "SELECT 'a;b'; DROP TABLE runs", + "SELECT 1;\n-- c\nUPDATE runs SET id = 0", + ] { + assert!(is_multi_statement(sql), "smuggled statement not caught: {sql}"); + } + assert!(!is_multi_statement("SELECT * FROM runs -- a trailing ; comment")); + } + + #[test] + fn sandbox_blocks_host_filesystem_reads() { + for sql in [ + "SELECT * FROM read_csv('/etc/passwd')", + "SELECT * FROM read_text('/etc/hostname')", + "SELECT * FROM read_json('/etc/passwd')", + "SELECT * FROM read_parquet('/tmp/x.parquet')", + "SELECT * FROM read_csv('/etc/*')", + "SELECT * FROM '/etc/passwd'", + ] { + assert!(run_hardened(sql).is_err(), "host file read NOT blocked: {sql}"); + } + let err = run_hardened("SELECT * FROM read_csv('/etc/passwd')").unwrap_err(); + assert!(err.to_string().to_lowercase().contains("disabled"), "{err}"); + } + + #[test] + fn sandbox_blocks_host_filesystem_writes() { + for sql in [ + "COPY (SELECT 1) TO '/tmp/pwned.csv'", + "COPY (SELECT 1) TO '/tmp/pwned.parquet' (FORMAT parquet)", + ] { + assert!(run_hardened(sql).is_err(), "host file write NOT blocked: {sql}"); + } + } + + #[test] + fn sandbox_blocks_configuration_escape() { + for sql in [ + "SET disabled_filesystems = ''", + "RESET disabled_filesystems", + "SET enable_external_access = true", + "SET allow_community_extensions = true", + "SET lock_configuration = false", + "PRAGMA disabled_filesystems=''", + ] { + assert!(run_hardened(sql).is_err(), "configuration escape NOT blocked: {sql}"); + } + } + + #[test] + fn sandbox_blocks_arbitrary_extension_loading() { + for sql in [ + "LOAD '/tmp/evil.duckdb_extension'", + "INSTALL some_untrusted_extension_xyz", + ] { + assert!(run_hardened(sql).is_err(), "extension load NOT blocked: {sql}"); + } + } + + #[test] + fn sandbox_blocks_network_ssrf_via_table_functions() { + for sql in [ + "SELECT * FROM read_csv('http://169.254.169.254/latest/meta-data/')", + "SELECT * FROM read_parquet('https://attacker.example/x.parquet')", + "SELECT * FROM read_csv('http://localhost:8080/internal')", + ] { + assert!(run_hardened(sql).is_err(), "SSRF NOT blocked: {sql}"); + } + } + + /// Polls MinIO's health endpoint over a raw socket until it answers 200. + fn wait_for_minio(port: u16) -> bool { + use std::io::{Read, Write}; + for _ in 0..60 { + if let Ok(mut stream) = std::net::TcpStream::connect(("127.0.0.1", port)) { + let _ = stream.set_read_timeout(Some(std::time::Duration::from_secs(2))); + let request = + "GET /minio/health/live HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"; + let mut response = String::new(); + if stream.write_all(request.as_bytes()).is_ok() + && stream.read_to_string(&mut response).is_ok() + && response.contains(" 200 ") + { + return true; + } + } + std::thread::sleep(std::time::Duration::from_millis(500)); + } + false + } + + /// The coverage the local-only tests can't give: it proves the hardening does + /// NOT break a real object-store Iceberg read (the reader uses S3FileSystem, + /// which the hardening leaves enabled) while local-filesystem access and + /// configuration changes stay blocked in the same session. Starts a MinIO + /// container via testcontainers and self-skips when no Docker daemon is + /// available, so a plain `cargo test` still passes without Docker. + #[test] + fn sandbox_holds_over_object_store_iceberg() { + use testcontainers::core::{IntoContainerPort, WaitFor}; + use testcontainers::runners::SyncRunner; + use testcontainers::{GenericImage, ImageExt}; + + let bucket = "warehouse"; + let image = GenericImage::new("minio/minio", "latest") + .with_wait_for(WaitFor::seconds(1)) + .with_exposed_port(9000.tcp()) + .with_entrypoint("sh") + .with_env_var("MINIO_ROOT_USER", "minioadmin") + .with_env_var("MINIO_ROOT_PASSWORD", "minioadmin") + .with_cmd([ + "-c".to_string(), + format!("mkdir -p /data/{bucket} && exec minio server /data"), + ]); + + let container = match image.start() { + Ok(container) => container, + Err(err) => { + eprintln!( + "skipping sandbox_holds_over_object_store_iceberg (no Docker daemon?): {err}" + ); + return; + } + }; + let port = container + .get_host_port_ipv4(9000.tcp()) + .expect("mapped MinIO port"); + assert!(wait_for_minio(port), "MinIO did not become healthy"); + + let secret = format!( + "CREATE SECRET s3sec (TYPE s3, KEY_ID 'minioadmin', SECRET 'minioadmin', \ + ENDPOINT '127.0.0.1:{port}', URL_STYLE 'path', USE_SSL false, REGION 'us-east-1')" + ); + let table = format!("s3://{bucket}/tbl"); + let extensions = || { + vec![ + "INSTALL httpfs".to_string(), + "LOAD httpfs".to_string(), + "INSTALL iceberg".to_string(), + "LOAD iceberg".to_string(), + ] + }; + + let seed = duckdb::Connection::open_in_memory().expect("open duckdb"); + for stmt in extensions() { + seed.execute_batch(&stmt).expect("load extension"); + } + seed.execute_batch(&secret).expect("create s3 secret"); + seed.execute_batch(&format!( + "COPY (SELECT * FROM (VALUES (1,'a'),(2,'b'),(3,NULL)) t(id,name)) \ + TO '{table}' (FORMAT iceberg)" + )) + .expect("seed iceberg table on object storage"); + + let mut setup = extensions(); + setup.push(secret); + setup.extend(hardening_statements()); + + let read = run_query( + &setup, + &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), + [], + None, + ) + .expect("hardening must not break object-store Iceberg reads"); + assert_eq!(read.columns, vec!["id", "name"]); + assert_eq!( + read.rows, + vec![ + vec![serde_json::json!(1), serde_json::json!("a")], + vec![serde_json::json!(2), serde_json::json!("b")], + vec![serde_json::json!(3), serde_json::Value::Null], + ] + ); + + let local = run_query(&setup, "SELECT * FROM read_csv('/etc/hostname')", [], None) + .expect_err("local filesystem reads must stay blocked"); + assert!( + local.to_string().to_lowercase().contains("disabled"), + "unexpected error: {local}" + ); + + let cfg = run_query(&setup, "SET memory_limit = '1GB'", [], None) + .expect_err("configuration must stay locked"); + let cfg = cfg.to_string().to_lowercase(); + assert!( + cfg.contains("lock") || cfg.contains("configuration"), + "unexpected error: {cfg}" + ); + } +} From 7d7d2af380d2d7ebfeb2c54d1bd7adb29fbd36c6 Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Fri, 24 Jul 2026 13:41:02 +0100 Subject: [PATCH 3/4] fix(tower-duckdb): gate read-only queries with DuckDB's parser, not a keyword denylist MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- crates/tower-cmd/src/catalogs.rs | 30 +++-- crates/tower-duckdb/src/guard.rs | 192 ++++++++++++------------------- crates/tower-duckdb/src/lib.rs | 89 +++++++++----- 3 files changed, 154 insertions(+), 157 deletions(-) diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index 0ab260b8..135de016 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -465,16 +465,28 @@ pub async fn do_query(out: &output::Out, config: Config, args: &ArgMatches) { } let write = cmd::get_bool_flag(args, "write"); - // Read mode runs untrusted SQL, so gate it before it reaches DuckDB: a - // smuggled second statement would otherwise execute as a side effect of - // `prepare`, and a write in read mode should fail with a clear message rather - // than a raw engine error. Write mode is the trusted power-user path. + // Read mode runs untrusted SQL, so gate it through DuckDB's parser before it + // reaches the query path: it must be exactly one SELECT. A smuggled second + // statement would otherwise execute as a side effect of `prepare`, and a + // write should fail with a clear message rather than a raw engine error. + // Write mode is the trusted power-user path and skips the gate. if !write { - if guard::is_multi_statement(&sql) { - out.die("Only a single SQL statement can be run at a time. Remove the extra statement(s)."); - } - if guard::is_write_statement(&sql) { - out.die("This command runs read-only queries. The statement looks like it writes or changes data; re-run with --write to modify the catalog."); + let sql_to_check = sql.clone(); + let verdict = + tokio::task::spawn_blocking(move || guard::classify_read_only(&sql_to_check)).await; + match verdict { + Ok(Ok(guard::ReadOnlyCheck::Allowed)) => {} + Ok(Ok(guard::ReadOnlyCheck::Empty)) => { + out.die("No SQL statement provided. Pass one with --sql or pipe it via stdin.") + } + Ok(Ok(guard::ReadOnlyCheck::Multiple)) => { + out.die("Only a single SQL statement can be run at a time. Remove the extra statement(s).") + } + Ok(Ok(guard::ReadOnlyCheck::NotReadOnly)) => out.die( + "This command runs read-only queries. Only a single SELECT statement is allowed; re-run with --write to modify the catalog.", + ), + Ok(Err(err)) => out.die(&format!("Could not validate the query: {err}")), + Err(err) => out.die(&format!("Could not validate the query: {err}")), } } let query_result = execute_catalog_query(out, &config, name, &env, sql, write).await; diff --git a/crates/tower-duckdb/src/guard.rs b/crates/tower-duckdb/src/guard.rs index 8e5b2f3a..0769c11a 100644 --- a/crates/tower-duckdb/src/guard.rs +++ b/crates/tower-duckdb/src/guard.rs @@ -1,135 +1,89 @@ -//! Static gates on untrusted SQL text, applied before a statement runs. They are +//! Read-only gate on untrusted SQL, applied before a statement runs. It is //! defence in depth on top of read-only credentials and the session hardening: -//! they reject write/DDL input and multi-statement input, and cap how many rows -//! an agent query may pull back. +//! it rejects anything that is not a single `SELECT`, and caps how many rows an +//! agent query may pull back. +//! +//! The check runs the SQL through DuckDB's own parser via `json_serialize_sql`, +//! which parses (but does not execute) a statement and serializes only `SELECT` +//! statements, erroring on everything else. Using the engine's parser rather +//! than scanning keywords is what makes this safe: a keyword denylist misses +//! comment tricks (a `--` comment ends at `\r` as well as `\n` in DuckDB, so a +//! `-- x\rDROP …` payload looks empty to a naive scanner but parses as a DROP) +//! and statements that start with an allowed keyword but still mutate (a `WITH …` +//! CTE, for one). The parser sees them the way the executor will. /// Row cap for agent-issued queries. Rows past this are dropped and the result /// is flagged truncated, so a model cannot pull an unbounded table into memory /// or its context. pub const AGENT_MAX_ROWS: usize = 1_000; -/// Leading keywords that write data or schema, or repoint the session. A query -/// starting with one of these is refused before it runs. -const WRITE_LEADING_KEYWORDS: &[&str] = &[ - "insert", "update", "delete", "merge", "create", "drop", "alter", "truncate", "replace", - "copy", "attach", "detach", -]; - -/// The leading SQL keyword, lowercased, after skipping leading whitespace and -/// `--` / `/* */` comments. -pub fn first_keyword(sql: &str) -> String { - let mut s = sql.trim_start(); - loop { - if let Some(rest) = s.strip_prefix("--") { - match rest.find('\n') { - Some(nl) => s = rest[nl + 1..].trim_start(), - None => return String::new(), - } - } else if let Some(rest) = s.strip_prefix("/*") { - match rest.find("*/") { - Some(end) => s = rest[end + 2..].trim_start(), - None => return String::new(), - } - } else { - break; - } - } - s.chars() - .take_while(|c| c.is_ascii_alphabetic()) - .collect::() - .to_lowercase() +/// The verdict for a piece of untrusted SQL. +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub enum ReadOnlyCheck { + /// Exactly one `SELECT` statement. Safe to run under a read-only session. + Allowed, + /// No statement at all (blank or comment-only input). + Empty, + /// More than one statement. Rejected: duckdb-rs `prepare` executes every + /// statement but the last as a side effect, so a second statement must never + /// reach it. + Multiple, + /// Parses as something other than a single `SELECT` (a write, DDL, `PRAGMA`, + /// `SET`, `COPY`, `ATTACH`, …) or does not parse at all. Rejected fail-closed. + NotReadOnly, } -/// True when `sql` starts with a write/DDL keyword. -pub fn is_write_statement(sql: &str) -> bool { - WRITE_LEADING_KEYWORDS.contains(&first_keyword(sql).as_str()) +/// Classify untrusted `sql` for the read-only path using DuckDB's parser. Opens a +/// throwaway in-memory connection with no catalog attached, so the parse cannot +/// touch customer data even in principle. Returns a `duckdb::Error` only if the +/// parser query itself cannot run; a response it cannot make sense of is treated +/// as [`ReadOnlyCheck::NotReadOnly`] (fail closed). +pub fn classify_read_only(sql: &str) -> Result { + let conn = duckdb::Connection::open_in_memory()?; + classify_read_only_on(&conn, sql) } -/// True when `sql` holds more than one statement. A `;` inside a string literal -/// or comment is data, not a separator, so those spans are skipped. This gate -/// matters because duckdb-rs `prepare` runs every statement but the last as a -/// side effect, so unguarded multi-statement SQL would execute its leading -/// statements even though only the final one is returned. -pub fn is_multi_statement(sql: &str) -> bool { - #[derive(PartialEq)] - enum State { - Normal, - Single, - Double, - Line, - Block, - } +/// [`classify_read_only`] against a caller-supplied connection, so a long-lived +/// caller (an MCP server) can reuse one connection instead of opening a fresh one +/// per query. The connection is only used to run `json_serialize_sql`, which does +/// not execute `sql`. +pub fn classify_read_only_on( + conn: &duckdb::Connection, + sql: &str, +) -> Result { + // `json_serialize_sql` parses and serializes SELECT statements to JSON and + // errors on anything else. The SQL is bound as a parameter, never spliced + // into this query. + let serialized: String = + conn.query_row("SELECT json_serialize_sql(CAST(? AS VARCHAR))", [sql], |row| { + row.get(0) + })?; - let mut state = State::Normal; - let mut statements = 0usize; - let mut current_has_content = false; - let mut chars = sql.chars().peekable(); + let parsed: serde_json::Value = match serde_json::from_str(&serialized) { + Ok(value) => value, + // DuckDB always returns valid JSON here; an unparseable response means + // something we don't understand, so refuse it rather than guess. + Err(_) => return Ok(ReadOnlyCheck::NotReadOnly), + }; - while let Some(c) = chars.next() { - match state { - State::Normal => match c { - '\'' => { - state = State::Single; - current_has_content = true; - } - '"' => { - state = State::Double; - current_has_content = true; - } - '-' if chars.peek() == Some(&'-') => { - chars.next(); - state = State::Line; - } - '/' if chars.peek() == Some(&'*') => { - chars.next(); - state = State::Block; - } - ';' => { - if current_has_content { - statements += 1; - if statements > 1 { - return true; - } - } - current_has_content = false; - } - c if c.is_whitespace() => {} - _ => current_has_content = true, - }, - State::Single => { - if c == '\'' { - if chars.peek() == Some(&'\'') { - chars.next(); - } else { - state = State::Normal; - } - } - } - State::Double => { - if c == '"' { - if chars.peek() == Some(&'"') { - chars.next(); - } else { - state = State::Normal; - } - } - } - State::Line => { - if c == '\n' { - state = State::Normal; - } - } - State::Block => { - if c == '*' && chars.peek() == Some(&'/') { - chars.next(); - state = State::Normal; - } - } - } + // `error: true` covers both non-SELECT statements ("Only SELECT statements + // can be serialized to json!") and malformed SQL. + if parsed + .get("error") + .and_then(serde_json::Value::as_bool) + .unwrap_or(true) + { + return Ok(ReadOnlyCheck::NotReadOnly); } - if current_has_content { - statements += 1; - } - statements > 1 + let statements = parsed + .get("statements") + .and_then(serde_json::Value::as_array) + .map_or(0, Vec::len); + + Ok(match statements { + 0 => ReadOnlyCheck::Empty, + 1 => ReadOnlyCheck::Allowed, + _ => ReadOnlyCheck::Multiple, + }) } diff --git a/crates/tower-duckdb/src/lib.rs b/crates/tower-duckdb/src/lib.rs index 5dffaa0f..f6a0bf6c 100644 --- a/crates/tower-duckdb/src/lib.rs +++ b/crates/tower-duckdb/src/lib.rs @@ -259,7 +259,7 @@ fn stringify_key(value: &serde_json::Value) -> String { #[cfg(test)] mod tests { - use super::guard::{first_keyword, is_multi_statement, is_write_statement}; + use super::guard::{classify_read_only, ReadOnlyCheck}; use super::{hardening_statements, run_query, value_to_json, QueryResult}; #[test] @@ -352,35 +352,60 @@ mod tests { assert!(!exact.truncated); } - // --- Guard gates ----------------------------------------------------- + // --- Read-only gate -------------------------------------------------- + + fn check(sql: &str) -> ReadOnlyCheck { + classify_read_only(sql).expect("parser should run") + } + + #[test] + fn read_only_gate_allows_only_single_selects() { + for sql in [ + "SELECT 1", + "SELECT 1;", + " SELECT 1 ; ", + "SELECT 'a;b'", + "SELECT 1 -- ; not a statement", + "SELECT /* ; */ 1", + "WITH x AS (SELECT 1) SELECT * FROM x", + "select COUNT(*) from runs", + ] { + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "should allow: {sql}"); + } + } #[test] - fn multi_statement_ignores_separators_in_strings_and_comments() { - assert!(!is_multi_statement("SELECT 1")); - assert!(!is_multi_statement("SELECT 1;")); - assert!(!is_multi_statement(" SELECT 1 ; ")); - assert!(!is_multi_statement("SELECT 'a;b'")); - assert!(!is_multi_statement("SELECT 1 -- ; not a statement")); - assert!(!is_multi_statement("SELECT 1; -- trailing comment")); - assert!(!is_multi_statement("SELECT /* ; */ 1")); - - assert!(is_multi_statement("SELECT 1; SELECT 2")); - assert!(is_multi_statement("SELECT 1; DROP TABLE t")); - assert!(is_multi_statement("SELECT 'a;b'; SELECT 2")); + fn read_only_gate_rejects_writes_and_ddl_whatever_the_leading_keyword() { + for sql in [ + "insert into t values (1)", + " DROP TABLE t", + "COPY t TO 'out.csv'", + "ATTACH 'x' AS y", + "SET memory_limit = '1GB'", + "PRAGMA version", + // Starts with an allowed keyword but still mutates: the parser sees + // the DELETE a first-keyword denylist would miss. + "WITH x AS (DELETE FROM t RETURNING *) SELECT * FROM x", + ] { + assert_eq!(check(sql), ReadOnlyCheck::NotReadOnly, "should reject: {sql}"); + } + } + + #[test] + fn read_only_gate_rejects_comment_terminator_smuggling() { + // A `--` comment ends at a carriage return in DuckDB, so this parses as a + // DROP even though it opens with what looks like a full-line comment. A + // text scanner that only breaks comments on `\n` would see it as empty. + assert_eq!(check("-- harmless\rDROP TABLE runs"), ReadOnlyCheck::NotReadOnly); + assert_eq!(check("-- harmless\r\nSELECT 1"), ReadOnlyCheck::Allowed); } #[test] - fn write_statements_are_detected_through_case_and_comments() { - assert_eq!(first_keyword(" SELECT 1"), "select"); - assert_eq!(first_keyword("/* c */ INSERT INTO t VALUES (1)"), "insert"); - assert_eq!(first_keyword("-- lead\nDELETE FROM t"), "delete"); - - assert!(!is_write_statement("SELECT * FROM t")); - assert!(!is_write_statement("WITH x AS (SELECT 1) SELECT * FROM x")); - assert!(is_write_statement("insert into t values (1)")); - assert!(is_write_statement(" DROP TABLE t")); - assert!(is_write_statement("COPY t TO 'out.csv'")); - assert!(is_write_statement("ATTACH 'x' AS y")); + fn read_only_gate_classifies_empty_and_multiple() { + assert_eq!(check(""), ReadOnlyCheck::Empty); + assert_eq!(check(" -- just a comment"), ReadOnlyCheck::Empty); + assert_eq!(check("SELECT 1; SELECT 2"), ReadOnlyCheck::Multiple); + assert_eq!(check("SELECT 'a;b'; SELECT 2"), ReadOnlyCheck::Multiple); } #[test] @@ -492,28 +517,34 @@ mod tests { " drop TABLE runs", "/* sneaky */ DELETE FROM runs", ] { - assert!(is_write_statement(sql), "not rejected as write/DDL: {sql}"); + assert_eq!(check(sql), ReadOnlyCheck::NotReadOnly, "not rejected: {sql}"); } for sql in [ "SELECT * FROM runs", "WITH t AS (SELECT 1) SELECT * FROM t", "SELECT count(*) FROM runs", ] { - assert!(!is_write_statement(sql), "legit read wrongly flagged: {sql}"); + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "legit read rejected: {sql}"); } } #[test] fn sandbox_gate_rejects_statement_smuggling() { + // A trailing statement of any kind is refused: an all-SELECT pair counts + // as Multiple, a mixed one trips the SELECT-only parser first. for sql in [ "SELECT 1; DROP TABLE runs", "SELECT 1; DELETE FROM runs", "SELECT 'a;b'; DROP TABLE runs", "SELECT 1;\n-- c\nUPDATE runs SET id = 0", + "SELECT 1; SELECT 2", ] { - assert!(is_multi_statement(sql), "smuggled statement not caught: {sql}"); + assert_ne!(check(sql), ReadOnlyCheck::Allowed, "smuggled statement allowed: {sql}"); } - assert!(!is_multi_statement("SELECT * FROM runs -- a trailing ; comment")); + assert_eq!( + check("SELECT * FROM runs -- a trailing ; comment"), + ReadOnlyCheck::Allowed + ); } #[test] From 085e39e513947dbfb0c21b087f2ca1d6e48758c4 Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Fri, 24 Jul 2026 14:08:25 +0100 Subject: [PATCH 4/4] 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. --- Cargo.toml | 2 +- crates/tower-duckdb/src/lib.rs | 28 ++++++++++++++++++---------- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index b737e493..3b92533e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,7 +28,7 @@ config = { path = "crates/config" } crypto = { path = "crates/crypto" } ctrlc = "3" dirs = "5" -duckdb = { version = "~1.10504.0", features = ["bundled"] } +duckdb = { version = "~1.10504.0", features = ["bundled", "json"] } flate2 = "1" fs2 = "0.4" futures = "0.3" diff --git a/crates/tower-duckdb/src/lib.rs b/crates/tower-duckdb/src/lib.rs index f6a0bf6c..0ba6f00f 100644 --- a/crates/tower-duckdb/src/lib.rs +++ b/crates/tower-duckdb/src/lib.rs @@ -4,11 +4,12 @@ //! //! The point of the crate is the hardened query path. Agent-issued SQL runs on //! a customer's machine with their catalog credentials, so before it runs we -//! reject write/DDL and multi-statement input (the [`guard`] module), lock the -//! session down so a query cannot read the local filesystem, load community -//! extensions, or unwind the settings ([`Session::harden`]), and cap the rows a -//! result can carry back ([`Session::query`]). The adversarial tests exercise -//! each of these invariants directly. +//! reject anything that is not a single read-only SELECT, using DuckDB's own +//! parser (the [`guard`] module), lock the session down so a query cannot read +//! the local filesystem, load community extensions, or unwind the settings +//! ([`Session::harden`]), and cap the rows a result can carry back +//! ([`Session::query`]). The adversarial tests exercise each of these invariants +//! directly. use std::time::Instant; @@ -437,18 +438,24 @@ mod tests { /// plain in-memory tables; this one writes a small Iceberg table with /// DuckDB's `COPY … (FORMAT iceberg)` (real metadata + manifests + parquet) /// and reads it back, so the iceberg reader and our column/row extraction are - /// exercised end to end, NULLs and the row cap included. Needs the `iceberg` - /// extension, fetched on first run and then cached. + /// exercised end to end, NULLs and the row cap included. The `iceberg` + /// extension is not bundled, so it is fetched on first run and cached; the + /// test self-skips where that install cannot happen (some CI has no writable + /// extension directory), the same way the object-store test skips without + /// Docker. #[test] fn iceberg_scan_reads_written_table_through_run_query() { + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + if conn.execute_batch("INSTALL iceberg; LOAD iceberg;").is_err() { + eprintln!("skipping iceberg_scan test (iceberg extension unavailable)"); + return; + } + let dir = std::env::temp_dir().join(format!("tower_iceberg_test_{}", std::process::id())); let _ = std::fs::remove_dir_all(&dir); std::fs::create_dir_all(&dir).expect("create temp dir"); let table = dir.join("events").to_string_lossy().replace('\'', "''"); - let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); - conn.execute_batch("INSTALL iceberg").expect("install iceberg"); - conn.execute_batch("LOAD iceberg").expect("load iceberg"); conn.execute_batch(&format!( "COPY (SELECT * FROM (VALUES (1, 'a'), (2, 'b'), (3, NULL)) AS t(id, name)) \ TO '{table}' (FORMAT iceberg)" @@ -729,3 +736,4 @@ mod tests { ); } } +