From 0821cf7870c901441e865424e04262bb7bc8cba0 Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Fri, 24 Jul 2026 09:23:33 +0100 Subject: [PATCH 1/9] 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/9] 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/9] 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/9] 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 { ); } } + From fc20da93bd71eb80eb5528263e51e211f7c03c46 Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Sun, 26 Jul 2026 14:35:02 +0100 Subject: [PATCH 5/9] test(tower-duckdb): co-locate read-only gate regression tests in guard.rs The gate's unit tests lived in lib.rs, away from the code they cover. This moves them next to `classify_read_only` in guard.rs and broadens them into a proper regression suite, so the parser-based read-only policy is pinned down where a reader of guard.rs will find it. The suite locks in the current DuckDB classification across the shapes that matter: the many forms of a single read (plain SELECT, CTEs, set operations, subqueries, VALUES/TABLE/FROM-first, DESCRIBE/SUMMARIZE/SHOW), that comments and whitespace and semicolons inside literals don't change the verdict, that writes and DDL and config/transaction/meta statements are rejected whatever the leading keyword, that a leading SELECT can't launder a trailing mutation (the `\r` comment-terminator and data-modifying-CTE bypasses), multi-statement and empty/comment-only classification, fail-closed rejection of unparseable input, and the shared-connection reuse path. It also documents one boundary explicitly: a SELECT that reads a local file or URL is Allowed by the gate because the gate classifies statement shape only; the session hardening is what refuses the read. The redundant gate tests are dropped from lib.rs; its adversarial suite keeps the shared `check` helper. --- crates/tower-duckdb/src/guard.rs | 193 +++++++++++++++++++++++++++++++ crates/tower-duckdb/src/lib.rs | 54 +-------- 2 files changed, 195 insertions(+), 52 deletions(-) diff --git a/crates/tower-duckdb/src/guard.rs b/crates/tower-duckdb/src/guard.rs index 0769c11a..92e51b3a 100644 --- a/crates/tower-duckdb/src/guard.rs +++ b/crates/tower-duckdb/src/guard.rs @@ -87,3 +87,196 @@ pub fn classify_read_only_on( _ => ReadOnlyCheck::Multiple, }) } + +#[cfg(test)] +mod tests { + use super::{classify_read_only, classify_read_only_on, ReadOnlyCheck}; + + fn check(sql: &str) -> ReadOnlyCheck { + classify_read_only(sql).expect("the parser query should run") + } + + #[test] + fn allows_a_single_select_in_its_many_shapes() { + // Everything DuckDB parses as one read: plain SELECTs, CTEs, set + // operations, subqueries, and the SELECT-sugar forms an agent might emit. + for sql in [ + "SELECT 1", + "select 1", + " SELECT 1 ", + "SELECT * FROM runs WHERE id > 0", + "SELECT count(*) FROM runs", + "WITH x AS (SELECT 1) SELECT * FROM x", + "WITH RECURSIVE t(n) AS (SELECT 1 UNION SELECT n + 1 FROM t WHERE n < 3) SELECT * FROM t", + "SELECT 1 UNION SELECT 2", + "(SELECT 1)", + "VALUES (1), (2)", + "TABLE runs", + "FROM runs", + "DESCRIBE SELECT 1", + "SUMMARIZE SELECT 1", + "SHOW TABLES", + "SHOW ALL TABLES", + ] { + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "should allow: {sql}"); + } + } + + #[test] + fn leading_comments_and_whitespace_do_not_change_the_verdict() { + for sql in [ + "-- a comment\nSELECT 1", + "/* a comment */ SELECT 1", + "\n\t SELECT 1", + "SELECT 1 -- trailing line comment", + "SELECT 1 /* trailing block comment */", + ] { + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "should allow: {sql}"); + } + } + + #[test] + fn semicolons_inside_literals_and_comments_are_not_separators() { + // A single statement whose text merely contains `;` in a string, a + // dollar-quoted string, or a comment stays a single statement. + for sql in [ + "SELECT 'a;b'", + "SELECT $$a;b$$", + "SELECT 1 -- ; not a separator", + "SELECT /* ; */ 1", + "SELECT 1;", + "SELECT 1 ; ", + "SELECT 1 \t ; \n ", + "SELECT 1;;", + ] { + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "should allow: {sql}"); + } + } + + #[test] + fn rejects_writes_and_ddl() { + for sql in [ + "INSERT INTO runs VALUES (1)", + "insert into runs values (1)", + "UPDATE runs SET id = 0", + "DELETE FROM runs", + "MERGE INTO runs USING x ON true WHEN MATCHED THEN DELETE", + "CREATE TABLE evil AS SELECT 1", + "CREATE OR REPLACE TABLE runs AS SELECT 1", + "DROP TABLE runs", + "ALTER TABLE runs ADD COLUMN x INTEGER", + "TRUNCATE runs", + "COPY runs TO '/tmp/exfil.csv'", + "ATTACH '/tmp/evil.db' AS e", + "DETACH runs", + ] { + assert_eq!(check(sql), ReadOnlyCheck::NotReadOnly, "should reject: {sql}"); + } + } + + #[test] + fn rejects_config_transaction_and_meta_statements() { + // Statements that are neither writes nor plain SELECTs still must not run + // on the read-only path: they can change session state or the plan. + for sql in [ + "SET memory_limit = '1GB'", + "RESET memory_limit", + "PRAGMA database_list", + "PRAGMA disable_verification", + "BEGIN", + "COMMIT", + "ROLLBACK", + "PREPARE p AS SELECT 1", + "EXECUTE p", + "CALL pragma_version()", + "EXPLAIN SELECT 1", + ] { + assert_eq!(check(sql), ReadOnlyCheck::NotReadOnly, "should reject: {sql}"); + } + } + + #[test] + fn a_leading_select_does_not_launder_a_trailing_mutation() { + // The bypass class the gate exists for: a statement that opens with an + // allowed keyword but carries a write. A first-keyword denylist would let + // these through; the parser does not. + for sql in [ + // A `--` comment ends at a carriage return in DuckDB, so this parses + // as a DROP even though it opens with a full-line comment. + "-- harmless\rDROP TABLE runs", + // A data-modifying CTE. + "WITH x AS (DELETE FROM runs RETURNING *) SELECT * FROM x", + ] { + assert_eq!(check(sql), ReadOnlyCheck::NotReadOnly, "should reject: {sql}"); + } + // The same shapes without the mutation stay allowed. + assert_eq!(check("-- harmless\r\nSELECT 1"), ReadOnlyCheck::Allowed); + assert_eq!( + check("WITH x AS (SELECT 1) SELECT * FROM x"), + ReadOnlyCheck::Allowed + ); + } + + #[test] + fn rejects_multiple_statements() { + // Two statements must never reach the query path: duckdb-rs `prepare` + // runs every statement but the last as a side effect. An all-SELECT pair + // is Multiple; a mix trips the SELECT-only parser first. + assert_eq!(check("SELECT 1; SELECT 2"), ReadOnlyCheck::Multiple); + assert_eq!(check("SELECT 'a;b'; SELECT 2"), ReadOnlyCheck::Multiple); + for sql in [ + "SELECT 1; DROP TABLE runs", + "SELECT 1; DELETE FROM runs", + "SELECT 1;\n-- c\nUPDATE runs SET id = 0", + ] { + assert_ne!(check(sql), ReadOnlyCheck::Allowed, "smuggled write allowed: {sql}"); + } + } + + #[test] + fn empty_and_comment_only_input_is_empty_not_allowed() { + for sql in ["", " ", "\n\t ", "-- just a comment", "/* only a block */"] { + assert_eq!(check(sql), ReadOnlyCheck::Empty, "should be empty: {sql:?}"); + } + } + + #[test] + fn unparseable_input_is_rejected_fail_closed() { + for sql in ["not valid sql at all", "SELECT FROM WHERE", "SELECT ((("] { + assert_eq!(check(sql), ReadOnlyCheck::NotReadOnly, "should reject: {sql}"); + } + } + + #[test] + fn table_functions_reaching_the_host_are_a_matter_for_the_hardening_not_the_gate() { + // The gate only classifies statement shape. A SELECT that reads a local + // file or a URL is a valid SELECT and is Allowed here; the session + // hardening (disabled_filesystems, no external access) is what refuses the + // actual read at execution time. This test documents that boundary so a + // future reader does not mistake the gate for the filesystem defense. + for sql in [ + "SELECT * FROM read_csv('/etc/passwd')", + "SELECT * FROM read_parquet('https://example.com/x.parquet')", + ] { + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "gate classifies shape only: {sql}"); + } + } + + #[test] + fn classify_on_a_shared_connection_matches_the_owning_variant() { + // The reuse path an MCP server would take: one connection, many queries. + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + assert_eq!( + classify_read_only_on(&conn, "SELECT 1").unwrap(), + ReadOnlyCheck::Allowed + ); + assert_eq!( + classify_read_only_on(&conn, "DROP TABLE runs").unwrap(), + ReadOnlyCheck::NotReadOnly + ); + assert_eq!( + classify_read_only_on(&conn, "SELECT 1; SELECT 2").unwrap(), + ReadOnlyCheck::Multiple + ); + } +} diff --git a/crates/tower-duckdb/src/lib.rs b/crates/tower-duckdb/src/lib.rs index 0ba6f00f..2b93b9f6 100644 --- a/crates/tower-duckdb/src/lib.rs +++ b/crates/tower-duckdb/src/lib.rs @@ -353,62 +353,12 @@ mod tests { assert!(!exact.truncated); } - // --- Read-only gate -------------------------------------------------- - + // The read-only gate is unit-tested in `guard.rs`. This helper backs the + // adversarial suite below, which exercises the gate alongside the hardening. 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 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 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] fn harden_allows_selects_but_blocks_local_files_and_config_changes() { let setup = hardening_statements(); From 6a107e5674091fb8663c90f8e49eb21d3e49464e Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Mon, 27 Jul 2026 09:00:07 +0100 Subject: [PATCH 6/9] feat(tower-duckdb): close SELECT-shaped holes, expand the lockdown, and bound queries Parsing as a SELECT is a statement shape, not a read-only property, and the gate was treating the two as the same thing. Verified against the DuckDB we ship (v1.5.4): `SELECT nextval('s')` classifies as a plain SELECT and really does advance the sequence, and `query()`/`query_table()` classify as SELECTs while handing a string to the execution pipeline. Those are now refused by name from the parsed tree, walking `function_name` nodes rather than matching text, so comments and quoting cannot hide them. The session lockdown gains the settings it was missing: no implicit extension install or load, no unsigned extensions, secrets kept redacted, and optional memory and temp-size ceilings, with `lock_configuration` still last because it freezes every later SET. `Hardening` is now a struct so a caller that needs no object storage can set `deny_external_access` and close network egress. The default cannot: an attached Iceberg catalog is made of S3 reads, so egress stays open on the catalog path and the docs now say so plainly rather than implying the lockdown is total. A row cap turned out to bound almost nothing, since one row can carry a whole column (`string_agg`, `list`, `to_json`). Results are now bounded by rows, total bytes, and wall-clock time together, via `Limits`; the byte ceiling is what actually holds, and the timeout exists because DuckDB has no statement timeout of its own, so it is enforced by interrupting the connection. `QueryResult.truncated` carries which ceiling was hit so the CLI can say which one. The docs now state that none of this is the security boundary: the read-only credential and the READ_ONLY attach are, and everything here is defence in depth in front of them. Tests pin each judgement against the shipped DuckDB, including canaries that fail loudly on upgrade if data-modifying CTEs start parsing as SELECTs or `prepare` stops executing leading statements. --- crates/tower-cmd/src/catalogs.rs | 51 ++-- crates/tower-duckdb/src/guard.rs | 242 +++++++++++++-- crates/tower-duckdb/src/lib.rs | 493 +++++++++++++++++++++++++++---- 3 files changed, 691 insertions(+), 95 deletions(-) diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index 135de016..d13f793b 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -7,7 +7,7 @@ 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_duckdb::{guard, params, run_query, Hardening, Limits, QueryResult, Session}; use tower_telemetry::debug; use crate::{api, beta, output, util::cmd}; @@ -407,7 +407,7 @@ async fn fetch_catalog_tables( &setup, "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", params![db_name], - None, + &Limits::none(), ) }) .await @@ -485,6 +485,9 @@ pub async fn do_query(out: &output::Out, config: Config, args: &ArgMatches) { 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(Ok(guard::ReadOnlyCheck::DeniedFunction(name))) => out.die(&format!( + "'{name}' is not allowed in a read-only query: it changes engine state or runs SQL built at runtime. Remove it, or re-run with --write." + )), Ok(Err(err)) => out.die(&format!("Could not validate the query: {err}")), Err(err) => out.die(&format!("Could not validate the query: {err}")), } @@ -525,18 +528,29 @@ async fn execute_catalog_query( let token = response.credentials.oauth_token.clone(); let setup = attach_statements(name, &response.credentials, mode); - // 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. + // Read mode is the sandboxed path: lock the session down after attach so the + // query cannot read the host or unwind the config, and bound the result so a + // terminal isn't flooded. No wall-clock ceiling here, because a legitimate + // analytical scan over a large catalog can take minutes and a person is + // driving; the agent path uses `Limits::agent()`, which adds one. Write mode + // is the trusted opt-in and runs unbounded. let harden = !write; - let max_rows = (!write).then_some(guard::AGENT_MAX_ROWS); + let limits = if write { + Limits::none() + } else { + Limits { + max_rows: Some(guard::AGENT_MAX_ROWS), + max_total_bytes: Some(guard::AGENT_MAX_RESULT_BYTES), + timeout: None, + } + }; let result = tokio::task::spawn_blocking(move || -> Result { let session = Session::open()?; session.run_setup(&setup)?; if harden { - session.harden()?; + session.harden(&Hardening::default())?; } - session.query(&sql, [], max_rows) + session.query(&sql, [], &limits) }) .await; @@ -713,7 +727,7 @@ async fn fetch_catalog_columns_via_rest( "column_types".to_string(), ], rows, - truncated: false, + truncated: None, }) } @@ -1001,13 +1015,16 @@ fn output_query_result(out: &output::Out, result: &QueryResult) { .collect(); out.table(result.columns.clone(), data, Some(&json_rows)); - if result.truncated { - out.note(&format!( - "\nShowing the first {} row(s); result truncated. Add a LIMIT or filter to narrow it.\n", + match result.truncated { + Some(tower_duckdb::Truncation::Rows) => out.note(&format!( + "\nShowing the first {} row(s); result truncated at the row limit. Add a LIMIT or filter to narrow it.\n", result.rows.len() - )); - } else { - out.note(&format!("\n{} row(s)\n", result.rows.len())); + )), + Some(tower_duckdb::Truncation::Bytes) => out.note(&format!( + "\nShowing {} row(s); result truncated at the size limit. Select fewer columns, or filter to narrow it.\n", + result.rows.len() + )), + None => out.note(&format!("\n{} row(s)\n", result.rows.len())), } } @@ -1284,7 +1301,7 @@ mod tests { token_export_command, }; use tower_api::models::{vend_catalog_credentials_body, CatalogCredentials}; - use tower_duckdb::{params, run_query}; + use tower_duckdb::{params, run_query, Limits}; #[test] fn list_defaults_to_default_environment() { @@ -1500,7 +1517,7 @@ mod tests { &setup, "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", params!["memory"], - None, + &Limits::none(), ) .expect("query should succeed"); diff --git a/crates/tower-duckdb/src/guard.rs b/crates/tower-duckdb/src/guard.rs index 92e51b3a..8395a19c 100644 --- a/crates/tower-duckdb/src/guard.rs +++ b/crates/tower-duckdb/src/guard.rs @@ -1,36 +1,82 @@ -//! 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: -//! it rejects anything that is not a single `SELECT`, and caps how many rows an -//! agent query may pull back. +//! Read-only gate on untrusted SQL, applied before a statement runs. +//! +//! This gate is defence in depth, not the security boundary. The load-bearing +//! control is the engine-enforced privilege: `catalogs query` vends read-only +//! credentials and attaches the catalog `READ_ONLY`, so a write is refused by the +//! catalog regardless of what this module concludes. The gate exists to refuse +//! bad input early, with a clear message, and to narrow the gap between "the +//! credential is read-only" and "the statement is a read". //! //! 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. +//! which parses (but does not execute) and serializes only `SELECT` statements, +//! erroring on everything else. Using the engine's parser rather than scanning +//! keywords is what makes this workable: 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 open with an allowed keyword but still mutate. The parser sees the +//! statement the way the executor will. +//! +//! Parsing as a `SELECT` is a statement *shape*, not a read-only property, so +//! shape alone is not enough. Two classes of SELECT-shaped statement still act: +//! functions that mutate engine state (`nextval` advances a sequence) and +//! functions that execute dynamically-built SQL (`query`, `query_table`). Those +//! are refused by name from the parsed tree, see [`MUTATING_OR_DYNAMIC_FUNCTIONS`]. +//! A `SELECT` that reads a file or a URL through a table function is still +//! allowed here; the session hardening is what refuses that read at execution. +//! +//! Every judgement in this module is pinned by tests against the DuckDB build we +//! ship, because both the grammar and the serialized JSON shape change between +//! versions. Re-run them on every DuckDB upgrade. /// 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. +/// is flagged truncated. +/// +/// A row cap alone bounds very little: `SELECT string_agg(email, ',') FROM users` +/// returns a whole column in one row. Pair it with [`AGENT_MAX_RESULT_BYTES`], +/// which is what actually bounds how much data a query carries back. pub const AGENT_MAX_ROWS: usize = 1_000; +/// Byte ceiling on a whole agent result set, measured as values are read. This is +/// the cap that survives aggregation tricks (`string_agg`, `list`, `to_json`) +/// that pack many rows into few. +pub const AGENT_MAX_RESULT_BYTES: usize = 1 << 20; + +/// Wall-clock ceiling on a single agent query. DuckDB has no statement timeout of +/// its own, so this is enforced from the host by interrupting the connection. +pub const AGENT_QUERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + +/// Functions refused even inside an otherwise valid `SELECT`. +/// +/// `nextval`/`currval` mutate sequence state, so they are writes wearing a +/// SELECT's clothes. `query`/`query_table` hand a string to DuckDB's execution +/// pipeline, which turns this gate's allowlist into an execution surface. None of +/// them have a legitimate use in a catalog data query. +/// +/// This deliberately excludes the `read_*`/`glob`/`*_scan` family: those are how a +/// catalog query legitimately reads Parquet and Iceberg over object storage. +/// Reaching the *local* filesystem through them is refused by the session +/// hardening instead. +pub const MUTATING_OR_DYNAMIC_FUNCTIONS: &[&str] = &["nextval", "currval", "query", "query_table"]; + /// The verdict for a piece of untrusted SQL. -#[derive(Debug, PartialEq, Eq, Clone, Copy)] +#[derive(Debug, PartialEq, Eq, Clone)] pub enum ReadOnlyCheck { - /// Exactly one `SELECT` statement. Safe to run under a read-only session. + /// Exactly one `SELECT` statement, calling nothing from the denied set. 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. + /// More than one statement. Rejected because duckdb-rs `prepare` executes + /// every statement but the last as a side effect: preparing + /// `DELETE FROM t; SELECT 1` deletes the rows and hands back the SELECT, so a + /// second statement must never reach it. Verified against the DuckDB build we + /// ship, in `prepare_executes_all_but_the_last_statement`. 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, + /// Parses as a `SELECT` but calls a function that mutates state or executes + /// dynamically-built SQL. Carries the offending name for the error message. + DeniedFunction(String), } /// Classify untrusted `sql` for the read-only path using DuckDB's parser. Opens a @@ -50,6 +96,18 @@ pub fn classify_read_only(sql: &str) -> Result { pub fn classify_read_only_on( conn: &duckdb::Connection, sql: &str, +) -> Result { + classify_read_only_with(conn, sql, MUTATING_OR_DYNAMIC_FUNCTIONS) +} + +/// [`classify_read_only_on`] with an explicit set of denied function names, for a +/// caller that wants to refuse more than the default (an agent path with no +/// business reading object storage might also deny the `read_*` family). Names are +/// matched case-insensitively against the parsed tree, never the raw SQL text. +pub fn classify_read_only_with( + conn: &duckdb::Connection, + sql: &str, + denied_functions: &[&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 @@ -81,11 +139,43 @@ pub fn classify_read_only_on( .and_then(serde_json::Value::as_array) .map_or(0, Vec::len); - Ok(match statements { - 0 => ReadOnlyCheck::Empty, - 1 => ReadOnlyCheck::Allowed, - _ => ReadOnlyCheck::Multiple, - }) + match statements { + 0 => return Ok(ReadOnlyCheck::Empty), + 1 => {} + _ => return Ok(ReadOnlyCheck::Multiple), + } + + if let Some(name) = find_denied_function(&parsed, denied_functions) { + return Ok(ReadOnlyCheck::DeniedFunction(name)); + } + Ok(ReadOnlyCheck::Allowed) +} + +/// Walks the serialized parse tree for a call to any denied function, scalar or +/// table-valued, at any depth. +/// +/// This reads `function_name` values out of the tree rather than pattern-matching +/// JSON text, because the serialized shape is not stable across DuckDB versions +/// and a substring match on raw SQL would be defeated by comments and quoting. If +/// a future DuckDB renames that key this walk silently stops matching, so each +/// denial is covered by a test that fails loudly when it stops being refused. +fn find_denied_function(node: &serde_json::Value, denied: &[&str]) -> Option { + match node { + serde_json::Value::Object(fields) => { + if let Some(serde_json::Value::String(name)) = fields.get("function_name") { + if denied.iter().any(|d| d.eq_ignore_ascii_case(name)) { + return Some(name.to_lowercase()); + } + } + fields + .values() + .find_map(|value| find_denied_function(value, denied)) + } + serde_json::Value::Array(items) => items + .iter() + .find_map(|item| find_denied_function(item, denied)), + _ => None, + } } #[cfg(test)] @@ -247,6 +337,110 @@ mod tests { } } + #[test] + fn rejects_functions_that_mutate_state_or_build_sql_dynamically() { + // These parse as perfectly good SELECTs, so the parser alone would let + // them through. `nextval` really does advance a sequence (see + // `nextval_is_a_select_that_mutates`), and `query`/`query_table` hand a + // string to the execution pipeline. + for (sql, expected) in [ + ("SELECT nextval('s')", "nextval"), + ("SELECT NEXTVAL('s')", "nextval"), + ("SELECT currval('s')", "currval"), + ("SELECT * FROM query('SELECT 1')", "query"), + ("SELECT * FROM query_table('t')", "query_table"), + // Nested well below the top level, to prove the walk is not shallow. + ( + "WITH x AS (SELECT nextval('s') AS n) SELECT sum(n) FROM x", + "nextval", + ), + ("SELECT (SELECT max(v) FROM (SELECT nextval('s') v))", "nextval"), + ] { + assert_eq!( + check(sql), + ReadOnlyCheck::DeniedFunction(expected.to_string()), + "should deny {expected} in: {sql}" + ); + } + } + + #[test] + fn a_denied_name_used_as_an_identifier_is_not_a_call() { + // The denial matches parsed function calls, not text, so a column or alias + // that happens to share the name is still a fine read. + for sql in [ + "SELECT query FROM runs", + "SELECT 1 AS nextval", + "SELECT 'nextval(x)' AS literal_text", + "SELECT * FROM runs -- nextval('s')", + ] { + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "should allow: {sql}"); + } + } + + /// Pins the reason `nextval` is on the denylist: it is a SELECT that writes. + /// If a DuckDB upgrade ever made this non-mutating the denial could be + /// revisited, and if the denial is dropped while this still mutates, the test + /// above starts failing. + #[test] + fn nextval_is_a_select_that_mutates() { + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + conn.execute_batch("CREATE SEQUENCE s START 1") + .expect("create sequence"); + let first: i64 = conn + .query_row("SELECT nextval('s')", [], |row| row.get(0)) + .expect("nextval"); + let second: i64 = conn + .query_row("SELECT nextval('s')", [], |row| row.get(0)) + .expect("nextval"); + assert_ne!( + first, second, + "nextval no longer mutates; revisit MUTATING_OR_DYNAMIC_FUNCTIONS" + ); + } + + /// Pins the reason multi-statement input is refused, and guards the claim in + /// `ReadOnlyCheck::Multiple`. duckdb-rs `prepare` runs every statement but the + /// last, so the leading statement takes effect even though only the final one + /// is returned. If a future duckdb-rs changes this to a hard error the test + /// fails and the doc comment needs updating, but the gate stays correct. + #[test] + fn prepare_executes_all_but_the_last_statement() { + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + conn.execute_batch("CREATE TABLE t(id INTEGER); INSERT INTO t VALUES (1), (2)") + .expect("seed table"); + + let _ = conn.prepare("DELETE FROM t; SELECT 1"); + + let remaining: i64 = conn + .query_row("SELECT count(*) FROM t", [], |row| row.get(0)) + .expect("count"); + assert_eq!( + remaining, 0, + "prepare no longer executes leading statements; revisit the Multiple rationale" + ); + } + + /// A DuckDB version canary. Data-modifying CTEs do not parse today ("A CTE + /// needs a SELECT"), so they land in `NotReadOnly`. Upstream has work to allow + /// DML as a CTE body; if that ships, these gain an outer SELECT and could + /// start classifying as `Allowed` while still deleting rows. This test fails + /// the moment that changes, which is the signal to add an explicit CTE check. + #[test] + fn data_modifying_ctes_do_not_parse_as_selects() { + for sql in [ + "WITH x AS (DELETE FROM runs RETURNING *) SELECT * FROM x", + "WITH x AS (INSERT INTO runs VALUES (9) RETURNING *) SELECT * FROM x", + "WITH x AS (UPDATE runs SET id = 5 RETURNING *) SELECT * FROM x", + ] { + assert_eq!( + check(sql), + ReadOnlyCheck::NotReadOnly, + "DML-in-CTE now parses as a SELECT; the gate needs an explicit check: {sql}" + ); + } + } + #[test] fn table_functions_reaching_the_host_are_a_matter_for_the_hardening_not_the_gate() { // The gate only classifies statement shape. A SELECT that reads a local @@ -280,3 +474,5 @@ mod tests { ); } } + + diff --git a/crates/tower-duckdb/src/lib.rs b/crates/tower-duckdb/src/lib.rs index 2b93b9f6..78531b6e 100644 --- a/crates/tower-duckdb/src/lib.rs +++ b/crates/tower-duckdb/src/lib.rs @@ -2,14 +2,28 @@ //! 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 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 +//! The point of the crate is the hardened query path. Agent-issued SQL runs on a +//! customer's machine with their catalog credentials, so it gets layered +//! treatment: 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, pull in extensions, or unwind the settings +//! ([`Hardening`]), and bound what a result may carry back in rows, bytes, and +//! wall-clock time ([`Limits`]). The adversarial tests exercise each of these //! directly. +//! +//! None of that is the security boundary. The boundary is the engine-enforced +//! privilege: the caller vends read-only credentials and attaches the catalog +//! `READ_ONLY`, so a write is refused by the catalog no matter what this crate +//! concludes. Everything here is defence in depth in front of that, worth having +//! because it turns a confusing engine error into a clear refusal and closes +//! SELECT-shaped holes the credential alone would not. +//! +//! Two gaps are deliberate and worth knowing. Network egress is open: an attached +//! Iceberg catalog is made of S3 reads, so `enable_external_access` cannot be +//! turned off on that path, and a query can still reach a URL through a table +//! function. And the ceilings bound what this process holds, not what DuckDB +//! spends internally on a materializing plan. Closing either one needs controls +//! outside this crate: a network boundary, and [`Hardening::memory_limit`]. use std::time::Instant; @@ -19,29 +33,141 @@ 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. +/// Why a result stopped short of everything the query would have returned. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Truncation { + /// Hit the row ceiling. + Rows, + /// Hit the byte ceiling. This is the one that catches a query which packs a + /// whole table into a handful of rows. + Bytes, +} + +/// A tabular query result: column names, rows as positional JSON values, and why +/// it stopped early if it did. #[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, + /// Set when a [`Limits`] ceiling cut the result short. + pub truncated: Option, } -/// 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. +impl QueryResult { + /// Whether a ceiling cut this result short. + pub fn is_truncated(&self) -> bool { + self.truncated.is_some() + } +} + +/// Ceilings applied while a result is read back. +/// +/// Rows alone are a weak bound, because a query can compact a whole column into +/// one row (`string_agg`, `list`, `to_json`). `max_total_bytes` is what actually +/// limits how much data a single query carries back, so set both for untrusted +/// callers. `timeout` bounds how long the query may run: DuckDB has no statement +/// timeout of its own, so it is enforced by interrupting the connection. +#[derive(Debug, Clone, Default)] +pub struct Limits { + pub max_rows: Option, + pub max_total_bytes: Option, + pub timeout: Option, +} + +impl Limits { + /// No ceilings. For trusted, caller-authored SQL. + pub fn none() -> Self { + Self::default() + } + + /// The ceilings for untrusted (agent-issued) SQL. + pub fn agent() -> Self { + Self { + max_rows: Some(guard::AGENT_MAX_ROWS), + max_total_bytes: Some(guard::AGENT_MAX_RESULT_BYTES), + timeout: Some(guard::AGENT_QUERY_TIMEOUT), + } + } + + /// Only a row ceiling, for a caller that wants a readable result without a + /// security budget. + pub fn rows(max_rows: usize) -> Self { + Self { + max_rows: Some(max_rows), + ..Self::default() + } + } +} + +/// How to lock a session down before untrusted SQL runs. +/// +/// The defaults are what a catalog query needs: local-filesystem access off (so +/// `read_csv('/etc/passwd')` and friends fail), no implicit extension +/// install/load, no community or unsigned extensions, secrets kept redacted, and +/// the configuration frozen so the query cannot unwind any of it. +/// +/// Note what is deliberately *not* set: `enable_external_access = false`. It is +/// DuckDB's master switch and would be the stronger control, but it also blocks +/// the S3 reads an attached Iceberg catalog is made of, so a catalog query cannot +/// use it. Network egress is therefore not closed by this lockdown; a query can +/// still reach a URL through a table function. The read-only credential is what +/// bounds what that query can *read*, and closing egress needs a network boundary +/// outside this process. Callers that do not need object storage should set +/// `deny_external_access` and get the stronger guarantee. +#[derive(Debug, Clone, Default)] +pub struct Hardening { + /// Refuse all external access, network and local files alike. Breaks attached + /// object-store catalogs, so it is off by default. + pub deny_external_access: bool, + /// Engine memory ceiling, e.g. `"2GB"`. `None` keeps DuckDB's default. + pub memory_limit: Option, + /// Ceiling on spill-to-disk, e.g. `"4GB"`. `None` keeps DuckDB's default. + pub max_temp_directory_size: Option, +} + +impl Hardening { + /// The lockdown statements, in the order they must run. + /// + /// `lock_configuration` is last because it freezes every later `SET`. The + /// whole sequence runs after setup, because attaching a catalog is what + /// installs the extensions and reaches the network that this then takes away. + pub fn statements(&self) -> Vec { + let mut statements = Vec::new(); + if self.deny_external_access { + statements.push("SET enable_external_access = false".to_string()); + } + statements.extend( + [ + "SET disabled_filesystems = 'LocalFileSystem'", + // Stop DuckDB reaching out for an extension mid-query. The ones a + // catalog needs are already loaded by setup. + "SET autoinstall_known_extensions = false", + "SET autoload_known_extensions = false", + "SET allow_community_extensions = false", + "SET allow_unsigned_extensions = false", + // Keep vended tokens redacted in duckdb_secrets(). + "SET allow_unredacted_secrets = false", + ] + .into_iter() + .map(str::to_string), + ); + if let Some(limit) = &self.memory_limit { + statements.push(format!("SET memory_limit = '{}'", limit.replace('\'', "''"))); + } + if let Some(limit) = &self.max_temp_directory_size { + statements.push(format!( + "SET max_temp_directory_size = '{}'", + limit.replace('\'', "''") + )); + } + statements.push("SET lock_configuration = true".to_string()); + statements + } +} + +/// The default lockdown statements, for callers that do not need to tune it. 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(), - ] + Hardening::default().statements() } /// An in-memory DuckDB connection Tower runs queries through. @@ -84,8 +210,8 @@ impl Session { /// attaching a catalog needs the access this removes. /// /// [`run_setup`]: Session::run_setup - pub fn harden(&self) -> Result<(), Error> { - for statement in hardening_statements() { + pub fn harden(&self, hardening: &Hardening) -> Result<(), Error> { + for statement in hardening.statements() { self.conn.execute_batch(&statement)?; } Ok(()) @@ -93,20 +219,33 @@ impl Session { /// 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. + /// text. + /// + /// `limits` bound the result as it is read: rows and total bytes are counted + /// as values come back and the read stops at either ceiling, so an untrusted + /// caller cannot pull an unbounded table into memory or a model's context. A + /// `timeout` interrupts the connection when the query outruns it. + /// + /// The ceilings bound what this process *holds*, which is not the same as + /// what the engine does: a query whose plan materializes (a large `ORDER BY` + /// or aggregate) spends that memory inside DuckDB before the first row is + /// handed over. Use [`Hardening::memory_limit`] for that. pub fn query( &self, sql: &str, params: P, - max_rows: Option, + limits: &Limits, ) -> Result { let query_start = Instant::now(); + // Armed before prepare so a statement that hangs while binding (a remote + // catalog read, say) is still interrupted. Disarms on drop. + let _deadline = limits.timeout.map(|timeout| Deadline::arm(&self.conn, timeout)); + let mut stmt = self.conn.prepare(sql)?; let mut columns: Vec = Vec::new(); let mut rows = Vec::new(); - let mut truncated = false; + let mut truncated = None; + let mut total_bytes = 0usize; { let mut result_rows = stmt.query(params)?; @@ -114,8 +253,8 @@ impl Session { if columns.is_empty() { columns = row.as_ref().column_names(); } - if max_rows.is_some_and(|max| rows.len() >= max) { - truncated = true; + if limits.max_rows.is_some_and(|max| rows.len() >= max) { + truncated = Some(Truncation::Rows); break; } let mut record = Vec::with_capacity(columns.len()); @@ -123,7 +262,14 @@ impl Session { let value: duckdb::types::Value = row.get(idx)?; record.push(value_to_json(value)); } + total_bytes += record.iter().map(json_size) .sum::(); rows.push(record); + // Checked after pushing, so a single oversized row is returned + // rather than silently yielding an empty result. + if limits.max_total_bytes.is_some_and(|max| total_bytes >= max) { + truncated = Some(Truncation::Bytes); + break; + } } } @@ -133,9 +279,10 @@ impl Session { } debug!( - "tower-duckdb: query took {:?} ({} rows): {}", + "tower-duckdb: query took {:?} ({} rows, {} bytes): {}", query_start.elapsed(), rows.len(), + total_bytes, sql ); @@ -147,6 +294,70 @@ impl Session { } } +/// Interrupts a connection if it is still working when the deadline passes. +/// +/// DuckDB has no statement timeout, so a wall-clock bound has to come from the +/// host. Interrupts are honoured at chunk boundaries, so this bounds a runaway +/// query rather than guaranteeing an exact deadline; a hard kill would need a +/// process boundary. +struct Deadline { + expired: std::sync::Arc, + watcher: Option>, +} + +impl Deadline { + fn arm(conn: &duckdb::Connection, timeout: std::time::Duration) -> Self { + let handle = conn.interrupt_handle(); + let expired = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let done = expired.clone(); + let watcher = std::thread::spawn(move || { + let deadline = Instant::now() + timeout; + // Woken in slices so a finished query disarms promptly instead of + // holding the thread for the whole timeout. + while Instant::now() < deadline { + if done.load(std::sync::atomic::Ordering::Relaxed) { + return; + } + std::thread::sleep(std::time::Duration::from_millis(25)); + } + if !done.load(std::sync::atomic::Ordering::Relaxed) { + handle.interrupt(); + } + }); + Self { + expired, + watcher: Some(watcher), + } + } +} + +impl Drop for Deadline { + fn drop(&mut self) { + self.expired + .store(true, std::sync::atomic::Ordering::Relaxed); + if let Some(watcher) = self.watcher.take() { + let _ = watcher.join(); + } + } +} + +/// Rough serialized size of a value, used only to bound how much a result may +/// carry back. Strings dominate real results, so they are measured exactly and +/// everything else is approximated. +fn json_size(value: &serde_json::Value) -> usize { + match value { + serde_json::Value::Null => 4, + serde_json::Value::Bool(_) => 5, + serde_json::Value::Number(n) => n.to_string().len(), + serde_json::Value::String(s) => s.len(), + serde_json::Value::Array(items) => items.iter().map(json_size).sum::() + 2, + serde_json::Value::Object(fields) => fields + .iter() + .map(|(key, value)| key.len() + json_size(value)) + .sum::(), + } +} + /// 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 @@ -155,11 +366,11 @@ pub fn run_query( setup: &[String], query: &str, params: P, - max_rows: Option, + limits: &Limits, ) -> Result { let session = Session::open()?; session.run_setup(setup)?; - session.query(query, params, max_rows) + session.query(query, params, limits) } /// Converts a DuckDB value into a `serde_json::Value`. Integers that overflow an @@ -261,7 +472,10 @@ fn stringify_key(value: &serde_json::Value) -> String { #[cfg(test)] mod tests { use super::guard::{classify_read_only, ReadOnlyCheck}; - use super::{hardening_statements, run_query, value_to_json, QueryResult}; + use super::{ + hardening_statements, run_query, value_to_json, Hardening, Limits, QueryResult, Session, + Truncation, + }; #[test] fn duckdb_values_convert_to_json() { @@ -296,7 +510,7 @@ mod tests { &[], "SELECT [1, 2] AS l, {'a': 1, 'b': 'x'} AS s, MAP {'k': 2} AS m", [], - None, + &Limits::none(), ) .expect("query should succeed"); @@ -317,7 +531,7 @@ mod tests { "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) + let result = run_query(&setup, "SELECT id, name FROM t ORDER BY id", [], &Limits::none()) .expect("query should succeed"); assert_eq!(result.columns, vec!["id", "name"]); @@ -328,13 +542,13 @@ mod tests { vec![serde_json::json!(2), serde_json::json!("b")], ] ); - assert!(!result.truncated); + assert!(!result.is_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"); + run_query(&[], "SELECT 1 AS x WHERE 1 = 0", [], &Limits::none()).expect("query should succeed"); assert_eq!(result.columns, vec!["x"]); assert!(result.rows.is_empty()); @@ -342,15 +556,184 @@ mod tests { #[test] fn run_query_caps_rows_and_flags_truncation() { - let capped = run_query(&[], "SELECT * FROM range(5) AS t(i)", [], Some(3)) + let capped = run_query(&[], "SELECT * FROM range(5) AS t(i)", [], &Limits::rows(3)) .expect("query should succeed"); assert_eq!(capped.rows.len(), 3); - assert!(capped.truncated); + assert_eq!(capped.truncated, Some(Truncation::Rows)); - let exact = run_query(&[], "SELECT * FROM range(3) AS t(i)", [], Some(3)) + let exact = run_query(&[], "SELECT * FROM range(3) AS t(i)", [], &Limits::rows(3)) .expect("query should succeed"); assert_eq!(exact.rows.len(), 3); - assert!(!exact.truncated); + assert!(!exact.is_truncated()); + } + + /// The reason a row cap is not enough on its own: one row can carry a whole + /// column. A row-only ceiling lets this through; the byte ceiling catches it. + #[test] + fn a_byte_ceiling_catches_what_a_row_ceiling_misses() { + let packed = "SELECT string_agg(i::VARCHAR, ',') AS all_rows FROM range(20000) AS t(i)"; + + let row_capped = run_query(&[], packed, [], &Limits::rows(1000)) + .expect("query should succeed"); + assert_eq!(row_capped.rows.len(), 1); + assert!( + !row_capped.is_truncated(), + "a row ceiling cannot see a single oversized row" + ); + let packed_bytes = row_capped.rows[0][0].as_str().map_or(0, str::len); + assert!( + packed_bytes > 64 * 1024, + "expected a large packed value, got {packed_bytes} bytes" + ); + + let byte_capped = run_query( + &[], + packed, + [], + &Limits { + max_rows: Some(1000), + max_total_bytes: Some(4096), + timeout: None, + }, + ) + .expect("query should succeed"); + assert_eq!( + byte_capped.truncated, + Some(Truncation::Bytes), + "the byte ceiling should have cut this short" + ); + } + + /// Many small rows trip the byte ceiling too, so the bound holds however the + /// result is shaped. + #[test] + fn the_byte_ceiling_also_bounds_many_small_rows() { + let result = run_query( + &[], + "SELECT repeat('x', 512) AS pad FROM range(10000)", + [], + &Limits { + max_rows: None, + max_total_bytes: Some(16 * 1024), + timeout: None, + }, + ) + .expect("query should succeed"); + + assert_eq!(result.truncated, Some(Truncation::Bytes)); + assert!( + result.rows.len() < 10_000, + "should have stopped early, got {} rows", + result.rows.len() + ); + } + + /// DuckDB has no statement timeout, so a runaway query is bounded by + /// interrupting the connection from the host. `range(1e12)` would run + /// effectively forever; it must come back as an error, not hang the test. + #[test] + fn a_runaway_query_is_interrupted_by_the_timeout() { + let started = std::time::Instant::now(); + let result = run_query( + &[], + "SELECT count(*) FROM range(1000000000000) AS t(i) WHERE i % 7 = 0", + [], + &Limits { + max_rows: None, + max_total_bytes: None, + timeout: Some(std::time::Duration::from_secs(2)), + }, + ); + + assert!(result.is_err(), "runaway query should have been interrupted"); + assert!( + started.elapsed() < std::time::Duration::from_secs(60), + "interrupt took too long: {:?}", + started.elapsed() + ); + } + + /// The timeout must not fire on a query that finishes inside it, and the + /// watcher must not outlive the call. + #[test] + fn a_fast_query_is_unaffected_by_an_armed_timeout() { + let result = run_query( + &[], + "SELECT 1 AS x", + [], + &Limits { + max_rows: None, + max_total_bytes: None, + timeout: Some(std::time::Duration::from_secs(30)), + }, + ) + .expect("a fast query should not be interrupted"); + assert_eq!(result.rows, vec![vec![serde_json::json!(1)]]); + } + + /// `lock_configuration` freezes every later `SET`, so it has to be issued + /// last or the rest of the lockdown silently fails to apply. + #[test] + fn lock_configuration_is_always_the_last_hardening_statement() { + for hardening in [ + Hardening::default(), + Hardening { + deny_external_access: true, + memory_limit: Some("2GB".to_string()), + max_temp_directory_size: Some("4GB".to_string()), + }, + ] { + let statements = hardening.statements(); + let last = statements.last().expect("hardening should not be empty"); + assert!( + last.contains("lock_configuration"), + "lock_configuration must be last, got: {last}" + ); + assert_eq!( + statements + .iter() + .filter(|s| s.contains("lock_configuration")) + .count(), + 1 + ); + } + } + + /// The stricter lockdown a caller with no object-store catalog can take. It + /// closes network egress, which the default cannot, because an attached + /// Iceberg catalog is made of S3 reads. + #[test] + fn deny_external_access_closes_the_network_the_default_leaves_open() { + let strict = Hardening { + deny_external_access: true, + ..Hardening::default() + }; + assert!(strict + .statements() + .iter() + .any(|s| s.contains("enable_external_access"))); + assert!( + !Hardening::default() + .statements() + .iter() + .any(|s| s.contains("enable_external_access")), + "the default must leave object-store reads working" + ); + + let session = Session::open().expect("open session"); + session.harden(&strict).expect("harden"); + let err = session + .query( + "SELECT * FROM read_csv('https://example.com/x.csv')", + [], + &Limits::none(), + ) + .expect_err("external access should be refused"); + let message = err.to_string().to_lowercase(); + assert!( + message.contains("disabled") || message.contains("permission"), + "unexpected error: {err}" + ); } // The read-only gate is unit-tested in `guard.rs`. This helper backs the @@ -363,19 +746,19 @@ mod tests { 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) + let ok = run_query(&setup, "SELECT 1 AS x", [], &Limits::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) + let fs_err = run_query(&setup, "SELECT * FROM read_csv('Cargo.toml')", [], &Limits::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) + let cfg_err = run_query(&setup, "SET memory_limit = '1GB'", [], &Limits::none()) .expect_err("configuration should be locked"); let cfg_msg = cfg_err.to_string().to_lowercase(); assert!( @@ -418,7 +801,7 @@ mod tests { &setup, &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), [], - None, + &Limits::none(), ) .expect("iceberg_scan should read the table"); assert_eq!(result.columns, vec!["id", "name"]); @@ -430,17 +813,17 @@ mod tests { vec![serde_json::json!(3), serde_json::Value::Null], ] ); - assert!(!result.truncated); + assert!(!result.is_truncated()); let capped = run_query( &setup, &format!("SELECT id FROM iceberg_scan('{table}') ORDER BY id"), [], - Some(2), + &Limits::rows(2), ) .expect("iceberg_scan should read the table"); assert_eq!(capped.rows.len(), 2); - assert!(capped.truncated); + assert_eq!(capped.truncated, Some(Truncation::Rows)); let _ = std::fs::remove_dir_all(&dir); } @@ -454,7 +837,7 @@ mod tests { /// 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) + run_query(&hardening_statements(), sql, [], &Limits::none()) } #[test] @@ -657,7 +1040,7 @@ mod tests { &setup, &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), [], - None, + &Limits::none(), ) .expect("hardening must not break object-store Iceberg reads"); assert_eq!(read.columns, vec!["id", "name"]); @@ -670,14 +1053,14 @@ mod tests { ] ); - let local = run_query(&setup, "SELECT * FROM read_csv('/etc/hostname')", [], None) + let local = run_query(&setup, "SELECT * FROM read_csv('/etc/hostname')", [], &Limits::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) + let cfg = run_query(&setup, "SET memory_limit = '1GB'", [], &Limits::none()) .expect_err("configuration must stay locked"); let cfg = cfg.to_string().to_lowercase(); assert!( From 58e266c47d6c21905e073ceeb2ef925fb87671dc Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Mon, 27 Jul 2026 09:03:35 +0100 Subject: [PATCH 7/9] feat(catalogs): add --max-rows to override the query result ceiling Read queries are bounded by default so a runaway result cannot flood a terminal or a model's context, but there was no way to ask for more. That makes the bound a wall rather than a default: an agent (or a person) with a legitimate need for a large extract had no option short of --write, which vends read-write credentials for what is still only a read. --max-rows sets the row ceiling and lifts the size ceiling with it, because a caller who asks for a million rows should not then be cut short by a byte budget they never set. --max-rows 0 removes the ceilings entirely. The help text says plainly that a large result can exhaust memory, since that is the trade being made. The read-only gate and the session hardening are untouched: this widens how much data a read may return, not what a query is allowed to do. The decision is factored into `query_limits` so it is unit-tested directly rather than only through argument parsing. --- crates/tower-cmd/src/catalogs.rs | 143 +++++++++++++++++++++++++++---- 1 file changed, 124 insertions(+), 19 deletions(-) diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index d13f793b..25c03023 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -160,6 +160,13 @@ pub fn catalogs_cmd() -> Command { .help("Allow write statements by vending read-write credentials; queries are read-only by default") .action(ArgAction::SetTrue), ) + .arg( + Arg::new("max_rows") + .long("max-rows") + .value_parser(value_parser!(usize)) + .help("Maximum rows to return; 0 for no limit. Also lifts the result size limit, so a large result can exhaust memory") + .action(ArgAction::Set), + ) .about(beta::STORAGE.short_about("Run a SQL query against a catalog using DuckDB")) .after_help( "Reference tables as .., e.g.:\n tower catalogs query default --sql 'SELECT * FROM \"default\".my_namespace.my_table LIMIT 10'", @@ -492,14 +499,40 @@ pub async fn do_query(out: &output::Out, config: Config, args: &ArgMatches) { Err(err) => out.die(&format!("Could not validate the query: {err}")), } } - let query_result = execute_catalog_query(out, &config, name, &env, sql, write).await; + let limits = query_limits(write, args.get_one::("max_rows").copied()); + let query_result = execute_catalog_query(out, &config, name, &env, sql, write, limits).await; output_query_result(out, &query_result); } +/// The result ceilings for a query. +/// +/// Read mode defaults to bounded, so a runaway query cannot flood a terminal or a +/// model's context. `--max-rows` is the escape hatch for a caller that knowingly +/// wants more: it sets the row ceiling and lifts the size ceiling, because a +/// caller who asked for a million rows should not then be cut off by a byte +/// budget they never saw. `--max-rows 0` removes the ceilings entirely. Write +/// mode is the trusted path and is unbounded unless a row count is asked for. +fn query_limits(write: bool, max_rows: Option) -> Limits { + match max_rows { + Some(0) => Limits::none(), + Some(rows) => Limits { + max_rows: Some(rows), + max_total_bytes: None, + timeout: None, + }, + None if write => Limits::none(), + None => Limits { + max_rows: Some(guard::AGENT_MAX_ROWS), + max_total_bytes: Some(guard::AGENT_MAX_RESULT_BYTES), + timeout: None, + }, + } +} + /// Vends credentials for the catalog, attaches it in an in-memory DuckDB, and /// 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 +/// after attach and the result is bounded by `limits`, 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( @@ -509,6 +542,7 @@ async fn execute_catalog_query( env: &str, sql: String, write: bool, + limits: Limits, ) -> QueryResult { let mode = if write { vend_catalog_credentials_body::Mode::ReadWrite @@ -529,21 +563,11 @@ async fn execute_catalog_query( let token = response.credentials.oauth_token.clone(); let setup = attach_statements(name, &response.credentials, mode); // Read mode is the sandboxed path: lock the session down after attach so the - // query cannot read the host or unwind the config, and bound the result so a - // terminal isn't flooded. No wall-clock ceiling here, because a legitimate - // analytical scan over a large catalog can take minutes and a person is - // driving; the agent path uses `Limits::agent()`, which adds one. Write mode - // is the trusted opt-in and runs unbounded. + // query cannot read the host or unwind the config. `limits` bounds the result + // (see `query_limits`). No wall-clock ceiling on this path, because a + // legitimate analytical scan over a large catalog can take minutes and a + // person is driving; the agent path uses `Limits::agent()`, which adds one. let harden = !write; - let limits = if write { - Limits::none() - } else { - Limits { - max_rows: Some(guard::AGENT_MAX_ROWS), - max_total_bytes: Some(guard::AGENT_MAX_RESULT_BYTES), - timeout: None, - } - }; let result = tokio::task::spawn_blocking(move || -> Result { let session = Session::open()?; session.run_setup(&setup)?; @@ -1297,8 +1321,8 @@ fn snippets( #[cfg(test)] mod tests { use super::{ - attach_statements, catalogs_cmd, is_storage_catalog_type, parse_mode, snippets, - token_export_command, + attach_statements, catalogs_cmd, is_storage_catalog_type, parse_mode, query_limits, + snippets, token_export_command, }; use tower_api::models::{vend_catalog_credentials_body, CatalogCredentials}; use tower_duckdb::{params, run_query, Limits}; @@ -1637,6 +1661,87 @@ mod tests { assert_eq!(query_args.get_one::("write").copied(), Some(true)); } + #[test] + fn query_accepts_a_max_rows_override() { + let matches = catalogs_cmd() + .try_get_matches_from(["catalogs", "query", "my-catalog", "--sql", "SELECT 1"]) + .expect("query should parse"); + let (_, query_args) = matches.subcommand().expect("expected query subcommand"); + assert_eq!(query_args.get_one::("max_rows").copied(), None); + + let matches = catalogs_cmd() + .try_get_matches_from([ + "catalogs", + "query", + "my-catalog", + "--sql", + "SELECT 1", + "--max-rows", + "50000", + ]) + .expect("query --max-rows should parse"); + let (_, query_args) = matches.subcommand().expect("expected query subcommand"); + assert_eq!(query_args.get_one::("max_rows").copied(), Some(50_000)); + + assert!( + catalogs_cmd() + .try_get_matches_from([ + "catalogs", + "query", + "my-catalog", + "--sql", + "SELECT 1", + "--max-rows", + "not-a-number", + ]) + .is_err(), + "a non-numeric --max-rows should be rejected" + ); + } + + #[test] + fn read_queries_are_bounded_unless_the_caller_asks_otherwise() { + // The default: bounded by rows and by size, so a runaway read cannot + // flood a terminal or a model's context. + let default = query_limits(false, None); + assert_eq!(default.max_rows, Some(tower_duckdb::guard::AGENT_MAX_ROWS)); + assert_eq!( + default.max_total_bytes, + Some(tower_duckdb::guard::AGENT_MAX_RESULT_BYTES) + ); + } + + #[test] + fn max_rows_override_raises_the_row_ceiling_and_lifts_the_size_ceiling() { + // A caller who asks for a million rows should not then be cut short by a + // byte budget they never set, so the size ceiling comes off with it. + let raised = query_limits(false, Some(1_000_000)); + assert_eq!(raised.max_rows, Some(1_000_000)); + assert_eq!( + raised.max_total_bytes, None, + "an explicit row count should not be second-guessed by the size cap" + ); + + // Lowering it is just as valid as raising it. + assert_eq!(query_limits(false, Some(10)).max_rows, Some(10)); + } + + #[test] + fn max_rows_zero_removes_every_result_ceiling() { + let unbounded = query_limits(false, Some(0)); + assert_eq!(unbounded.max_rows, None); + assert_eq!(unbounded.max_total_bytes, None); + } + + #[test] + fn write_mode_is_unbounded_but_still_honours_an_explicit_row_count() { + let write_default = query_limits(true, None); + assert_eq!(write_default.max_rows, None); + assert_eq!(write_default.max_total_bytes, None); + + assert_eq!(query_limits(true, Some(25)).max_rows, Some(25)); + } + #[test] fn token_export_command_fetches_token_without_printing_it() { let credentials = CatalogCredentials::new( From a23beed3e410b0a70ad311486414b0b8d496a17e Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Mon, 27 Jul 2026 10:33:34 +0100 Subject: [PATCH 8/9] fix(tower-duckdb): make the function gate fail-closed and the ceilings real Review found three ways the sandbox did less than it claimed, and two ways the tests hid it. The function denylist was not fail-closed. `json_execute_serialized_sql` runs whatever SQL it is handed and ships with the very `json` extension this crate enables to build the gate, so the gate supplied its own bypass: the outer statement is an ordinary SELECT and the effectful call hides in a string. `enable_logging`, `checkpoint`, `setseed` and friends were allowed for the same reason, that naming dangerous functions one at a time can never be complete. The gate now allowlists the table-function position, which is where the danger lives, so an unknown function is refused rather than waved through. Scalars are too numerous to allowlist, so it asks the engine instead: DuckDB's `has_side_effects` flags `nextval`, `setseed` and the rest, and keeps up with new versions on its own. It is NULL for every table function, which is why those are allowlisted rather than queried. An explicit list covers what neither reaches: the dynamic-SQL executors and the effectful table functions. This tightens the read path, since `read_csv`/`read_parquet` and the rest of that family are no longer permitted; a catalog query reads base tables, which are not functions. The byte ceiling measured a row after admitting it, so `SELECT repeat('x', 5e7)` returned fifty megabytes and labelled it truncated. It now measures first and withholds the row, which means a single over-budget row yields an empty truncated result. That is blunt but honest, and it is what the ceiling was for. Because that still only bounds what the caller is handed, the read path now hardens with `Hardening::agent()`, which sets an engine memory limit for what a query spends before the first row exists. The SSRF test passed for a reason production does not share: its helper never loaded `httpfs`, so there was no HTTP filesystem to block. It now loads the same extensions `attach_statements` does, and counts connections to a listener it owns rather than trusting an error string, with a control run so it cannot pass vacuously. Note the hardening blocks these today only because disabling LocalFileSystem breaks path resolution first; egress is still not closed by design, and that gap is documented rather than papered over. Also: the MinIO image is pinned to a release tag instead of `latest` and only skips when Docker is genuinely unreachable, so a registry problem fails loudly instead of silently disabling a security test; the iceberg skip is narrowed the same way; and the workspace `rust-version` is corrected to 1.88, which the toolchain and the dependency tree have required for some time. --- Cargo.toml | 6 +- crates/tower-cmd/src/catalogs.rs | 5 +- crates/tower-duckdb/src/guard.rs | 301 ++++++++++++++++++++++++++----- crates/tower-duckdb/src/lib.rs | 233 +++++++++++++++++++++--- 4 files changed, 473 insertions(+), 72 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3b92533e..d02601b7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,11 @@ resolver = "2" edition = "2021" version = "0.3.70-rc.1" description = "Tower is the best way to host Python data apps in production" -rust-version = "1.81" +# Matches rust-toolchain.toml. The two had drifted: the toolchain has been 1.88 +# for a while, and the dependency tree (testcontainers and its transitive deps, +# some on edition 2024) does not resolve on 1.81, so the older figure was a +# promise the workspace could not keep. +rust-version = "1.88" authors = ["Brad Heller ", "Ben Lovell "] license = "MIT" repository = "https://github.com/tower/tower-cli" diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index 25c03023..8b51b160 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -572,7 +572,10 @@ async fn execute_catalog_query( let session = Session::open()?; session.run_setup(&setup)?; if harden { - session.harden(&Hardening::default())?; + // `Hardening::agent()` adds an engine memory ceiling on top of the + // lockdown. The result `limits` bound what comes back; only the engine + // can bound what a query spends producing it. + session.harden(&Hardening::agent())?; } session.query(&sql, [], &limits) }) diff --git a/crates/tower-duckdb/src/guard.rs b/crates/tower-duckdb/src/guard.rs index 8395a19c..43ea9dff 100644 --- a/crates/tower-duckdb/src/guard.rs +++ b/crates/tower-duckdb/src/guard.rs @@ -45,18 +45,70 @@ pub const AGENT_MAX_RESULT_BYTES: usize = 1 << 20; /// its own, so this is enforced from the host by interrupting the connection. pub const AGENT_QUERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); -/// Functions refused even inside an otherwise valid `SELECT`. +/// Functions refused wherever they appear, whatever the engine says about them. /// -/// `nextval`/`currval` mutate sequence state, so they are writes wearing a -/// SELECT's clothes. `query`/`query_table` hand a string to DuckDB's execution -/// pipeline, which turns this gate's allowlist into an execution surface. None of -/// them have a legitimate use in a catalog data query. +/// Two kinds live here. The first executes SQL built at runtime, which would turn +/// this gate into an execution surface: `query`, `query_table`, and +/// `json_execute_serialized_sql` (the last one ships with the `json` extension +/// this crate enables for the parser, so the gate would otherwise supply its own +/// bypass). The second is effectful table functions, which the engine's +/// `has_side_effects` metadata does not cover because it is NULL for every table +/// function, so nothing but a list can catch them. +pub const ALWAYS_DENIED_FUNCTIONS: &[&str] = &[ + // Executes SQL assembled at runtime. + "query", + "query_table", + "json_execute_serialized_sql", + // Effectful table functions the engine cannot flag for us. + "enable_logging", + "disable_logging", + "truncate_duckdb_logs", + "checkpoint", + "force_checkpoint", + // Discloses vended credentials. + "duckdb_secrets", +]; + +/// Table functions a catalog query may use. Anything else in `FROM` position is +/// refused. /// -/// This deliberately excludes the `read_*`/`glob`/`*_scan` family: those are how a -/// catalog query legitimately reads Parquet and Iceberg over object storage. -/// Reaching the *local* filesystem through them is refused by the session -/// hardening instead. -pub const MUTATING_OR_DYNAMIC_FUNCTIONS: &[&str] = &["nextval", "currval", "query", "query_table"]; +/// This is the fail-closed half of the gate and the reason it does not rely on +/// naming every dangerous function. A catalog query reads base tables in the +/// attached catalog, which are not functions at all, so the legitimate +/// table-function surface is tiny. The `read_*`/`glob`/`*_scan` family is +/// deliberately absent: those are the file-read and network-egress vectors, and +/// they are not needed to query an attached catalog. A caller that genuinely +/// needs them can widen the set through [`classify_read_only_with`]. +pub const ALLOWED_TABLE_FUNCTIONS: &[&str] = &[ + "range", + "generate_series", + "unnest", + "iceberg_scan", + "iceberg_metadata", + "iceberg_snapshots", +]; + +/// What a caller will and will not allow in untrusted SQL. +#[derive(Debug, Clone)] +pub struct FunctionPolicy<'a> { + /// Refused wherever they appear. + pub always_denied: &'a [&'a str], + /// The only table functions permitted in `FROM` position. + pub allowed_table_functions: &'a [&'a str], + /// Ask the engine which scalar functions have side effects and refuse those. + /// Engine-sourced, so it keeps up with DuckDB rather than drifting. + pub deny_side_effecting_scalars: bool, +} + +impl Default for FunctionPolicy<'static> { + fn default() -> Self { + Self { + always_denied: ALWAYS_DENIED_FUNCTIONS, + allowed_table_functions: ALLOWED_TABLE_FUNCTIONS, + deny_side_effecting_scalars: true, + } + } +} /// The verdict for a piece of untrusted SQL. #[derive(Debug, PartialEq, Eq, Clone)] @@ -97,17 +149,16 @@ pub fn classify_read_only_on( conn: &duckdb::Connection, sql: &str, ) -> Result { - classify_read_only_with(conn, sql, MUTATING_OR_DYNAMIC_FUNCTIONS) + classify_read_only_with(conn, sql, &FunctionPolicy::default()) } -/// [`classify_read_only_on`] with an explicit set of denied function names, for a -/// caller that wants to refuse more than the default (an agent path with no -/// business reading object storage might also deny the `read_*` family). Names are -/// matched case-insensitively against the parsed tree, never the raw SQL text. +/// [`classify_read_only_on`] under an explicit [`FunctionPolicy`], for a caller +/// that needs to widen or narrow what functions are permitted. Names are matched +/// case-insensitively against the parsed tree, never the raw SQL text. pub fn classify_read_only_with( conn: &duckdb::Connection, sql: &str, - denied_functions: &[&str], + policy: &FunctionPolicy<'_>, ) -> 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 @@ -145,39 +196,116 @@ pub fn classify_read_only_with( _ => return Ok(ReadOnlyCheck::Multiple), } - if let Some(name) = find_denied_function(&parsed, denied_functions) { - return Ok(ReadOnlyCheck::DeniedFunction(name)); + let calls = collect_function_calls(&parsed); + + // 1. Names refused wherever they appear. + if let Some(name) = calls + .all + .iter() + .find(|name| policy.always_denied.iter().any(|d| d.eq_ignore_ascii_case(name))) + { + return Ok(ReadOnlyCheck::DeniedFunction(name.clone())); } + + // 2. Table functions are allowlisted, so an unknown one is refused rather + // than waved through. This is what keeps the gate fail-closed as DuckDB + // and its extensions add functions we have never heard of. + if let Some(name) = calls.table_functions.iter().find(|name| { + !policy + .allowed_table_functions + .iter() + .any(|a| a.eq_ignore_ascii_case(name)) + }) { + return Ok(ReadOnlyCheck::DeniedFunction(name.clone())); + } + + // 3. Scalars are too numerous to allowlist, so ask the engine which ones have + // side effects. `has_side_effects` is NULL for table functions, which is + // why those are handled above instead. + if policy.deny_side_effecting_scalars && !calls.all.is_empty() { + let effectful = side_effecting_functions(conn)?; + if let Some(name) = calls.all.iter().find(|name| effectful.contains(*name)) { + return Ok(ReadOnlyCheck::DeniedFunction(name.clone())); + } + } + Ok(ReadOnlyCheck::Allowed) } -/// Walks the serialized parse tree for a call to any denied function, scalar or -/// table-valued, at any depth. +/// Function names used by a statement, split by the position they appear in. +#[derive(Default)] +struct FunctionCalls { + /// Every function name in the statement, lowercased. + all: std::collections::BTreeSet, + /// Just the ones in `FROM` position. + table_functions: std::collections::BTreeSet, +} + +/// Walks the serialized parse tree collecting function names, noting which sit in +/// `FROM` position. /// -/// This reads `function_name` values out of the tree rather than pattern-matching -/// JSON text, because the serialized shape is not stable across DuckDB versions -/// and a substring match on raw SQL would be defeated by comments and quoting. If -/// a future DuckDB renames that key this walk silently stops matching, so each -/// denial is covered by a test that fails loudly when it stops being refused. -fn find_denied_function(node: &serde_json::Value, denied: &[&str]) -> Option { +/// This reads typed nodes rather than pattern-matching JSON text, because the +/// serialized shape is not stable across DuckDB versions and a substring match on +/// raw SQL would be defeated by comments and quoting. A table function appears as +/// a `from_table` node of type `TABLE_FUNCTION` carrying a nested `function`; a +/// scalar is any other node with a `function_name`. If a future DuckDB renames +/// those keys this walk quietly stops matching, so every denial below is covered +/// by a test that fails loudly the moment it stops being refused. +fn collect_function_calls(node: &serde_json::Value) -> FunctionCalls { + let mut calls = FunctionCalls::default(); + walk_function_calls(node, &mut calls); + calls +} + +fn walk_function_calls(node: &serde_json::Value, calls: &mut FunctionCalls) { match node { serde_json::Value::Object(fields) => { - if let Some(serde_json::Value::String(name)) = fields.get("function_name") { - if denied.iter().any(|d| d.eq_ignore_ascii_case(name)) { - return Some(name.to_lowercase()); + if fields.get("type").and_then(serde_json::Value::as_str) == Some("TABLE_FUNCTION") { + if let Some(name) = fields + .get("function") + .and_then(|f| f.get("function_name")) + .and_then(serde_json::Value::as_str) + { + calls.table_functions.insert(name.to_lowercase()); } } - fields - .values() - .find_map(|value| find_denied_function(value, denied)) + if let Some(name) = fields.get("function_name").and_then(serde_json::Value::as_str) { + calls.all.insert(name.to_lowercase()); + } + for value in fields.values() { + walk_function_calls(value, calls); + } } - serde_json::Value::Array(items) => items - .iter() - .find_map(|item| find_denied_function(item, denied)), - _ => None, + serde_json::Value::Array(items) => { + for item in items { + walk_function_calls(item, calls); + } + } + _ => {} } } +/// The scalar functions DuckDB reports as having side effects, lowercased. +/// +/// Sourcing this from the engine rather than a hand-kept list means a function +/// added by a future DuckDB is refused without us noticing it exists. It only +/// covers functions registered on `conn`, so a caller classifying against a +/// throwaway connection will not see extension-provided scalars; the +/// always-denied list carries the ones that matters for. +fn side_effecting_functions( + conn: &duckdb::Connection, +) -> Result, duckdb::Error> { + let mut stmt = conn.prepare( + "SELECT DISTINCT lower(function_name) FROM duckdb_functions() WHERE has_side_effects", + )?; + let mut names = std::collections::BTreeSet::new(); + let mut rows = stmt.query([])?; + while let Some(row) = rows.next()? { + names.insert(row.get::<_, String>(0)?); + } + Ok(names) +} + #[cfg(test)] mod tests { use super::{classify_read_only, classify_read_only_on, ReadOnlyCheck}; @@ -442,17 +570,100 @@ mod tests { } #[test] - fn table_functions_reaching_the_host_are_a_matter_for_the_hardening_not_the_gate() { - // The gate only classifies statement shape. A SELECT that reads a local - // file or a URL is a valid SELECT and is Allowed here; the session - // hardening (disabled_filesystems, no external access) is what refuses the - // actual read at execution time. This test documents that boundary so a - // future reader does not mistake the gate for the filesystem defense. + fn table_functions_that_reach_a_file_or_a_url_are_refused_by_the_allowlist() { + // These used to be Allowed here, on the reasoning that the session + // hardening would refuse the read at execution time. That left the gate + // depending on a filesystem setting to cover a network vector, so the + // table-function position is now allowlisted and these are refused + // outright, before anything runs. + for (sql, func) in [ + ("SELECT * FROM read_csv('/etc/passwd')", "read_csv"), + ( + "SELECT * FROM read_parquet('https://example.com/x.parquet')", + "read_parquet", + ), + ("SELECT * FROM read_text('/etc/hostname')", "read_text"), + ("SELECT * FROM read_blob('/etc/hostname')", "read_blob"), + ("SELECT * FROM glob('/etc/*')", "glob"), + ] { + assert_eq!( + check(sql), + ReadOnlyCheck::DeniedFunction(func.to_string()), + "should refuse {func} in: {sql}" + ); + } + } + + #[test] + fn the_table_function_allowlist_still_permits_catalog_reads() { + // The point of the allowlist is to refuse the unknown without getting in + // the way of an ordinary catalog query. Base tables are not functions at + // all, so the common case is untouched. for sql in [ - "SELECT * FROM read_csv('/etc/passwd')", - "SELECT * FROM read_parquet('https://example.com/x.parquet')", + "SELECT * FROM runs", + "SELECT * FROM \"default\".bronze.runs WHERE id > 0", + "SELECT count(*) FROM range(10)", + "SELECT * FROM generate_series(1, 5)", + "WITH x AS (SELECT * FROM runs) SELECT count(*) FROM x", + ] { + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "should allow: {sql}"); + } + } + + #[test] + fn an_unknown_table_function_is_refused_rather_than_waved_through() { + // The fail-closed property: a function nobody has heard of, which a future + // DuckDB or extension might add, is refused because it is not on the + // allowlist. A denylist would have admitted it. + assert_eq!( + check("SELECT * FROM some_brand_new_table_function('x')"), + ReadOnlyCheck::DeniedFunction("some_brand_new_table_function".to_string()) + ); + } + + #[test] + fn the_json_executor_bundled_for_the_parser_cannot_be_used_as_a_bypass() { + // This crate enables the `json` extension so the gate can parse SQL, and + // that same extension ships `json_execute_serialized_sql`, which runs + // whatever it is handed. The outer statement is a perfectly ordinary + // SELECT, so only an explicit denial catches it. + for sql in [ + "SELECT * FROM json_execute_serialized_sql(json_serialize_sql('SELECT 1'))", + "SELECT json_execute_serialized_sql(json_serialize_sql('SELECT 1'))", + ] { + assert_eq!( + check(sql), + ReadOnlyCheck::DeniedFunction("json_execute_serialized_sql".to_string()), + "the json executor must not be reachable: {sql}" + ); + } + } + + #[test] + fn effectful_functions_are_refused_whether_or_not_the_engine_flags_them() { + // `has_side_effects` is true for these scalars, so the engine tells us. + for (sql, func) in [ + ("SELECT nextval('s')", "nextval"), + ("SELECT currval('s')", "currval"), + ("SELECT setseed(0.5)", "setseed"), ] { - assert_eq!(check(sql), ReadOnlyCheck::Allowed, "gate classifies shape only: {sql}"); + assert_eq!( + check(sql), + ReadOnlyCheck::DeniedFunction(func.to_string()), + "engine-flagged effectful scalar should be refused: {sql}" + ); + } + // `has_side_effects` is NULL for every table function, so the engine + // cannot help here and the explicit list has to. + for (sql, func) in [ + ("SELECT * FROM enable_logging()", "enable_logging"), + ("SELECT * FROM duckdb_secrets()", "duckdb_secrets"), + ] { + assert_eq!( + check(sql), + ReadOnlyCheck::DeniedFunction(func.to_string()), + "effectful table function should be refused: {sql}" + ); } } diff --git a/crates/tower-duckdb/src/lib.rs b/crates/tower-duckdb/src/lib.rs index 78531b6e..00960cd1 100644 --- a/crates/tower-duckdb/src/lib.rs +++ b/crates/tower-duckdb/src/lib.rs @@ -126,6 +126,21 @@ pub struct Hardening { } impl Hardening { + /// The lockdown for untrusted (agent-issued) SQL. + /// + /// Adds an engine memory ceiling on top of the default lockdown. [`Limits`] + /// bounds what a result hands back, but a query can spend far more than that + /// inside DuckDB before the first row appears (a large `ORDER BY`, or a + /// `string_agg` over a whole column). Only the engine can bound that, so an + /// untrusted caller should set it. + pub fn agent() -> Self { + Self { + deny_external_access: false, + memory_limit: Some("1GB".to_string()), + max_temp_directory_size: Some("2GB".to_string()), + } + } + /// The lockdown statements, in the order they must run. /// /// `lock_configuration` is last because it freezes every later `SET`. The @@ -262,14 +277,25 @@ impl Session { let value: duckdb::types::Value = row.get(idx)?; record.push(value_to_json(value)); } - total_bytes += record.iter().map(json_size) .sum::(); - rows.push(record); - // Checked after pushing, so a single oversized row is returned - // rather than silently yielding an empty result. - if limits.max_total_bytes.is_some_and(|max| total_bytes >= max) { + // Measured before the row is kept, so a row that would blow the + // budget is discarded rather than returned and merely labelled + // truncated. A single `string_agg` can carry an entire table in + // one row, so admitting it and flagging it would leave the + // ceiling doing nothing at all. A first row that is already over + // budget yields an empty, truncated result, which is the honest + // answer. Note this bounds what the caller is handed, not what + // the engine allocated to produce it; that needs + // `Hardening::memory_limit`. + let record_bytes = record.iter().map(json_size).sum::(); + if limits + .max_total_bytes + .is_some_and(|max| total_bytes + record_bytes > max) + { truncated = Some(Truncation::Bytes); break; } + total_bytes += record_bytes; + rows.push(record); } } @@ -602,6 +628,48 @@ mod tests { Some(Truncation::Bytes), "the byte ceiling should have cut this short" ); + // The ceiling has to *withhold* the oversized row, not hand it over with a + // label on it. Returning it and calling it truncated would leave the bound + // doing nothing. + let returned: usize = byte_capped + .rows + .iter() + .flat_map(|row| row.iter()) + .map(|value| value.as_str().map_or(0, str::len)) + .sum(); + assert!( + returned <= 4096, + "budget was 4096 bytes but {returned} bytes came back" + ); + } + + /// A single row larger than the whole budget is refused outright, leaving an + /// empty truncated result. That is deliberately blunt: the alternative is + /// admitting an arbitrarily large row, which is what the ceiling exists to + /// prevent. + #[test] + fn one_row_bigger_than_the_whole_budget_is_withheld() { + let result = run_query( + &[], + "SELECT repeat('x', 5000000) AS pad", + [], + &Limits { + max_rows: None, + max_total_bytes: Some(1024 * 1024), + timeout: None, + }, + ) + .expect("query should succeed"); + + assert_eq!(result.truncated, Some(Truncation::Bytes)); + assert!( + result.rows.is_empty(), + "a row over the whole budget must not be returned, got {} row(s)", + result.rows.len() + ); + // Columns are still reported so the caller can see the shape of what it + // asked for. + assert_eq!(result.columns, vec!["pad"]); } /// Many small rows trip the byte ceiling too, so the bound holds however the @@ -779,8 +847,22 @@ mod tests { #[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)"); + // The iceberg extension is fetched at runtime, so a machine that cannot + // install it (no writable extension directory, no network) is a legitimate + // skip. Anything else is a real failure: matching on the install error + // keeps a broken build from turning this into a silent pass. + if let Err(err) = conn.execute_batch("INSTALL iceberg; LOAD iceberg;") { + let message = err.to_string().to_lowercase(); + let unavailable = message.contains("install") + || message.contains("download") + || message.contains("network") + || message.contains("access is denied") + || message.contains("no such file"); + assert!( + unavailable, + "iceberg failed for a reason other than being unavailable: {err}" + ); + eprintln!("skipping iceberg_scan test (iceberg extension unavailable): {err}"); return; } @@ -937,15 +1019,90 @@ mod tests { } } + /// Proves no request leaves the process, by counting connections to a listener + /// we control rather than trusting an error message. + /// + /// The earlier version of this test ran through a helper that never loaded + /// `httpfs`, so it passed for a reason production does not share: with no HTTP + /// filesystem registered there was nothing to block. This one loads the same + /// extensions `attach_statements` does before hardening, so it fails if the + /// sandbox ever stops covering the real configuration. #[test] - fn sandbox_blocks_network_ssrf_via_table_functions() { + fn no_query_can_reach_the_network_under_production_hardening() { + use std::io::{Read, Write}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + let connections = Arc::new(AtomicUsize::new(0)); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind probe listener"); + let port = listener.local_addr().expect("listener addr").port(); + let counter = connections.clone(); + std::thread::spawn(move || { + for stream in listener.incoming() { + counter.fetch_add(1, Ordering::SeqCst); + if let Ok(mut stream) = stream { + let mut buf = [0u8; 1024]; + let _ = stream.read(&mut buf); + let body = "col\nreached\n"; + let _ = stream.write_all( + format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nAccept-Ranges: bytes\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ) + .as_bytes(), + ); + } + } + }); + + // A control run proves the listener and the URL are reachable at all, so a + // later "no connections" result means the sandbox stopped it rather than + // the probe being broken. + let control = duckdb::Connection::open_in_memory().expect("open duckdb"); + let extensions_available = control.execute_batch("INSTALL httpfs; LOAD httpfs;").is_ok(); + if !extensions_available { + eprintln!("skipping SSRF test: httpfs extension unavailable"); + return; + } + let url = format!("http://127.0.0.1:{port}/probe.csv"); + let reached_without_hardening = control + .query_row( + &format!("SELECT count(*) FROM read_csv('{url}')"), + [], + |row| row.get::<_, i64>(0), + ) + .is_ok(); + assert!( + reached_without_hardening && connections.load(Ordering::SeqCst) > 0, + "control failed: the probe listener was never reached even unhardened, so this test would pass vacuously" + ); + + let baseline = connections.load(Ordering::SeqCst); 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')", + "SELECT count(*) FROM read_csv('URL')", + "SELECT count(*) FROM read_text('URL')", + "SELECT count(*) FROM read_blob('URL')", + "SELECT count(*) FROM read_json_auto('URL')", + "SELECT count(*) FROM 'URL'", + "SELECT count(*) FROM read_csv('http://169.254.169.254/latest/meta-data/')", ] { - assert!(run_hardened(sql).is_err(), "SSRF NOT blocked: {sql}"); + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + // Mirror production: extensions loaded first, then the lockdown. + conn.execute_batch("INSTALL httpfs; LOAD httpfs;") + .expect("load httpfs"); + let _ = conn.execute_batch("INSTALL iceberg; LOAD iceberg;"); + for statement in hardening_statements() { + conn.execute_batch(&statement).expect("harden session"); + } + let _ = conn.query_row(&sql.replace("URL", &url), [], |row| row.get::<_, i64>(0)); } + + assert_eq!( + connections.load(Ordering::SeqCst), + baseline, + "a hardened query reached the network" + ); } /// Polls MinIO's health endpoint over a raw socket until it answers 200. @@ -969,20 +1126,49 @@ mod tests { false } + /// MinIO pinned to an immutable release tag rather than `latest`. + /// + /// A floating tag lets a new upstream image change this test's behaviour, or + /// break it into a skip, without anyone choosing that. A security regression + /// test should only change when someone means it. Bump deliberately and re-run. + const MINIO_IMAGE_TAG: &str = "RELEASE.2025-09-07T16-13-09Z"; + + /// Whether a container runtime is actually reachable. + /// + /// Used to tell "no Docker on this machine", which is a legitimate skip, from + /// "the container failed to start", which is a failure worth surfacing. + /// Treating both as a skip is how a security test quietly becomes a no-op. + fn docker_is_available() -> bool { + std::process::Command::new("docker") + .args(["info", "--format", "{{.ServerVersion}}"]) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .map(|status| status.success()) + .unwrap_or(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. + /// configuration changes stay blocked in the same session. + /// + /// Skips only when no container runtime is reachable. If Docker is present and + /// the container will not start, that is a failure, not a skip, so a broken + /// image or a registry problem cannot turn this into a silent pass. #[test] fn sandbox_holds_over_object_store_iceberg() { use testcontainers::core::{IntoContainerPort, WaitFor}; use testcontainers::runners::SyncRunner; use testcontainers::{GenericImage, ImageExt}; + if !docker_is_available() { + eprintln!("skipping sandbox_holds_over_object_store_iceberg: no container runtime"); + return; + } + let bucket = "warehouse"; - let image = GenericImage::new("minio/minio", "latest") + let image = GenericImage::new("minio/minio", MINIO_IMAGE_TAG) .with_wait_for(WaitFor::seconds(1)) .with_exposed_port(9000.tcp()) .with_entrypoint("sh") @@ -993,15 +1179,12 @@ mod tests { 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; - } - }; + // Docker answered `info` above, so a failure here is a real problem (a bad + // image, a registry outage, a broken container config) and must fail the + // test rather than quietly skip it. + let container = image + .start() + .expect("MinIO container failed to start despite a reachable Docker daemon"); let port = container .get_host_port_ipv4(9000.tcp()) .expect("mapped MinIO port"); From e337d0895312679e8f912a635b21e2008d27e374 Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Mon, 27 Jul 2026 11:19:16 +0100 Subject: [PATCH 9/9] fix(tower-duckdb): close the replacement-scan SSRF path CI found on Linux MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI caught what my local testing could not. The SSRF I reported as non-reproducible does reproduce on Linux: three hardened queries reached the probe listener there. The block I measured on macOS was incidental, disabling LocalFileSystem breaks HTTP path resolution before the HTTP filesystem is consulted, and Linux does not do that. Relying on it was the mistake; the reviewer was right that the hardening is not an egress control. Two holes are closed at the gate, which is platform-independent. `SELECT * FROM 'http://…'` is a replacement scan: it parses as an ordinary base table with the URL as its name, so the table-function allowlist never saw it. Base-table references that name a file or a URL are now refused outright, with their own verdict so the error can say what is actually wrong. `iceberg_scan` is out of the allowlist. It takes a location, so allowing it left a way to reach internal services even with the rest of the family refused. Every remaining entry is a pure generator that accepts no path and no URL, which is the property to preserve: adding anything there that takes a location reopens this. The test now mirrors production by running the gate first and executing only what it allows, because asserting that hardening alone blocks the network was asserting something this crate's own documentation says is untrue. A companion test records that hardening does not close egress, so the next person does not rediscover it in CI. Separately, the MinIO test failed on Windows because MinIO publishes Linux images only, and my previous change had turned that into a hard failure. It now skips when no Linux container runtime is available, decided from the platform rather than inferred from a pull error, and still fails loudly for anything else. --- crates/tower-cmd/src/catalogs.rs | 5 +- crates/tower-duckdb/src/guard.rs | 110 +++++++++++++++++++++++++++---- crates/tower-duckdb/src/lib.rs | 103 ++++++++++++++++++++++++----- 3 files changed, 190 insertions(+), 28 deletions(-) diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index 8b51b160..b0bd8799 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -493,7 +493,10 @@ pub async fn do_query(out: &output::Out, config: Config, args: &ArgMatches) { "This command runs read-only queries. Only a single SELECT statement is allowed; re-run with --write to modify the catalog.", ), Ok(Ok(guard::ReadOnlyCheck::DeniedFunction(name))) => out.die(&format!( - "'{name}' is not allowed in a read-only query: it changes engine state or runs SQL built at runtime. Remove it, or re-run with --write." + "'{name}' is not allowed in a read-only query: it changes engine state, runs SQL built at runtime, or reads outside the catalog. Remove it, or re-run with --write." + )), + Ok(Ok(guard::ReadOnlyCheck::DeniedTableReference(reference))) => out.die(&format!( + "'{reference}' is a file or URL, not a table in this catalog. Read-only queries can only read the catalog's own tables; re-run with --write to read elsewhere." )), Ok(Err(err)) => out.die(&format!("Could not validate the query: {err}")), Err(err) => out.die(&format!("Could not validate the query: {err}")), diff --git a/crates/tower-duckdb/src/guard.rs b/crates/tower-duckdb/src/guard.rs index 43ea9dff..8df75a34 100644 --- a/crates/tower-duckdb/src/guard.rs +++ b/crates/tower-duckdb/src/guard.rs @@ -75,18 +75,17 @@ pub const ALWAYS_DENIED_FUNCTIONS: &[&str] = &[ /// This is the fail-closed half of the gate and the reason it does not rely on /// naming every dangerous function. A catalog query reads base tables in the /// attached catalog, which are not functions at all, so the legitimate -/// table-function surface is tiny. The `read_*`/`glob`/`*_scan` family is -/// deliberately absent: those are the file-read and network-egress vectors, and -/// they are not needed to query an attached catalog. A caller that genuinely -/// needs them can widen the set through [`classify_read_only_with`]. -pub const ALLOWED_TABLE_FUNCTIONS: &[&str] = &[ - "range", - "generate_series", - "unnest", - "iceberg_scan", - "iceberg_metadata", - "iceberg_snapshots", -]; +/// table-function surface is tiny. +/// +/// Every entry here is a pure generator that takes no path and no URL. That is +/// the property worth preserving: the session hardening does not close network +/// egress, and on Linux it does not incidentally block it either, so any table +/// function that accepts a location is a way to reach internal services. That is +/// why the `read_*`/`glob` family is absent, and why `iceberg_scan` is too even +/// though it reads the same data an attached catalog does. **Adding anything here +/// that accepts a path or URL reopens SSRF.** A caller that needs one can widen +/// the set through [`classify_read_only_with`] and take that on knowingly. +pub const ALLOWED_TABLE_FUNCTIONS: &[&str] = &["range", "generate_series", "unnest"]; /// What a caller will and will not allow in untrusted SQL. #[derive(Debug, Clone)] @@ -129,6 +128,11 @@ pub enum ReadOnlyCheck { /// Parses as a `SELECT` but calls a function that mutates state or executes /// dynamically-built SQL. Carries the offending name for the error message. DeniedFunction(String), + /// Names a file or URL where a table belongs, which DuckDB resolves through a + /// replacement scan. `SELECT * FROM 'http://…'` reads as an ordinary base + /// table in the parse tree, so nothing about the function allowlist catches + /// it. Carries the offending reference. + DeniedTableReference(String), } /// Classify untrusted `sql` for the read-only path using DuckDB's parser. Opens a @@ -198,6 +202,18 @@ pub fn classify_read_only_with( let calls = collect_function_calls(&parsed); + // 0. A table reference that names a file or a URL. DuckDB turns these into a + // replacement scan, so `SELECT * FROM 'http://internal/x.csv'` reaches the + // network while parsing as a plain base table. The function allowlist never + // sees it, which is why this is checked separately. + if let Some(reference) = calls + .table_references + .iter() + .find(|name| looks_like_a_location(name)) + { + return Ok(ReadOnlyCheck::DeniedTableReference(reference.clone())); + } + // 1. Names refused wherever they appear. if let Some(name) = calls .all @@ -239,6 +255,29 @@ struct FunctionCalls { all: std::collections::BTreeSet, /// Just the ones in `FROM` position. table_functions: std::collections::BTreeSet, + /// Base-table names, which is where a replacement scan hides a file or URL. + /// Kept with original case, since these are paths rather than identifiers. + table_references: std::collections::BTreeSet, +} + +/// Whether a table reference is really a file or URL, and so a replacement scan +/// rather than a table. +/// +/// A DuckDB identifier does not contain a path separator or a scheme, so those +/// are decisive. The extension check catches a bare relative name like +/// `data.csv`. Erring toward refusal is right here: the cost of rejecting an +/// oddly-named table is a clear error, and the cost of missing one is egress. +fn looks_like_a_location(reference: &str) -> bool { + let lowered = reference.to_lowercase(); + if lowered.contains("://") || lowered.contains('/') || lowered.contains('\\') { + return true; + } + [ + ".csv", ".tsv", ".parquet", ".json", ".ndjson", ".jsonl", ".avro", ".arrow", ".xlsx", + ".gz", ".zst", + ] + .iter() + .any(|extension| lowered.ends_with(extension)) } /// Walks the serialized parse tree collecting function names, noting which sit in @@ -269,6 +308,14 @@ fn walk_function_calls(node: &serde_json::Value, calls: &mut FunctionCalls) { calls.table_functions.insert(name.to_lowercase()); } } + if fields.get("type").and_then(serde_json::Value::as_str) == Some("BASE_TABLE") { + if let Some(name) = fields + .get("table_name") + .and_then(serde_json::Value::as_str) + { + calls.table_references.insert(name.to_string()); + } + } if let Some(name) = fields.get("function_name").and_then(serde_json::Value::as_str) { calls.all.insert(name.to_lowercase()); } @@ -610,6 +657,45 @@ mod tests { } } + #[test] + fn a_replacement_scan_naming_a_file_or_url_is_refused() { + // `FROM ''` parses as an ordinary base table, so the function + // allowlist never sees it. Missing this was a live SSRF path: it reached + // the network on Linux, where disabling LocalFileSystem does not + // incidentally block HTTP the way it does on macOS. + for reference in [ + "http://169.254.169.254/latest/meta-data/", + "https://attacker.example/x.parquet", + "s3://bucket/object.parquet", + "/etc/passwd", + "./local.csv", + "data.parquet", + "C:\\Windows\\System32\\drivers\\etc\\hosts", + ] { + let sql = format!("SELECT * FROM '{reference}'"); + assert_eq!( + check(&sql), + ReadOnlyCheck::DeniedTableReference(reference.to_string()), + "replacement scan not refused: {sql}" + ); + } + } + + #[test] + fn ordinary_table_names_are_not_mistaken_for_locations() { + // The location check must not get in the way of real tables, including + // qualified ones, where the catalog and schema are separate fields. + for sql in [ + "SELECT * FROM runs", + "SELECT * FROM \"default\".bronze.runs", + "SELECT * FROM my_schema.events_2026", + "SELECT * FROM \"weird name\"", + "WITH staging AS (SELECT 1) SELECT * FROM staging", + ] { + assert_eq!(check(sql), ReadOnlyCheck::Allowed, "should allow: {sql}"); + } + } + #[test] fn an_unknown_table_function_is_refused_rather_than_waved_through() { // The fail-closed property: a function nobody has heard of, which a future diff --git a/crates/tower-duckdb/src/lib.rs b/crates/tower-duckdb/src/lib.rs index 00960cd1..57ab1b11 100644 --- a/crates/tower-duckdb/src/lib.rs +++ b/crates/tower-duckdb/src/lib.rs @@ -1085,23 +1085,85 @@ mod tests { "SELECT count(*) FROM read_blob('URL')", "SELECT count(*) FROM read_json_auto('URL')", "SELECT count(*) FROM 'URL'", + "SELECT count(*) FROM iceberg_scan('URL')", "SELECT count(*) FROM read_csv('http://169.254.169.254/latest/meta-data/')", ] { + let sql = sql.replace("URL", &url); + + // Production runs the gate first and never executes what it refuses, + // so the test has to as well. Executing regardless would measure the + // hardening alone, which is not a layer that closes egress: see + // `hardening_alone_does_not_close_network_egress`. + let verdict = classify_read_only(&sql).expect("gate should run"); + assert_ne!( + verdict, + ReadOnlyCheck::Allowed, + "the gate must refuse a query that can reach a URL: {sql}" + ); + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); - // Mirror production: extensions loaded first, then the lockdown. conn.execute_batch("INSTALL httpfs; LOAD httpfs;") .expect("load httpfs"); let _ = conn.execute_batch("INSTALL iceberg; LOAD iceberg;"); for statement in hardening_statements() { conn.execute_batch(&statement).expect("harden session"); } - let _ = conn.query_row(&sql.replace("URL", &url), [], |row| row.get::<_, i64>(0)); + if verdict == ReadOnlyCheck::Allowed { + let _ = conn.query_row(&sql, [], |row| row.get::<_, i64>(0)); + } } assert_eq!( connections.load(Ordering::SeqCst), baseline, - "a hardened query reached the network" + "a query that survived the gate reached the network" + ); + } + + /// The layer that is *not* load-bearing here, recorded so nobody assumes it is. + /// + /// The session hardening does not close network egress: it cannot, because an + /// attached Iceberg catalog is made of S3 reads, so `enable_external_access` + /// has to stay on. On macOS it happens to block HTTP reads anyway, because + /// disabling LocalFileSystem breaks path resolution before the HTTP filesystem + /// is consulted. On Linux it does not, and the request goes out. Relying on + /// that accident is exactly the mistake this test exists to prevent: the + /// read-only gate's table-function allowlist is what actually stops these. + #[test] + fn hardening_alone_does_not_close_network_egress() { + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + if conn.execute_batch("INSTALL httpfs; LOAD httpfs;").is_err() { + eprintln!("skipping: httpfs extension unavailable"); + return; + } + for statement in hardening_statements() { + conn.execute_batch(&statement).expect("harden session"); + } + + // Whether this reaches the network is platform-dependent, so assert only + // the part that holds everywhere: hardening does not *refuse* the + // statement the way the gate does. The error, if any, comes from the read + // failing, not from the sandbox declining to try. + let hardened_error = conn + .query_row( + "SELECT count(*) FROM read_csv('http://127.0.0.1:9/nothing')", + [], + |row| row.get::<_, i64>(0), + ) + .err() + .map(|err| err.to_string().to_lowercase()); + if let Some(message) = &hardened_error { + assert!( + !message.contains("not allowed") && !message.contains("denied function"), + "hardening is not the layer that refuses this; the gate is: {message}" + ); + } + + // The gate is. That is the invariant worth depending on. + assert_ne!( + classify_read_only("SELECT count(*) FROM read_csv('http://127.0.0.1:9/nothing')") + .expect("gate should run"), + ReadOnlyCheck::Allowed ); } @@ -1133,18 +1195,27 @@ mod tests { /// test should only change when someone means it. Bump deliberately and re-run. const MINIO_IMAGE_TAG: &str = "RELEASE.2025-09-07T16-13-09Z"; - /// Whether a container runtime is actually reachable. + /// Whether a container runtime that can run this image is reachable. /// - /// Used to tell "no Docker on this machine", which is a legitimate skip, from - /// "the container failed to start", which is a failure worth surfacing. - /// Treating both as a skip is how a security test quietly becomes a no-op. - fn docker_is_available() -> bool { + /// Used to tell an environment that cannot run the test, which is a legitimate + /// skip, from a container that failed to start, which is a failure worth + /// surfacing. Treating both as a skip is how a security test quietly becomes a + /// no-op. + /// + /// Windows is excluded outright: MinIO publishes Linux images only, and + /// Windows CI runs Windows containers, so the pull can never succeed there. + /// That is a property of the image, not a transient problem, so it is decided + /// here rather than inferred from a pull error. + fn can_run_linux_containers() -> bool { + if cfg!(windows) { + return false; + } std::process::Command::new("docker") - .args(["info", "--format", "{{.ServerVersion}}"]) - .stdout(std::process::Stdio::null()) - .stderr(std::process::Stdio::null()) - .status() - .map(|status| status.success()) + .args(["info", "--format", "{{.OSType}}"]) + .output() + .map(|out| { + out.status.success() && String::from_utf8_lossy(&out.stdout).trim() == "linux" + }) .unwrap_or(false) } @@ -1162,8 +1233,10 @@ mod tests { use testcontainers::runners::SyncRunner; use testcontainers::{GenericImage, ImageExt}; - if !docker_is_available() { - eprintln!("skipping sandbox_holds_over_object_store_iceberg: no container runtime"); + if !can_run_linux_containers() { + eprintln!( + "skipping sandbox_holds_over_object_store_iceberg: no Linux container runtime" + ); return; }