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..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" @@ -28,7 +32,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 +72,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..b0bd8799 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, Hardening, Limits, QueryResult, Session}; use tower_telemetry::debug; use crate::{api, beta, output, util::cmd}; @@ -159,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'", @@ -402,10 +410,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], + &Limits::none(), ) }) .await @@ -463,13 +472,71 @@ pub async fn do_query(out: &output::Out, config: Config, args: &ArgMatches) { } let write = cmd::get_bool_flag(args, "write"); - let query_result = execute_catalog_query(out, &config, name, &env, sql, write).await; + // 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(Ok(guard::ReadOnlyCheck::DeniedFunction(name))) => out.die(&format!( + "'{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}")), + } + } + 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. 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 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( out: &output::Out, @@ -478,6 +545,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 @@ -497,7 +565,24 @@ 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 so the + // 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 result = tokio::task::spawn_blocking(move || -> Result { + let session = Session::open()?; + session.run_setup(&setup)?; + if harden { + // `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) + }) + .await; match result { Ok(Ok(query_result)) => { @@ -533,11 +618,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 +651,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 +757,7 @@ async fn fetch_catalog_columns_via_rest( "column_types".to_string(), ], rows, + truncated: None, }) } @@ -999,89 +1024,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 +1045,17 @@ 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())); + 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() + )), + 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())), + } } fn json_value_to_cell(value: &serde_json::Value) -> String { @@ -1375,10 +1327,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, query_limits, + snippets, token_export_command, }; use tower_api::models::{vend_catalog_credentials_body, CatalogCredentials}; + use tower_duckdb::{params, run_query, Limits}; #[test] fn list_defaults_to_default_environment() { @@ -1590,10 +1543,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"], + &Limits::none(), ) .expect("query should succeed"); @@ -1714,76 +1668,84 @@ mod tests { } #[test] - fn duckdb_values_convert_to_json() { - use duckdb::types::{TimeUnit, Value}; + 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); - 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") + 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 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); + 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!( - 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] + default.max_total_bytes, + Some(tower_duckdb::guard::AGENT_MAX_RESULT_BYTES) ); } #[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"]); + 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!( - result.rows[0], - vec![ - serde_json::json!([1, 2]), - serde_json::json!({"a": 1, "b": "x"}), - serde_json::json!({"k": 2}), - ] + 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 run_duckdb_query_reports_columns_for_empty_results() { - let result = - run_duckdb_query(&[], "SELECT 1 AS x WHERE 1 = 0", []).expect("query should succeed"); + 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!(result.columns, vec!["x"]); - assert!(result.rows.is_empty()); + assert_eq!(query_limits(true, Some(25)).max_rows, Some(25)); } #[test] 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..8df75a34 --- /dev/null +++ b/crates/tower-duckdb/src/guard.rs @@ -0,0 +1,775 @@ +//! 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) 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. +/// +/// 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 wherever they appear, whatever the engine says about them. +/// +/// 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 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. +/// +/// 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)] +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)] +pub enum ReadOnlyCheck { + /// 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 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), + /// 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 +/// 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 { + classify_read_only_with(conn, sql, &FunctionPolicy::default()) +} + +/// [`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, + 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 + // 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); + + match statements { + 0 => return Ok(ReadOnlyCheck::Empty), + 1 => {} + _ => return Ok(ReadOnlyCheck::Multiple), + } + + 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 + .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) +} + +/// 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, + /// 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 +/// `FROM` position. +/// +/// 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 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()); + } + } + 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()); + } + for value in fields.values() { + walk_function_calls(value, calls); + } + } + 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}; + + 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 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_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 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 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 + // 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::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}" + ); + } + } + + #[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 new file mode 100644 index 00000000..57ab1b11 --- /dev/null +++ b/crates/tower-duckdb/src/lib.rs @@ -0,0 +1,1328 @@ +//! 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 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; + +use tower_telemetry::debug; + +pub use duckdb::{params, Error, Params}; + +pub mod guard; + +/// 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>, + /// Set when a [`Limits`] ceiling cut the result short. + pub truncated: Option, +} + +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 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 + /// 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 { + Hardening::default().statements() +} + +/// 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, hardening: &Hardening) -> 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. + /// + /// `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, + 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 = None; + let mut total_bytes = 0usize; + + { + 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 limits.max_rows.is_some_and(|max| rows.len() >= max) { + truncated = Some(Truncation::Rows); + 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)); + } + // 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); + } + } + + // A query with no result rows never populates columns above. + if columns.is_empty() { + columns = stmt.column_names(); + } + + debug!( + "tower-duckdb: query took {:?} ({} rows, {} bytes): {}", + query_start.elapsed(), + rows.len(), + total_bytes, + sql + ); + + Ok(QueryResult { + columns, + rows, + truncated, + }) + } +} + +/// 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 +/// query instead. +pub fn run_query( + setup: &[String], + query: &str, + params: P, + limits: &Limits, +) -> Result { + let session = Session::open()?; + session.run_setup(setup)?; + session.query(query, params, limits) +} + +/// 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, Hardening, Limits, QueryResult, Session, + Truncation, + }; + + #[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", + [], + &Limits::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", [], &Limits::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.is_truncated()); + } + + #[test] + fn run_query_reports_columns_for_empty_results() { + let result = + 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()); + } + + #[test] + fn run_query_caps_rows_and_flags_truncation() { + 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_eq!(capped.truncated, Some(Truncation::Rows)); + + 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.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" + ); + // 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 + /// 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 + // 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 harden_allows_selects_but_blocks_local_files_and_config_changes() { + let setup = hardening_statements(); + + 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')", [], &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'", [], &Limits::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"); + // 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; + } + + 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"), + [], + &Limits::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.is_truncated()); + + let capped = run_query( + &setup, + &format!("SELECT id FROM iceberg_scan('{table}') ORDER BY id"), + [], + &Limits::rows(2), + ) + .expect("iceberg_scan should read the table"); + assert_eq!(capped.rows.len(), 2); + assert_eq!(capped.truncated, Some(Truncation::Rows)); + + 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, [], &Limits::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}"); + } + } + + /// 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 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 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 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"); + 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"); + } + if verdict == ReadOnlyCheck::Allowed { + let _ = conn.query_row(&sql, [], |row| row.get::<_, i64>(0)); + } + } + + assert_eq!( + connections.load(Ordering::SeqCst), + baseline, + "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 + ); + } + + /// 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 + } + + /// 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 that can run this image is reachable. + /// + /// 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", "{{.OSType}}"]) + .output() + .map(|out| { + out.status.success() && String::from_utf8_lossy(&out.stdout).trim() == "linux" + }) + .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. + /// + /// 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 !can_run_linux_containers() { + eprintln!( + "skipping sandbox_holds_over_object_store_iceberg: no Linux container runtime" + ); + return; + } + + let bucket = "warehouse"; + let image = GenericImage::new("minio/minio", MINIO_IMAGE_TAG) + .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"), + ]); + + // 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"); + 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"), + [], + &Limits::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')", [], &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'", [], &Limits::none()) + .expect_err("configuration must stay locked"); + let cfg = cfg.to_string().to_lowercase(); + assert!( + cfg.contains("lock") || cfg.contains("configuration"), + "unexpected error: {cfg}" + ); + } +} +