diff --git a/Cargo.lock b/Cargo.lock index b445be24..f728341f 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" @@ -4571,7 +4817,6 @@ dependencies = [ "crypto", "ctrlc", "cucumber", - "duckdb", "futures", "futures-util", "http", @@ -4597,6 +4842,7 @@ dependencies = [ "toml", "toml_edit 0.22.27", "tower-api", + "tower-duckdb", "tower-package", "tower-runtime", "tower-telemetry", @@ -4606,13 +4852,24 @@ 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" 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 +5404,7 @@ checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" dependencies = [ "windows-implement", "windows-interface", - "windows-link", + "windows-link 0.1.3", "windows-result", "windows-strings", ] @@ -5180,13 +5437,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 +5458,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 +5506,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 +5778,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/Cargo.toml b/Cargo.toml index a005bb5a..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" @@ -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 6d5bcf5e..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 } diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index 4dcf6fd7..135de016 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -7,6 +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_telemetry::debug; use crate::{api, beta, output, util::cmd}; @@ -402,10 +403,11 @@ 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, ) }) .await @@ -463,13 +465,39 @@ 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 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 { + 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; 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, @@ -497,7 +525,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, [])).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)) => { @@ -533,11 +574,6 @@ fn read_sql_from_stdin(out: &output::Out) -> String { sql } -struct QueryResult { - columns: Vec, - rows: Vec>, -} - /// 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 @@ -571,62 +607,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. -fn run_duckdb_query( - setup: &[String], - query: &str, - params: P, -) -> 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 result_rows = stmt.query(params)?; - while let Some(row) = result_rows.next()? { - if columns.is_empty() { - columns = row.as_ref().column_names(); - } - 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 }) -} - /// 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 @@ -733,6 +713,7 @@ async fn fetch_catalog_columns_via_rest( "column_types".to_string(), ], rows, + truncated: false, }) } @@ -999,89 +980,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 @@ -1103,7 +1001,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 { @@ -1375,10 +1280,11 @@ fn snippets( #[cfg(test)] mod tests { 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() { @@ -1590,10 +1496,11 @@ 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"); @@ -1713,79 +1620,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", []) - .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", - [], - ) - .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", []).expect("query should succeed"); - - assert_eq!(result.columns, vec!["x"]); - assert!(result.rows.is_empty()); - } - #[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..0769c11a --- /dev/null +++ b/crates/tower-duckdb/src/guard.rs @@ -0,0 +1,89 @@ +//! 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. +//! +//! 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; + +/// 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, +} + +/// 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) +} + +/// [`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 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), + }; + + // `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); + } + + 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 new file mode 100644 index 00000000..0ba6f00f --- /dev/null +++ b/crates/tower-duckdb/src/lib.rs @@ -0,0 +1,739 @@ +//! 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 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; + +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::{classify_read_only, ReadOnlyCheck}; + 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); + } + + // --- 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 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(); + + 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. 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('\'', "''"); + + 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_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_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_ne!(check(sql), ReadOnlyCheck::Allowed, "smuggled statement allowed: {sql}"); + } + assert_eq!( + check("SELECT * FROM runs -- a trailing ; comment"), + ReadOnlyCheck::Allowed + ); + } + + #[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}" + ); + } +} +