From 635780b95a23d6be2c3fc1e82d1cfce7276d4354 Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Mon, 19 Oct 2026 17:43:14 -0400 Subject: [PATCH 1/7] feat(mcp): keep SQL access within explicit column permissions Signed-off-by: Yordis Prieto --- rsworkspace/Cargo.lock | 1211 ++++++++++++++++- .../crates/mcp/mcp-postgres/Cargo.toml | 29 + .../mcp/mcp-postgres/examples/config.toml | 9 + .../mcp-postgres/sql/install-ddl-guard.sql | 25 + .../crates/mcp/mcp-postgres/sql/schema.zed | 20 + .../crates/mcp/mcp-postgres/src/config.rs | 79 ++ .../crates/mcp/mcp-postgres/src/constants.rs | 1 + .../mcp/mcp-postgres/src/database_options.rs | 21 + .../mcp-postgres/src/environment_variable.rs | 21 + .../crates/mcp/mcp-postgres/src/error.rs | 29 + .../crates/mcp/mcp-postgres/src/identifier.rs | 18 + .../crates/mcp/mcp-postgres/src/lib.rs | 45 + .../crates/mcp/mcp-postgres/src/main.rs | 37 + .../crates/mcp/mcp-postgres/src/object_id.rs | 18 + .../crates/mcp/mcp-postgres/src/parameter.rs | 35 + .../crates/mcp/mcp-postgres/src/parser.rs | 293 ++++ .../crates/mcp/mcp-postgres/src/postgres.rs | 383 ++++++ .../crates/mcp/mcp-postgres/src/row_limit.rs | 23 + .../crates/mcp/mcp-postgres/src/server.rs | 94 ++ .../crates/mcp/mcp-postgres/src/spicedb.rs | 106 ++ .../mcp/mcp-postgres/src/spicedb_endpoint.rs | 25 + .../crates/mcp/mcp-postgres/src/telemetry.rs | 71 + .../crates/mcp/mcp-postgres/tests/config.rs | 49 + .../crates/mcp/mcp-postgres/tests/parser.rs | 47 + .../crates/mcp/mcp-postgres/tests/postgres.rs | 683 ++++++++++ .../crates/mcp/mcp-postgres/tests/spicedb.rs | 228 ++++ 26 files changed, 3564 insertions(+), 36 deletions(-) create mode 100644 rsworkspace/crates/mcp/mcp-postgres/Cargo.toml create mode 100644 rsworkspace/crates/mcp/mcp-postgres/examples/config.toml create mode 100644 rsworkspace/crates/mcp/mcp-postgres/sql/install-ddl-guard.sql create mode 100644 rsworkspace/crates/mcp/mcp-postgres/sql/schema.zed create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/config.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/constants.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/database_options.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/environment_variable.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/error.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/identifier.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/lib.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/main.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/object_id.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/parameter.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/parser.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/row_limit.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/server.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/spicedb.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/spicedb_endpoint.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/config.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/parser.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/spicedb.rs diff --git a/rsworkspace/Cargo.lock b/rsworkspace/Cargo.lock index 0f88c493d..b8ac160cb 100644 --- a/rsworkspace/Cargo.lock +++ b/rsworkspace/Cargo.lock @@ -66,6 +66,52 @@ dependencies = [ "wiremock", ] +[[package]] +name = "a2a-gateway" +version = "0.1.0" +dependencies = [ + "a2a-auth-callout", + "a2a-lf", + "a2a-nats", + "a2a-pack", + "a2a-redaction", + "async-nats", + "async-trait", + "axum", + "base64 0.23.1", + "bytes", + "cel-interpreter", + "clap", + "filetime", + "futures", + "jsonwebtoken 10.4.0", + "nkeys", + "p256", + "rand_core 0.6.4", + "reqwest 0.12.28", + "serde", + "serde_json", + "sha2 0.11.0", + "spicedb-grpc-tonic", + "tempfile", + "thiserror 2.0.21", + "time", + "tokio", + "tokio-util", + "toml 1.1.6+spec-1.1.0", + "tonic", + "tracing", + "tracing-subscriber", + "trogon-aauth-person", + "trogon-aauth-sdk", + "trogon-aauth-verify", + "trogon-identity-types", + "trogon-jwks-publisher", + "trogon-nats", + "trogon-std", + "uuid", +] + [[package]] name = "a2a-identity-types" version = "0.1.0" @@ -125,6 +171,71 @@ dependencies = [ "wiremock", ] +[[package]] +name = "a2a-nats-http" +version = "0.1.0" +dependencies = [ + "a2a-identity-types", + "a2a-lf", + "a2a-nats", + "async-compat", + "async-nats", + "async-trait", + "axum", + "base64 0.23.1", + "bytes", + "futures", + "futures-util", + "jsonrpc-nats", + "serde", + "serde_json", + "thiserror 2.0.21", + "tokio", + "tokio-tungstenite", + "tower", + "tower-http 0.7.1", + "tracing", + "tracing-subscriber", + "trogon-nats", + "trogon-std", +] + +[[package]] +name = "a2a-nats-server" +version = "0.1.0" +dependencies = [ + "a2a-lf", + "a2a-nats", + "async-nats", + "async-trait", + "thiserror 2.0.21", + "tokio", + "tokio-util", + "tracing", + "tracing-subscriber", + "trogon-nats", + "trogon-std", +] + +[[package]] +name = "a2a-nats-stdio" +version = "0.1.0" +dependencies = [ + "a2a-lf", + "a2a-nats", + "async-nats", + "bytes", + "futures", + "jsonrpc-nats", + "serde", + "serde_json", + "thiserror 2.0.21", + "tokio", + "tracing", + "trogon-nats", + "trogon-std", +] + [[package]] name = "a2a-pack" version = "0.1.0" @@ -179,6 +290,75 @@ dependencies = [ "uuid", ] +[[package]] +name = "acp-nats-agent" +version = "0.0.1" +dependencies = [ + "acp-nats", + "agent-client-protocol", + "async-nats", + "async-trait", + "bytes", + "futures", + "jsonrpc-nats", + "serde", + "serde_json", + "thiserror 2.0.21", + "tokio", + "tracing", + "tracing-subscriber", + "trogon-nats", + "trogon-std", +] + +[[package]] +name = "acp-nats-server" +version = "0.1.0" +dependencies = [ + "acp-nats", + "agent-client-protocol", + "agent-client-protocol-http", + "anyhow", + "async-nats", + "axum", + "bytes", + "clap", + "futures-util", + "opentelemetry", + "reqwest 0.12.28", + "serde_json", + "thiserror 2.0.21", + "tokio", + "tokio-tungstenite", + "tower", + "tracing", + "tracing-subscriber", + "trogon-nats", + "trogon-std", + "trogon-telemetry", +] + +[[package]] +name = "acp-nats-stdio" +version = "0.0.1" +dependencies = [ + "acp-nats", + "agent-client-protocol", + "anyhow", + "async-compat", + "async-nats", + "bytes", + "clap", + "futures", + "opentelemetry", + "tokio", + "tracing", + "tracing-subscriber", + "trogon-nats", + "trogon-std", + "trogon-telemetry", +] + [[package]] name = "addr2line" version = "0.26.1" @@ -188,6 +368,12 @@ dependencies = [ "gimli", ] +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "aead" version = "0.5.2" @@ -211,7 +397,7 @@ dependencies = [ "blocking", "futures", "futures-concurrency", - "rustc-hash", + "rustc-hash 2.1.2", "rustix 1.1.4", "schemars 1.2.1", "serde", @@ -232,6 +418,23 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "agent-client-protocol-http" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b4d8db045bc66b84526dfe4ef2ffb87d651b81c4ff945f91c63a2bc677a582c" +dependencies = [ + "agent-client-protocol", + "async-stream", + "axum", + "futures", + "serde_json", + "tokio", + "tower-http 0.7.1", + "tracing", + "uuid", +] + [[package]] name = "agent-client-protocol-schema" version = "1.5.0" @@ -271,6 +474,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -336,6 +554,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "antlr4rust" +version = "0.3.0-rc2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d240d49ee89063f90fa0cb18aead41a5893cd544a1785983dc3bf5c3d5faa58b" +dependencies = [ + "better_any", + "bit-set", + "byteorder", + "lazy_static", + "murmur3", + "once_cell", + "parking_lot", + "typed-arena", + "uuid", +] + [[package]] name = "anyhow" version = "1.0.104" @@ -470,7 +705,7 @@ dependencies = [ "futures-core", "libc", "portable-atomic", - "rustc-hash", + "rustc-hash 2.1.2", "rustix 0.38.44", "tokio", "tokio-stream", @@ -502,6 +737,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "async-compression" +version = "0.4.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "515a1f282e33d55983c499d7e9e87082e81cbc32974825bf9032f928392d5844" +dependencies = [ + "compression-codecs", + "compression-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "async-io" version = "2.6.0" @@ -560,7 +807,7 @@ dependencies = [ "tokio-rustls", "tokio-stream", "tokio-util", - "tokio-websockets", + "tokio-websockets 0.10.1", "tracing", "tryhard", "url", @@ -641,6 +888,15 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -683,6 +939,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" dependencies = [ "axum-core", + "axum-macros", "base64 0.22.1", "bytes", "form_urlencoded", @@ -702,7 +959,7 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_urlencoded", - "sha1", + "sha1 0.10.6", "sync_wrapper", "tokio", "tokio-tungstenite", @@ -731,6 +988,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum-macros" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7aa268c23bfbbd2c4363b9cd302a4f504fb2a9dfe7e3451d66f35dd392e20aca" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "base16ct" version = "0.2.0" @@ -755,6 +1023,35 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "better_any" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4372b9543397a4b86050cc5e7ee36953edf4bac9518e8a774c2da694977fb6e4" + +[[package]] +name = "bindgen" +version = "0.66.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2b84e06fc203107bfbad243f4aba2af864eb7db3b1cf46ea0a023b0b433d2a7" +dependencies = [ + "bitflags 2.13.0", + "cexpr", + "clang-sys", + "lazy_static", + "lazycell", + "log", + "peeking_take_while", + "prettyplease 0.2.37", + "proc-macro2", + "quote", + "regex", + "rustc-hash 1.1.0", + "shlex", + "syn 2.0.118", + "which", +] + [[package]] name = "bit-set" version = "0.8.0" @@ -790,6 +1087,9 @@ name = "bitflags" version = "2.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +dependencies = [ + "serde_core", +] [[package]] name = "block-buffer" @@ -873,8 +1173,8 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "85a885520bf6249ab931a764ffdb87b0ceef48e6e7d807cfdb21b751e086e1ad" dependencies = [ - "prost", - "prost-types", + "prost 0.14.3", + "prost-types 0.14.3", "tonic", "tonic-prost", "ureq", @@ -889,7 +1189,7 @@ dependencies = [ "base64 0.22.1", "bollard-buildkit-proto", "bytes", - "prost", + "prost 0.14.3", "serde", "serde_json", "serde_repr", @@ -902,6 +1202,27 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc0b364ead1874514c8c2855ab558056ebfeb775653e7ae45ff72f28f8f3166c" +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "bs58" version = "0.5.1" @@ -964,6 +1285,12 @@ version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.12.1" @@ -991,6 +1318,40 @@ dependencies = [ "shlex", ] +[[package]] +name = "cel-interpreter" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a76c07820046cc8239526fceec6df147a979ae48644dbad274fc3ce38ab0973b" +dependencies = [ + "cel-parser", + "chrono", + "nom", + "paste", + "regex", + "serde", + "thiserror 1.0.69", +] + +[[package]] +name = "cel-parser" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "546fb134998490c5c47fc7a29c7535e725d2e403f172040e8f263d0b318bff5f" +dependencies = [ + "antlr4rust", + "lazy_static", +] + +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -1072,6 +1433,17 @@ dependencies = [ "zeroize", ] +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + [[package]] name = "clap" version = "4.6.7" @@ -1161,6 +1533,26 @@ dependencies = [ "memchr", ] +[[package]] +name = "compression-codecs" +version = "0.4.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe67f2944eef52fc7b106b8c9450d243a88701a0c065f7f57235e76abaed7df" +dependencies = [ + "brotli", + "compression-core", + "flate2", + "memchr", + "zstd", + "zstd-safe", +] + +[[package]] +name = "compression-core" +version = "0.4.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ebdda89edc534bf8d5dd2857562633b6495557e5846552ade994b99ed13c36a" + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -1330,7 +1722,7 @@ dependencies = [ "postcard", "pulley-interpreter", "regalloc2", - "rustc-hash", + "rustc-hash 2.1.2", "serde", "serde_derive", "sha2 0.10.9", @@ -1415,6 +1807,21 @@ version = "0.134.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f6977c2a71ab1e0d1e62f966b411a498aa04c4dce47d93d52f8a360a06058922" +[[package]] +name = "crc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" + [[package]] name = "crc32fast" version = "1.5.0" @@ -1464,6 +1871,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-queue" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03e8bd762f7479489c70ed6c768ddca99d7296857de437a68dcb2a94365b3fae" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -1876,6 +2292,12 @@ dependencies = [ "serde_json", ] +[[package]] +name = "dotenvy" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" + [[package]] name = "dptree" version = "0.5.1" @@ -1943,6 +2365,9 @@ name = "either" version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +dependencies = [ + "serde", +] [[package]] name = "elliptic-curve" @@ -1956,7 +2381,7 @@ dependencies = [ "ff", "generic-array", "group", - "hkdf", + "hkdf 0.12.4", "pem-rfc7468", "pkcs8", "rand_core 0.6.4", @@ -2096,6 +2521,16 @@ version = "0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" +[[package]] +name = "filetime" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -2114,6 +2549,17 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flate2" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + [[package]] name = "fluent-uri" version = "0.4.1" @@ -2125,6 +2571,17 @@ dependencies = [ "serde", ] +[[package]] +name = "flume" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" +dependencies = [ + "futures-core", + "futures-sink", + "spin", +] + [[package]] name = "fnv" version = "1.0.7" @@ -2223,6 +2680,17 @@ dependencies = [ "futures-util", ] +[[package]] +name = "futures-intrusive" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d930c203dd0b6ff06e0201a4a2fe9149b43c684fd4420555b26d21b1a02956f" +dependencies = [ + "futures-core", + "lock_api", + "parking_lot", +] + [[package]] name = "futures-io" version = "0.3.34" @@ -2290,7 +2758,7 @@ checksum = "25234f20a3ec0a962a61770cfe39ecf03cb529a6e474ad8cff025ed497eda557" dependencies = [ "bitflags 2.13.0", "debugid", - "rustc-hash", + "rustc-hash 2.1.2", "serde", "serde_derive", "serde_json", @@ -2438,6 +2906,11 @@ name = "hashbrown" version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] [[package]] name = "hashbrown" @@ -2452,6 +2925,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "hashlink" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824e001ac4f3012dd16a264bec811403a67ca9deb6c102fc5049b32c4574b35f" +dependencies = [ + "hashbrown 0.16.1", +] + [[package]] name = "heck" version = "0.5.0" @@ -2479,6 +2961,15 @@ dependencies = [ "hmac 0.12.1", ] +[[package]] +name = "hkdf" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018" +dependencies = [ + "hmac 0.13.0", +] + [[package]] name = "hmac" version = "0.12.1" @@ -3141,6 +3632,12 @@ dependencies = [ "spin", ] +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + [[package]] name = "leb128fmt" version = "0.1.0" @@ -3153,6 +3650,16 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "libm" version = "0.2.16" @@ -3168,6 +3675,16 @@ dependencies = [ "libc", ] +[[package]] +name = "libsqlite3-sys" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f111c8c41e7c61a49cd34e44c7619462967221a6443b0ec299e0ac30cfb9b1" +dependencies = [ + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.4.15" @@ -3298,6 +3815,39 @@ dependencies = [ "trogon-telemetry", ] +[[package]] +name = "mcp-postgres" +version = "0.0.1" +dependencies = [ + "async-trait", + "futures-util", + "pg_query", + "rmcp", + "serde", + "serde_json", + "spicedb-grpc-tonic", + "sqlx", + "testcontainers-modules", + "thiserror 2.0.21", + "tokio", + "toml 1.1.6+spec-1.1.0", + "tonic", + "tracing", + "tracing-subscriber", + "trogon-std", + "url", +] + +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.2", +] + [[package]] name = "memchr" version = "2.8.0" @@ -3341,6 +3891,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.0" @@ -3378,6 +3938,15 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" +[[package]] +name = "murmur3" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a198f9589efc03f544388dfc4a19fe8af4323662b62f598b8dcfdac62c14771c" +dependencies = [ + "byteorder", +] + [[package]] name = "nats-jwt-rs" version = "0.1.1" @@ -3639,7 +4208,7 @@ dependencies = [ "opentelemetry-http", "opentelemetry-proto", "opentelemetry_sdk", - "prost", + "prost 0.14.3", "reqwest 0.13.4", "serde_json", "thiserror 2.0.21", @@ -3655,7 +4224,7 @@ dependencies = [ "const-hex", "opentelemetry", "opentelemetry_sdk", - "prost", + "prost 0.14.3", "serde", ] @@ -3677,6 +4246,15 @@ dependencies = [ "tokio-stream", ] +[[package]] +name = "ordered-float" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" +dependencies = [ + "num-traits", +] + [[package]] name = "outref" version = "0.5.2" @@ -3761,12 +4339,24 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pastey" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c5a797f0e07bdf071d15742978fc3128ec6c22891c31a3a931513263904c982a" +[[package]] +name = "peeking_take_while" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" + [[package]] name = "pem" version = "3.0.6" @@ -3813,6 +4403,24 @@ dependencies = [ "indexmap 2.14.0", ] +[[package]] +name = "pg_query" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da43168013d137b3ca49b0044c1108ec09ff596a0626d4db89209c0a1d450cd2" +dependencies = [ + "bindgen", + "cc", + "fs_extra", + "glob", + "itertools 0.10.5", + "prost 0.13.5", + "prost-build 0.13.5", + "serde", + "serde_json", + "thiserror 1.0.69", +] + [[package]] name = "phf" version = "0.11.3" @@ -4091,6 +4699,16 @@ dependencies = [ "unarray", ] +[[package]] +name = "prost" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +dependencies = [ + "bytes", + "prost-derive 0.13.5", +] + [[package]] name = "prost" version = "0.14.3" @@ -4098,7 +4716,27 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" dependencies = [ "bytes", - "prost-derive", + "prost-derive 0.14.3", +] + +[[package]] +name = "prost-build" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" +dependencies = [ + "heck", + "itertools 0.14.0", + "log", + "multimap", + "once_cell", + "petgraph 0.6.5", + "prettyplease 0.2.37", + "prost 0.13.5", + "prost-types 0.13.5", + "regex", + "syn 2.0.118", + "tempfile", ] [[package]] @@ -4113,8 +4751,8 @@ dependencies = [ "multimap", "petgraph 0.8.3", "prettyplease 0.2.37", - "prost", - "prost-types", + "prost 0.14.3", + "prost-types 0.14.3", "pulldown-cmark", "pulldown-cmark-to-cmark", "regex", @@ -4122,6 +4760,19 @@ dependencies = [ "tempfile", ] +[[package]] +name = "prost-derive" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn 2.0.118", +] + [[package]] name = "prost-derive" version = "0.14.3" @@ -4135,13 +4786,22 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "prost-types" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" +dependencies = [ + "prost 0.13.5", +] + [[package]] name = "prost-types" version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" dependencies = [ - "prost", + "prost 0.14.3", ] [[package]] @@ -4214,7 +4874,7 @@ dependencies = [ "pin-project-lite", "quinn-proto", "quinn-udp", - "rustc-hash", + "rustc-hash 2.1.2", "rustls", "socket2", "thiserror 2.0.21", @@ -4236,7 +4896,7 @@ dependencies = [ "rand 0.10.1", "rand_pcg", "ring", - "rustc-hash", + "rustc-hash 2.1.2", "rustls", "rustls-pki-types", "slab", @@ -4485,7 +5145,7 @@ dependencies = [ "bumpalo", "hashbrown 0.17.1", "log", - "rustc-hash", + "rustc-hash 2.1.2", "serde", "smallvec", ] @@ -4701,6 +5361,12 @@ version = "0.1.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b50b8869d9fc858ce7266cce0194bd74df58b9d0e3f6df3a9fc8eb470d95c09d" +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + [[package]] name = "rustc-hash" version = "2.1.2" @@ -4989,6 +5655,16 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-value" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3a1a3341211875ef120e117ea7fd5228530ae7e7036a779fdc9117be6b3282c" +dependencies = [ + "ordered-float", + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -5142,6 +5818,17 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.2", +] + [[package]] name = "sha1_smol" version = "1.0.1" @@ -5223,6 +5910,12 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + [[package]] name = "simd_cesu8" version = "1.1.1" @@ -5297,8 +5990,8 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06f14ba9151ca67327f59d4b5d7bbd21dc7f3df6e1da94790f1156b513cec7c1" dependencies = [ - "prost", - "prost-types", + "prost 0.14.3", + "prost-types 0.14.3", "tonic", "tonic-prost", "tonic-prost-build", @@ -5306,19 +5999,199 @@ dependencies = [ ] [[package]] -name = "spin" -version = "0.9.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" - -[[package]] -name = "spki" -version = "0.7.3" +name = "spin" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +dependencies = [ + "lock_api", +] + +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + +[[package]] +name = "sqlx" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "378620ccc25c62c89d8be1c819e76a88d59bdcc3304733330788948e619bfd71" +dependencies = [ + "sqlx-core", + "sqlx-macros", + "sqlx-mysql", + "sqlx-postgres", + "sqlx-sqlite", +] + +[[package]] +name = "sqlx-core" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05b44e85bf579a8eeb4ceaa77a3a523baf2bf0e9bac7e40f405d537b5d2d5ccb" +dependencies = [ + "base64 0.22.1", + "bytes", + "cfg-if", + "chrono", + "crc", + "crossbeam-queue", + "either", + "event-listener", + "futures-core", + "futures-intrusive", + "futures-io", + "futures-util", + "hashbrown 0.16.1", + "hashlink", + "indexmap 2.14.0", + "log", + "memchr", + "percent-encoding", + "rustls", + "serde", + "serde_json", + "sha2 0.10.9", + "smallvec", + "thiserror 2.0.21", + "tokio", + "tokio-stream", + "tracing", + "url", + "webpki-roots 1.0.6", +] + +[[package]] +name = "sqlx-macros" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd2b84f2bc39a5705ef27ec785a11c934a41bbd4a24941e257927cddc26b60bf" +dependencies = [ + "proc-macro2", + "quote", + "sqlx-core", + "sqlx-macros-core", + "syn 2.0.118", +] + +[[package]] +name = "sqlx-macros-core" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb8d96de5fdc85a5c4ec813432b523ec637e80ba98f046555f75f7908ddac7c3" +dependencies = [ + "cfg-if", + "dotenvy", + "either", + "heck", + "hex", + "proc-macro2", + "quote", + "serde", + "serde_json", + "sha2 0.10.9", + "sqlx-core", + "sqlx-mysql", + "sqlx-postgres", + "sqlx-sqlite", + "syn 2.0.118", + "thiserror 2.0.21", + "tokio", + "url", +] + +[[package]] +name = "sqlx-mysql" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90b8020fe17c5f2c245bfa2505d7ef59c5604839527c740266ad2214acebea27" +dependencies = [ + "bitflags 2.13.0", + "byteorder", + "bytes", + "chrono", + "crc", + "digest 0.11.2", + "dotenvy", + "either", + "futures-core", + "futures-util", + "generic-array", + "log", + "percent-encoding", + "serde", + "sha1 0.11.0", + "sha2 0.11.0", + "sqlx-core", + "thiserror 2.0.21", + "tracing", +] + +[[package]] +name = "sqlx-postgres" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87a2bdd6e83f6b3ea525ca9fee568030508b58355a43d0b2c1674d5f79dcd65e" +dependencies = [ + "atoi", + "base64 0.22.1", + "bitflags 2.13.0", + "byteorder", + "chrono", + "crc", + "dotenvy", + "etcetera", + "futures-channel", + "futures-core", + "futures-util", + "hex", + "hkdf 0.13.0", + "hmac 0.13.0", + "itoa", + "log", + "md-5", + "memchr", + "rand 0.10.1", + "serde", + "serde_json", + "sha2 0.11.0", + "smallvec", + "sqlx-core", + "stringprep", + "thiserror 2.0.21", + "tracing", + "whoami", +] + +[[package]] +name = "sqlx-sqlite" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +checksum = "488e99c397a62007e4229aec669a179816339afc6d2620ca6fa420dbee2e982c" dependencies = [ - "base64ct", - "der", + "atoi", + "chrono", + "flume", + "form_urlencoded", + "futures-channel", + "futures-core", + "futures-executor", + "futures-intrusive", + "futures-util", + "libsqlite3-sys", + "log", + "percent-encoding", + "serde", + "sqlx-core", + "thiserror 2.0.21", + "tracing", + "url", ] [[package]] @@ -5353,6 +6226,17 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "stringprep" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4df3d392d81bd458a8a621b8bffbd2302a12ffe288a9d931670948749463b1" +dependencies = [ + "unicode-bidi", + "unicode-normalization", + "unicode-properties", +] + [[package]] name = "strsim" version = "0.11.1" @@ -5730,6 +6614,7 @@ dependencies = [ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", @@ -5777,7 +6662,11 @@ checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c" dependencies = [ "futures-util", "log", + "rustls", + "rustls-native-certs", + "rustls-pki-types", "tokio", + "tokio-rustls", "tungstenite", ] @@ -5816,6 +6705,28 @@ dependencies = [ "webpki-roots 0.26.11", ] +[[package]] +name = "tokio-websockets" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d52efb639344a7c6adb8e62c6f3d2c19c001ff1b79a5041ba1c6ed42e19c6aa5" +dependencies = [ + "base64 0.22.1", + "bytes", + "fastrand", + "futures-core", + "futures-sink", + "http", + "httparse", + "rustls-native-certs", + "rustls-pki-types", + "sha1_smol", + "simdutf8", + "tokio", + "tokio-rustls", + "tokio-util", +] + [[package]] name = "toml" version = "0.9.12+spec-1.1.0" @@ -5898,9 +6809,11 @@ dependencies = [ "hyper-util", "percent-encoding", "pin-project", + "rustls-native-certs", "socket2", "sync_wrapper", "tokio", + "tokio-rustls", "tokio-stream", "tower", "tower-layer", @@ -5927,7 +6840,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" dependencies = [ "bytes", - "prost", + "prost 0.14.3", "tonic", ] @@ -5939,8 +6852,8 @@ checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27" dependencies = [ "prettyplease 0.2.37", "proc-macro2", - "prost-build", - "prost-types", + "prost-build 0.14.3", + "prost-types 0.14.3", "quote", "syn 2.0.118", "tempfile", @@ -5991,12 +6904,16 @@ version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c" dependencies = [ + "async-compression", "bitflags 2.13.0", "bytes", + "futures-core", "http", "http-body", "percent-encoding", "pin-project-lite", + "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", @@ -6105,6 +7022,50 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "trogon-aauth-as" +version = "0.0.1" +dependencies = [ + "async-trait", + "axum", + "base64 0.23.1", + "jsonwebtoken 10.4.0", + "p256", + "pkcs8", + "rand_core 0.6.4", + "serde", + "serde_json", + "thiserror 2.0.21", + "tokio", + "tower", + "trogon-aauth-verify", + "trogon-identity-types", + "uuid", +] + +[[package]] +name = "trogon-aauth-person" +version = "0.0.1" +dependencies = [ + "async-trait", + "axum", + "base64 0.23.1", + "jsonwebtoken 10.4.0", + "p256", + "pkcs8", + "rand_core 0.6.4", + "serde", + "serde_json", + "sha2 0.11.0", + "thiserror 2.0.21", + "tokio", + "tower", + "trogon-aauth-verify", + "trogon-identity-types", + "uuid", + "wiremock", +] + [[package]] name = "trogon-aauth-sdk" version = "0.0.1" @@ -6372,6 +7333,46 @@ dependencies = [ "wit-bindgen 0.62.0", ] +[[package]] +name = "trogon-gateway" +version = "0.1.0" +dependencies = [ + "anyhow", + "async-nats", + "axum", + "base64 0.22.1", + "bytes", + "clap", + "confique", + "form_urlencoded", + "futures-core", + "futures-util", + "hex", + "hmac 0.13.0", + "reqwest 0.12.28", + "rustls", + "serde", + "serde_json", + "sha2 0.11.0", + "subtle", + "tempfile", + "thiserror 2.0.21", + "time", + "tokio", + "tokio-tungstenite", + "tower", + "tracing", + "tracing-subscriber", + "trogon-nats", + "trogon-semconv", + "trogon-service-config", + "trogon-std", + "trogon-telemetry", + "twilight-gateway", + "twilight-model", + "url", +] + [[package]] name = "trogon-identity-types" version = "0.1.0" @@ -6423,6 +7424,43 @@ dependencies = [ "uuid", ] +[[package]] +name = "trogon-scheduler" +version = "0.1.0" +dependencies = [ + "async-nats", + "buffa", + "buffa-types", + "bytes", + "chrono", + "chrono-tz", + "cron", + "futures", + "opentelemetry", + "opentelemetry_sdk", + "proptest", + "rrule", + "serde", + "serde_json", + "sqlx", + "testcontainers-modules", + "thiserror 2.0.21", + "time", + "tokio", + "tracing", + "tracing-opentelemetry", + "trogon-decider", + "trogon-decider-nats", + "trogon-decider-runtime", + "trogon-nats", + "trogon-scheduler-domain", + "trogon-semconv", + "trogon-std", + "trogon-telemetry", + "trogonai-proto", + "uuid", +] + [[package]] name = "trogon-scheduler-domain" version = "0.1.0" @@ -6570,10 +7608,60 @@ dependencies = [ "httparse", "log", "rand 0.9.2", - "sha1", + "rustls", + "rustls-pki-types", + "sha1 0.10.6", "thiserror 2.0.21", ] +[[package]] +name = "twilight-gateway" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59267541c31f888c1587da5e7cbab182ae6efd98faedd9ea2e35a4eef43ff204" +dependencies = [ + "bitflags 2.13.0", + "fastrand", + "futures-core", + "futures-sink", + "serde", + "serde_json", + "tokio", + "tokio-websockets 0.13.3", + "tracing", + "twilight-gateway-queue", + "twilight-model", +] + +[[package]] +name = "twilight-gateway-queue" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "366a73fe47f61a3d522c3aaf70475e60634b0ae59e7b94272ed7496fffa7ceb7" +dependencies = [ + "tokio", + "tracing", +] + +[[package]] +name = "twilight-model" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bf6bb7b93a7f765d89b3388cc710c0ae16104579e06bb30ea1ee6bd41420a8b" +dependencies = [ + "bitflags 2.13.0", + "serde", + "serde-value", + "serde_repr", + "time", +] + +[[package]] +name = "typed-arena" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" + [[package]] name = "typenum" version = "1.20.0" @@ -6592,6 +7680,12 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" +[[package]] +name = "unicode-bidi" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c1cb5db39152898a79168971543b1cb5020dff7fe43c8dc468b0885f5e29df5" + [[package]] name = "unicode-general-category" version = "1.1.0" @@ -6604,6 +7698,21 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization" +version = "0.1.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75cd3fb03d5288b3ba9a1ee103ac7e9cbb59f0b03d64f8ddfcc04dd42758a85d" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-properties" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" + [[package]] name = "unicode-segmentation" version = "1.13.2" @@ -6730,6 +7839,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + [[package]] name = "version_check" version = "0.9.5" @@ -7290,6 +8405,24 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix 0.38.44", +] + +[[package]] +name = "whoami" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c" + [[package]] name = "winapi" version = "0.3.9" @@ -7933,6 +9066,12 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "zlib-rs" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b268e58e7c693d7c271f93ffc4ba3b380412554231c85bf61ca7af91042a4112" + [[package]] name = "zmij" version = "1.0.21" diff --git a/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml b/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml new file mode 100644 index 000000000..f18942ea6 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "mcp-postgres" +version = "0.0.1" +edition = "2024" +license = "Apache-2.0" + +[lints] +workspace = true + +[dependencies] +authzed = { workspace = true } +async-trait = { workspace = true } +futures-util = { workspace = true } +pg_query = "=6.2.1" +rmcp = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +sqlx = { workspace = true, features = ["tls-rustls-aws-lc-rs"] } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-std", "sync", "time"] } +trogon-std = { workspace = true } +toml = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true, features = ["json"] } +tonic = { workspace = true, features = ["tls-aws-lc", "tls-native-roots"] } +url = { workspace = true } + +[dev-dependencies] +testcontainers-modules = { workspace = true, features = ["postgres"] } diff --git a/rsworkspace/crates/mcp/mcp-postgres/examples/config.toml b/rsworkspace/crates/mcp/mcp-postgres/examples/config.toml new file mode 100644 index 000000000..dcd7714db --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/examples/config.toml @@ -0,0 +1,9 @@ +user = "example-user" +client = "example-client" +spicedb_endpoint = "http://127.0.0.1:50051" +spicedb_token_env = "MCP_POSTGRES_SPICEDB_TOKEN" + +[[connections]] +name = "analytics-readonly" +id = "analytics-local-v1" +dsn_env = "MCP_POSTGRES_ANALYTICS_DSN" diff --git a/rsworkspace/crates/mcp/mcp-postgres/sql/install-ddl-guard.sql b/rsworkspace/crates/mcp/mcp-postgres/sql/install-ddl-guard.sql new file mode 100644 index 000000000..e1ae888c0 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/sql/install-ddl-guard.sql @@ -0,0 +1,25 @@ +BEGIN; + +SELECT pg_catalog.pg_advisory_xact_lock(1879445091, 1347635021); + +CREATE SCHEMA mcp_postgres_guard; +REVOKE ALL ON SCHEMA mcp_postgres_guard FROM PUBLIC; + +CREATE FUNCTION mcp_postgres_guard.lock_ddl() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog +AS $guard$BEGIN + PERFORM pg_catalog.pg_advisory_xact_lock(1879445091, 1347635021); +END;$guard$; + +REVOKE ALL ON FUNCTION mcp_postgres_guard.lock_ddl() FROM PUBLIC; + +CREATE EVENT TRIGGER mcp_postgres_ddl_guard +ON ddl_command_start +EXECUTE FUNCTION mcp_postgres_guard.lock_ddl(); + +ALTER EVENT TRIGGER mcp_postgres_ddl_guard ENABLE ALWAYS; + +COMMIT; diff --git a/rsworkspace/crates/mcp/mcp-postgres/sql/schema.zed b/rsworkspace/crates/mcp/mcp-postgres/sql/schema.zed new file mode 100644 index 000000000..604f9b252 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/sql/schema.zed @@ -0,0 +1,20 @@ +definition sql_user {} +definition sql_client {} + +definition sql_connection { + relation reader: sql_user | sql_client + relation forbidden: sql_user | sql_user:* | sql_client | sql_client:* + permission query = reader - forbidden +} + +definition sql_table { + relation reader: sql_user | sql_client + relation forbidden: sql_user | sql_user:* | sql_client | sql_client:* + permission query = reader - forbidden +} + +definition sql_column { + relation reader: sql_user | sql_client + relation forbidden: sql_user | sql_user:* | sql_client | sql_client:* + permission query = reader - forbidden +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/config.rs b/rsworkspace/crates/mcp/mcp-postgres/src/config.rs new file mode 100644 index 000000000..b9ffd6d9d --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/config.rs @@ -0,0 +1,79 @@ +use crate::{Error, ObjectId}; +use serde::Deserialize; +use std::collections::BTreeMap; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ConfigInput { + pub user: String, + pub client: String, + pub spicedb_endpoint: String, + pub spicedb_token_env: String, + pub connections: Vec, +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ConnectionInput { + pub name: String, + pub id: String, + pub dsn_env: String, +} +pub struct RegisteredConnection { + pub(crate) id: ObjectId, + pub(crate) options: crate::DatabaseOptions, +} +pub struct RuntimeConfig { + pub(crate) user: ObjectId, + pub(crate) client: ObjectId, + pub(crate) connections: BTreeMap, + pub(crate) spicedb_endpoint: crate::SpiceDbEndpoint, + pub(crate) spicedb_token: String, +} +impl RuntimeConfig { + pub fn from_input(input: ConfigInput, env: impl Fn(&str) -> Option) -> Result { + let mut connections = BTreeMap::new(); + let mut ids = std::collections::BTreeSet::new(); + for connection in input.connections { + let id = ObjectId::new(connection.id)?; + if !ids.insert(id.clone()) { + return Err(Error::Configuration); + } + let name = ObjectId::new(connection.name)?; + let dsn = env(crate::EnvironmentVariable::new(connection.dsn_env)?.as_str()) + .filter(|v| !v.is_empty()) + .ok_or(Error::Configuration)?; + if connections + .insert( + name.as_str().to_owned(), + RegisteredConnection { + id, + options: crate::DatabaseOptions::new(&dsn)?, + }, + ) + .is_some() + { + return Err(Error::Configuration); + } + } + if connections.is_empty() { + return Err(Error::Configuration); + } + Ok(Self { + user: ObjectId::new(input.user)?, + client: ObjectId::new(input.client)?, + connections, + spicedb_endpoint: crate::SpiceDbEndpoint::new(input.spicedb_endpoint)?, + spicedb_token: env(crate::EnvironmentVariable::new(input.spicedb_token_env)?.as_str()) + .filter(|v| !v.is_empty()) + .ok_or(Error::Configuration)?, + }) + } +} +impl RuntimeConfig { + pub fn spicedb_endpoint(&self) -> &str { + self.spicedb_endpoint.as_str() + } + pub fn spicedb_token(&self) -> &str { + &self.spicedb_token + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs b/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs new file mode 100644 index 000000000..2eed7f912 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs @@ -0,0 +1 @@ +pub(crate) const MAX_BYTES: usize = 1_048_576; diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/database_options.rs b/rsworkspace/crates/mcp/mcp-postgres/src/database_options.rs new file mode 100644 index 000000000..4f3930725 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/database_options.rs @@ -0,0 +1,21 @@ +use crate::Error; +use sqlx::{ + ConnectOptions, + postgres::{PgConnectOptions, PgSslMode}, +}; +use std::str::FromStr; + +pub struct DatabaseOptions(PgConnectOptions); +impl DatabaseOptions { + pub fn new(dsn: &str) -> Result { + let options = PgConnectOptions::from_str(dsn).map_err(|_| Error::Configuration)?; + let local = matches!(options.get_host(), "localhost" | "127.0.0.1" | "::1"); + if !local && !matches!(options.get_ssl_mode(), PgSslMode::VerifyFull) { + return Err(Error::Configuration); + } + Ok(Self(options.disable_statement_logging())) + } + pub(crate) fn into_inner(self) -> PgConnectOptions { + self.0 + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/environment_variable.rs b/rsworkspace/crates/mcp/mcp-postgres/src/environment_variable.rs new file mode 100644 index 000000000..1704a4df5 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/environment_variable.rs @@ -0,0 +1,21 @@ +use crate::Error; + +#[derive(Debug, Clone)] +pub struct EnvironmentVariable(String); +impl EnvironmentVariable { + pub fn new(value: String) -> Result { + if value.is_empty() + || value.len() > 200 + || !value + .bytes() + .all(|c| c.is_ascii_uppercase() || c.is_ascii_digit() || c == b'_') + || value.as_bytes()[0].is_ascii_digit() + { + return Err(Error::Configuration); + } + Ok(Self(value)) + } + pub fn as_str(&self) -> &str { + &self.0 + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/error.rs b/rsworkspace/crates/mcp/mcp-postgres/src/error.rs new file mode 100644 index 000000000..024c01996 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/error.rs @@ -0,0 +1,29 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid configuration")] + Configuration, + #[error("unsupported or invalid SQL")] + UnsupportedSql, + #[error("SQL parsing failed")] + Parse(#[from] pg_query::Error), + #[error("unknown connection")] + UnknownConnection, + #[error("access denied")] + Denied, + #[error("database operation failed")] + Database(#[from] sqlx::Error), + #[error("authorization service failed")] + Authorization(#[from] tonic::Status), + #[error("authorization connection failed")] + AuthorizationConnection(#[from] tonic::transport::Error), + #[error("resource limit exceeded")] + Limit, + #[error("request timed out")] + Timeout, + #[error("unsupported database object")] + UnsupportedObject, + #[error("invalid identifier")] + Identifier, + #[error("invalid object ID")] + ObjectId, +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/identifier.rs b/rsworkspace/crates/mcp/mcp-postgres/src/identifier.rs new file mode 100644 index 000000000..6ebf4b2ea --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/identifier.rs @@ -0,0 +1,18 @@ +use crate::Error; + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct Identifier(String); +impl Identifier { + pub fn new(value: String) -> Result { + if value.is_empty() || value.len() > 63 || value.contains('\0') { + return Err(Error::Identifier); + } + Ok(Self(value)) + } + pub fn as_str(&self) -> &str { + &self.0 + } + pub fn quoted(&self) -> String { + format!("\"{}\"", self.0.replace('"', "\"\"")) + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs b/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs new file mode 100644 index 000000000..6fe3fb10f --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs @@ -0,0 +1,45 @@ +//! Governed PostgreSQL queries over MCP for a trusted local stdio launcher. +//! +//! Named connections and user/client identities come exclusively from runtime configuration. +//! The bounded SELECT grammar rejects unsupported syntax and catalog features. +//! Every referenced column needs an explicit grant for both identities in one SpiceDB snapshot. +//! The required administrator-installed DDL guard fences catalog changes until the query ends. +//! Production access and deployment are outside this crate's local implementation contract. +//! +//! The launcher and database administrators are trusted. Callers must not control the configuration +//! or credentials, and the query role must have only the required SELECT privileges. +//! Install `sql/install-ddl-guard.sql` in each registered database. Migrations must take the session +//! advisory lock `pg_advisory_lock(1879445091, 1347635021)` before touching relations and retain it +//! through all transactions. This is essential for concurrent index operations, guard maintenance, +//! and shared-object changes that the event trigger cannot fence on its own. +//! Replacing or restoring a database target requires a new connection ID when catalog IDs can repeat. +//! Permissions are checked at admission; revocation does not retroactively cancel an admitted query. + +mod config; +mod constants; +mod database_options; +mod environment_variable; +mod error; +mod identifier; +mod object_id; +mod parameter; +mod parser; +mod postgres; +mod row_limit; +mod server; +mod spicedb; +mod spicedb_endpoint; +mod telemetry; + +pub use config::{ConfigInput, RegisteredConnection, RuntimeConfig}; +pub use database_options::DatabaseOptions; +pub use environment_variable::EnvironmentVariable; +pub use error::Error; +pub use identifier::Identifier; +pub use object_id::ObjectId; +pub use parameter::{Parameter, ParameterInput}; +pub use parser::Query; +pub use postgres::{QueryResult, QueryService}; +pub use server::PostgresMcp; +pub use spicedb::{PermissionCheck, PermissionRequest, Resource, SpiceDb}; +pub use spicedb_endpoint::SpiceDbEndpoint; diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/main.rs b/rsworkspace/crates/mcp/mcp-postgres/src/main.rs new file mode 100644 index 000000000..c34c284dc --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/main.rs @@ -0,0 +1,37 @@ +use mcp_postgres::{ConfigInput, PostgresMcp, QueryService, RuntimeConfig, SpiceDb}; +use rmcp::ServiceExt; +use std::{process::ExitCode, sync::Arc}; +use tracing_subscriber::{Layer, layer::SubscriberExt, util::SubscriberInitExt}; +use trogon_std::env::{ReadEnv, SystemEnv}; + +#[tokio::main] +async fn main() -> ExitCode { + let audit = tracing_subscriber::fmt::layer() + .json() + .with_ansi(false) + .with_writer(std::io::stderr) + .with_filter(tracing_subscriber::filter::Targets::new().with_target("mcp_postgres", tracing::Level::INFO)); + if tracing_subscriber::registry().with(audit).try_init().is_err() { + return ExitCode::FAILURE; + } + if run().await.is_err() { + tracing::error!("mcp-postgres startup or transport failed"); + return ExitCode::FAILURE; + } + ExitCode::SUCCESS +} +async fn run() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let path = args.next().ok_or(mcp_postgres::Error::Configuration)?; + if args.next().is_some() { + return Err(mcp_postgres::Error::Configuration.into()); + } + let input: ConfigInput = toml::from_str(&std::fs::read_to_string(path)?)?; + let config = RuntimeConfig::from_input(input, |key| SystemEnv.var(key).ok())?; + let permissions = + Arc::new(SpiceDb::connect(config.spicedb_endpoint().into(), config.spicedb_token().into()).await?); + let service = QueryService::connect(config, permissions).await?; + let server = PostgresMcp::new(service).serve(rmcp::transport::stdio()).await?; + server.waiting().await?; + Ok(()) +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/object_id.rs b/rsworkspace/crates/mcp/mcp-postgres/src/object_id.rs new file mode 100644 index 000000000..55c377e7b --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/object_id.rs @@ -0,0 +1,18 @@ +use crate::Error; + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct ObjectId(String); +impl ObjectId { + pub fn new(value: String) -> Result { + if value.is_empty() + || value.len() > 200 + || !value.bytes().all(|c| c.is_ascii_alphanumeric() || b"_-./".contains(&c)) + { + return Err(Error::ObjectId); + } + Ok(Self(value)) + } + pub fn as_str(&self) -> &str { + &self.0 + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/parameter.rs b/rsworkspace/crates/mcp/mcp-postgres/src/parameter.rs new file mode 100644 index 000000000..3aaf97dc3 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/parameter.rs @@ -0,0 +1,35 @@ +use crate::Error; +use rmcp::schemars; +use serde::Deserialize; + +#[derive(Deserialize, schemars::JsonSchema)] +#[serde(untagged)] +pub enum ParameterInput { + Boolean(bool), + Integer(i64), + Text(String), +} +#[derive(Debug, Clone)] +pub struct Parameter(ParameterValue); +#[derive(Debug, Clone)] +pub(crate) enum ParameterValue { + Boolean(bool), + Integer(i64), + Text(String), +} +impl TryFrom for Parameter { + type Error = Error; + fn try_from(input: ParameterInput) -> Result { + Ok(Self(match input { + ParameterInput::Boolean(v) => ParameterValue::Boolean(v), + ParameterInput::Integer(v) => ParameterValue::Integer(v), + ParameterInput::Text(v) if v.len() <= 16_384 && !v.contains('\0') => ParameterValue::Text(v), + ParameterInput::Text(_) => return Err(Error::Limit), + })) + } +} +impl Parameter { + pub(crate) fn value(&self) -> &ParameterValue { + &self.0 + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/parser.rs b/rsworkspace/crates/mcp/mcp-postgres/src/parser.rs new file mode 100644 index 000000000..dc0857ecd --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/parser.rs @@ -0,0 +1,293 @@ +use crate::{Error, Identifier}; +use pg_query::{ + NodeEnum, + protobuf::{self, a_const}, +}; + +#[derive(Debug)] +pub struct Query { + pub(crate) select: Vec<(Column, Option)>, + pub(crate) tables: Vec, + pub(crate) from: String, + pub(crate) predicates: Vec, + pub(crate) order: Vec<(Column, &'static str)>, + pub(crate) limit: crate::row_limit::RowLimit, +} +#[derive(Debug, Clone)] +pub(crate) struct Table { + pub schema: Identifier, + pub name: Identifier, + pub alias: Identifier, +} +#[derive(Debug, Clone)] +pub(crate) struct Column { + pub qualifier: Option, + pub name: Identifier, +} +#[derive(Debug)] +pub(crate) enum Expr { + Column(Column), + Literal(String), + Parameter(usize), + Compare(Box, String, Box), + Boolean(Vec, &'static str), +} +impl Query { + pub fn parse(sql: &str) -> Result { + if sql.len() > 16_384 { + return Err(Error::Limit); + } + let parsed = pg_query::parse(sql)?; + if parsed.protobuf.stmts.len() != 1 { + return Err(Error::UnsupportedSql); + } + let statement = parsed + .protobuf + .stmts + .first() + .and_then(|s| s.stmt.as_ref()) + .ok_or(Error::UnsupportedSql)?; + let NodeEnum::SelectStmt(stmt) = node(statement)? else { + return Err(Error::UnsupportedSql); + }; + let mut normalized = stmt.as_ref().clone(); + normalized.target_list.clear(); + normalized.from_clause.clear(); + normalized.where_clause = None; + normalized.sort_clause.clear(); + normalized.limit_count = None; + normalized.limit_option = 0; + normalized.op = 0; + if normalized != protobuf::SelectStmt::default() + || stmt.op != protobuf::SetOperation::SetopNone as i32 + || stmt.from_clause.len() != 1 + || stmt.target_list.is_empty() + || stmt.target_list.len() > 64 + { + return Err(Error::UnsupportedSql); + } + if stmt.limit_option != protobuf::LimitOption::Default as i32 + && stmt.limit_option != protobuf::LimitOption::Count as i32 + { + return Err(Error::UnsupportedSql); + } + let mut query = Self { + select: Vec::new(), + tables: Vec::new(), + from: String::new(), + predicates: Vec::new(), + order: Vec::new(), + limit: crate::row_limit::RowLimit::bounded_default(), + }; + query.from = query.render_from_node(&stmt.from_clause[0], 0)?; + for target in &stmt.target_list { + let NodeEnum::ResTarget(target) = node(target)? else { + return Err(Error::UnsupportedSql); + }; + if !target.indirection.is_empty() { + return Err(Error::UnsupportedSql); + } + query.select.push(( + column(target.val.as_deref().ok_or(Error::UnsupportedSql)?)?, + if target.name.is_empty() { + None + } else { + Some(Identifier::new(target.name.clone())?) + }, + )); + } + if let Some(predicate) = stmt.where_clause.as_deref() { + query.predicates.push(expr(predicate, 0)?); + } + for sort in &stmt.sort_clause { + let NodeEnum::SortBy(sort) = node(sort)? else { + return Err(Error::UnsupportedSql); + }; + if !sort.use_op.is_empty() || sort.sortby_nulls != protobuf::SortByNulls::SortbyNullsDefault as i32 { + return Err(Error::UnsupportedSql); + } + let direction = match protobuf::SortByDir::try_from(sort.sortby_dir) { + Ok(protobuf::SortByDir::SortbyDefault | protobuf::SortByDir::SortbyAsc) => "ASC", + Ok(protobuf::SortByDir::SortbyDesc) => "DESC", + _ => return Err(Error::UnsupportedSql), + }; + query + .order + .push((column(sort.node.as_deref().ok_or(Error::UnsupportedSql)?)?, direction)); + } + if let Some(limit) = stmt.limit_count.as_deref() { + let NodeEnum::AConst(value) = node(limit)? else { + return Err(Error::UnsupportedSql); + }; + let Some(a_const::Val::Ival(value)) = value.val.as_ref() else { + return Err(Error::UnsupportedSql); + }; + query.limit = + crate::row_limit::RowLimit::requested(usize::try_from(value.ival).map_err(|_| Error::UnsupportedSql)?)?; + } + Ok(query) + } + fn render_from_node(&mut self, value: &protobuf::Node, depth: usize) -> Result { + if depth > 8 || self.tables.len() > 8 { + return Err(Error::Limit); + } + match node(value)? { + NodeEnum::RangeVar(table) => { + if self.tables.len() >= 8 { + return Err(Error::Limit); + } + if !table.catalogname.is_empty() + || table.schemaname.is_empty() + || table.alias.as_ref().is_some_and(|a| !a.colnames.is_empty()) + { + return Err(Error::UnsupportedSql); + } + let schema = Identifier::new(table.schemaname.clone())?; + if schema.as_str() == "pg_catalog" + || schema.as_str() == "information_schema" + || schema.as_str().starts_with("pg_") + { + return Err(Error::UnsupportedSql); + } + let name = Identifier::new(table.relname.clone())?; + let alias = Identifier::new( + table + .alias + .as_ref() + .map_or_else(|| table.relname.clone(), |a| a.aliasname.clone()), + )?; + if self.tables.iter().any(|t| t.alias == alias) { + return Err(Error::UnsupportedSql); + } + let from = format!("ONLY {}.{} AS {}", schema.quoted(), name.quoted(), alias.quoted()); + self.tables.push(Table { schema, name, alias }); + Ok(from) + } + NodeEnum::JoinExpr(join) => { + if join.is_natural + || !join.using_clause.is_empty() + || join.alias.is_some() + || join.join_using_alias.is_some() + { + return Err(Error::UnsupportedSql); + } + let kind = match protobuf::JoinType::try_from(join.jointype) { + Ok(protobuf::JoinType::JoinInner) => "INNER", + Ok(protobuf::JoinType::JoinLeft) => "LEFT", + _ => return Err(Error::UnsupportedSql), + }; + let scope_start = self.tables.len(); + let left = self.render_from_node(join.larg.as_deref().ok_or(Error::UnsupportedSql)?, depth + 1)?; + let right = self.render_from_node(join.rarg.as_deref().ok_or(Error::UnsupportedSql)?, depth + 1)?; + // Join predicates are rendered later after every column has been resolved. + let predicate = expr(join.quals.as_deref().ok_or(Error::UnsupportedSql)?, 0)?; + check_join_scope(&predicate, &self.tables[scope_start..])?; + self.predicates.push(predicate); + if kind == "LEFT" { + return Err(Error::UnsupportedSql); + } + Ok(format!("({left} {kind} JOIN {right} ON true)")) + } + _ => Err(Error::UnsupportedSql), + } + } +} +fn node(value: &protobuf::Node) -> Result<&NodeEnum, Error> { + value.node.as_ref().ok_or(Error::UnsupportedSql) +} +fn string(value: &protobuf::Node) -> Result { + match node(value)? { + NodeEnum::String(value) => Ok(value.sval.clone()), + _ => Err(Error::UnsupportedSql), + } +} +fn column(value: &protobuf::Node) -> Result { + let NodeEnum::ColumnRef(value) = node(value)? else { + return Err(Error::UnsupportedSql); + }; + match value.fields.as_slice() { + [name] => Ok(Column { + qualifier: None, + name: Identifier::new(string(name)?)?, + }), + [qualifier, name] => Ok(Column { + qualifier: Some(Identifier::new(string(qualifier)?)?), + name: Identifier::new(string(name)?)?, + }), + _ => Err(Error::UnsupportedSql), + } +} +fn expr(value: &protobuf::Node, depth: usize) -> Result { + if depth > 16 { + return Err(Error::Limit); + } + match node(value)? { + NodeEnum::ColumnRef(_) => Ok(Expr::Column(column(value)?)), + NodeEnum::ParamRef(value) if (1..=32).contains(&value.number) => Ok(Expr::Parameter(value.number as usize)), + NodeEnum::AConst(value) if !value.isnull => match value.val.as_ref() { + Some(a_const::Val::Ival(value)) => Ok(Expr::Literal(value.ival.to_string())), + Some(a_const::Val::Sval(value)) => Ok(Expr::Literal(format!("'{}'", value.sval.replace('\'', "''")))), + Some(a_const::Val::Boolval(value)) => { + Ok(Expr::Literal(if value.boolval { "true" } else { "false" }.into())) + } + _ => Err(Error::UnsupportedSql), + }, + NodeEnum::AExpr(value) if value.kind == protobuf::AExprKind::AexprOp as i32 && value.name.len() == 1 => { + let operator = string(&value.name[0])?; + if !["=", "<>", "<", ">", "<=", ">="].contains(&operator.as_str()) { + return Err(Error::UnsupportedSql); + } + Ok(Expr::Compare( + Box::new(expr(value.lexpr.as_deref().ok_or(Error::UnsupportedSql)?, depth + 1)?), + operator, + Box::new(expr(value.rexpr.as_deref().ok_or(Error::UnsupportedSql)?, depth + 1)?), + )) + } + NodeEnum::BoolExpr(value) => { + let operator = match protobuf::BoolExprType::try_from(value.boolop) { + Ok(protobuf::BoolExprType::AndExpr) => "AND", + Ok(protobuf::BoolExprType::OrExpr) => "OR", + _ => return Err(Error::UnsupportedSql), + }; + if value.args.len() < 2 { + return Err(Error::UnsupportedSql); + } + Ok(Expr::Boolean( + value + .args + .iter() + .map(|v| expr(v, depth + 1)) + .collect::>()?, + operator, + )) + } + _ => Err(Error::UnsupportedSql), + } +} + +fn check_join_scope(expr: &Expr, tables: &[Table]) -> Result<(), Error> { + match expr { + Expr::Column(column) => { + if column + .qualifier + .as_ref() + .is_some_and(|q| tables.iter().any(|t| t.alias == *q)) + { + Ok(()) + } else { + Err(Error::UnsupportedSql) + } + } + Expr::Literal(_) | Expr::Parameter(_) => Ok(()), + Expr::Compare(left, _, right) => { + check_join_scope(left, tables)?; + check_join_scope(right, tables) + } + Expr::Boolean(args, _) => { + for expr in args { + check_join_scope(expr, tables)?; + } + Ok(()) + } + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs new file mode 100644 index 000000000..dda89f144 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs @@ -0,0 +1,383 @@ +use crate::{ + Error, ObjectId, PermissionCheck, Query, Resource, RuntimeConfig, + parser::{Column, Expr}, +}; +use futures_util::TryStreamExt; +use serde::Serialize; +use sqlx::{ + AssertSqlSafe, Column as _, Connection, Executor, PgConnection, PgPool, Row, SqlSafeStr, Statement, TypeInfo, + postgres::{PgPoolOptions, PgRow}, +}; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::Arc, + time::Duration, +}; +use tokio::sync::Semaphore; + +use crate::constants::MAX_BYTES; +#[derive(Serialize, Debug)] +pub struct QueryResult { + pub columns: Vec, + pub rows: Vec>, +} +pub struct QueryService { + user: ObjectId, + client: ObjectId, + connections: BTreeMap, + permissions: Arc, + permits: Semaphore, +} +#[derive(Debug)] +struct CatalogColumn { + table: usize, + name: String, + attnum: i16, + type_oid: u32, +} +impl QueryService { + pub async fn connect(config: RuntimeConfig, permissions: Arc) -> Result { + let mut connections = BTreeMap::new(); + for (name, connection) in config.connections { + let pool = PgPoolOptions::new() + .max_connections(4) + .acquire_timeout(Duration::from_secs(2)) + .connect_with(connection.options.into_inner()) + .await?; + connections.insert(name, (connection.id, pool)); + } + Ok(Self { + user: config.user, + client: config.client, + connections, + permissions, + permits: Semaphore::new(4), + }) + } + pub async fn query(&self, connection: &str, sql: &str) -> Result { + self.query_with_parameters(connection, sql, &[]).await + } + pub async fn query_with_parameters( + &self, + connection: &str, + sql: &str, + parameters: &[crate::Parameter], + ) -> Result { + let mut audit = crate::telemetry::QueryAudit::start( + &self.user, + &self.client, + self.connections.get(connection).map(|(id, _)| id), + ); + let result = async { + if parameters.len() > 32 { + return Err(Error::Limit); + } + let _permit = self.permits.try_acquire().map_err(|_| Error::Limit)?; + tokio::time::timeout( + Duration::from_secs(10), + self.execute(connection, Query::parse(sql)?, parameters), + ) + .await + .map_err(|_| Error::Timeout)? + } + .await; + audit.complete(&result); + result + } + async fn authorize(&self, resource: &Resource) -> Result<(), Error> { + self.permissions.check(resource, "sql_user", &self.user).await?; + self.permissions.check(resource, "sql_client", &self.client).await + } + async fn execute( + &self, + connection: &str, + query: Query, + parameters: &[crate::Parameter], + ) -> Result { + let (connection_id, pool) = self.connections.get(connection).ok_or(Error::UnknownConnection)?; + self.authorize(&Resource { + kind: "sql_connection", + id: connection_id.clone(), + }) + .await?; + let mut connection = pool.acquire().await?; + connection.close_on_drop(); + let mut tx = connection.begin().await?; + tx.execute("SET TRANSACTION ISOLATION LEVEL READ COMMITTED READ ONLY") + .await?; + for setting in [ + "SET LOCAL search_path = pg_catalog", + "SET LOCAL statement_timeout = '5s'", + "SET LOCAL lock_timeout = '1s'", + "SET LOCAL idle_in_transaction_session_timeout = '5s'", + "SET LOCAL standard_conforming_strings = on", + "SET LOCAL row_security = on", + "SET LOCAL enable_indexscan = off", + "SET LOCAL enable_indexonlyscan = off", + "SET LOCAL enable_bitmapscan = off", + "SET LOCAL constraint_exclusion = off", + "SET LOCAL jit = off", + "SET LOCAL client_connection_check_interval = '100ms'", + ] { + tx.execute(setting).await?; + } + let unsafe_role: bool = sqlx::query_scalar("SELECT rolsuper OR rolbypassrls OR rolcreaterole OR rolcreatedb OR rolreplication FROM pg_catalog.pg_roles WHERE rolname = current_user").fetch_one(&mut *tx).await?; + if unsafe_role { + return Err(Error::UnsupportedObject); + } + tx.execute("SELECT pg_catalog.pg_advisory_xact_lock_shared(1879445091,1347635021)") + .await?; + let guard: bool = sqlx::query_scalar("SELECT EXISTS (SELECT 1 FROM pg_catalog.pg_event_trigger e JOIN pg_catalog.pg_proc p ON p.oid=e.evtfoid JOIN pg_catalog.pg_namespace n ON n.oid=p.pronamespace WHERE e.evtname='mcp_postgres_ddl_guard' AND e.evtevent='ddl_command_start' AND e.evtenabled='A' AND e.evttags IS NULL AND n.nspname='mcp_postgres_guard' AND p.proname='lock_ddl' AND p.prosecdef AND p.proconfig = ARRAY['search_path=pg_catalog']::text[] AND p.pronargs=0 AND p.prorettype='pg_catalog.event_trigger'::regtype AND p.prolang=(SELECT oid FROM pg_catalog.pg_language WHERE lanname='plpgsql') AND p.proowner=n.nspowner AND EXISTS (SELECT 1 FROM pg_catalog.pg_roles r WHERE r.oid=p.proowner AND r.rolsuper) AND NOT EXISTS (SELECT 1 FROM pg_catalog.aclexplode(COALESCE(n.nspacl,pg_catalog.acldefault('n',n.nspowner))) a WHERE a.privilege_type='CREATE' AND a.grantee<>n.nspowner) AND p.prosrc=$1)").bind("BEGIN\n PERFORM pg_catalog.pg_advisory_xact_lock(1879445091, 1347635021);\nEND;").fetch_one(&mut *tx).await?; + if !guard { + return Err(Error::UnsupportedObject); + } + let database: i64 = + sqlx::query_scalar("SELECT oid::bigint FROM pg_catalog.pg_database WHERE datname=current_database()") + .fetch_one(&mut *tx) + .await?; + let mut catalog = Vec::new(); + let mut resources = BTreeSet::new(); + resources.insert(Resource { + kind: "sql_connection", + id: connection_id.clone(), + }); + for (index, table) in query.tables.iter().enumerate() { + let lock = format!( + "LOCK TABLE ONLY {}.{} IN ACCESS SHARE MODE", + table.schema.quoted(), + table.name.quoted() + ); + sqlx::query(AssertSqlSafe(lock)).execute(&mut *tx).await?; + let relation = sqlx::query("SELECT c.oid::bigint AS oid FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace JOIN pg_catalog.pg_am a ON a.oid=c.relam WHERE n.nspname=$1 AND c.relname=$2 AND c.relkind='r' AND a.amname='heap' AND a.oid=2 AND NOT c.relrowsecurity AND NOT c.relforcerowsecurity AND NOT c.relhasrules AND NOT c.relispartition AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits i WHERE i.inhrelid=c.oid OR i.inhparent=c.oid) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i WHERE i.indrelid=c.oid AND (i.indexprs IS NOT NULL OR i.indpred IS NOT NULL)) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i JOIN pg_catalog.pg_class ic ON ic.oid=i.indexrelid JOIN pg_catalog.pg_am ia ON ia.oid=ic.relam WHERE i.indrelid=c.oid AND ia.amname NOT IN ('btree','hash')) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i, unnest(i.indcollation) co WHERE i.indrelid=c.oid AND co<>0 AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_collation cl WHERE cl.oid=co AND cl.collnamespace='pg_catalog'::regnamespace)) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i, unnest(i.indclass) op WHERE i.indrelid=c.oid AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_opclass oc WHERE oc.oid=op AND oc.opcnamespace='pg_catalog'::regnamespace))") + .bind(table.schema.as_str()).bind(table.name.as_str()).fetch_optional(&mut *tx).await?.ok_or(Error::UnsupportedObject)?; + let oid: i64 = relation.try_get("oid")?; + let table_id = format!("{}/{database}/{oid}", connection_id.as_str()); + resources.insert(Resource { + kind: "sql_table", + id: ObjectId::new(table_id.clone())?, + }); + let columns = sqlx::query("SELECT a.attname,a.attnum,a.atttypid::bigint AS type_oid,a.attgenerated::text AS generated,COALESCE(n.nspname='pg_catalog',true) AS safe_collation FROM pg_catalog.pg_attribute a LEFT JOIN pg_catalog.pg_collation c ON c.oid=a.attcollation LEFT JOIN pg_catalog.pg_namespace n ON n.oid=c.collnamespace WHERE a.attrelid=$1::bigint::oid AND a.attnum>0 AND NOT a.attisdropped").bind(oid).fetch_all(&mut *tx).await?; + for row in columns { + let type_oid = + u32::try_from(row.try_get::("type_oid")?).map_err(|_| Error::UnsupportedObject)?; + let generated: String = row.try_get("generated")?; + let safe_collation: bool = row.try_get("safe_collation")?; + // Reject the whole table: even unselected generated attributes can introduce planner dependencies. + if !generated.is_empty() { + return Err(Error::UnsupportedObject); + } + let name: String = row.try_get("attname")?; + let attnum: i16 = row.try_get("attnum")?; + if !safe_collation || ![16, 20, 21, 23, 25, 1043].contains(&type_oid) { + return Err(Error::UnsupportedObject); + } + catalog.push(CatalogColumn { + table: index, + name, + attnum, + type_oid, + }); + } + } + let mut used = BTreeSet::new(); + let mut projections = Vec::new(); + let mut sizes = Vec::new(); + for (column, alias) in &query.select { + let rendered = resolve(column, &query, &catalog, &mut used)?; + if catalog.iter().any(|c| { + c.name == column.name.as_str() + && column + .qualifier + .as_ref() + .is_none_or(|q| query.tables[c.table].alias == *q) + && [25, 1043].contains(&c.type_oid) + }) { + sizes.push(format!( + "COALESCE(pg_catalog.octet_length({rendered}),0)::pg_catalog.int8" + )); + } + projections.push((rendered, alias.as_ref().unwrap_or(&column.name).quoted())); + } + let predicates = query + .predicates + .iter() + .map(|expr| render_expr(expr, &query, &catalog, &mut used, parameters)) + .collect::, _>>()?; + let mut orders = Vec::new(); + for (column, direction) in &query.order { + if column.qualifier.is_none() + && query + .select + .iter() + .any(|(_, alias)| alias.as_ref() == Some(&column.name)) + { + return Err(Error::UnsupportedSql); + } + orders.push(format!("{} {direction}", resolve(column, &query, &catalog, &mut used)?)); + } + for (table, attnum) in used { + let relation_oid: i64 = sqlx::query_scalar("SELECT c.oid::bigint FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace WHERE n.nspname=$1 AND c.relname=$2").bind(query.tables[table].schema.as_str()).bind(query.tables[table].name.as_str()).fetch_one(&mut *tx).await?; + resources.insert(Resource { + kind: "sql_column", + id: ObjectId::new(format!("{}/{database}/{relation_oid}/{attnum}", connection_id.as_str()))?, + }); + } + let overflow = if sizes.is_empty() { + "false".to_owned() + } else { + format!("({}) > {}", sizes.join(" + "), MAX_BYTES / 8) + }; + let projections = projections + .into_iter() + .map(|(value, label)| format!("CASE WHEN ({overflow}) THEN NULL ELSE {value} END AS {label}")) + .collect::>(); + let mut statement = format!( + "SELECT ({overflow}) AS __mcp_overflow, {} FROM {}", + projections.join(", "), + query.from + ); + if !predicates.is_empty() { + statement.push_str(&format!(" WHERE {}", predicates.join(" AND "))); + } + if !orders.is_empty() { + statement.push_str(&format!(" ORDER BY {}", orders.join(", "))); + } + statement.push_str(&format!(" LIMIT {}", query.limit.statement_count())); + let prepared = (&mut *tx as &mut PgConnection) + .prepare(AssertSqlSafe(statement).into_sql_str()) + .await?; + crate::telemetry::references(&resources, &self.user, &self.client, connection_id); + let requests = resources + .into_iter() + .flat_map(|resource| { + [ + crate::PermissionRequest { + resource: resource.clone(), + subject_kind: "sql_user", + subject: self.user.clone(), + }, + crate::PermissionRequest { + resource, + subject_kind: "sql_client", + subject: self.client.clone(), + }, + ] + }) + .collect::>(); + self.permissions.check_all(&requests).await?; + let mut bound = prepared.query(); + for parameter in parameters { + bound = match parameter.value() { + crate::parameter::ParameterValue::Boolean(v) => bound.bind(*v), + crate::parameter::ParameterValue::Integer(v) => bound.bind(*v), + crate::parameter::ParameterValue::Text(v) => bound.bind(v), + }; + } + let columns = prepared.columns().iter().skip(1).map(|c| c.name().to_owned()).collect(); + let mut stream = bound.fetch(&mut *tx); + let mut result = QueryResult { + columns, + rows: Vec::new(), + }; + let mut bytes = serde_json::to_vec(&result).map_err(|_| Error::Limit)?.len(); + while let Some(row) = stream.try_next().await? { + if row.try_get::(0)? { + return Err(Error::Limit); + } + if result.rows.len() == query.limit.result_count() { + return Err(Error::Limit); + } + + let values = decode(&row)?; + bytes += serde_json::to_vec(&values).map_err(|_| Error::Limit)?.len() + 1; + if bytes > MAX_BYTES { + return Err(Error::Limit); + } + result.rows.push(values); + } + drop(stream); + tx.rollback().await?; + if serde_json::to_vec(&result).map_err(|_| Error::Limit)?.len() > MAX_BYTES { + return Err(Error::Limit); + } + Ok(result) + } +} +fn resolve( + column: &Column, + query: &Query, + catalog: &[CatalogColumn], + used: &mut BTreeSet<(usize, i16)>, +) -> Result { + let matches: Vec<_> = catalog + .iter() + .filter(|c| { + c.name == column.name.as_str() + && column + .qualifier + .as_ref() + .is_none_or(|q| query.tables[c.table].alias == *q) + }) + .collect(); + if matches.len() != 1 { + return Err(Error::UnsupportedObject); + } + let found = matches[0]; + used.insert((found.table, found.attnum)); + Ok(format!( + "{}.{}", + query.tables[found.table].alias.quoted(), + column.name.quoted() + )) +} +fn render_expr( + expr: &Expr, + query: &Query, + catalog: &[CatalogColumn], + used: &mut BTreeSet<(usize, i16)>, + parameters: &[crate::Parameter], +) -> Result { + match expr { + Expr::Column(c) => resolve(c, query, catalog, used), + Expr::Literal(value) => Ok(value.clone()), + Expr::Parameter(number) => { + let parameter = parameters.get(number - 1).ok_or(Error::UnsupportedSql)?; + let kind = match parameter.value() { + crate::parameter::ParameterValue::Boolean(_) => "bool", + crate::parameter::ParameterValue::Integer(_) => "int8", + crate::parameter::ParameterValue::Text(_) => "text", + }; + Ok(format!("CAST(${number} AS pg_catalog.{kind})")) + } + Expr::Compare(left, operator, right) => Ok(format!( + "({} OPERATOR(pg_catalog.{operator}) {})", + render_expr(left, query, catalog, used, parameters)?, + render_expr(right, query, catalog, used, parameters)? + )), + Expr::Boolean(args, operator) => Ok(format!( + "({})", + args.iter() + .map(|a| render_expr(a, query, catalog, used, parameters)) + .collect::, _>>()? + .join(&format!(" {operator} ")) + )), + } +} +fn decode(row: &PgRow) -> Result, Error> { + row.columns() + .iter() + .enumerate() + .skip(1) + .map(|(index, column)| { + match column.type_info().name() { + "BOOL" => serde_json::to_value(row.try_get::, _>(index)?), + "INT2" => serde_json::to_value(row.try_get::, _>(index)?), + "INT4" => serde_json::to_value(row.try_get::, _>(index)?), + "INT8" => serde_json::to_value(row.try_get::, _>(index)?), + "TEXT" | "VARCHAR" => serde_json::to_value(row.try_get::, _>(index)?), + _ => return Err(Error::UnsupportedObject), + } + .map_err(|_| Error::Limit) + }) + .collect() +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/row_limit.rs b/rsworkspace/crates/mcp/mcp-postgres/src/row_limit.rs new file mode 100644 index 000000000..39fe3e2c2 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/row_limit.rs @@ -0,0 +1,23 @@ +use crate::Error; + +#[derive(Debug)] +pub(crate) struct RowLimit { + requested: Option, +} +impl RowLimit { + pub(crate) fn bounded_default() -> Self { + Self { requested: None } + } + pub(crate) fn requested(value: usize) -> Result { + if value > 1000 { + return Err(Error::Limit); + } + Ok(Self { requested: Some(value) }) + } + pub(crate) fn statement_count(&self) -> usize { + self.requested.unwrap_or(1001) + } + pub(crate) fn result_count(&self) -> usize { + self.requested.unwrap_or(1000) + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/server.rs b/rsworkspace/crates/mcp/mcp-postgres/src/server.rs new file mode 100644 index 000000000..78504aeb7 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/server.rs @@ -0,0 +1,94 @@ +use crate::QueryService; +use rmcp::schemars; +use rmcp::{ + ErrorData, ServerHandler, + model::*, + service::{RequestContext, RoleServer}, +}; +use serde::Deserialize; +use std::sync::Arc; + +#[derive(Clone)] +pub struct PostgresMcp { + service: Arc, +} +impl PostgresMcp { + pub fn new(service: QueryService) -> Self { + Self { + service: Arc::new(service), + } + } +} +#[derive(Deserialize, schemars::JsonSchema)] +#[serde(deny_unknown_fields)] +struct QueryRequest { + connection: String, + #[schemars(length(max = 16384))] + sql: String, + #[serde(default)] + #[schemars(length(max = 32))] + parameters: Vec, +} +impl ServerHandler for PostgresMcp { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()).with_instructions("Queries use registered connections and the launcher's trusted identity. Only the documented bounded SELECT subset is supported.") + } + async fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + let input_schema = rmcp::handler::server::common::schema_for_type::(); + Ok(ListToolsResult { + tools: vec![Tool::new( + "query", + "Query a registered PostgreSQL connection with column authorization", + input_schema.as_ref().clone(), + )], + ..Default::default() + }) + } + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + if request.name != "query" { + return Err(ErrorData::invalid_params("unknown tool", None)); + } + let args: QueryRequest = + serde_json::from_value(serde_json::Value::Object(request.arguments.unwrap_or_default())) + .map_err(|_| ErrorData::invalid_params("invalid query arguments", None))?; + let parameters = args + .parameters + .into_iter() + .map(crate::Parameter::try_from) + .collect::, _>>() + .map_err(|_| ErrorData::invalid_params("invalid parameters", None))?; + let outcome = tokio::select! { + biased; + _ = context.ct.cancelled() => return Err(ErrorData::internal_error("request cancelled",None)), + outcome = self.service.query_with_parameters(&args.connection,&args.sql,¶meters) => outcome, + }; + let result = match outcome { + Ok(result) => CallToolResult::success(vec![ContentBlock::text( + serde_json::to_string(&result) + .map_err(|_| ErrorData::internal_error("result encoding failed", None))?, + )]), + Err(error) => { + let message = match error { + crate::Error::Denied => "access denied", + crate::Error::Limit => "resource limit exceeded", + crate::Error::Timeout => "request timed out", + crate::Error::UnsupportedSql | crate::Error::Parse(_) | crate::Error::Identifier => { + "unsupported or invalid SQL" + } + crate::Error::UnknownConnection => "unknown connection", + _ => "query unavailable", + }; + CallToolResult::error(vec![ContentBlock::text(message)]) + } + }; + Ok(CallToolResponse::Complete(result)) + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/spicedb.rs b/rsworkspace/crates/mcp/mcp-postgres/src/spicedb.rs new file mode 100644 index 000000000..cdce2b13c --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/spicedb.rs @@ -0,0 +1,106 @@ +use crate::{Error, ObjectId}; +use async_trait::async_trait; +use authzed::v1::{ + CheckBulkPermissionsRequest, CheckBulkPermissionsRequestItem, Consistency, ObjectReference, SubjectReference, + check_bulk_permissions_pair::Response, check_permission_response::Permissionship, consistency::Requirement, + permissions_service_client::PermissionsServiceClient, +}; +use tonic::{ + metadata::{Ascii, MetadataValue}, + transport::Channel, +}; + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct Resource { + pub kind: &'static str, + pub id: ObjectId, +} +#[derive(Debug, Clone)] +pub struct PermissionRequest { + pub resource: Resource, + pub subject_kind: &'static str, + pub subject: ObjectId, +} +#[async_trait] +pub trait PermissionCheck: Send + Sync { + async fn check(&self, resource: &Resource, subject_kind: &'static str, subject: &ObjectId) -> Result<(), Error> { + self.check_all(&[PermissionRequest { + resource: resource.clone(), + subject_kind, + subject: subject.clone(), + }]) + .await + } + /// Evaluate the complete admission set at one coherent authorization revision. + async fn check_all(&self, requests: &[PermissionRequest]) -> Result<(), Error>; +} +#[derive(Clone)] +pub struct SpiceDb { + client: PermissionsServiceClient, + token: MetadataValue, +} +impl SpiceDb { + pub async fn connect(endpoint: String, token: String) -> Result { + let token = MetadataValue::try_from(format!("Bearer {token}")).map_err(|_| Error::Configuration)?; + let endpoint = crate::SpiceDbEndpoint::new(endpoint)?; + let mut endpoint = Channel::from_shared(endpoint.as_str().to_owned()).map_err(|_| Error::Configuration)?; + if endpoint.uri().scheme_str() == Some("https") { + endpoint = endpoint.tls_config(tonic::transport::ClientTlsConfig::new().with_native_roots())?; + } + let client = PermissionsServiceClient::new(endpoint.connect().await?); + Ok(Self { client, token }) + } +} +#[async_trait] +impl PermissionCheck for SpiceDb { + async fn check_all(&self, requests: &[PermissionRequest]) -> Result<(), Error> { + if requests.is_empty() || requests.len() > 2048 { + return Err(Error::Denied); + } + let mut items = requests + .iter() + .map(|request| CheckBulkPermissionsRequestItem { + resource: Some(ObjectReference { + object_type: request.resource.kind.into(), + object_id: request.resource.id.as_str().into(), + }), + permission: "query".into(), + subject: Some(SubjectReference { + object: Some(ObjectReference { + object_type: request.subject_kind.into(), + object_id: request.subject.as_str().into(), + }), + optional_relation: String::new(), + }), + context: None, + }) + .collect::>(); + let mut request = tonic::Request::new(CheckBulkPermissionsRequest { + consistency: Some(Consistency { + requirement: Some(Requirement::FullyConsistent(true)), + }), + items: items.clone(), + with_tracing: false, + }); + request.metadata_mut().insert("authorization", self.token.clone()); + request.set_timeout(std::time::Duration::from_secs(2)); + let response = self.client.clone().check_bulk_permissions(request).await?.into_inner(); + if response.checked_at.as_ref().is_none_or(|token| token.token.is_empty()) + || response.pairs.len() != items.len() + { + return Err(Error::Denied); + } + for pair in response.pairs { + let echoed = pair.request.ok_or(Error::Denied)?; + let index = items.iter().position(|item| *item == echoed).ok_or(Error::Denied)?; + items.swap_remove(index); + match pair.response { + Some(Response::Item(item)) + if item.permissionship == Permissionship::HasPermission as i32 + && item.partial_caveat_info.is_none() => {} + _ => return Err(Error::Denied), + } + } + if items.is_empty() { Ok(()) } else { Err(Error::Denied) } + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/spicedb_endpoint.rs b/rsworkspace/crates/mcp/mcp-postgres/src/spicedb_endpoint.rs new file mode 100644 index 000000000..09d6dc520 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/spicedb_endpoint.rs @@ -0,0 +1,25 @@ +use crate::Error; + +#[derive(Clone)] +pub struct SpiceDbEndpoint(String); +impl SpiceDbEndpoint { + pub fn new(value: String) -> Result { + let url = url::Url::parse(&value).map_err(|_| Error::Configuration)?; + if !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || url.path() != "/" + { + return Err(Error::Configuration); + } + let local = matches!(url.host_str(), Some("localhost" | "127.0.0.1" | "[::1]" | "::1")); + if !(url.scheme() == "https" || local && url.scheme() == "http") { + return Err(Error::Configuration); + } + Ok(Self(value)) + } + pub fn as_str(&self) -> &str { + &self.0 + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs b/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs new file mode 100644 index 000000000..2e98a2100 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs @@ -0,0 +1,71 @@ +use crate::{Error, ObjectId, Resource}; +use std::{collections::BTreeSet, time::Instant}; + +#[derive(Debug)] +enum QueryOutcome { + Cancelled, + Permitted, + Denied, + Timeout, + Limited, + Rejected, + Unavailable, +} + +pub(crate) struct QueryAudit<'a> { + user: &'a ObjectId, + client: &'a ObjectId, + connection: Option<&'a ObjectId>, + started: Instant, + outcome: QueryOutcome, +} +impl<'a> QueryAudit<'a> { + pub(crate) fn start(user: &'a ObjectId, client: &'a ObjectId, connection: Option<&'a ObjectId>) -> Self { + tracing::info!( + event = "sql_query_admission", + user = user.as_str(), + client = client.as_str(), + connection = connection.map(ObjectId::as_str) + ); + Self { + user, + client, + connection, + started: Instant::now(), + outcome: QueryOutcome::Cancelled, + } + } + pub(crate) fn complete(&mut self, result: &Result) { + self.outcome = match result { + Ok(_) => QueryOutcome::Permitted, + Err(Error::Denied) => QueryOutcome::Denied, + Err(Error::Timeout) => QueryOutcome::Timeout, + Err(Error::Limit) => QueryOutcome::Limited, + Err( + Error::UnsupportedSql + | Error::UnsupportedObject + | Error::Parse(_) + | Error::UnknownConnection + | Error::Identifier, + ) => QueryOutcome::Rejected, + Err(_) => QueryOutcome::Unavailable, + }; + } +} +impl Drop for QueryAudit<'_> { + fn drop(&mut self) { + tracing::info!(event="sql_query_completion",user=self.user.as_str(),client=self.client.as_str(),connection=self.connection.map(ObjectId::as_str),outcome=?self.outcome,elapsed_ms=self.started.elapsed().as_millis() as u64); + } +} +pub(crate) fn references(resources: &BTreeSet, user: &ObjectId, client: &ObjectId, connection: &ObjectId) { + for resource in resources { + tracing::info!( + event = "sql_query_reference", + user = user.as_str(), + client = client.as_str(), + connection = connection.as_str(), + resource_kind = resource.kind, + resource_id = resource.id.as_str() + ); + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs new file mode 100644 index 000000000..e569c99ec --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs @@ -0,0 +1,49 @@ +use mcp_postgres::{ConfigInput, DatabaseOptions, EnvironmentVariable, ObjectId, RuntimeConfig, SpiceDbEndpoint}; +#[test] +fn rejects_duplicate_connection_names_and_ids() -> Result<(), Box> { + for connections in [ + "[{name='a',id='one',dsn_env='DSN'},{name='a',id='two',dsn_env='DSN'}]", + "[{name='a',id='one',dsn_env='DSN'},{name='b',id='one',dsn_env='DSN'}]", + ] { + let input: ConfigInput = toml::from_str(&format!( + "user='alice'\nclient='desktop'\nspicedb_endpoint='http://localhost:50051'\nspicedb_token_env='TOKEN'\nconnections={connections}" + ))?; + assert!( + RuntimeConfig::from_input(input, |key| Some( + if key == "DSN" { + "postgres://reader@localhost/example" + } else { + "token" + } + .into() + )) + .is_err() + ); + } + Ok(()) +} +#[test] +fn unknown_settings_fail_closed() { + let valid = include_str!("../examples/config.toml"); + assert!(toml::from_str::(valid).is_ok()); + assert!(toml::from_str::(&format!("permissions_disabled=true\n{valid}")).is_err()); +} + +#[test] +fn secrets_and_remote_transport_settings_are_validated_at_the_boundary() -> Result<(), Box> { + assert!(DatabaseOptions::new("postgres://reader@localhost/example").is_ok()); + assert!(DatabaseOptions::new("postgres://reader@db.example.test/example?sslmode=verify-full").is_ok()); + assert!(DatabaseOptions::new("postgres://reader@db.example.test/example?sslmode=disable").is_err()); + assert!(SpiceDbEndpoint::new("https://policy.example.test:443".into()).is_ok()); + assert!(SpiceDbEndpoint::new("http://policy.example.test:50051".into()).is_err()); + assert!(SpiceDbEndpoint::new("https://user:password@policy.example.test".into()).is_err()); + assert!(SpiceDbEndpoint::new("https://policy.example.test/path".into()).is_err()); + assert!(EnvironmentVariable::new("MCP_POSTGRES_DSN".into()).is_ok()); + assert!(EnvironmentVariable::new("not-an-env-reference".into()).is_err()); + assert!(ObjectId::new(String::new()).is_err()); + assert!(ObjectId::new("untrusted:*".into()).is_err()); + + let input: ConfigInput = toml::from_str(include_str!("../examples/config.toml"))?; + assert!(RuntimeConfig::from_input(input, |_| None).is_err()); + Ok(()) +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/parser.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/parser.rs new file mode 100644 index 000000000..e4830806b --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/parser.rs @@ -0,0 +1,47 @@ +use mcp_postgres::Query; + +#[test] +fn supports_explicit_columns_and_dependencies() -> Result<(), Box> { + for sql in [ + "SELECT id FROM public.customers", + "SELECT c.id AS customer FROM public.customers c WHERE c.email = 'example' ORDER BY c.id DESC LIMIT 10", + "SELECT c.id, o.id FROM public.customers c INNER JOIN public.orders o ON c.id = o.customer_id WHERE o.id > 3 AND c.id < 100", + "SELECT \"odd\"\"name\" FROM \"business schema\".\"table name\"", + ] { + Query::parse(sql)?; + } + Ok(()) +} +#[test] +fn rejects_every_unsupported_construct() { + for sql in [ + "SELECT a.id FROM public.a a JOIN public.b b ON c.id=b.id JOIN public.c c ON c.id=a.id", + "SELECT a.id FROM public.a a JOIN public.b b ON id=b.id", + "SELECT * FROM public.customers", + "SELECT c.* FROM public.customers c", + "SELECT 1", + "SELECT id FROM customers", + "SELECT id FROM public.customers; SELECT id FROM public.customers", + "DELETE FROM public.customers", + "SELECT id INTO public.copy FROM public.customers", + "SELECT pg_sleep(1) FROM public.customers", + "SELECT id::text FROM public.customers", + "SELECT id FROM public.customers FOR UPDATE", + "WITH x AS (SELECT id FROM public.customers) SELECT id FROM x", + "SELECT id FROM public.customers WHERE EXISTS (SELECT 1)", + "SELECT id FROM public.customers GROUP BY id", + "SELECT DISTINCT id FROM public.customers", + "SELECT id FROM public.customers UNION SELECT id FROM public.orders", + "SELECT id FROM public.customers WHERE id OPERATOR(public.=) 1", + "SELECT id FROM public.customers ORDER BY id USING OPERATOR(public.<)", + "SELECT id FROM public.customers TABLESAMPLE SYSTEM (10)", + "SELECT id FROM public.customers OFFSET 5", + "SELECT id FROM public.customers c NATURAL JOIN public.orders o", + "SELECT id FROM public.customers c JOIN public.orders o USING(id)", + "SELECT id FROM public.customers c LEFT JOIN public.orders o ON c.id=o.id", + "SELECT id FROM public.customers LIMIT ALL", + "SELECT id FROM public.customers WHERE email COLLATE public.custom = 'x'", + ] { + assert!(Query::parse(sql).is_err(), "accepted unsupported SQL: {sql}"); + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs new file mode 100644 index 000000000..e964e2fe7 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs @@ -0,0 +1,683 @@ +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +use std::collections::BTreeSet; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use async_trait::async_trait; +use authzed::v1::{ + ObjectReference, Relationship, RelationshipUpdate, SubjectReference, WriteRelationshipsRequest, WriteSchemaRequest, + permissions_service_client::PermissionsServiceClient, relationship_update::Operation, + schema_service_client::SchemaServiceClient, +}; +use mcp_postgres::{ + ConfigInput, Error, Parameter, ParameterInput, PermissionCheck, PermissionRequest, PostgresMcp, QueryService, + RuntimeConfig, +}; +use rmcp::{ + ServiceExt, + model::{CallToolRequest, CallToolRequestParams, ClientRequest}, + service::PeerRequestOptions, +}; +use serde_json::json; +use sqlx::{Executor, PgPool}; +use testcontainers_modules::postgres::Postgres; +use testcontainers_modules::testcontainers::{ + ContainerAsync, GenericImage, ImageExt, + core::{IntoContainerPort, WaitFor}, + runners::AsyncRunner, +}; +use tokio::sync::Notify; + +#[derive(Default)] +struct Permissions { + denied: Mutex>, + checked: Mutex>, +} + +#[async_trait] +impl PermissionCheck for Permissions { + async fn check_all(&self, requests: &[PermissionRequest]) -> Result<(), Error> { + let denied = self.denied.lock().unwrap(); + let mut checked = self.checked.lock().unwrap(); + for request in requests { + checked.push(( + request.resource.kind.to_owned(), + request.resource.id.as_str().to_owned(), + request.subject.as_str().to_owned(), + )); + if denied.contains(&(request.resource.id.as_str().to_owned(), request.subject_kind.to_owned())) { + return Err(Error::Denied); + } + } + Ok(()) + } +} + +struct Fixture { + _container: ContainerAsync, + admin: PgPool, + dsn: String, +} + +impl Fixture { + async fn new() -> Self { + let container = Postgres::default().with_tag("18.3-alpine").start().await.unwrap(); + let host = container.get_host().await.unwrap(); + let port = container.get_host_port_ipv4(5432).await.unwrap(); + let admin = PgPool::connect(&format!("postgres://postgres:postgres@{host}:{port}/postgres")) + .await + .unwrap(); + sqlx::raw_sql( + "CREATE ROLE query_role LOGIN PASSWORD 'query-only'; + CREATE TABLE public.customers (id integer, email text, active boolean); + INSERT INTO public.customers VALUES (1, 'one@example.test', true), (2, 'two@example.test', false); + CREATE TABLE public.orders (id integer, customer_id integer, note text); + INSERT INTO public.orders VALUES (9, 1, 'one@example.test'); + GRANT USAGE ON SCHEMA public TO query_role; + GRANT SELECT ON ALL TABLES IN SCHEMA public TO query_role; + ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO query_role;", + ) + .execute(&admin) + .await + .unwrap(); + Self { + _container: container, + admin, + dsn: format!("postgres://query_role:query-only@{host}:{port}/postgres"), + } + } + + async fn install_guard(&self) { + sqlx::raw_sql(include_str!("../sql/install-ddl-guard.sql")) + .execute(&self.admin) + .await + .unwrap(); + } + + async fn service(&self, permissions: Arc) -> QueryService { + let input: ConfigInput = toml::from_str(include_str!("../examples/config.toml")).unwrap(); + let config = RuntimeConfig::from_input(input, |key| match key { + "MCP_POSTGRES_SPICEDB_TOKEN" => Some("test-only-token".into()), + "MCP_POSTGRES_ANALYTICS_DSN" => Some(self.dsn.clone()), + _ => None, + }) + .unwrap(); + QueryService::connect(config, permissions).await.unwrap() + } + + async fn column_id(&self, table: &str, column: &str) -> String { + let (database, relation, attribute): (i64, i64, i16) = sqlx::query_as("SELECT d.oid::bigint,c.oid::bigint,a.attnum FROM pg_catalog.pg_database d,pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace JOIN pg_catalog.pg_attribute a ON a.attrelid=c.oid WHERE d.datname=current_database() AND n.nspname='public' AND c.relname=$1 AND a.attname=$2") + .bind(table).bind(column).fetch_one(&self.admin).await.unwrap(); + format!("analytics-local-v1/{database}/{relation}/{attribute}") + } +} + +#[tokio::test] +async fn actual_postgres_returns_authorized_results_and_preserves_limits() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + let permissions = Arc::new(Permissions::default()); + let service = fixture.service(permissions.clone()).await; + let result = service + .query( + "analytics-readonly", + "SELECT c.id, c.active FROM public.customers c ORDER BY c.id LIMIT 1", + ) + .await + .unwrap(); + assert_eq!(result.columns, ["id", "active"]); + assert_eq!(result.rows, vec![vec![json!(1), json!(true)]]); + let empty = service + .query("analytics-readonly", "SELECT id FROM public.customers LIMIT 0") + .await + .unwrap(); + assert_eq!(empty.columns, ["id"]); + assert!(empty.rows.is_empty()); + let absent = service + .query("analytics-readonly", "SELECT id FROM public.customers WHERE id = 999") + .await + .unwrap(); + assert_eq!(absent.columns, ["id"]); + assert!(absent.rows.is_empty()); + let joined = service + .query( + "analytics-readonly", + "SELECT c.id, o.id AS order_id FROM public.customers c INNER JOIN public.orders o ON c.id = o.customer_id", + ) + .await + .unwrap(); + assert_eq!(joined.rows, vec![vec![json!(1), json!(9)]]); + let checks = permissions.checked.lock().unwrap(); + for kind in ["sql_connection", "sql_table", "sql_column"] { + for subject in ["example-user", "example-client"] { + assert!(checks.iter().any(|(k, _, s)| k == kind && s == subject)); + } + } +} + +#[tokio::test] +async fn forbidden_columns_block_projections_predicates_joins_and_ordering_for_either_subject() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + let permissions = Arc::new(Permissions::default()); + let email = fixture.column_id("customers", "email").await; + let service = fixture.service(permissions.clone()).await; + for subject in ["sql_user", "sql_client"] { + permissions.denied.lock().unwrap().clear(); + permissions + .denied + .lock() + .unwrap() + .insert((email.clone(), subject.into())); + for sql in [ + "SELECT email FROM public.customers", + "SELECT id FROM public.customers WHERE email = 'one@example.test'", + "SELECT c.id FROM public.customers c INNER JOIN public.orders o ON c.email = o.note", + "SELECT id FROM public.customers ORDER BY email", + ] { + assert!( + matches!(service.query("analytics-readonly", sql).await, Err(Error::Denied)), + "{sql}" + ); + } + assert!( + service + .query("analytics-readonly", "SELECT id FROM public.customers") + .await + .is_ok() + ); + } +} + +#[tokio::test] +async fn missing_disabled_or_replaced_guard_fails_closed() { + let fixture = Fixture::new().await; + let service = fixture.service(Arc::new(Permissions::default())).await; + let query = "SELECT id FROM public.customers"; + assert!(matches!( + service.query("analytics-readonly", query).await, + Err(Error::UnsupportedObject) + )); + fixture.install_guard().await; + assert!(service.query("analytics-readonly", query).await.is_ok()); + fixture + .admin + .execute("ALTER EVENT TRIGGER mcp_postgres_ddl_guard DISABLE") + .await + .unwrap(); + assert!(matches!( + service.query("analytics-readonly", query).await, + Err(Error::UnsupportedObject) + )); + fixture + .admin + .execute("ALTER EVENT TRIGGER mcp_postgres_ddl_guard ENABLE ALWAYS") + .await + .unwrap(); + fixture.admin.execute("CREATE OR REPLACE FUNCTION mcp_postgres_guard.lock_ddl() RETURNS event_trigger LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog AS 'BEGIN RETURN; END;'").await.unwrap(); + assert!(matches!( + service.query("analytics-readonly", query).await, + Err(Error::UnsupportedObject) + )); +} + +#[tokio::test] +async fn catalog_features_and_unsupported_sql_cannot_expose_hidden_dependencies() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + sqlx::raw_sql("CREATE TABLE public.rls_table (id integer); ALTER TABLE public.rls_table ENABLE ROW LEVEL SECURITY; + CREATE TABLE public.generated_table (id integer, generated integer GENERATED ALWAYS AS (id+1) STORED); + CREATE TABLE public.parent_table (id integer); CREATE TABLE public.child_table () INHERITS (public.parent_table); + CREATE DOMAIN public.custom_text AS text; CREATE TABLE public.custom_type_table (id integer, secret public.custom_text); + CREATE VIEW public.customer_view AS SELECT id FROM public.customers; + CREATE TABLE public.partitioned_table (id integer) PARTITION BY RANGE (id); + CREATE TABLE public.partition_child PARTITION OF public.partitioned_table FOR VALUES FROM (0) TO (10); + CREATE TABLE public.index_expression_table (id integer); CREATE INDEX ON public.index_expression_table ((id+1));") + .execute(&fixture.admin).await.unwrap(); + let service = fixture.service(Arc::new(Permissions::default())).await; + for table in [ + "rls_table", + "generated_table", + "parent_table", + "child_table", + "custom_type_table", + "customer_view", + "partitioned_table", + "partition_child", + "index_expression_table", + ] { + assert!( + service + .query("analytics-readonly", &format!("SELECT id FROM public.{table}")) + .await + .is_err(), + "{table}" + ); + } + for sql in [ + "SELECT * FROM public.customers", + "SELECT c.* FROM public.customers c", + "SELECT c FROM public.customers c", + "SELECT ctid FROM public.customers", + "SELECT tableoid FROM public.customers", + "SELECT id::text FROM public.customers", + "SELECT length(email) FROM public.customers", + "SELECT id FROM public.customers WHERE email OPERATOR(public.=) 'x'", + "SELECT id FROM public.customers FOR UPDATE", + "SELECT id INTO public.copied FROM public.customers", + "WITH c AS (SELECT id FROM public.customers) SELECT id FROM c", + "SELECT id FROM public.customers; SELECT email FROM public.customers", + "SELECT id FROM public.customers WHERE id IN (SELECT id FROM public.orders)", + ] { + assert!(service.query("analytics-readonly", sql).await.is_err(), "{sql}"); + } +} + +struct PausingPermissions { + entered: Notify, + release: Notify, +} + +#[async_trait] +impl PermissionCheck for PausingPermissions { + async fn check_all(&self, requests: &[PermissionRequest]) -> Result<(), Error> { + for request in requests { + if request.resource.kind == "sql_column" { + self.entered.notify_one(); + self.release.notified().await; + } + } + Ok(()) + } +} + +#[tokio::test] +async fn ddl_guard_blocks_schema_rename_through_authorization_and_execution() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + let gate = Arc::new(PausingPermissions { + entered: Notify::new(), + release: Notify::new(), + }); + let service = Arc::new(fixture.service(gate.clone()).await); + let task = tokio::spawn({ + let service = service.clone(); + async move { + service + .query("analytics-readonly", "SELECT id FROM public.customers") + .await + } + }); + tokio::time::timeout(Duration::from_secs(3), gate.entered.notified()) + .await + .unwrap(); + let mut migration = fixture.admin.acquire().await.unwrap(); + migration.execute("SET lock_timeout='150ms'").await.unwrap(); + let error = migration + .execute("ALTER SCHEMA public RENAME TO moved") + .await + .unwrap_err(); + assert_eq!( + error.as_database_error().and_then(|e| e.code()).as_deref(), + Some("55P03") + ); + gate.release.notify_one(); + tokio::time::timeout(Duration::from_secs(3), gate.entered.notified()) + .await + .unwrap(); + gate.release.notify_one(); + assert_eq!(task.await.unwrap().unwrap().rows.len(), 2); + migration.execute("ALTER SCHEMA public RENAME TO moved").await.unwrap(); +} + +#[tokio::test] +async fn resource_limits_abort_without_returning_partial_results() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + sqlx::raw_sql( + "CREATE TABLE public.many_rows (id integer); INSERT INTO public.many_rows SELECT generate_series(1,1001); + CREATE TABLE public.large_rows (content text); INSERT INTO public.large_rows VALUES (repeat('x',1048577));", + ) + .execute(&fixture.admin) + .await + .unwrap(); + let service = fixture.service(Arc::new(Permissions::default())).await; + assert!(matches!( + service + .query("analytics-readonly", "SELECT id FROM public.many_rows") + .await, + Err(Error::Limit) + )); + assert!(matches!( + service + .query("analytics-readonly", "SELECT content FROM public.large_rows") + .await, + Err(Error::Limit) + )); + assert_eq!( + service + .query("analytics-readonly", "SELECT id FROM public.customers LIMIT 1") + .await + .unwrap() + .rows + .len(), + 1 + ); +} + +#[tokio::test] +async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_widths() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + fixture + .admin + .execute("CREATE TABLE public.widths (small smallint, regular integer, large bigint)") + .await + .unwrap(); + fixture + .admin + .execute("INSERT INTO public.widths VALUES (7,7,7)") + .await + .unwrap(); + let service = fixture.service(Arc::new(Permissions::default())).await; + for column in ["small", "regular", "large"] { + let parameters = [Parameter::try_from(ParameterInput::Integer(7)).unwrap()]; + let result = service + .query_with_parameters( + "analytics-readonly", + &format!("SELECT {column} FROM public.widths WHERE {column} = $1"), + ¶meters, + ) + .await + .unwrap(); + assert_eq!(result.rows, vec![vec![json!(7)]]); + } + let parameters = [Parameter::try_from(ParameterInput::Text("' OR true --".into())).unwrap()]; + let result = service + .query_with_parameters( + "analytics-readonly", + "SELECT id FROM public.customers WHERE email = $1", + ¶meters, + ) + .await + .unwrap(); + assert!(result.rows.is_empty()); + let parameters = [Parameter::try_from(ParameterInput::Boolean(true)).unwrap()]; + let result = service + .query_with_parameters( + "analytics-readonly", + "SELECT id FROM public.customers WHERE active = $1", + ¶meters, + ) + .await + .unwrap(); + assert_eq!(result.rows, vec![vec![json!(1)]]); + assert!( + service + .query("analytics-readonly", "SELECT id FROM public.customers WHERE id = $1") + .await + .is_err() + ); +} + +#[tokio::test] +async fn mcp_client_initializes_discovers_queries_and_cannot_spoof_identity() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + let service = fixture.service(Arc::new(Permissions::default())).await; + let (client_io, server_io) = tokio::io::duplex(65536); + let server_task = tokio::spawn(async move { + PostgresMcp::new(service) + .serve(server_io) + .await + .unwrap() + .waiting() + .await + .unwrap(); + }); + let client = ().serve(client_io).await.unwrap(); + let tools = client.list_tools(None).await.unwrap(); + assert_eq!(tools.tools.len(), 1); + assert_eq!(tools.tools[0].name, "query"); + let result = client + .call_tool( + CallToolRequestParams::new("query").with_arguments( + json!({"connection":"analytics-readonly","sql":"SELECT id FROM public.customers LIMIT 1"}) + .as_object() + .unwrap() + .clone(), + ), + ) + .await + .unwrap(); + assert_ne!(result.is_error, Some(true)); + let spoofed = client.call_tool(CallToolRequestParams::new("query").with_arguments(json!({"connection":"analytics-readonly","sql":"SELECT id FROM public.customers","user":"administrator"}).as_object().unwrap().clone())).await; + assert!(spoofed.is_err()); + let invalid = client + .call_tool( + CallToolRequestParams::new("query").with_arguments( + json!({"connection":"analytics-readonly","sql":"SELECT secret-marker-839 '"}) + .as_object() + .unwrap() + .clone(), + ), + ) + .await + .unwrap(); + assert_eq!(invalid.is_error, Some(true)); + assert!(!serde_json::to_string(&invalid).unwrap().contains("secret-marker-839")); + client.cancel().await.unwrap(); + tokio::time::timeout(Duration::from_secs(3), server_task) + .await + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn cancelling_an_mcp_request_releases_its_database_guard() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + let gate = Arc::new(PausingPermissions { + entered: Notify::new(), + release: Notify::new(), + }); + let service = fixture.service(gate.clone()).await; + let (client_io, server_io) = tokio::io::duplex(65536); + let server = tokio::spawn(async move { + PostgresMcp::new(service) + .serve(server_io) + .await + .unwrap() + .waiting() + .await + .unwrap(); + }); + let client = ().serve(client_io).await.unwrap(); + let params = CallToolRequestParams::new("query").with_arguments( + json!({"connection":"analytics-readonly","sql":"SELECT id FROM public.customers"}) + .as_object() + .unwrap() + .clone(), + ); + let pending = client + .peer() + .send_request_with_option( + ClientRequest::CallToolRequest(CallToolRequest::new(params)), + PeerRequestOptions::no_options(), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(3), gate.entered.notified()) + .await + .unwrap(); + let unlocked: bool = sqlx::query_scalar("SELECT pg_catalog.pg_try_advisory_xact_lock(1879445091,1347635021)") + .fetch_one(&fixture.admin) + .await + .unwrap(); + assert!(!unlocked); + pending.cancel(Some("test cancellation".into())).await.unwrap(); + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let unlocked: bool = + sqlx::query_scalar("SELECT pg_catalog.pg_try_advisory_xact_lock(1879445091,1347635021)") + .fetch_one(&fixture.admin) + .await + .unwrap(); + if unlocked { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .unwrap(); + client.cancel().await.unwrap(); + tokio::time::timeout(Duration::from_secs(3), server) + .await + .unwrap() + .unwrap(); +} + +fn spice_request(value: T) -> tonic::Request { + let mut request = tonic::Request::new(value); + request + .metadata_mut() + .insert("authorization", "Bearer test-only-token".parse().unwrap()); + request +} + +#[tokio::test] +async fn mcp_spicedb_and_postgres_enforce_hard_column_denial_together() { + let fixture = Fixture::new().await; + fixture.install_guard().await; + let spice = GenericImage::new("authzed/spicedb", "v1.53.0") + .with_exposed_port(50051.tcp()) + .with_wait_for(WaitFor::message_on_either_std("grpc server started serving")) + .with_startup_timeout(Duration::from_secs(90)) + .with_cmd([ + "serve", + "--datastore-engine", + "memory", + "--grpc-preshared-key", + "test-only-token", + ]) + .start() + .await + .unwrap(); + let endpoint = format!( + "http://{}:{}", + spice.get_host().await.unwrap(), + spice.get_host_port_ipv4(50051.tcp()).await.unwrap() + ); + let channel = tonic::transport::Channel::from_shared(endpoint.clone()) + .unwrap() + .connect() + .await + .unwrap(); + SchemaServiceClient::new(channel.clone()) + .write_schema(spice_request(WriteSchemaRequest { + schema: include_str!("../sql/schema.zed").into(), + })) + .await + .unwrap(); + let mut writer = PermissionsServiceClient::new(channel); + let id_column = fixture.column_id("customers", "id").await; + let email_column = fixture.column_id("customers", "email").await; + let table = id_column.rsplit_once('/').unwrap().0; + let update = |kind: &str, id: &str, relation: &str, subject_kind: &str, subject: &str| RelationshipUpdate { + operation: Operation::Touch as i32, + relationship: Some(Relationship { + resource: Some(ObjectReference { + object_type: kind.into(), + object_id: id.into(), + }), + relation: relation.into(), + subject: Some(SubjectReference { + object: Some(ObjectReference { + object_type: subject_kind.into(), + object_id: subject.into(), + }), + optional_relation: String::new(), + }), + ..Default::default() + }), + }; + let mut grants = Vec::new(); + for (kind, id) in [ + ("sql_connection", "analytics-local-v1"), + ("sql_table", table), + ("sql_column", id_column.as_str()), + ("sql_column", email_column.as_str()), + ] { + grants.push(update(kind, id, "reader", "sql_user", "example-user")); + grants.push(update(kind, id, "reader", "sql_client", "example-client")); + } + writer + .write_relationships(spice_request(WriteRelationshipsRequest { + updates: grants, + ..Default::default() + })) + .await + .unwrap(); + let adapter = Arc::new( + mcp_postgres::SpiceDb::connect(endpoint, "test-only-token".into()) + .await + .unwrap(), + ); + let service = fixture.service(adapter).await; + let (client_io, server_io) = tokio::io::duplex(65536); + let server = tokio::spawn(async move { + PostgresMcp::new(service) + .serve(server_io) + .await + .unwrap() + .waiting() + .await + .unwrap(); + }); + let client = ().serve(client_io).await.unwrap(); + let query = |sql: &str| { + CallToolRequestParams::new("query").with_arguments( + json!({"connection":"analytics-readonly","sql":sql}) + .as_object() + .unwrap() + .clone(), + ) + }; + let allowed = client + .call_tool(query("SELECT email FROM public.customers ORDER BY id")) + .await + .unwrap(); + assert_ne!(allowed.is_error, Some(true)); + let wire = serde_json::to_value(&allowed).unwrap(); + let data: serde_json::Value = serde_json::from_str(wire["content"][0]["text"].as_str().unwrap()).unwrap(); + assert_eq!(data["rows"][0][0], "one@example.test"); + writer + .write_relationships(spice_request(WriteRelationshipsRequest { + updates: vec![update("sql_column", &email_column, "forbidden", "sql_client", "*")], + ..Default::default() + })) + .await + .unwrap(); + for sql in [ + "SELECT email FROM public.customers", + "SELECT id FROM public.customers WHERE email = 'one@example.test'", + ] { + let denied = client.call_tool(query(sql)).await.unwrap(); + assert_eq!(denied.is_error, Some(true)); + let wire = serde_json::to_string(&denied).unwrap(); + assert!(wire.contains("access denied")); + assert!(!wire.contains("one@example.test")); + } + let still_allowed = client + .call_tool(query("SELECT id FROM public.customers")) + .await + .unwrap(); + assert_ne!(still_allowed.is_error, Some(true)); + client.cancel().await.unwrap(); + tokio::time::timeout(Duration::from_secs(3), server) + .await + .unwrap() + .unwrap(); +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/spicedb.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/spicedb.rs new file mode 100644 index 000000000..5bc130417 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/spicedb.rs @@ -0,0 +1,228 @@ +use std::{error::Error as StdError, time::Duration}; + +use authzed::v1::{ + ContextualizedCaveat, ObjectReference, Relationship, RelationshipUpdate, SubjectReference, + WriteRelationshipsRequest, WriteSchemaRequest, permissions_service_client::PermissionsServiceClient, + relationship_update::Operation, schema_service_client::SchemaServiceClient, +}; +use mcp_postgres::{Error, ObjectId, PermissionCheck, Resource, SpiceDb}; +use testcontainers_modules::testcontainers::{ + GenericImage, ImageExt, + core::{IntoContainerPort, WaitFor}, + runners::AsyncRunner, +}; + +const TOKEN: &str = "test-only-token"; + +fn authorized(value: T) -> Result, Box> { + let mut request = tonic::Request::new(value); + request + .metadata_mut() + .insert("authorization", format!("Bearer {TOKEN}").parse()?); + Ok(request) +} + +fn relationship(kind: &str, id: &str, relation: &str, subject_kind: &str, subject: &str) -> RelationshipUpdate { + RelationshipUpdate { + operation: Operation::Touch as i32, + relationship: Some(Relationship { + resource: Some(ObjectReference { + object_type: kind.into(), + object_id: id.into(), + }), + relation: relation.into(), + subject: Some(SubjectReference { + object: Some(ObjectReference { + object_type: subject_kind.into(), + object_id: subject.into(), + }), + optional_relation: String::new(), + }), + ..Default::default() + }), + } +} + +#[tokio::test] +async fn actual_spicedb_enforces_schema_grants_denials_and_subject_types() -> Result<(), Box> { + let container = GenericImage::new("authzed/spicedb", "v1.53.0") + .with_exposed_port(50051.tcp()) + .with_wait_for(WaitFor::message_on_either_std("grpc server started serving")) + .with_startup_timeout(Duration::from_secs(90)) + .with_cmd(["serve", "--datastore-engine", "memory", "--grpc-preshared-key", TOKEN]) + .start() + .await?; + let endpoint = format!( + "http://{}:{}", + container.get_host().await?, + container.get_host_port_ipv4(50051.tcp()).await? + ); + let channel = tonic::transport::Channel::from_shared(endpoint.clone())? + .connect() + .await?; + let mut schema = SchemaServiceClient::new(channel.clone()); + schema + .write_schema(authorized(WriteSchemaRequest { + schema: include_str!("../sql/schema.zed").into(), + })?) + .await?; + let mut writer = PermissionsServiceClient::new(channel); + let adapter = SpiceDb::connect(endpoint.clone(), TOKEN.into()).await?; + let alice = ObjectId::new("alice".into())?; + let desktop = ObjectId::new("desktop".into())?; + let stranger = ObjectId::new("stranger".into())?; + + for kind in ["sql_connection", "sql_table", "sql_column"] { + let resource = Resource { + kind, + id: ObjectId::new("registered/123/456/1".into())?, + }; + assert!(matches!( + adapter.check(&resource, "sql_user", &alice).await, + Err(Error::Denied) + )); + writer + .write_relationships(authorized(WriteRelationshipsRequest { + updates: vec![ + relationship(kind, resource.id.as_str(), "reader", "sql_user", "alice"), + relationship(kind, resource.id.as_str(), "reader", "sql_client", "desktop"), + ], + ..Default::default() + })?) + .await?; + adapter.check(&resource, "sql_user", &alice).await?; + adapter.check(&resource, "sql_client", &desktop).await?; + assert!(matches!( + adapter.check(&resource, "sql_user", &stranger).await, + Err(Error::Denied) + )); + assert!(matches!( + adapter.check(&resource, "sql_client", &alice).await, + Err(Error::Denied) + )); + assert!(matches!( + adapter.check(&resource, "sql_user", &desktop).await, + Err(Error::Denied) + )); + + writer + .write_relationships(authorized(WriteRelationshipsRequest { + updates: vec![relationship( + kind, + resource.id.as_str(), + "forbidden", + "sql_user", + "alice", + )], + ..Default::default() + })?) + .await?; + assert!(matches!( + adapter.check(&resource, "sql_user", &alice).await, + Err(Error::Denied) + )); + adapter.check(&resource, "sql_client", &desktop).await?; + + writer + .write_relationships(authorized(WriteRelationshipsRequest { + updates: vec![relationship(kind, resource.id.as_str(), "forbidden", "sql_client", "*")], + ..Default::default() + })?) + .await?; + assert!(matches!( + adapter.check(&resource, "sql_client", &desktop).await, + Err(Error::Denied) + )); + + let other_connection = Resource { + kind, + id: ObjectId::new("other/123/456/1".into())?, + }; + assert!(matches!( + adapter.check(&other_connection, "sql_user", &alice).await, + Err(Error::Denied) + )); + } + + let global = Resource { + kind: "sql_column", + id: ObjectId::new("registered/123/456/2".into())?, + }; + writer + .write_relationships(authorized(WriteRelationshipsRequest { + updates: vec![ + relationship("sql_column", global.id.as_str(), "reader", "sql_user", "alice"), + relationship("sql_column", global.id.as_str(), "reader", "sql_client", "desktop"), + relationship("sql_column", global.id.as_str(), "forbidden", "sql_user", "*"), + relationship("sql_column", global.id.as_str(), "forbidden", "sql_client", "*"), + ], + ..Default::default() + })?) + .await?; + assert!(matches!( + adapter.check(&global, "sql_user", &alice).await, + Err(Error::Denied) + )); + assert!(matches!( + adapter.check(&global, "sql_client", &desktop).await, + Err(Error::Denied) + )); + + let invalid = writer + .write_relationships(authorized(WriteRelationshipsRequest { + updates: vec![relationship("sql_column", "invalid", "reader", "sql_user", "*")], + ..Default::default() + })?) + .await; + assert!( + invalid.is_err(), + "reader wildcard must not be accepted by production schema" + ); + + let resource = Resource { + kind: "sql_column", + id: ObjectId::new("registered/123/456/1".into())?, + }; + let unauthenticated = SpiceDb::connect(endpoint, "wrong-test-token".into()).await?; + assert!(matches!( + unauthenticated.check(&resource, "sql_user", &alice).await, + Err(Error::Authorization(_)) + )); + assert!(matches!( + adapter.check(&resource, "missing_subject_type", &alice).await, + Err(Error::Denied) + )); + + let extended = format!( + "{}\ncaveat request_approved(approved bool) {{ approved }}\ndefinition conditional_resource {{\n relation reader: sql_user with request_approved\n permission query = reader\n}}", + include_str!("../sql/schema.zed") + ); + schema + .write_schema(authorized(WriteSchemaRequest { schema: extended })?) + .await?; + let mut caveated = relationship("conditional_resource", "conditional", "reader", "sql_user", "alice"); + if let Some(relationship) = caveated.relationship.as_mut() { + relationship.optional_caveat = Some(ContextualizedCaveat { + caveat_name: "request_approved".into(), + context: None, + }); + } + writer + .write_relationships(authorized(WriteRelationshipsRequest { + updates: vec![caveated], + ..Default::default() + })?) + .await?; + let conditional = Resource { + kind: "conditional_resource", + id: ObjectId::new("conditional".into())?, + }; + assert!( + matches!( + adapter.check(&conditional, "sql_user", &alice).await, + Err(Error::Denied) + ), + "conditional permissions must fail closed" + ); + Ok(()) +} From 5a590694a6a5b9f830ac9d6a94b3a3445a9ed79f Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Mon, 19 Oct 2026 17:58:45 -0400 Subject: [PATCH 2/7] fix(mcp): keep query authorization outside target databases Signed-off-by: Yordis Prieto --- .../mcp-postgres/sql/install-ddl-guard.sql | 25 -- .../crates/mcp/mcp-postgres/src/lib.rs | 11 +- .../crates/mcp/mcp-postgres/src/postgres.rs | 71 +++-- .../mcp/mcp-postgres/tests/cursor_binding.rs | 63 ++++ .../crates/mcp/mcp-postgres/tests/postgres.rs | 289 ++++++++++++++---- 5 files changed, 343 insertions(+), 116 deletions(-) delete mode 100644 rsworkspace/crates/mcp/mcp-postgres/sql/install-ddl-guard.sql create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/cursor_binding.rs diff --git a/rsworkspace/crates/mcp/mcp-postgres/sql/install-ddl-guard.sql b/rsworkspace/crates/mcp/mcp-postgres/sql/install-ddl-guard.sql deleted file mode 100644 index e1ae888c0..000000000 --- a/rsworkspace/crates/mcp/mcp-postgres/sql/install-ddl-guard.sql +++ /dev/null @@ -1,25 +0,0 @@ -BEGIN; - -SELECT pg_catalog.pg_advisory_xact_lock(1879445091, 1347635021); - -CREATE SCHEMA mcp_postgres_guard; -REVOKE ALL ON SCHEMA mcp_postgres_guard FROM PUBLIC; - -CREATE FUNCTION mcp_postgres_guard.lock_ddl() -RETURNS event_trigger -LANGUAGE plpgsql -SECURITY DEFINER -SET search_path = pg_catalog -AS $guard$BEGIN - PERFORM pg_catalog.pg_advisory_xact_lock(1879445091, 1347635021); -END;$guard$; - -REVOKE ALL ON FUNCTION mcp_postgres_guard.lock_ddl() FROM PUBLIC; - -CREATE EVENT TRIGGER mcp_postgres_ddl_guard -ON ddl_command_start -EXECUTE FUNCTION mcp_postgres_guard.lock_ddl(); - -ALTER EVENT TRIGGER mcp_postgres_ddl_guard ENABLE ALWAYS; - -COMMIT; diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs b/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs index 6fe3fb10f..facd92215 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs @@ -3,15 +3,16 @@ //! Named connections and user/client identities come exclusively from runtime configuration. //! The bounded SELECT grammar rejects unsupported syntax and catalog features. //! Every referenced column needs an explicit grant for both identities in one SpiceDB snapshot. -//! The required administrator-installed DDL guard fences catalog changes until the query ends. +//! Transactional relation locks and cursor metadata bind authorization to the executed relations. //! Production access and deployment are outside this crate's local implementation contract. //! //! The launcher and database administrators are trusted. Callers must not control the configuration //! or credentials, and the query role must have only the required SELECT privileges. -//! Install `sql/install-ddl-guard.sql` in each registered database. Migrations must take the session -//! advisory lock `pg_advisory_lock(1879445091, 1347635021)` before touching relations and retain it -//! through all transactions. This is essential for concurrent index operations, guard maintenance, -//! and shared-object changes that the event trigger cannot fence on its own. +//! The service installs no database objects and requires no migration coordination. A read-only +//! transaction locks supported base tables, verifies cursor relation identities, and authorizes +//! before fetching any rows. Database administrators, server hooks, and catalog integrity remain +//! trusted; concurrent privileged changes to executable objects, indexes, or shared objects are +//! outside this boundary. //! Replacing or restoring a database target requires a new connection ID when catalog IDs can repeat. //! Permissions are checked at admission; revocation does not retroactively cancel an admitted query. diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs index dda89f144..10fa365a2 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs @@ -6,7 +6,7 @@ use futures_util::TryStreamExt; use serde::Serialize; use sqlx::{ AssertSqlSafe, Column as _, Connection, Executor, PgConnection, PgPool, Row, SqlSafeStr, Statement, TypeInfo, - postgres::{PgPoolOptions, PgRow}, + postgres::{PgPoolOptions, PgRow, types::Oid}, }; use std::{ collections::{BTreeMap, BTreeSet}, @@ -125,17 +125,12 @@ impl QueryService { if unsafe_role { return Err(Error::UnsupportedObject); } - tx.execute("SELECT pg_catalog.pg_advisory_xact_lock_shared(1879445091,1347635021)") - .await?; - let guard: bool = sqlx::query_scalar("SELECT EXISTS (SELECT 1 FROM pg_catalog.pg_event_trigger e JOIN pg_catalog.pg_proc p ON p.oid=e.evtfoid JOIN pg_catalog.pg_namespace n ON n.oid=p.pronamespace WHERE e.evtname='mcp_postgres_ddl_guard' AND e.evtevent='ddl_command_start' AND e.evtenabled='A' AND e.evttags IS NULL AND n.nspname='mcp_postgres_guard' AND p.proname='lock_ddl' AND p.prosecdef AND p.proconfig = ARRAY['search_path=pg_catalog']::text[] AND p.pronargs=0 AND p.prorettype='pg_catalog.event_trigger'::regtype AND p.prolang=(SELECT oid FROM pg_catalog.pg_language WHERE lanname='plpgsql') AND p.proowner=n.nspowner AND EXISTS (SELECT 1 FROM pg_catalog.pg_roles r WHERE r.oid=p.proowner AND r.rolsuper) AND NOT EXISTS (SELECT 1 FROM pg_catalog.aclexplode(COALESCE(n.nspacl,pg_catalog.acldefault('n',n.nspowner))) a WHERE a.privilege_type='CREATE' AND a.grantee<>n.nspowner) AND p.prosrc=$1)").bind("BEGIN\n PERFORM pg_catalog.pg_advisory_xact_lock(1879445091, 1347635021);\nEND;").fetch_one(&mut *tx).await?; - if !guard { - return Err(Error::UnsupportedObject); - } let database: i64 = sqlx::query_scalar("SELECT oid::bigint FROM pg_catalog.pg_database WHERE datname=current_database()") .fetch_one(&mut *tx) .await?; let mut catalog = Vec::new(); + let mut relations = Vec::new(); let mut resources = BTreeSet::new(); resources.insert(Resource { kind: "sql_connection", @@ -148,9 +143,10 @@ impl QueryService { table.name.quoted() ); sqlx::query(AssertSqlSafe(lock)).execute(&mut *tx).await?; - let relation = sqlx::query("SELECT c.oid::bigint AS oid FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace JOIN pg_catalog.pg_am a ON a.oid=c.relam WHERE n.nspname=$1 AND c.relname=$2 AND c.relkind='r' AND a.amname='heap' AND a.oid=2 AND NOT c.relrowsecurity AND NOT c.relforcerowsecurity AND NOT c.relhasrules AND NOT c.relispartition AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits i WHERE i.inhrelid=c.oid OR i.inhparent=c.oid) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i WHERE i.indrelid=c.oid AND (i.indexprs IS NOT NULL OR i.indpred IS NOT NULL)) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i JOIN pg_catalog.pg_class ic ON ic.oid=i.indexrelid JOIN pg_catalog.pg_am ia ON ia.oid=ic.relam WHERE i.indrelid=c.oid AND ia.amname NOT IN ('btree','hash')) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i, unnest(i.indcollation) co WHERE i.indrelid=c.oid AND co<>0 AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_collation cl WHERE cl.oid=co AND cl.collnamespace='pg_catalog'::regnamespace)) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i, unnest(i.indclass) op WHERE i.indrelid=c.oid AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_opclass oc WHERE oc.oid=op AND oc.opcnamespace='pg_catalog'::regnamespace))") + let relation = sqlx::query("SELECT c.oid::bigint AS oid FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace JOIN pg_catalog.pg_am a ON a.oid=c.relam WHERE n.nspname=$1 AND c.relname=$2 AND EXISTS (SELECT 1 FROM pg_catalog.pg_locks l WHERE l.pid=pg_catalog.pg_backend_pid() AND l.locktype='relation' AND l.relation=c.oid AND l.mode='AccessShareLock' AND l.granted) AND c.relkind='r' AND a.amname='heap' AND a.oid=2 AND NOT c.relrowsecurity AND NOT c.relforcerowsecurity AND NOT c.relhasrules AND NOT c.relispartition AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits i WHERE i.inhrelid=c.oid OR i.inhparent=c.oid) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i WHERE i.indrelid=c.oid AND (i.indexprs IS NOT NULL OR i.indpred IS NOT NULL)) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i JOIN pg_catalog.pg_class ic ON ic.oid=i.indexrelid JOIN pg_catalog.pg_am ia ON ia.oid=ic.relam WHERE i.indrelid=c.oid AND ia.amname NOT IN ('btree','hash')) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i, unnest(i.indcollation) co WHERE i.indrelid=c.oid AND co<>0 AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_collation cl WHERE cl.oid=co AND cl.collnamespace='pg_catalog'::regnamespace)) AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_index i, unnest(i.indclass) op WHERE i.indrelid=c.oid AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_opclass oc WHERE oc.oid=op AND oc.opcnamespace='pg_catalog'::regnamespace))") .bind(table.schema.as_str()).bind(table.name.as_str()).fetch_optional(&mut *tx).await?.ok_or(Error::UnsupportedObject)?; let oid: i64 = relation.try_get("oid")?; + relations.push(Oid(u32::try_from(oid).map_err(|_| Error::UnsupportedObject)?)); let table_id = format!("{}/{database}/{oid}", connection_id.as_str()); resources.insert(Resource { kind: "sql_table", @@ -216,7 +212,7 @@ impl QueryService { orders.push(format!("{} {direction}", resolve(column, &query, &catalog, &mut used)?)); } for (table, attnum) in used { - let relation_oid: i64 = sqlx::query_scalar("SELECT c.oid::bigint FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid=c.relnamespace WHERE n.nspname=$1 AND c.relname=$2").bind(query.tables[table].schema.as_str()).bind(query.tables[table].name.as_str()).fetch_one(&mut *tx).await?; + let relation_oid = relations[table].0; resources.insert(Resource { kind: "sql_column", id: ObjectId::new(format!("{}/{database}/{relation_oid}/{attnum}", connection_id.as_str()))?, @@ -231,9 +227,16 @@ impl QueryService { .into_iter() .map(|(value, label)| format!("CASE WHEN ({overflow}) THEN NULL ELSE {value} END AS {label}")) .collect::>(); + let markers = query + .tables + .iter() + .enumerate() + .map(|(index, table)| format!("{}.tableoid AS __mcp_relation_{index}", table.alias.quoted())) + .collect::>(); let mut statement = format!( - "SELECT ({overflow}) AS __mcp_overflow, {} FROM {}", + "SELECT ({overflow}) AS __mcp_overflow, {}, {} FROM {}", projections.join(", "), + markers.join(", "), query.from ); if !predicates.is_empty() { @@ -243,9 +246,30 @@ impl QueryService { statement.push_str(&format!(" ORDER BY {}", orders.join(", "))); } statement.push_str(&format!(" LIMIT {}", query.limit.statement_count())); - let prepared = (&mut *tx as &mut PgConnection) - .prepare(AssertSqlSafe(statement).into_sql_str()) + let mut declaration = sqlx::query(AssertSqlSafe(format!( + "DECLARE mcp_query NO SCROLL CURSOR WITHOUT HOLD FOR {statement}" + ))); + for parameter in parameters { + declaration = match parameter.value() { + crate::parameter::ParameterValue::Boolean(v) => declaration.bind(*v), + crate::parameter::ParameterValue::Integer(v) => declaration.bind(*v), + crate::parameter::ParameterValue::Text(v) => declaration.bind(v), + }; + } + declaration.execute(&mut *tx).await?; + // A cursor retains its bound plan even if a schema name changes during authorization. + let metadata = (&mut *tx as &mut PgConnection) + .prepare("FETCH FORWARD 0 FROM mcp_query".into_sql_str()) .await?; + let visible_columns = query.select.len(); + if metadata.columns().len() != 1 + visible_columns + relations.len() { + return Err(Error::UnsupportedObject); + } + for (column, relation) in metadata.columns().iter().skip(1 + visible_columns).zip(&relations) { + if column.relation_id() != Some(*relation) || column.relation_attribute_no() != Some(-6) { + return Err(Error::UnsupportedObject); + } + } crate::telemetry::references(&resources, &self.user, &self.client, connection_id); let requests = resources .into_iter() @@ -265,16 +289,14 @@ impl QueryService { }) .collect::>(); self.permissions.check_all(&requests).await?; - let mut bound = prepared.query(); - for parameter in parameters { - bound = match parameter.value() { - crate::parameter::ParameterValue::Boolean(v) => bound.bind(*v), - crate::parameter::ParameterValue::Integer(v) => bound.bind(*v), - crate::parameter::ParameterValue::Text(v) => bound.bind(v), - }; - } - let columns = prepared.columns().iter().skip(1).map(|c| c.name().to_owned()).collect(); - let mut stream = bound.fetch(&mut *tx); + let columns = metadata + .columns() + .iter() + .skip(1) + .take(visible_columns) + .map(|c| c.name().to_owned()) + .collect(); + let mut stream = sqlx::query("FETCH FORWARD ALL FROM mcp_query").fetch(&mut *tx); let mut result = QueryResult { columns, rows: Vec::new(), @@ -288,7 +310,7 @@ impl QueryService { return Err(Error::Limit); } - let values = decode(&row)?; + let values = decode(&row, visible_columns)?; bytes += serde_json::to_vec(&values).map_err(|_| Error::Limit)?.len() + 1; if bytes > MAX_BYTES { return Err(Error::Limit); @@ -363,11 +385,12 @@ fn render_expr( )), } } -fn decode(row: &PgRow) -> Result, Error> { +fn decode(row: &PgRow, visible_columns: usize) -> Result, Error> { row.columns() .iter() .enumerate() .skip(1) + .take(visible_columns) .map(|(index, column)| { match column.type_info().name() { "BOOL" => serde_json::to_value(row.try_get::, _>(index)?), diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/cursor_binding.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/cursor_binding.rs new file mode 100644 index 000000000..c746dcfc9 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/cursor_binding.rs @@ -0,0 +1,63 @@ +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +use sqlx::{Connection, Executor, PgConnection, Row, SqlSafeStr, Statement}; +use testcontainers_modules::{ + postgres::Postgres, + testcontainers::{ImageExt, runners::AsyncRunner}, +}; + +#[tokio::test] +async fn cursor_metadata_binds_original_relation_before_namespace_changes() { + let container = Postgres::default().with_tag("18.3-alpine").start().await.unwrap(); + let host = container.get_host().await.unwrap(); + let port = container.get_host_port_ipv4(5432).await.unwrap(); + let mut admin = PgConnection::connect(&format!("postgres://postgres:postgres@{host}:{port}/postgres")) + .await + .unwrap(); + admin + .execute("CREATE ROLE query_role LOGIN PASSWORD 'query-only'") + .await + .unwrap(); + admin + .execute("CREATE TABLE public.customers (id integer)") + .await + .unwrap(); + admin.execute("INSERT INTO public.customers VALUES (1)").await.unwrap(); + admin + .execute("GRANT SELECT ON public.customers TO query_role") + .await + .unwrap(); + let expected: u32 = sqlx::query_scalar::<_, i64>("SELECT 'public.customers'::regclass::oid::bigint") + .fetch_one(&mut admin) + .await + .unwrap() + .try_into() + .unwrap(); + let mut reader = PgConnection::connect(&format!("postgres://query_role:query-only@{host}:{port}/postgres")) + .await + .unwrap(); + let mut tx = reader.begin().await.unwrap(); + tx.execute("SET TRANSACTION READ ONLY").await.unwrap(); + sqlx::query("DECLARE mcp_query NO SCROLL CURSOR WITHOUT HOLD FOR SELECT id, customers.tableoid FROM public.customers WHERE id=$1") + .bind(1_i32).execute(&mut *tx).await.unwrap(); + let metadata = (&mut *tx) + .prepare("FETCH FORWARD 0 FROM mcp_query".into_sql_str()) + .await + .unwrap(); + assert_eq!(metadata.columns()[1].relation_id().unwrap().0, expected); + assert_eq!(metadata.columns()[1].relation_attribute_no(), Some(-6)); + admin.execute("ALTER SCHEMA public RENAME TO previous").await.unwrap(); + admin.execute("CREATE SCHEMA public").await.unwrap(); + admin + .execute("CREATE TABLE public.customers (id integer)") + .await + .unwrap(); + admin.execute("INSERT INTO public.customers VALUES (9)").await.unwrap(); + let rows = sqlx::query("FETCH FORWARD ALL FROM mcp_query") + .fetch_all(&mut *tx) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::(0), 1); + assert_eq!(rows[0].columns()[1].relation_id().unwrap().0, expected); +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs index e964e2fe7..a51164ed9 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs @@ -1,6 +1,7 @@ #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] use std::collections::BTreeSet; +use std::io; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -28,6 +29,10 @@ use testcontainers_modules::testcontainers::{ runners::AsyncRunner, }; use tokio::sync::Notify; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, +}; #[derive(Default)] struct Permissions { @@ -70,6 +75,7 @@ impl Fixture { .unwrap(); sqlx::raw_sql( "CREATE ROLE query_role LOGIN PASSWORD 'query-only'; + REVOKE TEMPORARY ON DATABASE postgres FROM PUBLIC; CREATE TABLE public.customers (id integer, email text, active boolean); INSERT INTO public.customers VALUES (1, 'one@example.test', true), (2, 'two@example.test', false); CREATE TABLE public.orders (id integer, customer_id integer, note text); @@ -88,18 +94,15 @@ impl Fixture { } } - async fn install_guard(&self) { - sqlx::raw_sql(include_str!("../sql/install-ddl-guard.sql")) - .execute(&self.admin) - .await - .unwrap(); + async fn service(&self, permissions: Arc) -> QueryService { + self.service_with_dsn(permissions, &self.dsn).await } - async fn service(&self, permissions: Arc) -> QueryService { + async fn service_with_dsn(&self, permissions: Arc, dsn: &str) -> QueryService { let input: ConfigInput = toml::from_str(include_str!("../examples/config.toml")).unwrap(); let config = RuntimeConfig::from_input(input, |key| match key { "MCP_POSTGRES_SPICEDB_TOKEN" => Some("test-only-token".into()), - "MCP_POSTGRES_ANALYTICS_DSN" => Some(self.dsn.clone()), + "MCP_POSTGRES_ANALYTICS_DSN" => Some(dsn.to_owned()), _ => None, }) .unwrap(); @@ -113,10 +116,39 @@ impl Fixture { } } +#[tokio::test] +async fn ordinary_readonly_database_requires_no_installed_objects() { + let fixture = Fixture::new().await; + let service = fixture.service(Arc::new(Permissions::default())).await; + let result = service + .query( + "analytics-readonly", + "SELECT id FROM public.customers ORDER BY id LIMIT 1", + ) + .await + .unwrap(); + assert_eq!(result.rows, vec![vec![json!(1)]]); + let installed: i64 = + sqlx::query_scalar("SELECT count(*) FROM pg_catalog.pg_namespace WHERE nspname='mcp_postgres_guard'") + .fetch_one(&fixture.admin) + .await + .unwrap(); + assert_eq!(installed, 0); + let triggers: i64 = + sqlx::query_scalar("SELECT count(*) FROM pg_catalog.pg_event_trigger WHERE evtname LIKE 'mcp_postgres%'") + .fetch_one(&fixture.admin) + .await + .unwrap(); + assert_eq!(triggers, 0); + let reader = PgPool::connect(&fixture.dsn).await.unwrap(); + let can_install: bool = sqlx::query_scalar("SELECT pg_catalog.has_schema_privilege(current_user,'public','CREATE') OR pg_catalog.has_database_privilege(current_user,current_database(),'CREATE') OR pg_catalog.has_database_privilege(current_user,current_database(),'TEMP')") + .fetch_one(&reader).await.unwrap(); + assert!(!can_install); +} + #[tokio::test] async fn actual_postgres_returns_authorized_results_and_preserves_limits() { let fixture = Fixture::new().await; - fixture.install_guard().await; let permissions = Arc::new(Permissions::default()); let service = fixture.service(permissions.clone()).await; let result = service @@ -148,6 +180,16 @@ async fn actual_postgres_returns_authorized_results_and_preserves_limits() { .await .unwrap(); assert_eq!(joined.rows, vec![vec![json!(1), json!(9)]]); + let self_join = service.query("analytics-readonly", "SELECT c.id, other.id FROM public.customers c INNER JOIN public.customers other ON c.id = other.id ORDER BY c.id").await.unwrap(); + assert_eq!(self_join.rows, vec![vec![json!(1), json!(1)], vec![json!(2), json!(2)]]); + let unused_join = service + .query( + "analytics-readonly", + "SELECT c.id FROM public.customers c INNER JOIN public.orders o ON true ORDER BY c.id", + ) + .await + .unwrap(); + assert_eq!(unused_join.rows, vec![vec![json!(1)], vec![json!(2)]]); let checks = permissions.checked.lock().unwrap(); for kind in ["sql_connection", "sql_table", "sql_column"] { for subject in ["example-user", "example-client"] { @@ -159,7 +201,6 @@ async fn actual_postgres_returns_authorized_results_and_preserves_limits() { #[tokio::test] async fn forbidden_columns_block_projections_predicates_joins_and_ordering_for_either_subject() { let fixture = Fixture::new().await; - fixture.install_guard().await; let permissions = Arc::new(Permissions::default()); let email = fixture.column_id("customers", "email").await; let service = fixture.service(permissions.clone()).await; @@ -190,42 +231,9 @@ async fn forbidden_columns_block_projections_predicates_joins_and_ordering_for_e } } -#[tokio::test] -async fn missing_disabled_or_replaced_guard_fails_closed() { - let fixture = Fixture::new().await; - let service = fixture.service(Arc::new(Permissions::default())).await; - let query = "SELECT id FROM public.customers"; - assert!(matches!( - service.query("analytics-readonly", query).await, - Err(Error::UnsupportedObject) - )); - fixture.install_guard().await; - assert!(service.query("analytics-readonly", query).await.is_ok()); - fixture - .admin - .execute("ALTER EVENT TRIGGER mcp_postgres_ddl_guard DISABLE") - .await - .unwrap(); - assert!(matches!( - service.query("analytics-readonly", query).await, - Err(Error::UnsupportedObject) - )); - fixture - .admin - .execute("ALTER EVENT TRIGGER mcp_postgres_ddl_guard ENABLE ALWAYS") - .await - .unwrap(); - fixture.admin.execute("CREATE OR REPLACE FUNCTION mcp_postgres_guard.lock_ddl() RETURNS event_trigger LANGUAGE plpgsql SECURITY DEFINER SET search_path=pg_catalog AS 'BEGIN RETURN; END;'").await.unwrap(); - assert!(matches!( - service.query("analytics-readonly", query).await, - Err(Error::UnsupportedObject) - )); -} - #[tokio::test] async fn catalog_features_and_unsupported_sql_cannot_expose_hidden_dependencies() { let fixture = Fixture::new().await; - fixture.install_guard().await; sqlx::raw_sql("CREATE TABLE public.rls_table (id integer); ALTER TABLE public.rls_table ENABLE ROW LEVEL SECURITY; CREATE TABLE public.generated_table (id integer, generated integer GENERATED ALWAYS AS (id+1) STORED); CREATE TABLE public.parent_table (id integer); CREATE TABLE public.child_table () INHERITS (public.parent_table); @@ -293,9 +301,163 @@ impl PermissionCheck for PausingPermissions { } #[tokio::test] -async fn ddl_guard_blocks_schema_rename_through_authorization_and_execution() { +async fn namespace_replacement_during_authorization_cannot_redirect_the_bound_query() { + let fixture = Fixture::new().await; + let gate = Arc::new(PausingPermissions { + entered: Notify::new(), + release: Notify::new(), + }); + let service = Arc::new(fixture.service(gate.clone()).await); + let task = tokio::spawn({ + let service = service.clone(); + async move { + service + .query("analytics-readonly", "SELECT id FROM public.customers") + .await + } + }); + tokio::time::timeout(Duration::from_secs(3), gate.entered.notified()) + .await + .unwrap(); + let mut migration = fixture.admin.acquire().await.unwrap(); + migration.execute("SET lock_timeout='150ms'").await.unwrap(); + migration.execute("ALTER SCHEMA public RENAME TO moved").await.unwrap(); + sqlx::raw_sql( + "CREATE SCHEMA public; + CREATE TABLE public.customers (id integer, email text, active boolean); + INSERT INTO public.customers VALUES (999, 'replacement@example.test', true); + GRANT USAGE ON SCHEMA public TO query_role; + GRANT SELECT ON public.customers TO query_role;", + ) + .execute(&mut *migration) + .await + .unwrap(); + gate.release.notify_one(); + tokio::time::timeout(Duration::from_secs(3), gate.entered.notified()) + .await + .unwrap(); + gate.release.notify_one(); + assert_eq!(task.await.unwrap().unwrap().rows, vec![vec![json!(1)], vec![json!(2)]]); +} + +#[tokio::test] +async fn namespace_replacement_before_binding_is_rejected_before_column_authorization() { + let fixture = Fixture::new().await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let upstream_url = url::Url::parse(&fixture.dsn).unwrap(); + let upstream = format!("{}:{}", upstream_url.host_str().unwrap(), upstream_url.port().unwrap()); + let reached = Arc::new(Notify::new()); + let proceed = Arc::new(Notify::new()); + let proxy = tokio::spawn({ + let reached = reached.clone(); + let proceed = proceed.clone(); + async move { + while let Ok((client, _)) = listener.accept().await { + let upstream = upstream.clone(); + let reached = reached.clone(); + let proceed = proceed.clone(); + tokio::spawn(async move { + let _ = intercept_cursor_parse(client, &upstream, &reached, &proceed).await; + }); + } + } + }); + let permissions = Arc::new(Permissions::default()); + let service = Arc::new( + fixture + .service_with_dsn( + permissions.clone(), + &format!("postgres://query_role:query-only@127.0.0.1:{port}/postgres?sslmode=disable"), + ) + .await, + ); + let query = tokio::spawn({ + let service = service.clone(); + async move { + service + .query("analytics-readonly", "SELECT id FROM public.customers LIMIT 1") + .await + } + }); + tokio::time::timeout(Duration::from_secs(3), reached.notified()) + .await + .unwrap(); + sqlx::raw_sql( + "ALTER SCHEMA public RENAME TO moved; + CREATE SCHEMA public; + CREATE TABLE public.customers (id integer, email text, active boolean); + INSERT INTO public.customers VALUES (999, 'replacement@example.test', true); + GRANT USAGE ON SCHEMA public TO query_role; + GRANT SELECT ON public.customers TO query_role;", + ) + .execute(&fixture.admin) + .await + .unwrap(); + proceed.notify_one(); + assert!(matches!(query.await.unwrap(), Err(Error::UnsupportedObject))); + assert!( + !permissions + .checked + .lock() + .unwrap() + .iter() + .any(|(kind, _, _)| kind == "sql_column") + ); + proxy.abort(); + let _ = proxy.await; +} + +async fn intercept_cursor_parse( + client: TcpStream, + upstream: &str, + reached: &Notify, + proceed: &Notify, +) -> io::Result<()> { + let server = TcpStream::connect(upstream).await?; + let (mut client_read, mut client_write) = client.into_split(); + let (mut server_read, mut server_write) = server.into_split(); + let responses = async { tokio::io::copy(&mut server_read, &mut client_write).await.map(|_| ()) }; + let requests = async { + let startup_length = client_read.read_u32().await?; + let mut startup = vec![0; startup_length as usize - 4]; + client_read.read_exact(&mut startup).await?; + server_write.write_u32(startup_length).await?; + server_write.write_all(&startup).await?; + server_write.flush().await?; + loop { + let tag = match client_read.read_u8().await { + Ok(tag) => tag, + Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => { + server_write.shutdown().await?; + return Ok(()); + } + Err(error) => return Err(error), + }; + let length = client_read.read_u32().await?; + let mut body = vec![0; length as usize - 4]; + client_read.read_exact(&mut body).await?; + if tag == b'P' + && body + .windows(b"DECLARE mcp_query".len()) + .any(|bytes| bytes == b"DECLARE mcp_query") + { + reached.notify_one(); + proceed.notified().await; + } + server_write.write_u8(tag).await?; + server_write.write_u32(length).await?; + server_write.write_all(&body).await?; + server_write.flush().await?; + } + }; + tokio::try_join!(requests, responses)?; + Ok(()) +} + +#[tokio::test] +async fn column_ddl_waits_for_the_bound_query_without_a_database_guard() { let fixture = Fixture::new().await; - fixture.install_guard().await; let gate = Arc::new(PausingPermissions { entered: Notify::new(), release: Notify::new(), @@ -315,7 +477,7 @@ async fn ddl_guard_blocks_schema_rename_through_authorization_and_execution() { let mut migration = fixture.admin.acquire().await.unwrap(); migration.execute("SET lock_timeout='150ms'").await.unwrap(); let error = migration - .execute("ALTER SCHEMA public RENAME TO moved") + .execute("ALTER TABLE public.customers RENAME COLUMN id TO renamed_id") .await .unwrap_err(); assert_eq!( @@ -328,13 +490,15 @@ async fn ddl_guard_blocks_schema_rename_through_authorization_and_execution() { .unwrap(); gate.release.notify_one(); assert_eq!(task.await.unwrap().unwrap().rows.len(), 2); - migration.execute("ALTER SCHEMA public RENAME TO moved").await.unwrap(); + migration + .execute("ALTER TABLE public.customers RENAME COLUMN id TO renamed_id") + .await + .unwrap(); } #[tokio::test] async fn resource_limits_abort_without_returning_partial_results() { let fixture = Fixture::new().await; - fixture.install_guard().await; sqlx::raw_sql( "CREATE TABLE public.many_rows (id integer); INSERT INTO public.many_rows SELECT generate_series(1,1001); CREATE TABLE public.large_rows (content text); INSERT INTO public.large_rows VALUES (repeat('x',1048577));", @@ -369,7 +533,6 @@ async fn resource_limits_abort_without_returning_partial_results() { #[tokio::test] async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_widths() { let fixture = Fixture::new().await; - fixture.install_guard().await; fixture .admin .execute("CREATE TABLE public.widths (small smallint, regular integer, large bigint)") @@ -424,7 +587,6 @@ async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_width #[tokio::test] async fn mcp_client_initializes_discovers_queries_and_cannot_spoof_identity() { let fixture = Fixture::new().await; - fixture.install_guard().await; let service = fixture.service(Arc::new(Permissions::default())).await; let (client_io, server_io) = tokio::io::duplex(65536); let server_task = tokio::spawn(async move { @@ -475,9 +637,8 @@ async fn mcp_client_initializes_discovers_queries_and_cannot_spoof_identity() { } #[tokio::test] -async fn cancelling_an_mcp_request_releases_its_database_guard() { +async fn cancelling_an_mcp_request_releases_its_relation_locks() { let fixture = Fixture::new().await; - fixture.install_guard().await; let gate = Arc::new(PausingPermissions { entered: Notify::new(), release: Notify::new(), @@ -511,20 +672,25 @@ async fn cancelling_an_mcp_request_releases_its_database_guard() { tokio::time::timeout(Duration::from_secs(3), gate.entered.notified()) .await .unwrap(); - let unlocked: bool = sqlx::query_scalar("SELECT pg_catalog.pg_try_advisory_xact_lock(1879445091,1347635021)") - .fetch_one(&fixture.admin) + let mut probe = fixture.admin.begin().await.unwrap(); + let blocked = probe + .execute("LOCK TABLE public.customers IN ACCESS EXCLUSIVE MODE NOWAIT") .await - .unwrap(); - assert!(!unlocked); + .unwrap_err(); + assert_eq!( + blocked.as_database_error().and_then(|e| e.code()).as_deref(), + Some("55P03") + ); + probe.rollback().await.unwrap(); pending.cancel(Some("test cancellation".into())).await.unwrap(); tokio::time::timeout(Duration::from_secs(2), async { loop { - let unlocked: bool = - sqlx::query_scalar("SELECT pg_catalog.pg_try_advisory_xact_lock(1879445091,1347635021)") - .fetch_one(&fixture.admin) - .await - .unwrap(); - if unlocked { + let mut probe = fixture.admin.begin().await.unwrap(); + let result = probe + .execute("LOCK TABLE public.customers IN ACCESS EXCLUSIVE MODE NOWAIT") + .await; + probe.rollback().await.unwrap(); + if result.is_ok() { break; } tokio::time::sleep(Duration::from_millis(20)).await; @@ -550,7 +716,6 @@ fn spice_request(value: T) -> tonic::Request { #[tokio::test] async fn mcp_spicedb_and_postgres_enforce_hard_column_denial_together() { let fixture = Fixture::new().await; - fixture.install_guard().await; let spice = GenericImage::new("authzed/spicedb", "v1.53.0") .with_exposed_port(50051.tcp()) .with_wait_for(WaitFor::message_on_either_std("grpc server started serving")) From c764d26f4988e9d17f57a697e3eb9c38ffdfe2ee Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Mon, 19 Oct 2026 18:03:27 -0400 Subject: [PATCH 3/7] fix(mcp): prevent unusable authorization identities Signed-off-by: Yordis Prieto --- .../crates/mcp/mcp-postgres/src/config.rs | 6 +-- .../mcp/mcp-postgres/src/connection_id.rs | 21 ++++++++ .../crates/mcp/mcp-postgres/src/constants.rs | 2 + .../crates/mcp/mcp-postgres/src/lib.rs | 2 + .../crates/mcp/mcp-postgres/src/object_id.rs | 4 +- .../crates/mcp/mcp-postgres/src/postgres.rs | 2 +- .../crates/mcp/mcp-postgres/tests/config.rs | 50 ++++++++++++++++++- 7 files changed, 80 insertions(+), 7 deletions(-) create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/connection_id.rs diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/config.rs b/rsworkspace/crates/mcp/mcp-postgres/src/config.rs index b9ffd6d9d..fd988661d 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/config.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/config.rs @@ -1,4 +1,4 @@ -use crate::{Error, ObjectId}; +use crate::{ConnectionId, Error, ObjectId}; use serde::Deserialize; use std::collections::BTreeMap; @@ -19,7 +19,7 @@ pub struct ConnectionInput { pub dsn_env: String, } pub struct RegisteredConnection { - pub(crate) id: ObjectId, + pub(crate) id: ConnectionId, pub(crate) options: crate::DatabaseOptions, } pub struct RuntimeConfig { @@ -34,7 +34,7 @@ impl RuntimeConfig { let mut connections = BTreeMap::new(); let mut ids = std::collections::BTreeSet::new(); for connection in input.connections { - let id = ObjectId::new(connection.id)?; + let id = ConnectionId::new(connection.id)?; if !ids.insert(id.clone()) { return Err(Error::Configuration); } diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/connection_id.rs b/rsworkspace/crates/mcp/mcp-postgres/src/connection_id.rs new file mode 100644 index 000000000..9e4c982d5 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/connection_id.rs @@ -0,0 +1,21 @@ +use crate::{Error, ObjectId}; + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct ConnectionId(ObjectId); + +impl ConnectionId { + pub fn new(value: String) -> Result { + if value.len() > crate::constants::MAX_CONNECTION_ID_BYTES { + return Err(Error::ObjectId); + } + ObjectId::new(value).map(Self) + } + + pub fn as_str(&self) -> &str { + self.0.as_str() + } + + pub fn into_object_id(self) -> ObjectId { + self.0 + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs b/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs index 2eed7f912..1e300b767 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs @@ -1 +1,3 @@ pub(crate) const MAX_BYTES: usize = 1_048_576; +pub(crate) const MAX_OBJECT_ID_BYTES: usize = 200; +pub(crate) const MAX_CONNECTION_ID_BYTES: usize = MAX_OBJECT_ID_BYTES - "/4294967295/4294967295/32767".len(); diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs b/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs index facd92215..f0ede8ea3 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs @@ -17,6 +17,7 @@ //! Permissions are checked at admission; revocation does not retroactively cancel an admitted query. mod config; +mod connection_id; mod constants; mod database_options; mod environment_variable; @@ -33,6 +34,7 @@ mod spicedb_endpoint; mod telemetry; pub use config::{ConfigInput, RegisteredConnection, RuntimeConfig}; +pub use connection_id::ConnectionId; pub use database_options::DatabaseOptions; pub use environment_variable::EnvironmentVariable; pub use error::Error; diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/object_id.rs b/rsworkspace/crates/mcp/mcp-postgres/src/object_id.rs index 55c377e7b..339e39f89 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/object_id.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/object_id.rs @@ -5,8 +5,8 @@ pub struct ObjectId(String); impl ObjectId { pub fn new(value: String) -> Result { if value.is_empty() - || value.len() > 200 - || !value.bytes().all(|c| c.is_ascii_alphanumeric() || b"_-./".contains(&c)) + || value.len() > crate::constants::MAX_OBJECT_ID_BYTES + || !value.bytes().all(|c| c.is_ascii_alphanumeric() || b"_-/".contains(&c)) { return Err(Error::ObjectId); } diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs index 10fa365a2..d588f7080 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs @@ -44,7 +44,7 @@ impl QueryService { .acquire_timeout(Duration::from_secs(2)) .connect_with(connection.options.into_inner()) .await?; - connections.insert(name, (connection.id, pool)); + connections.insert(name, (connection.id.into_object_id(), pool)); } Ok(Self { user: config.user, diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs index e569c99ec..4ab863b98 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs @@ -1,4 +1,6 @@ -use mcp_postgres::{ConfigInput, DatabaseOptions, EnvironmentVariable, ObjectId, RuntimeConfig, SpiceDbEndpoint}; +use mcp_postgres::{ + ConfigInput, ConnectionId, DatabaseOptions, EnvironmentVariable, ObjectId, RuntimeConfig, SpiceDbEndpoint, +}; #[test] fn rejects_duplicate_connection_names_and_ids() -> Result<(), Box> { for connections in [ @@ -47,3 +49,49 @@ fn secrets_and_remote_transport_settings_are_validated_at_the_boundary() -> Resu assert!(RuntimeConfig::from_input(input, |_| None).is_err()); Ok(()) } + +#[test] +fn authorization_ids_reject_dots_at_configuration_boundary() -> Result<(), Box> { + assert!(ObjectId::new("invalid.id".into()).is_err()); + for field in ["user", "client", "connection"] { + let mut input: ConfigInput = toml::from_str(include_str!("../examples/config.toml"))?; + match field { + "user" => input.user = "invalid.id".into(), + "client" => input.client = "invalid.id".into(), + _ => input.connections[0].id = "invalid.id".into(), + } + assert!(RuntimeConfig::from_input(input, config_environment).is_err(), "{field}"); + } + Ok(()) +} + +#[test] +fn connection_ids_reserve_room_for_catalog_resource_suffixes() -> Result<(), Box> { + let suffix = format!("/{}/{}/{}", u32::MAX, u32::MAX, i16::MAX); + let maximum = 200 - suffix.len(); + let id = ConnectionId::new("a".repeat(maximum))?; + assert_eq!(id.as_str().len(), maximum); + assert_eq!(id.into_object_id(), ObjectId::new("a".repeat(maximum))?); + assert!(ConnectionId::new("a".repeat(maximum + 1)).is_err()); + assert!(ConnectionId::new("invalid.id".into()).is_err()); + let mut input: ConfigInput = toml::from_str(include_str!("../examples/config.toml"))?; + input.connections[0].id = "a".repeat(maximum); + assert!(RuntimeConfig::from_input(input, config_environment).is_ok()); + assert!(ObjectId::new(format!("{}{suffix}", "a".repeat(maximum))).is_ok()); + let mut input: ConfigInput = toml::from_str(include_str!("../examples/config.toml"))?; + input.connections[0].id = "a".repeat(maximum + 1); + assert!(RuntimeConfig::from_input(input, config_environment).is_err()); + assert!(ObjectId::new("a".repeat(200)).is_ok()); + Ok(()) +} + +fn config_environment(key: &str) -> Option { + Some( + if key.contains("DSN") { + "postgres://reader@localhost/example" + } else { + "token" + } + .into(), + ) +} From 165b82196fa5d5655905fb53ca750e1098439d68 Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Mon, 19 Oct 2026 18:47:31 -0400 Subject: [PATCH 4/7] test(mcp): protect query safety at failure boundaries Signed-off-by: Yordis Prieto --- .../crates/mcp/mcp-postgres/src/main.rs | 3 + .../crates/mcp/mcp-postgres/src/parser.rs | 9 +- .../mcp/mcp-postgres/src/parser_tests.rs | 79 ++++++++++ .../crates/mcp/mcp-postgres/src/postgres.rs | 52 +++++-- .../mcp/mcp-postgres/src/postgres_tests.rs | 55 +++++++ .../crates/mcp/mcp-postgres/src/server.rs | 5 +- .../crates/mcp/mcp-postgres/src/telemetry.rs | 33 +++-- .../crates/mcp/mcp-postgres/tests/config.rs | 61 +++++++- .../crates/mcp/mcp-postgres/tests/parser.rs | 85 ++++++++++- .../crates/mcp/mcp-postgres/tests/postgres.rs | 50 +++++++ .../crates/mcp/mcp-postgres/tests/server.rs | 135 ++++++++++++++++++ .../mcp/mcp-postgres/tests/spicedb_tls.rs | 23 +++ .../crates/mcp/mcp-postgres/tests/startup.rs | 90 ++++++++++++ 13 files changed, 649 insertions(+), 31 deletions(-) create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/parser_tests.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/postgres_tests.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/server.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/spicedb_tls.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/startup.rs diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/main.rs b/rsworkspace/crates/mcp/mcp-postgres/src/main.rs index c34c284dc..e2bdc4549 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/main.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/main.rs @@ -1,3 +1,5 @@ +#![cfg_attr(coverage_nightly, feature(coverage_attribute))] + use mcp_postgres::{ConfigInput, PostgresMcp, QueryService, RuntimeConfig, SpiceDb}; use rmcp::ServiceExt; use std::{process::ExitCode, sync::Arc}; @@ -5,6 +7,7 @@ use tracing_subscriber::{Layer, layer::SubscriberExt, util::SubscriberInitExt}; use trogon_std::env::{ReadEnv, SystemEnv}; #[tokio::main] +#[cfg_attr(coverage_nightly, coverage(off))] async fn main() -> ExitCode { let audit = tracing_subscriber::fmt::layer() .json() diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/parser.rs b/rsworkspace/crates/mcp/mcp-postgres/src/parser.rs index dc0857ecd..a3e0ddd5a 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/parser.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/parser.rs @@ -50,7 +50,10 @@ impl Query { let NodeEnum::SelectStmt(stmt) = node(statement)? else { return Err(Error::UnsupportedSql); }; - let mut normalized = stmt.as_ref().clone(); + Self::from_statement(stmt) + } + fn from_statement(stmt: &protobuf::SelectStmt) -> Result { + let mut normalized = stmt.clone(); normalized.target_list.clear(); normalized.from_clause.clear(); normalized.where_clause = None; @@ -291,3 +294,7 @@ fn check_join_scope(expr: &Expr, tables: &[Table]) -> Result<(), Error> { } } } + +#[cfg(test)] +#[path = "parser_tests.rs"] +mod tests; diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/parser_tests.rs b/rsworkspace/crates/mcp/mcp-postgres/src/parser_tests.rs new file mode 100644 index 000000000..5e65d73f3 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/parser_tests.rs @@ -0,0 +1,79 @@ +use super::{Query, expr}; +use crate::Error; +use pg_query::{NodeEnum, protobuf}; + +fn select() -> Result> { + let parsed = pg_query::parse("SELECT id FROM public.customers ORDER BY id")?; + match parsed + .protobuf + .stmts + .into_iter() + .next() + .and_then(|s| s.stmt) + .and_then(|s| s.node) + { + Some(NodeEnum::SelectStmt(stmt)) => Ok(*stmt), + _ => Err("expected SELECT fixture".into()), + } +} + +fn wrapped(node: NodeEnum) -> protobuf::Node { + protobuf::Node { node: Some(node) } +} + +#[test] +fn rejects_unexpected_projection_and_sort_ast_shapes() -> Result<(), Box> { + let mut statement = select()?; + statement.target_list = vec![wrapped(NodeEnum::Integer(protobuf::Integer { ival: 1 }))]; + assert!(matches!(Query::from_statement(&statement), Err(Error::UnsupportedSql))); + + let mut statement = select()?; + statement.target_list = vec![wrapped(NodeEnum::ResTarget(Box::new(protobuf::ResTarget { + indirection: vec![wrapped(NodeEnum::Integer(protobuf::Integer { ival: 1 }))], + ..Default::default() + })))]; + assert!(matches!(Query::from_statement(&statement), Err(Error::UnsupportedSql))); + + let mut statement = select()?; + statement.sort_clause = vec![wrapped(NodeEnum::Integer(protobuf::Integer { ival: 1 }))]; + assert!(matches!(Query::from_statement(&statement), Err(Error::UnsupportedSql))); + + let mut statement = select()?; + statement.sort_clause = vec![wrapped(NodeEnum::SortBy(Box::new(protobuf::SortBy { + sortby_dir: i32::MAX, + sortby_nulls: protobuf::SortByNulls::SortbyNullsDefault as i32, + ..Default::default() + })))]; + assert!(matches!(Query::from_statement(&statement), Err(Error::UnsupportedSql))); + Ok(()) +} + +#[test] +fn rejects_invalid_relation_aliases_from_the_parser_boundary() -> Result<(), Box> { + for alias in [String::new(), "x".repeat(64), "a\0b".into()] { + let mut statement = select()?; + let Some(NodeEnum::RangeVar(table)) = statement.from_clause[0].node.as_mut() else { + return Err("expected table fixture".into()); + }; + table.alias = Some(protobuf::Alias { + aliasname: alias, + ..Default::default() + }); + assert!(matches!(Query::from_statement(&statement), Err(Error::Identifier))); + } + Ok(()) +} + +#[test] +fn rejects_incomplete_boolean_and_missing_ast_nodes() { + let incomplete_boolean = wrapped(NodeEnum::BoolExpr(Box::new(protobuf::BoolExpr { + boolop: protobuf::BoolExprType::AndExpr as i32, + args: Vec::new(), + ..Default::default() + }))); + assert!(matches!(expr(&incomplete_boolean, 0), Err(Error::UnsupportedSql))); + assert!(matches!( + expr(&protobuf::Node::default(), 0), + Err(Error::UnsupportedSql) + )); +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs index d588f7080..240f9f28f 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs @@ -6,7 +6,7 @@ use futures_util::TryStreamExt; use serde::Serialize; use sqlx::{ AssertSqlSafe, Column as _, Connection, Executor, PgConnection, PgPool, Row, SqlSafeStr, Statement, TypeInfo, - postgres::{PgPoolOptions, PgRow, types::Oid}, + postgres::{PgColumn, PgPoolOptions, PgRow, types::Oid}, }; use std::{ collections::{BTreeMap, BTreeSet}, @@ -21,6 +21,27 @@ pub struct QueryResult { pub columns: Vec, pub rows: Vec>, } +impl QueryResult { + pub(crate) fn into_json(self) -> serde_json::Value { + serde_json::Value::Object(serde_json::Map::from_iter([ + ( + "columns".into(), + serde_json::Value::Array(self.columns.into_iter().map(serde_json::Value::String).collect()), + ), + ( + "rows".into(), + serde_json::Value::Array(self.rows.into_iter().map(serde_json::Value::Array).collect()), + ), + ])) + } + + fn check_size(&self) -> Result<(), Error> { + if serde_json::to_vec(self).map_err(|_| Error::Limit)?.len() > MAX_BYTES { + return Err(Error::Limit); + } + Ok(()) + } +} pub struct QueryService { user: ObjectId, client: ObjectId, @@ -262,14 +283,7 @@ impl QueryService { .prepare("FETCH FORWARD 0 FROM mcp_query".into_sql_str()) .await?; let visible_columns = query.select.len(); - if metadata.columns().len() != 1 + visible_columns + relations.len() { - return Err(Error::UnsupportedObject); - } - for (column, relation) in metadata.columns().iter().skip(1 + visible_columns).zip(&relations) { - if column.relation_id() != Some(*relation) || column.relation_attribute_no() != Some(-6) { - return Err(Error::UnsupportedObject); - } - } + verify_bindings(metadata.columns(), visible_columns, &relations)?; crate::telemetry::references(&resources, &self.user, &self.client, connection_id); let requests = resources .into_iter() @@ -319,12 +333,26 @@ impl QueryService { } drop(stream); tx.rollback().await?; - if serde_json::to_vec(&result).map_err(|_| Error::Limit)?.len() > MAX_BYTES { - return Err(Error::Limit); - } + result.check_size()?; Ok(result) } } + +#[cfg(test)] +#[path = "postgres_tests.rs"] +mod tests; +fn verify_bindings(columns: &[PgColumn], visible_columns: usize, relations: &[Oid]) -> Result<(), Error> { + if columns.len() != 1 + visible_columns + relations.len() { + return Err(Error::UnsupportedObject); + } + for (column, relation) in columns.iter().skip(1 + visible_columns).zip(relations) { + if column.relation_id() != Some(*relation) || column.relation_attribute_no() != Some(-6) { + return Err(Error::UnsupportedObject); + } + } + Ok(()) +} + fn resolve( column: &Column, query: &Query, diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/postgres_tests.rs b/rsworkspace/crates/mcp/mcp-postgres/src/postgres_tests.rs new file mode 100644 index 000000000..d3ed505f5 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/postgres_tests.rs @@ -0,0 +1,55 @@ +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +use super::{QueryResult, decode, verify_bindings}; +use crate::{Error, constants::MAX_BYTES}; +use sqlx::{Connection, PgConnection, Row, postgres::types::Oid}; +use testcontainers_modules::{ + postgres::Postgres, + testcontainers::{ImageExt, runners::AsyncRunner}, +}; + +#[test] +fn final_result_budget_includes_column_labels_and_json_escaping() { + let valid = QueryResult { + columns: vec!["content".into()], + rows: vec![vec![serde_json::Value::String("small".into())]], + }; + assert!(valid.check_size().is_ok()); + let oversized_label = QueryResult { + columns: vec!["x".repeat(MAX_BYTES)], + rows: Vec::new(), + }; + assert!(matches!(oversized_label.check_size(), Err(Error::Limit))); + let escaped = QueryResult { + columns: vec!["content".into()], + rows: vec![vec![serde_json::Value::String("\u{1}".repeat(MAX_BYTES / 4))]], + }; + assert!(matches!(escaped.check_size(), Err(Error::Limit))); +} + +#[tokio::test] +async fn cursor_metadata_and_decoder_reject_incompatible_server_shapes() { + let container = Postgres::default().with_tag("18.3-alpine").start().await.unwrap(); + let host = container.get_host().await.unwrap(); + let port = container.get_host_port_ipv4(5432).await.unwrap(); + let mut db = PgConnection::connect(&format!("postgres://postgres:postgres@{host}:{port}/postgres")) + .await + .unwrap(); + let row = sqlx::query("SELECT false AS overflow, 1::float8 AS unexpected") + .fetch_one(&mut db) + .await + .unwrap(); + assert!(matches!(decode(&row, 1), Err(Error::UnsupportedObject))); + assert!(matches!( + verify_bindings(row.columns(), 1, &[Oid(123)]), + Err(Error::UnsupportedObject) + )); + let row = sqlx::query("SELECT false AS overflow, 1 AS visible, 123::oid AS spoofed_relation") + .fetch_one(&mut db) + .await + .unwrap(); + assert!(matches!( + verify_bindings(row.columns(), 1, &[Oid(123)]), + Err(Error::UnsupportedObject) + )); +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/server.rs b/rsworkspace/crates/mcp/mcp-postgres/src/server.rs index 78504aeb7..925de8510 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/server.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/server.rs @@ -71,10 +71,7 @@ impl ServerHandler for PostgresMcp { outcome = self.service.query_with_parameters(&args.connection,&args.sql,¶meters) => outcome, }; let result = match outcome { - Ok(result) => CallToolResult::success(vec![ContentBlock::text( - serde_json::to_string(&result) - .map_err(|_| ErrorData::internal_error("result encoding failed", None))?, - )]), + Ok(result) => CallToolResult::success(vec![ContentBlock::text(result.into_json().to_string())]), Err(error) => { let message = match error { crate::Error::Denied => "access denied", diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs b/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs index 2e98a2100..4fb2f29e2 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs @@ -21,12 +21,12 @@ pub(crate) struct QueryAudit<'a> { } impl<'a> QueryAudit<'a> { pub(crate) fn start(user: &'a ObjectId, client: &'a ObjectId, connection: Option<&'a ObjectId>) -> Self { - tracing::info!( - event = "sql_query_admission", - user = user.as_str(), - client = client.as_str(), - connection = connection.map(ObjectId::as_str) - ); + { + let user = user.as_str(); + let client = client.as_str(); + let connection = connection.map(ObjectId::as_str); + tracing::info!(event = "sql_query_admission", user, client, connection); + } Self { user, client, @@ -54,18 +54,27 @@ impl<'a> QueryAudit<'a> { } impl Drop for QueryAudit<'_> { fn drop(&mut self) { - tracing::info!(event="sql_query_completion",user=self.user.as_str(),client=self.client.as_str(),connection=self.connection.map(ObjectId::as_str),outcome=?self.outcome,elapsed_ms=self.started.elapsed().as_millis() as u64); + let user = self.user.as_str(); + let client = self.client.as_str(); + let connection = self.connection.map(ObjectId::as_str); + let elapsed_ms = self.started.elapsed().as_millis() as u64; + tracing::info!(event = "sql_query_completion", user, client, connection, outcome = ?self.outcome, elapsed_ms); } } pub(crate) fn references(resources: &BTreeSet, user: &ObjectId, client: &ObjectId, connection: &ObjectId) { + let user = user.as_str(); + let client = client.as_str(); + let connection = connection.as_str(); for resource in resources { + let resource_kind = resource.kind; + let resource_id = resource.id.as_str(); tracing::info!( event = "sql_query_reference", - user = user.as_str(), - client = client.as_str(), - connection = connection.as_str(), - resource_kind = resource.kind, - resource_id = resource.id.as_str() + user, + client, + connection, + resource_kind, + resource_id ); } } diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs index 4ab863b98..1a56a4c2a 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs @@ -1,5 +1,6 @@ use mcp_postgres::{ - ConfigInput, ConnectionId, DatabaseOptions, EnvironmentVariable, ObjectId, RuntimeConfig, SpiceDbEndpoint, + ConfigInput, ConnectionId, DatabaseOptions, EnvironmentVariable, Error, Identifier, ObjectId, Parameter, + ParameterInput, RuntimeConfig, SpiceDbEndpoint, }; #[test] fn rejects_duplicate_connection_names_and_ids() -> Result<(), Box> { @@ -95,3 +96,61 @@ fn config_environment(key: &str) -> Option { .into(), ) } + +#[test] +fn runtime_configuration_requires_connections_and_nonempty_secrets() -> Result<(), Box> { + let valid = include_str!("../examples/config.toml"); + let config = RuntimeConfig::from_input(toml::from_str(valid)?, config_environment)?; + assert_eq!(config.spicedb_endpoint(), "http://127.0.0.1:50051"); + assert_eq!(config.spicedb_token(), "token"); + + let mut empty: ConfigInput = toml::from_str(valid)?; + empty.connections.clear(); + assert!(matches!( + RuntimeConfig::from_input(empty, config_environment), + Err(Error::Configuration) + )); + for missing_token in [false, true] { + let input: ConfigInput = toml::from_str(valid)?; + let token_key = input.spicedb_token_env.clone(); + assert!(matches!( + RuntimeConfig::from_input(input, |key| { + if key == token_key { + if missing_token { None } else { Some(String::new()) } + } else { + config_environment(key) + } + }), + Err(Error::Configuration) + )); + } + assert!(matches!( + RuntimeConfig::from_input(toml::from_str(valid)?, |_| Some(String::new())), + Err(Error::Configuration) + )); + Ok(()) +} + +#[test] +fn identifiers_and_parameters_enforce_postgres_and_request_boundaries() -> Result<(), Box> { + for invalid in [String::new(), "x".repeat(64), "contains\0nul".into(), "é".repeat(32)] { + assert!(matches!(Identifier::new(invalid), Err(Error::Identifier))); + } + assert_eq!(Identifier::new("x".repeat(63))?.as_str().len(), 63); + assert_eq!(Identifier::new("odd\"name".into())?.quoted(), "\"odd\"\"name\""); + for input in [ + ParameterInput::Boolean(false), + ParameterInput::Integer(i64::MIN), + ParameterInput::Text("x".repeat(16_384)), + ] { + Parameter::try_from(input)?; + } + for input in [ + ParameterInput::Text("x".repeat(16_385)), + ParameterInput::Text("a\0b".into()), + ParameterInput::Text("é".repeat(8193)), + ] { + assert!(matches!(Parameter::try_from(input), Err(Error::Limit))); + } + Ok(()) +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/parser.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/parser.rs index e4830806b..22e6fee9d 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/parser.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/parser.rs @@ -1,4 +1,4 @@ -use mcp_postgres::Query; +use mcp_postgres::{Error, Query}; #[test] fn supports_explicit_columns_and_dependencies() -> Result<(), Box> { @@ -12,6 +12,89 @@ fn supports_explicit_columns_and_dependencies() -> Result<(), Box Result<(), Box> { + let valid = "SELECT id FROM public.customers"; + Query::parse(&format!("{valid}{}", " ".repeat(16_384 - valid.len())))?; + assert!(matches!(Query::parse(&" ".repeat(16_385)), Err(Error::Limit))); + Query::parse("SELECT id FROM public.customers LIMIT 0")?; + Query::parse("SELECT id FROM public.customers LIMIT 1000")?; + assert!(matches!( + Query::parse("SELECT id FROM public.customers LIMIT 1001"), + Err(Error::Limit) + )); + + let joined = (1..8).fold("public.customers c0".to_owned(), |left, index| { + format!("{left} JOIN public.customers c{index} ON true") + }); + Query::parse(&format!("SELECT c0.id FROM {joined}"))?; + assert!(matches!( + Query::parse(&format!("SELECT c0.id FROM {joined} JOIN public.customers c8 ON true")), + Err(Error::Limit) + )); + let too_deep = (1..11).fold("public.customers c0".to_owned(), |left, index| { + format!("{left} JOIN public.customers c{index} ON true") + }); + assert!(matches!( + Query::parse(&format!("SELECT c0.id FROM {too_deep}")), + Err(Error::Limit) + )); + + let predicate = (0..18).fold("id = 1".to_owned(), |inner, index| { + if index % 2 == 0 { + format!("id = 2 AND ({inner})") + } else { + format!("id = 3 OR ({inner})") + } + }); + assert!(matches!( + Query::parse(&format!("{valid} WHERE {predicate}")), + Err(Error::Limit) + )); + Ok(()) +} + +#[test] +fn accepts_scoped_boolean_join_predicates_and_supported_literals() -> Result<(), Box> { + for sql in [ + "SELECT a.id FROM public.customers a JOIN public.orders b ON (a.id=b.id AND b.id=$1) OR false", + "SELECT id FROM public.customers WHERE id=1 OR id=2 ORDER BY id ASC", + "SELECT id FROM public.customers WHERE true AND false", + "SELECT id FROM public.customers WHERE email='can''t'", + ] { + Query::parse(sql)?; + } + Ok(()) +} #[test] fn rejects_every_unsupported_construct() { for sql in [ diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs index a51164ed9..235fe343e 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs @@ -530,6 +530,56 @@ async fn resource_limits_abort_without_returning_partial_results() { ); } +#[tokio::test] +async fn admission_rejects_privileged_roles_and_excess_parameters() { + let fixture = Fixture::new().await; + let service = fixture.service(Arc::new(Permissions::default())).await; + let parameters = vec![Parameter::try_from(ParameterInput::Boolean(true)).unwrap(); 33]; + assert!(matches!( + service + .query_with_parameters("analytics-readonly", "SELECT id FROM public.customers", ¶meters) + .await, + Err(Error::Limit) + )); + fixture.admin.execute("ALTER ROLE query_role SUPERUSER").await.unwrap(); + assert!(matches!( + service + .query("analytics-readonly", "SELECT id FROM public.customers") + .await, + Err(Error::UnsupportedObject) + )); +} + +#[tokio::test] +async fn boolean_predicates_alias_rejection_and_aggregate_byte_budget() { + let fixture = Fixture::new().await; + let service = fixture.service(Arc::new(Permissions::default())).await; + let result = service + .query( + "analytics-readonly", + "SELECT id FROM public.customers WHERE (id = 1 OR id = 2) AND active = true ORDER BY id", + ) + .await + .unwrap(); + assert_eq!(result.rows, vec![vec![json!(1)]]); + assert!(matches!( + service + .query( + "analytics-readonly", + "SELECT id AS renamed FROM public.customers ORDER BY renamed" + ) + .await, + Err(Error::UnsupportedSql) + )); + sqlx::raw_sql("CREATE TABLE public.aggregate_large (content text); INSERT INTO public.aggregate_large SELECT repeat('x',100000) FROM generate_series(1,11); GRANT SELECT ON public.aggregate_large TO query_role;").execute(&fixture.admin).await.unwrap(); + assert!(matches!( + service + .query("analytics-readonly", "SELECT content FROM public.aggregate_large") + .await, + Err(Error::Limit) + )); +} + #[tokio::test] async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_widths() { let fixture = Fixture::new().await; diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs new file mode 100644 index 000000000..08ffdd6ca --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs @@ -0,0 +1,135 @@ +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +use async_trait::async_trait; +use mcp_postgres::{ConfigInput, Error, PermissionCheck, PermissionRequest, PostgresMcp, QueryService, RuntimeConfig}; +use rmcp::{ServiceExt, model::CallToolRequestParams}; +use serde_json::json; +use std::sync::{Arc, Mutex}; +use testcontainers_modules::{ + postgres::Postgres, + testcontainers::{ImageExt, runners::AsyncRunner}, +}; +use tracing_subscriber::{Layer, layer::SubscriberExt}; +#[derive(Clone, Default)] +struct LogBuffer(Arc>>); +impl std::io::Write for LogBuffer { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} +impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for LogBuffer { + type Writer = Self; + fn make_writer(&'a self) -> Self::Writer { + self.clone() + } +} +struct Permissions(Arc>>); +#[async_trait] +impl PermissionCheck for Permissions { + async fn check_all(&self, _: &[PermissionRequest]) -> Result<(), Error> { + match self.0.lock().unwrap().take() { + Some(error) => Err(error), + None => Ok(()), + } + } +} +#[tokio::test(flavor = "current_thread")] +async fn mcp_errors_are_safe_and_audits_exclude_sql_parameters_and_rows() { + let container = Postgres::default().with_tag("18.3-alpine").start().await.unwrap(); + let dsn = format!( + "postgres://postgres:postgres@{}:{}/postgres", + container.get_host().await.unwrap(), + container.get_host_port_ipv4(5432).await.unwrap() + ); + let pool = sqlx::PgPool::connect(&dsn).await.unwrap(); + sqlx::query("CREATE TABLE public.customers (id integer, email text)") + .execute(&pool) + .await + .unwrap(); + sqlx::query("INSERT INTO public.customers VALUES (1, 'private-row@example.test')") + .execute(&pool) + .await + .unwrap(); + sqlx::raw_sql("CREATE ROLE query_role LOGIN PASSWORD 'query-only'; GRANT USAGE ON SCHEMA public TO query_role; GRANT SELECT ON public.customers TO query_role;").execute(&pool).await.unwrap(); + let dsn = dsn.replace("postgres:postgres@", "query_role:query-only@"); + let input: ConfigInput = toml::from_str(include_str!("../examples/config.toml")).unwrap(); + let config = RuntimeConfig::from_input(input, |key| { + Some(if key.ends_with("DSN") { + dsn.clone() + } else { + "private-token".into() + }) + }) + .unwrap(); + let next_error = Arc::new(Mutex::new(None)); + let service = QueryService::connect(config, Arc::new(Permissions(next_error.clone()))) + .await + .unwrap(); + let logs = LogBuffer::default(); + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .json() + .with_writer(logs.clone()) + .with_filter(tracing_subscriber::filter::Targets::new().with_target("mcp_postgres", tracing::Level::INFO)), + ); + let _guard = tracing::subscriber::set_default(subscriber); + let (client_io, server_io) = tokio::io::duplex(65536); + let server = tokio::spawn(async move { + PostgresMcp::new(service) + .serve(server_io) + .await + .unwrap() + .waiting() + .await + .unwrap(); + }); + let client = ().serve(client_io).await.unwrap(); + assert!(client.call_tool(CallToolRequestParams::new("missing")).await.is_err()); + for (connection, failure, expected) in [ + ("missing", None, "unknown connection"), + ("analytics-readonly", Some(Error::Limit), "resource limit exceeded"), + ("analytics-readonly", Some(Error::Timeout), "request timed out"), + ("analytics-readonly", Some(Error::Configuration), "query unavailable"), + ] { + *next_error.lock().unwrap() = failure; + let response = client + .call_tool( + CallToolRequestParams::new("query").with_arguments( + json!({"connection":connection,"sql":"SELECT email FROM public.customers"}) + .as_object() + .unwrap() + .clone(), + ), + ) + .await + .unwrap(); + assert_eq!(response.is_error, Some(true)); + assert_eq!(response.content[0].as_text().unwrap().text, expected); + } + let response = client.call_tool(CallToolRequestParams::new("query").with_arguments(json!({"connection":"analytics-readonly","sql":"SELECT email FROM public.customers WHERE email = $1", "parameters":["private-row@example.test"]}).as_object().unwrap().clone())).await.unwrap(); + assert_ne!(response.is_error, Some(true)); + let result: serde_json::Value = serde_json::from_str(&response.content[0].as_text().unwrap().text).unwrap(); + assert_eq!(result["columns"], json!(["email"])); + assert_eq!(result["rows"], json!([["private-row@example.test"]])); + client.cancel().await.unwrap(); + server.await.unwrap(); + let output = String::from_utf8(logs.0.lock().unwrap().clone()).unwrap(); + for expected in [ + "sql_query_admission", + "sql_query_completion", + "sql_query_reference", + "Timeout", + "Unavailable", + "resource_id", + "example-user", + "example-client", + ] { + assert!(output.contains(expected), "missing {expected}: {output}"); + } + for secret in ["SELECT", "private-row@example.test", "private-token", &dsn] { + assert!(!output.contains(secret), "audit leaked {secret}"); + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/spicedb_tls.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/spicedb_tls.rs new file mode 100644 index 000000000..bb51e3d8f --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/spicedb_tls.rs @@ -0,0 +1,23 @@ +use mcp_postgres::{Error, SpiceDb}; +use tokio::{io::AsyncReadExt, net::TcpListener}; + +#[tokio::test] +async fn https_authorization_endpoint_requires_tls_and_never_sends_a_plaintext_token() +-> Result<(), Box> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("https://{}", listener.local_addr()?); + let received = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await?; + let mut bytes = [0; 1024]; + let count = socket.read(&mut bytes).await?; + Ok::<_, std::io::Error>(bytes[..count].to_vec()) + }); + assert!(matches!( + SpiceDb::connect(endpoint, "private-transport-token".into()).await, + Err(Error::AuthorizationConnection(_)) + )); + let bytes = received.await??; + assert_eq!(bytes.first(), Some(&22), "the first record must be a TLS handshake"); + assert!(!String::from_utf8_lossy(&bytes).contains("private-transport-token")); + Ok(()) +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/startup.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/startup.rs new file mode 100644 index 000000000..ae3560bef --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/startup.rs @@ -0,0 +1,90 @@ +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +use rmcp::ServiceExt; +use std::{ + process::{Command, Stdio}, + time::Duration, +}; +use testcontainers_modules::{ + postgres::Postgres, + testcontainers::{ + GenericImage, ImageExt, + core::{IntoContainerPort, WaitFor}, + runners::AsyncRunner, + }, +}; +use tokio::io::AsyncReadExt; + +#[test] +fn invalid_startup_reports_only_safe_failure() { + for arguments in [ + vec![], + vec!["missing-private-config", "extra"], + vec!["missing-private-config"], + ] { + let output = Command::new(env!("CARGO_BIN_EXE_mcp-postgres")) + .args(arguments) + .output() + .unwrap(); + assert!(!output.status.success()); + assert!(output.stdout.is_empty()); + let stderr = String::from_utf8(output.stderr).unwrap(); + assert!(stderr.contains("startup or transport failed")); + assert!(!stderr.contains("missing-private-config")); + } +} + +#[tokio::test] +async fn configured_binary_initializes_stdio_and_exits_when_client_disconnects() { + let postgres = Postgres::default().with_tag("18.3-alpine").start().await.unwrap(); + let spicedb = GenericImage::new("authzed/spicedb", "v1.53.0") + .with_exposed_port(50051.tcp()) + .with_wait_for(WaitFor::message_on_either_std("grpc server started serving")) + .with_cmd([ + "serve", + "--datastore-engine", + "memory", + "--grpc-preshared-key", + "startup-test-token", + ]) + .start() + .await + .unwrap(); + let endpoint = format!( + "http://{}:{}", + spicedb.get_host().await.unwrap(), + spicedb.get_host_port_ipv4(50051.tcp()).await.unwrap() + ); + let dsn = format!( + "postgres://postgres:postgres@{}:{}/postgres", + postgres.get_host().await.unwrap(), + postgres.get_host_port_ipv4(5432).await.unwrap() + ); + let config = include_str!("../examples/config.toml").replace("http://127.0.0.1:50051", &endpoint); + let path = std::env::temp_dir().join(format!("mcp-startup-{}.toml", std::process::id())); + std::fs::write(&path, config).unwrap(); + let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_mcp-postgres")) + .arg(&path) + .env("MCP_POSTGRES_SPICEDB_TOKEN", "startup-test-token") + .env("MCP_POSTGRES_ANALYTICS_DSN", &dsn) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .unwrap(); + let client = ().serve((child.stdout.take().unwrap(), child.stdin.take().unwrap())).await.unwrap(); + let tools = client.list_tools(None).await.unwrap(); + assert_eq!(tools.tools.len(), 1); + assert_eq!(tools.tools[0].name, "query"); + client.cancel().await.unwrap(); + let status = tokio::time::timeout(Duration::from_secs(5), child.wait()) + .await + .unwrap() + .unwrap(); + assert!(status.success()); + let mut stderr = String::new(); + child.stderr.take().unwrap().read_to_string(&mut stderr).await.unwrap(); + assert!(!stderr.contains("startup-test-token")); + assert!(!stderr.contains(&dsn)); + std::fs::remove_file(path).unwrap(); +} From 83bab9909e68311dc52493905350ec5d43cc6c1d Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Mon, 19 Oct 2026 19:36:31 -0400 Subject: [PATCH 5/7] test(mcp): keep audit failures from exposing sensitive data Signed-off-by: Yordis Prieto --- .../crates/mcp/mcp-postgres/tests/server.rs | 31 +++++++++++++++++-- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs index 08ffdd6ca..c2b5af397 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs @@ -117,6 +117,10 @@ async fn mcp_errors_are_safe_and_audits_exclude_sql_parameters_and_rows() { client.cancel().await.unwrap(); server.await.unwrap(); let output = String::from_utf8(logs.0.lock().unwrap().clone()).unwrap(); + assert_audit_output(&output, &["SELECT", "private-row@example.test", "private-token", &dsn]); +} + +fn assert_audit_output(output: &str, forbidden: &[&str]) { for expected in [ "sql_query_admission", "sql_query_completion", @@ -127,9 +131,30 @@ async fn mcp_errors_are_safe_and_audits_exclude_sql_parameters_and_rows() { "example-user", "example-client", ] { - assert!(output.contains(expected), "missing {expected}: {output}"); + assert!(output.contains(expected), "audit metadata is missing"); + } + for secret in forbidden { + assert!(!output.contains(secret), "audit contains a forbidden value"); } - for secret in ["SELECT", "private-row@example.test", "private-token", &dsn] { - assert!(!output.contains(secret), "audit leaked {secret}"); +} + +#[test] +fn audit_assertion_failures_do_not_disclose_captured_output_or_forbidden_values() { + let canary = "audit-privacy-regression-canary"; + let complete_metadata = "sql_query_admission sql_query_completion sql_query_reference Timeout Unavailable resource_id example-user example-client"; + let failures = [canary.to_owned(), format!("{complete_metadata} {canary}")] + .map(|output| std::panic::catch_unwind(|| assert_audit_output(&output, &[canary]))); + for failure in failures { + let payload = failure.expect_err("audit assertion accepted invalid output"); + let message = payload + .downcast_ref::() + .map(String::as_str) + .or_else(|| payload.downcast_ref::<&str>().copied()) + .expect("audit assertion did not provide a text diagnostic"); + assert!(!message.contains(canary), "audit assertion disclosed a forbidden value"); + assert!( + !message.contains(complete_metadata), + "audit assertion disclosed captured output" + ); } } From 001ffd1dba312a19fcbe7b839561fe3cd7208f5f Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Mon, 19 Oct 2026 23:30:52 -0400 Subject: [PATCH 6/7] test(mcp): protect PostgreSQL certificate trust boundaries Signed-off-by: Yordis Prieto --- rsworkspace/Cargo.lock | 3 + .../crates/mcp/mcp-postgres/Cargo.toml | 3 + .../mcp/mcp-postgres/src/database_options.rs | 2 +- .../crates/mcp/mcp-postgres/tests/config.rs | 39 ++++ .../mcp/mcp-postgres/tests/postgres_tls.rs | 183 ++++++++++++++++++ 5 files changed, 229 insertions(+), 1 deletion(-) create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/postgres_tls.rs diff --git a/rsworkspace/Cargo.lock b/rsworkspace/Cargo.lock index b8ac160cb..156e8e2c3 100644 --- a/rsworkspace/Cargo.lock +++ b/rsworkspace/Cargo.lock @@ -3822,11 +3822,14 @@ dependencies = [ "async-trait", "futures-util", "pg_query", + "rcgen", "rmcp", + "rustls", "serde", "serde_json", "spicedb-grpc-tonic", "sqlx", + "tempfile", "testcontainers-modules", "thiserror 2.0.21", "tokio", diff --git a/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml b/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml index f18942ea6..518da11b2 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml +++ b/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml @@ -26,4 +26,7 @@ tonic = { workspace = true, features = ["tls-aws-lc", "tls-native-roots"] } url = { workspace = true } [dev-dependencies] +rcgen = "=0.14.8" +rustls = { workspace = true } +tempfile = { workspace = true } testcontainers-modules = { workspace = true, features = ["postgres"] } diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/database_options.rs b/rsworkspace/crates/mcp/mcp-postgres/src/database_options.rs index 4f3930725..a3f20a12f 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/database_options.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/database_options.rs @@ -9,7 +9,7 @@ pub struct DatabaseOptions(PgConnectOptions); impl DatabaseOptions { pub fn new(dsn: &str) -> Result { let options = PgConnectOptions::from_str(dsn).map_err(|_| Error::Configuration)?; - let local = matches!(options.get_host(), "localhost" | "127.0.0.1" | "::1"); + let local = matches!(options.get_host(), "localhost" | "127.0.0.1" | "::1" | "[::1]"); if !local && !matches!(options.get_ssl_mode(), PgSslMode::VerifyFull) { return Err(Error::Configuration); } diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs index 1a56a4c2a..ac8d973dc 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs @@ -32,6 +32,45 @@ fn unknown_settings_fail_closed() { assert!(toml::from_str::(&format!("permissions_disabled=true\n{valid}")).is_err()); } +#[test] +fn remote_postgres_connections_require_full_certificate_verification() { + for mode in [ + "", + "?sslmode=disable", + "?sslmode=allow", + "?sslmode=prefer", + "?sslmode=require", + "?sslmode=verify-ca", + ] { + assert!( + matches!( + DatabaseOptions::new(&format!("postgres://reader@db.example.test/example{mode}")), + Err(Error::Configuration) + ), + "rejected TLS mode: {mode}" + ); + } + assert!(DatabaseOptions::new("postgres://reader@db.example.test/example?sslmode=verify-full").is_ok()); +} + +#[test] +fn loopback_postgres_connections_allow_local_plaintext_fixtures() { + for host in ["localhost", "127.0.0.1", "[::1]"] { + for mode in ["", "?sslmode=disable", "?sslmode=verify-full"] { + assert!( + DatabaseOptions::new(&format!("postgres://reader@{host}/example{mode}")).is_ok(), + "host={host}, mode={mode}" + ); + } + } + for host in ["localhost.example.test", "127.0.0.2", "db.example.test"] { + assert!(matches!( + DatabaseOptions::new(&format!("postgres://reader@{host}/example?sslmode=disable")), + Err(Error::Configuration) + )); + } +} + #[test] fn secrets_and_remote_transport_settings_are_validated_at_the_boundary() -> Result<(), Box> { assert!(DatabaseOptions::new("postgres://reader@localhost/example").is_ok()); diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres_tls.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres_tls.rs new file mode 100644 index 000000000..547c94435 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres_tls.rs @@ -0,0 +1,183 @@ +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] + +use std::{io, sync::Arc}; + +use async_trait::async_trait; +use mcp_postgres::{ConfigInput, Error, PermissionCheck, PermissionRequest, QueryService, RuntimeConfig}; +use rcgen::{BasicConstraints, CertificateParams, DnType, IsCa, Issuer, KeyPair}; +use rustls::{CertificateError, Error as TlsError}; +use serde_json::json; +use sqlx::PgPool; +use tempfile::TempDir; +use testcontainers_modules::{ + postgres::Postgres, + testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner}, +}; +use url::Url; + +struct Permissions; + +#[async_trait] +impl PermissionCheck for Permissions { + async fn check_all(&self, _: &[PermissionRequest]) -> Result<(), Error> { + Ok(()) + } +} + +fn authority(name: &str) -> (CertificateParams, KeyPair) { + let mut params = CertificateParams::default(); + params.distinguished_name.push(DnType::CommonName, name); + params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + (params, KeyPair::generate().unwrap()) +} + +struct Fixture { + _container: ContainerAsync, + certificates: TempDir, + admin: PgPool, + port: u16, +} + +impl Fixture { + async fn new() -> Self { + let certificates = tempfile::tempdir().unwrap(); + let (ca_params, ca_key) = authority("trusted-test-ca"); + let ca = ca_params.self_signed(&ca_key).unwrap(); + std::fs::write(certificates.path().join("ca.pem"), ca.pem()).unwrap(); + let (other_params, other_key) = authority("untrusted-test-ca"); + let other_ca = other_params.self_signed(&other_key).unwrap(); + std::fs::write(certificates.path().join("other-ca.pem"), other_ca.pem()).unwrap(); + let server_key = KeyPair::generate().unwrap(); + let server = CertificateParams::new(vec!["localhost".into()]) + .unwrap() + .signed_by(&server_key, &Issuer::from_params(&ca_params, &ca_key)) + .unwrap(); + let container = Postgres::default() + .with_tag("18.3-alpine") + .with_copy_to("/tmp/server.crt", server.pem().into_bytes()) + .with_copy_to("/tmp/server.key", server_key.serialize_pem().into_bytes()) + .with_cmd([ + "sh", + "-c", + "chown postgres:postgres /tmp/server.key && chmod 600 /tmp/server.key && exec docker-entrypoint.sh postgres -c ssl=on -c ssl_cert_file=/tmp/server.crt -c ssl_key_file=/tmp/server.key", + ]) + .start() + .await + .unwrap(); + let port = container.get_host_port_ipv4(5432).await.unwrap(); + let host = container.get_host().await.unwrap(); + let admin = PgPool::connect(&format!( + "postgres://postgres:postgres@{host}:{port}/postgres?sslmode=disable" + )) + .await + .unwrap(); + sqlx::raw_sql( + "CREATE ROLE query_role LOGIN PASSWORD 'query-only'; + REVOKE TEMPORARY ON DATABASE postgres FROM PUBLIC; + CREATE TABLE public.customers (id integer); + INSERT INTO public.customers VALUES (1); + GRANT USAGE ON SCHEMA public TO query_role; + GRANT SELECT ON public.customers TO query_role;", + ) + .execute(&admin) + .await + .unwrap(); + Self { + _container: container, + certificates, + admin, + port, + } + } + + fn dsn(&self, host: &str, ca: &str) -> String { + let mut dsn = Url::parse(&format!( + "postgres://query_role:query-only@{host}:{}/postgres", + self.port + )) + .unwrap(); + dsn.query_pairs_mut() + .append_pair("sslmode", "verify-full") + .append_pair("sslrootcert", self.certificates.path().join(ca).to_str().unwrap()) + .append_pair("application_name", "mcp-postgres-tls-test"); + dsn.into() + } + + async fn trusted_service(&self) -> QueryService { + service(&self.dsn("localhost", "ca.pem")).await.unwrap() + } +} + +async fn service(dsn: &str) -> Result { + let input: ConfigInput = toml::from_str(include_str!("../examples/config.toml")).unwrap(); + let config = RuntimeConfig::from_input(input, |key| match key { + "MCP_POSTGRES_SPICEDB_TOKEN" => Some("test-only-token".into()), + "MCP_POSTGRES_ANALYTICS_DSN" => Some(dsn.to_owned()), + _ => None, + }) + .unwrap(); + QueryService::connect(config, Arc::new(Permissions)).await +} + +fn certificate_error(result: Result) -> CertificateError { + let Err(Error::Database(sqlx::Error::Io(error))) = result else { + panic!("expected a TLS certificate validation failure"); + }; + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + let Some(TlsError::InvalidCertificate(error)) = error.get_ref().and_then(|error| error.downcast_ref::()) + else { + panic!("expected a typed certificate validation cause"); + }; + error.clone() +} + +#[tokio::test] +async fn trusted_ca_and_matching_hostname_protect_the_query_connection() { + let fixture = Fixture::new().await; + let service = fixture.trusted_service().await; + let encrypted: bool = sqlx::query_scalar( + "SELECT count(*) > 0 AND bool_and(s.ssl) FROM pg_catalog.pg_stat_activity a + JOIN pg_catalog.pg_stat_ssl s USING (pid) + WHERE a.usename='query_role' AND a.application_name='mcp-postgres-tls-test'", + ) + .fetch_one(&fixture.admin) + .await + .unwrap(); + assert!(encrypted); + let result = service + .query("analytics-readonly", "SELECT id FROM public.customers") + .await + .unwrap(); + assert_eq!(result.rows, vec![vec![json!(1)]]); +} + +#[tokio::test] +async fn an_untrusted_ca_is_rejected_by_the_query_service() { + let fixture = Fixture::new().await; + let _trusted = fixture.trusted_service().await; + let error = certificate_error(service(&fixture.dsn("localhost", "other-ca.pem")).await); + assert!(matches!(error, CertificateError::UnknownIssuer)); +} + +#[tokio::test] +async fn a_trusted_certificate_for_another_hostname_is_rejected() { + let fixture = Fixture::new().await; + let _trusted = fixture.trusted_service().await; + let error = certificate_error(service(&fixture.dsn("127.0.0.1", "ca.pem")).await); + assert!(matches!(error, CertificateError::NotValidForNameContext { .. })); +} + +#[tokio::test] +async fn verify_full_refuses_a_server_without_tls() { + let container = Postgres::default().with_tag("18.3-alpine").start().await.unwrap(); + let host = container.get_host().await.unwrap(); + let port = container.get_host_port_ipv4(5432).await.unwrap(); + let dsn = format!("postgres://postgres:postgres@{host}:{port}/postgres"); + let admin = PgPool::connect(&format!("{dsn}?sslmode=disable")).await.unwrap(); + let tls_enabled: String = sqlx::query_scalar("SHOW ssl").fetch_one(&admin).await.unwrap(); + assert_eq!(tls_enabled, "off"); + assert!(matches!( + service(&format!("{dsn}?sslmode=verify-full")).await, + Err(Error::Database(sqlx::Error::Tls(_))) + )); +} From 448b101a377e017dab30013d3b5cd648e49db627 Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Tue, 20 Oct 2026 02:00:54 -0400 Subject: [PATCH 7/7] feat(mcp): keep SQL access tied to authenticated callers Signed-off-by: Yordis Prieto --- rsworkspace/Cargo.lock | 27 +- .../crates/mcp/mcp-postgres/Cargo.toml | 10 +- .../mcp/mcp-postgres/examples/config.toml | 23 +- .../crates/mcp/mcp-postgres/src/auth_error.rs | 11 + .../mcp-postgres/src/authentication_config.rs | 123 ++++++ .../mcp/mcp-postgres/src/authenticator.rs | 248 +++++++++++ .../mcp-postgres/src/authenticator_tests.rs | 173 ++++++++ .../crates/mcp/mcp-postgres/src/config.rs | 10 +- .../crates/mcp/mcp-postgres/src/constants.rs | 6 + .../mcp/mcp-postgres/src/http_config.rs | 81 ++++ .../mcp/mcp-postgres/src/http_transport.rs | 127 ++++++ .../crates/mcp/mcp-postgres/src/lib.rs | 18 +- .../crates/mcp/mcp-postgres/src/main.rs | 29 +- .../crates/mcp/mcp-postgres/src/postgres.rs | 44 +- .../mcp/mcp-postgres/src/request_identity.rs | 30 ++ .../crates/mcp/mcp-postgres/src/server.rs | 19 +- .../crates/mcp/mcp-postgres/src/telemetry.rs | 9 +- .../mcp/mcp-postgres/tests/authentication.rs | 402 ++++++++++++++++++ .../mcp/mcp-postgres/tests/common/mod.rs | 131 ++++++ .../crates/mcp/mcp-postgres/tests/config.rs | 17 +- .../mcp/mcp-postgres/tests/http_config.rs | 58 +++ .../mcp/mcp-postgres/tests/http_identity.rs | 217 ++++++++++ .../crates/mcp/mcp-postgres/tests/postgres.rs | 238 ++++++++--- .../mcp/mcp-postgres/tests/postgres_tls.rs | 10 +- .../crates/mcp/mcp-postgres/tests/server.rs | 161 ++++++- .../crates/mcp/mcp-postgres/tests/startup.rs | 126 ++++-- 26 files changed, 2187 insertions(+), 161 deletions(-) create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/auth_error.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/authentication_config.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/authenticator.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/authenticator_tests.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/http_config.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/http_transport.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/src/request_identity.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/authentication.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/common/mod.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/http_config.rs create mode 100644 rsworkspace/crates/mcp/mcp-postgres/tests/http_identity.rs diff --git a/rsworkspace/Cargo.lock b/rsworkspace/Cargo.lock index 156e8e2c3..f8878549f 100644 --- a/rsworkspace/Cargo.lock +++ b/rsworkspace/Cargo.lock @@ -3820,10 +3820,16 @@ name = "mcp-postgres" version = "0.0.1" dependencies = [ "async-trait", + "axum", + "base64 0.23.1", "futures-util", + "http", + "jsonwebtoken 11.1.0", "pg_query", "rcgen", + "reqwest 0.12.28", "rmcp", + "rsa", "rustls", "serde", "serde_json", @@ -3835,6 +3841,7 @@ dependencies = [ "tokio", "toml 1.1.6+spec-1.1.0", "tonic", + "tower-http 0.7.1", "tracing", "tracing-subscriber", "trogon-std", @@ -5219,7 +5226,7 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams", + "wasm-streams 0.4.2", "web-sys", "webpki-roots 1.0.6", ] @@ -5249,15 +5256,19 @@ dependencies = [ "rustls", "rustls-pki-types", "rustls-platform-verifier", + "serde", + "serde_json", "sync_wrapper", "tokio", "tokio-rustls", + "tokio-util", "tower", "tower-http 0.6.11", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams 0.5.0", "web-sys", ] @@ -5312,6 +5323,7 @@ dependencies = [ "pastey", "pin-project-lite", "rand 0.10.1", + "reqwest 0.13.4", "schemars 1.2.1", "serde", "serde_json", @@ -8051,6 +8063,19 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" diff --git a/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml b/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml index 518da11b2..95fbccb1e 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml +++ b/rsworkspace/crates/mcp/mcp-postgres/Cargo.toml @@ -10,23 +10,31 @@ workspace = true [dependencies] authzed = { workspace = true } async-trait = { workspace = true } +axum = { workspace = true } +http = { workspace = true } +jsonwebtoken = { workspace = true } futures-util = { workspace = true } pg_query = "=6.2.1" rmcp = { workspace = true } +reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } sqlx = { workspace = true, features = ["tls-rustls-aws-lc-rs"] } thiserror = { workspace = true } -tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-std", "sync", "time"] } +tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "signal", "sync", "time"] } trogon-std = { workspace = true } toml = { workspace = true } +tower-http = { workspace = true, features = ["cors"] } tracing = { workspace = true } tracing-subscriber = { workspace = true, features = ["json"] } tonic = { workspace = true, features = ["tls-aws-lc", "tls-native-roots"] } url = { workspace = true } [dev-dependencies] +base64 = { workspace = true } +rmcp = { workspace = true, features = ["transport-streamable-http-client-reqwest"] } rcgen = "=0.14.8" +rsa = { version = "=0.9.10", features = ["pem"] } rustls = { workspace = true } tempfile = { workspace = true } testcontainers-modules = { workspace = true, features = ["postgres"] } diff --git a/rsworkspace/crates/mcp/mcp-postgres/examples/config.toml b/rsworkspace/crates/mcp/mcp-postgres/examples/config.toml index dcd7714db..eabee60fd 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/examples/config.toml +++ b/rsworkspace/crates/mcp/mcp-postgres/examples/config.toml @@ -1,5 +1,3 @@ -user = "example-user" -client = "example-client" spicedb_endpoint = "http://127.0.0.1:50051" spicedb_token_env = "MCP_POSTGRES_SPICEDB_TOKEN" @@ -7,3 +5,24 @@ spicedb_token_env = "MCP_POSTGRES_SPICEDB_TOKEN" name = "analytics-readonly" id = "analytics-local-v1" dsn_env = "MCP_POSTGRES_ANALYTICS_DSN" + +[http] +bind = "0.0.0.0:8080" +resource_url = "https://sql.example.com/mcp" +allowed_origins = ["https://admin.example.com"] + +[[authentication.issuers]] +issuer = "https://identity.example.com" +jwks_uri = "https://identity.example.com/.well-known/jwks.json" +audience = "https://sql.example.com/mcp" +user_claim = "sub" +client_claim = "client_id" +scope_claim = "scope" +required_scopes = ["sql:query"] +delegation = "optional" + +[authentication.issuers.users] +"external-user-id" = "example-user" + +[authentication.issuers.clients] +"external-client-id" = "example-client" diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/auth_error.rs b/rsworkspace/crates/mcp/mcp-postgres/src/auth_error.rs new file mode 100644 index 000000000..c471b7927 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/auth_error.rs @@ -0,0 +1,11 @@ +#[derive(Debug, thiserror::Error)] +pub enum AuthError { + #[error("invalid authentication configuration")] + Configuration, + #[error("authentication required")] + Unauthenticated, + #[error("required scope is missing")] + InsufficientScope, + #[error("authentication service unavailable")] + Unavailable, +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/authentication_config.rs b/rsworkspace/crates/mcp/mcp-postgres/src/authentication_config.rs new file mode 100644 index 000000000..03c7a411d --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/authentication_config.rs @@ -0,0 +1,123 @@ +use crate::{AuthError, ObjectId}; +use serde::Deserialize; +use std::collections::BTreeMap; +use url::Url; + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AuthenticationConfigInput { + pub issuers: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct IssuerInput { + pub issuer: String, + pub jwks_uri: String, + pub audience: String, + #[serde(default = "default_user_claim")] + pub user_claim: String, + #[serde(default = "default_client_claim")] + pub client_claim: String, + #[serde(default = "default_scope_claim")] + pub scope_claim: String, + #[serde(default)] + pub required_scopes: Vec, + pub users: BTreeMap, + pub clients: BTreeMap, + #[serde(default)] + pub delegation: DelegationPolicy, +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DelegationPolicy { + #[default] + Optional, + Required, + Forbidden, +} + +pub(crate) struct Issuer { + pub(crate) issuer: String, + pub(crate) jwks_uri: Url, + pub(crate) audience: String, + pub(crate) user_claim: String, + pub(crate) client_claim: String, + pub(crate) scope_claim: String, + pub(crate) required_scopes: Vec, + pub(crate) users: BTreeMap, + pub(crate) clients: BTreeMap, + pub(crate) delegation: DelegationPolicy, +} + +impl TryFrom for Issuer { + type Error = AuthError; + + fn try_from(input: IssuerInput) -> Result { + let issuer_url = endpoint(&input.issuer)?; + let jwks_uri = endpoint(&input.jwks_uri)?; + if issuer_url.query().is_some() + || input.audience.is_empty() + || [&input.user_claim, &input.client_claim, &input.scope_claim] + .iter() + .any(|claim| claim.is_empty() || claim.len() > 128) + || input + .required_scopes + .iter() + .any(|scope| scope.is_empty() || scope.split_whitespace().count() != 1) + || input.users.is_empty() + || input.clients.is_empty() + { + return Err(AuthError::Configuration); + } + Ok(Self { + issuer: input.issuer, + jwks_uri, + audience: input.audience, + user_claim: input.user_claim, + client_claim: input.client_claim, + scope_claim: input.scope_claim, + required_scopes: input.required_scopes, + users: mappings(input.users)?, + clients: mappings(input.clients)?, + delegation: input.delegation, + }) + } +} + +fn mappings(input: BTreeMap) -> Result, AuthError> { + input + .into_iter() + .map(|(external, local)| { + if external.is_empty() { + return Err(AuthError::Configuration); + } + Ok((external, ObjectId::new(local).map_err(|_| AuthError::Configuration)?)) + }) + .collect() +} + +fn endpoint(input: &str) -> Result { + let url = Url::parse(input).map_err(|_| AuthError::Configuration)?; + let loopback = matches!(url.host_str(), Some("localhost" | "127.0.0.1" | "[::1]")); + if !(url.scheme() == "https" || (url.scheme() == "http" && loopback)) + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.fragment().is_some() + { + return Err(AuthError::Configuration); + } + Ok(url) +} + +fn default_user_claim() -> String { + "sub".to_owned() +} +fn default_client_claim() -> String { + "client_id".to_owned() +} +fn default_scope_claim() -> String { + "scope".to_owned() +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/authenticator.rs b/rsworkspace/crates/mcp/mcp-postgres/src/authenticator.rs new file mode 100644 index 000000000..37d9dce4e --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/authenticator.rs @@ -0,0 +1,248 @@ +use crate::authentication_config::Issuer; +use crate::{AuthError, AuthenticationConfigInput, DelegationPolicy, RequestIdentity}; +use jsonwebtoken::jwk::{Jwk, JwkSet, KeyAlgorithm, KeyOperations, PublicKeyUse}; +use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header}; +use serde_json::Value; +use std::collections::{BTreeMap, BTreeSet}; +use std::time::{Duration, Instant}; +use tokio::sync::Mutex; + +pub struct Authenticator { + issuers: BTreeMap, + client: reqwest::Client, +} + +struct IssuerState { + config: Issuer, + keys: Mutex, +} + +struct KeyCache { + keys: JwkSet, + fetched: Instant, + attempted: Option, +} + +impl Authenticator { + pub async fn new(input: AuthenticationConfigInput) -> Result { + if input.issuers.is_empty() { + return Err(AuthError::Configuration); + } + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .timeout(Duration::from_secs(5)) + .build() + .map_err(|_| AuthError::Configuration)?; + let mut issuers = BTreeMap::new(); + for input in input.issuers { + let config = Issuer::try_from(input)?; + if issuers.contains_key(&config.issuer) { + return Err(AuthError::Configuration); + } + let keys = fetch_keys(&client, &config).await?; + let now = Instant::now(); + issuers.insert( + config.issuer.clone(), + IssuerState { + config, + keys: Mutex::new(KeyCache { + keys, + fetched: now, + attempted: None, + }), + }, + ); + } + Ok(Self { issuers, client }) + } + + pub fn issuer_urls(&self) -> Vec { + self.issuers.keys().cloned().collect() + } + + pub async fn authenticate(&self, token: &str) -> Result { + if token.len() > crate::constants::MAX_TOKEN_BYTES { + return Err(AuthError::Unauthenticated); + } + let unverified = + jsonwebtoken::dangerous::insecure_decode::(token).map_err(|_| AuthError::Unauthenticated)?; + let issuer = string_claim(&unverified.claims, "iss")?; + let state = self.issuers.get(issuer).ok_or(AuthError::Unauthenticated)?; + let header = decode_header(token).map_err(|_| AuthError::Unauthenticated)?; + if header.crit.is_some() || header.enc.is_some() || header.zip.is_some() { + return Err(AuthError::Unauthenticated); + } + if !matches!(header.alg, Algorithm::RS256 | Algorithm::ES256 | Algorithm::EdDSA) { + return Err(AuthError::Unauthenticated); + } + let kid = header + .kid + .as_deref() + .filter(|id| !id.is_empty()) + .ok_or(AuthError::Unauthenticated)?; + let key = self.key(state, kid, header.alg).await?; + let mut validation = Validation::new(header.alg); + validation.set_required_spec_claims(&["exp", "iss", "aud", "sub"]); + validation.set_issuer(&[&state.config.issuer]); + validation.set_audience(&[&state.config.audience]); + validation.validate_nbf = true; + validation.leeway = 0; + let claims = decode::(token, &key, &validation) + .map_err(|_| AuthError::Unauthenticated)? + .claims; + principal(&state.config, &claims) + } + + async fn key(&self, state: &IssuerState, kid: &str, algorithm: Algorithm) -> Result { + let mut cache = state.keys.lock().await; + let expired = cache.fetched.elapsed() >= Duration::from_secs(crate::constants::JWKS_CACHE_SECONDS); + let missing = !cache + .keys + .keys + .iter() + .any(|key| key.common.key_id.as_deref() == Some(kid)); + if expired || missing { + if cache.attempted.is_some_and(|attempted| { + attempted.elapsed() < Duration::from_secs(crate::constants::JWKS_REFRESH_SECONDS) + }) { + return Err(if expired { + AuthError::Unavailable + } else { + AuthError::Unauthenticated + }); + } + cache.attempted = Some(Instant::now()); + cache.keys = fetch_keys(&self.client, &state.config).await?; + cache.fetched = Instant::now(); + } + let mut matches = cache + .keys + .keys + .iter() + .filter(|key| key.common.key_id.as_deref() == Some(kid)); + let jwk = matches.next().ok_or(AuthError::Unauthenticated)?; + if matches.next().is_some() || !compatible(jwk, algorithm) { + return Err(AuthError::Unauthenticated); + } + DecodingKey::from_jwk(jwk).map_err(|_| AuthError::Unauthenticated) + } +} + +async fn fetch_keys(client: &reqwest::Client, issuer: &Issuer) -> Result { + let mut response = client + .get(issuer.jwks_uri.clone()) + .send() + .await + .map_err(|_| AuthError::Unavailable)?; + if !response.status().is_success() + || response + .content_length() + .is_some_and(|size| size > crate::constants::MAX_JWKS_BYTES as u64) + { + return Err(AuthError::Unavailable); + } + let mut bytes = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(|_| AuthError::Unavailable)? { + if bytes.len().saturating_add(chunk.len()) > crate::constants::MAX_JWKS_BYTES { + return Err(AuthError::Unavailable); + } + bytes.extend_from_slice(&chunk); + } + let keys: JwkSet = serde_json::from_slice(&bytes).map_err(|_| AuthError::Unavailable)?; + if keys.keys.is_empty() { + return Err(AuthError::Unavailable); + } + Ok(keys) +} + +fn compatible(key: &Jwk, algorithm: Algorithm) -> bool { + key.common + .public_key_use + .as_ref() + .is_none_or(|usage| *usage == PublicKeyUse::Signature) + && key + .common + .key_operations + .as_ref() + .is_none_or(|operations| operations.contains(&KeyOperations::Verify)) + && key + .common + .key_algorithm + .is_none_or(|alg| alg == KeyAlgorithm::from(algorithm)) +} + +fn principal(config: &Issuer, claims: &Value) -> Result { + string_claim(claims, "sub")?; + if claims.get("chain_act").is_some() { + return Err(AuthError::Unauthenticated); + } + let user = config + .users + .get(string_claim(claims, &config.user_claim)?) + .ok_or(AuthError::Unauthenticated)? + .clone(); + let scopes: BTreeSet<&str> = match claims.get(&config.scope_claim) { + Some(Value::String(scope)) => scope.split_whitespace().collect(), + Some(Value::Array(scopes)) => scopes + .iter() + .map(|scope| scope.as_str().ok_or(AuthError::Unauthenticated)) + .collect::>()?, + None if config.required_scopes.is_empty() => BTreeSet::new(), + _ => return Err(AuthError::Unauthenticated), + }; + if config + .required_scopes + .iter() + .any(|scope| !scopes.contains(scope.as_str())) + { + return Err(AuthError::InsufficientScope); + } + let act = claims.get("act"); + if (act.is_some() && config.delegation == DelegationPolicy::Forbidden) + || (act.is_none() && config.delegation == DelegationPolicy::Required) + { + return Err(AuthError::Unauthenticated); + } + let mut actors = Vec::new(); + let mut seen = BTreeSet::new(); + let mut current = act; + while let Some(actor) = current { + if actors.len() >= crate::constants::MAX_DELEGATION_DEPTH || !actor.is_object() { + return Err(AuthError::Unauthenticated); + } + let external = string_claim(actor, "sub")?; + if actor + .get("iss") + .is_some_and(|issuer| issuer.as_str() != Some(config.issuer.as_str())) + { + return Err(AuthError::Unauthenticated); + } + let mapped = config.clients.get(external).ok_or(AuthError::Unauthenticated)?.clone(); + if !seen.insert(mapped.clone()) { + return Err(AuthError::Unauthenticated); + } + actors.push(mapped); + current = actor.get("act"); + } + let client = match actors.first() { + Some(actor) => actor.clone(), + None => config + .clients + .get(string_claim(claims, &config.client_claim)?) + .ok_or(AuthError::Unauthenticated)? + .clone(), + }; + Ok(RequestIdentity::new(user, client, actors)) +} + +fn string_claim<'a>(claims: &'a Value, name: &str) -> Result<&'a str, AuthError> { + claims + .get(name) + .and_then(Value::as_str) + .filter(|claim| !claim.is_empty()) + .ok_or(AuthError::Unauthenticated) +} + +#[cfg(test)] +#[path = "authenticator_tests.rs"] +mod tests; diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/authenticator_tests.rs b/rsworkspace/crates/mcp/mcp-postgres/src/authenticator_tests.rs new file mode 100644 index 000000000..f5d47957d --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/authenticator_tests.rs @@ -0,0 +1,173 @@ +use super::*; +use crate::IssuerInput; +use serde_json::json; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +type TestResult = Result<(), Box>; + +fn input() -> Result { + serde_json::from_value(json!({ + "issuer":"http://127.0.0.1:1", "jwks_uri":"http://127.0.0.1:1/jwks", "audience":"mcp", + "users":{"user":"local-user"},"clients":{"agent":"local-agent"} + })) +} + +#[test] +fn validates_configuration_without_accepting_unsafe_endpoints_or_mappings() -> TestResult { + for value in [ + "not a url", + "ftp://localhost/jwks", + "https://user:password@example.test/jwks", + "https://example.test/#fragment", + ] { + let mut config = input()?; + config.jwks_uri = value.to_owned(); + assert!(matches!(Issuer::try_from(config), Err(AuthError::Configuration))); + } + for field in [ + "issuer", + "user_claim", + "client_claim", + "scope_claim", + "required_scopes", + "users", + "clients", + ] { + let mut value = serde_json::to_value(json!({ + "issuer":"https://example.test", "jwks_uri":"https://example.test/jwks", "audience":"mcp", + "users":{"user":"local-user"},"clients":{"agent":"local-agent"} + }))?; + value[field] = match field { + "issuer" => json!("https://example.test?query"), + "required_scopes" => json!(["two scopes"]), + "users" | "clients" => json!({}), + _ => json!(""), + }; + assert!(matches!( + Issuer::try_from(serde_json::from_value::(value)?), + Err(AuthError::Configuration) + )); + } + for mapping in [json!({"":"valid"}), json!({"user":"invalid@local"})] { + let mut config = input()?; + config.users = serde_json::from_value(mapping)?; + assert!(matches!(Issuer::try_from(config), Err(AuthError::Configuration))); + } + Ok(()) +} + +#[tokio::test] +async fn refuses_empty_configuration_and_unreachable_key_source() -> TestResult { + assert!(matches!( + Authenticator::new(AuthenticationConfigInput { issuers: vec![] }).await, + Err(AuthError::Configuration) + )); + assert!(matches!( + Authenticator::new(AuthenticationConfigInput { + issuers: vec![input()?] + }) + .await, + Err(AuthError::Unavailable) + )); + Ok(()) +} + +#[test] +fn bounds_delegation_and_accepts_array_scopes() -> TestResult { + let mut config = Issuer::try_from(input()?)?; + config.required_scopes = vec!["read".to_owned()]; + let mut claims = json!({"sub":"user","client_id":"agent","scope":["read"]}); + assert!(principal(&config, &claims).is_ok()); + claims["scope"] = json!([1]); + assert!(matches!(principal(&config, &claims), Err(AuthError::Unauthenticated))); + claims["scope"] = json!(null); + assert!(matches!(principal(&config, &claims), Err(AuthError::Unauthenticated))); + claims.as_object_mut().ok_or("claims")?.remove("scope"); + assert!(matches!(principal(&config, &claims), Err(AuthError::Unauthenticated))); + config.required_scopes.clear(); + assert!(principal(&config, &claims).is_ok()); + let mut act = json!({"sub":"actor-8"}); + for index in (0..8).rev() { + act = json!({"sub":format!("actor-{index}"),"act":act}); + } + for index in 0..9 { + config.clients.insert( + format!("actor-{index}"), + crate::ObjectId::new(format!("actor-{index}"))?, + ); + } + claims["act"] = act; + assert!(matches!(principal(&config, &claims), Err(AuthError::Unauthenticated))); + claims["act"] = json!(null); + assert!(matches!(principal(&config, &claims), Err(AuthError::Unauthenticated))); + Ok(()) +} + +#[tokio::test] +async fn expires_cached_keys_and_throttles_failed_refresh() -> TestResult { + let config = Issuer::try_from(input()?)?; + let state = IssuerState { + config, + keys: Mutex::new(KeyCache { + keys: JwkSet { keys: vec![] }, + fetched: Instant::now() - Duration::from_secs(400), + attempted: Some(Instant::now()), + }), + }; + let auth = Authenticator { + issuers: BTreeMap::new(), + client: reqwest::Client::new(), + }; + assert!(matches!( + auth.key(&state, "missing", Algorithm::ES256).await, + Err(AuthError::Unavailable) + )); + state.keys.lock().await.attempted = None; + assert!(matches!( + auth.key(&state, "missing", Algorithm::ES256).await, + Err(AuthError::Unavailable) + )); + assert!(state.keys.lock().await.attempted.is_some()); + Ok(()) +} + +async fn source(response: Vec) -> Result<(Issuer, tokio::task::JoinHandle<()>), Box> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let mut config = input()?; + config.jwks_uri = format!("http://{}/jwks", listener.local_addr()?); + let task = tokio::spawn(async move { + if let Ok((mut socket, _)) = listener.accept().await { + let mut bytes = [0; 4096]; + let _ = socket.read(&mut bytes).await; + let _ = socket.write_all(&response).await; + } + }); + Ok((Issuer::try_from(config)?, task)) +} + +#[tokio::test] +async fn rejects_failed_oversized_truncated_and_malformed_key_responses() -> TestResult { + let oversized = "a".repeat(crate::constants::MAX_JWKS_BYTES + 1); + let responses = [ + b"HTTP/1.1 503 Unavailable\r\nContent-Length: 0\r\n\r\n".to_vec(), + b"HTTP/1.1 200 OK\r\nContent-Length: 999999\r\n\r\n".to_vec(), + b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n{".to_vec(), + b"HTTP/1.1 200 OK\r\nContent-Length: 1\r\n\r\n{".to_vec(), + b"HTTP/1.1 200 OK\r\nContent-Length: 11\r\n\r\n{\"keys\":[]}".to_vec(), + format!( + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n{:x}\r\n{}\r\n0\r\n\r\n", + oversized.len(), + oversized + ) + .into_bytes(), + ]; + for response in responses { + let (issuer, task) = source(response).await?; + assert!(matches!( + fetch_keys(&reqwest::Client::new(), &issuer).await, + Err(AuthError::Unavailable) + )); + task.await?; + } + Ok(()) +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/config.rs b/rsworkspace/crates/mcp/mcp-postgres/src/config.rs index fd988661d..54b3fd989 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/config.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/config.rs @@ -5,8 +5,10 @@ use std::collections::BTreeMap; #[derive(Deserialize)] #[serde(deny_unknown_fields)] pub struct ConfigInput { - pub user: String, - pub client: String, + #[serde(default)] + pub http: Option, + #[serde(default)] + pub authentication: Option, pub spicedb_endpoint: String, pub spicedb_token_env: String, pub connections: Vec, @@ -23,8 +25,6 @@ pub struct RegisteredConnection { pub(crate) options: crate::DatabaseOptions, } pub struct RuntimeConfig { - pub(crate) user: ObjectId, - pub(crate) client: ObjectId, pub(crate) connections: BTreeMap, pub(crate) spicedb_endpoint: crate::SpiceDbEndpoint, pub(crate) spicedb_token: String, @@ -59,8 +59,6 @@ impl RuntimeConfig { return Err(Error::Configuration); } Ok(Self { - user: ObjectId::new(input.user)?, - client: ObjectId::new(input.client)?, connections, spicedb_endpoint: crate::SpiceDbEndpoint::new(input.spicedb_endpoint)?, spicedb_token: env(crate::EnvironmentVariable::new(input.spicedb_token_env)?.as_str()) diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs b/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs index 1e300b767..77660ffea 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/constants.rs @@ -1,3 +1,9 @@ pub(crate) const MAX_BYTES: usize = 1_048_576; pub(crate) const MAX_OBJECT_ID_BYTES: usize = 200; pub(crate) const MAX_CONNECTION_ID_BYTES: usize = MAX_OBJECT_ID_BYTES - "/4294967295/4294967295/32767".len(); +pub(crate) const MAX_TOKEN_BYTES: usize = 16_384; +pub(crate) const MAX_JWKS_BYTES: usize = 262_144; +pub(crate) const MAX_DELEGATION_DEPTH: usize = 8; +pub(crate) const JWKS_CACHE_SECONDS: u64 = 300; +pub(crate) const JWKS_REFRESH_SECONDS: u64 = 30; +pub(crate) const MCP_MAX_REQUEST_BODY_BYTES: usize = 131_072; diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/http_config.rs b/rsworkspace/crates/mcp/mcp-postgres/src/http_config.rs new file mode 100644 index 000000000..1c41102aa --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/http_config.rs @@ -0,0 +1,81 @@ +use crate::Error; +use serde::Deserialize; +use std::net::SocketAddr; +use url::Url; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub struct HttpTransportConfigInput { + pub bind: SocketAddr, + pub resource_url: String, + #[serde(default)] + pub allowed_origins: Vec, +} + +pub struct HttpTransportConfig { + bind: SocketAddr, + resource_url: Url, + allowed_origins: Vec, + transport_origins: Vec, +} + +impl HttpTransportConfig { + pub fn new(input: HttpTransportConfigInput) -> Result { + let resource_url = Url::parse(&input.resource_url).map_err(|_| Error::Configuration)?; + if !secure_url(&resource_url) || resource_url.path() != "/mcp" || resource_url.as_str() != input.resource_url { + return Err(Error::Configuration); + } + let mut allowed_origins = Vec::new(); + let mut transport_origins = Vec::new(); + for origin in input.allowed_origins { + let url = Url::parse(&origin).map_err(|_| Error::Configuration)?; + if !secure_url(&url) || url.path() != "/" || url.origin().ascii_serialization() != origin { + return Err(Error::Configuration); + } + transport_origins.push(format!( + "{}://{}:{}", + url.scheme(), + &url[url::Position::BeforeHost..url::Position::AfterHost], + url.port_or_known_default().ok_or(Error::Configuration)? + )); + allowed_origins.push(origin); + } + Ok(Self { + bind: input.bind, + resource_url, + allowed_origins, + transport_origins, + }) + } + + pub fn bind(&self) -> SocketAddr { + self.bind + } + pub fn resource_url(&self) -> &str { + self.resource_url.as_str() + } + pub(crate) fn allowed_origins(&self) -> &[String] { + &self.allowed_origins + } + pub(crate) fn transport_origins(&self) -> &[String] { + &self.transport_origins + } + pub(crate) fn authority(&self) -> String { + self.resource_url[url::Position::BeforeHost..url::Position::AfterPort].to_owned() + } + pub(crate) fn metadata_url(&self) -> String { + let mut url = self.resource_url.clone(); + url.set_path("/.well-known/oauth-protected-resource/mcp"); + url.to_string() + } +} + +fn secure_url(url: &Url) -> bool { + let local = matches!(url.host_str(), Some("localhost" | "127.0.0.1" | "[::1]")); + (url.scheme() == "https" || (url.scheme() == "http" && local)) + && url.host_str().is_some() + && url.username().is_empty() + && url.password().is_none() + && url.query().is_none() + && url.fragment().is_none() +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/http_transport.rs b/rsworkspace/crates/mcp/mcp-postgres/src/http_transport.rs new file mode 100644 index 000000000..f8ad6c29c --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/http_transport.rs @@ -0,0 +1,127 @@ +use crate::{AuthError, Authenticator, HttpTransportConfig, PostgresMcp, QueryService}; +use axum::{ + Json, Router, + extract::{Request, State}, + http::{HeaderName, HeaderValue, Method, StatusCode, header}, + middleware::{self, Next}, + response::{IntoResponse, Response}, + routing::get, +}; +use rmcp::transport::streamable_http_server::{ + StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager, +}; +use serde::Serialize; +use std::sync::Arc; +use tower_http::cors::{AllowOrigin, CorsLayer}; + +#[derive(Clone)] +struct Admission { + authenticator: Arc, + challenge: HeaderValue, + insufficient_scope_challenge: HeaderValue, +} + +#[derive(Clone, Serialize)] +struct ProtectedResourceMetadata { + resource: String, + authorization_servers: Vec, + bearer_methods_supported: Vec, +} + +pub fn http_router(service: QueryService, authenticator: Arc, config: HttpTransportConfig) -> Router { + let origins = config.allowed_origins().to_vec(); + let cors = CorsLayer::new() + .allow_origin(AllowOrigin::predicate(move |origin, _| { + origins.iter().any(|allowed| origin.as_bytes() == allowed.as_bytes()) + })) + .allow_methods([Method::GET, Method::POST, Method::DELETE]) + .allow_headers([ + header::AUTHORIZATION, + header::CONTENT_TYPE, + header::ACCEPT, + HeaderName::from_static("mcp-protocol-version"), + ]) + .expose_headers([header::WWW_AUTHENTICATE]); + let mcp = PostgresMcp::new(service); + let http_config = StreamableHttpServerConfig::default() + .with_legacy_session_mode(false) + .with_json_response(true) + .with_allowed_hosts(vec![config.authority()]) + .with_allowed_origins(config.transport_origins().to_vec()) + .enforce_origin_validation() + .with_max_request_body_bytes(crate::constants::MCP_MAX_REQUEST_BODY_BYTES); + let transport: StreamableHttpService = + StreamableHttpService::new(move || Ok(mcp.clone()), Default::default(), http_config); + let metadata = ProtectedResourceMetadata { + resource: config.resource_url().to_owned(), + authorization_servers: authenticator.issuer_urls(), + bearer_methods_supported: vec!["header".into()], + }; + let challenge = HeaderValue::from_str(&format!("Bearer resource_metadata=\"{}\"", config.metadata_url())) + .unwrap_or_else(|_| HeaderValue::from_static("Bearer")); + let insufficient_scope_challenge = HeaderValue::from_str(&format!( + "Bearer error=\"insufficient_scope\", resource_metadata=\"{}\"", + config.metadata_url() + )) + .unwrap_or_else(|_| HeaderValue::from_static("Bearer error=\"insufficient_scope\"")); + let admission = Admission { + authenticator, + challenge, + insufficient_scope_challenge, + }; + Router::new() + .route_service("/mcp", transport) + .route_layer(middleware::from_fn_with_state(admission, authenticate)) + .route( + "/.well-known/oauth-protected-resource/mcp", + get(move || async move { Json(metadata) }), + ) + .layer(cors) +} + +async fn authenticate(State(admission): State, mut request: Request, next: Next) -> Response { + let token = bearer_token(request.headers()); + let Some(token) = token else { + return rejected(&admission, StatusCode::UNAUTHORIZED); + }; + match admission.authenticator.authenticate(token).await { + Ok(identity) => { + request.extensions_mut().insert(identity); + next.run(request).await + } + Err(AuthError::InsufficientScope) => rejected(&admission, StatusCode::FORBIDDEN), + Err(AuthError::Unavailable) => rejected(&admission, StatusCode::SERVICE_UNAVAILABLE), + Err(AuthError::Configuration | AuthError::Unauthenticated) => rejected(&admission, StatusCode::UNAUTHORIZED), + } +} + +fn bearer_token(headers: &axum::http::HeaderMap) -> Option<&str> { + let mut values = headers.get_all(header::AUTHORIZATION).iter(); + let value = values.next()?.to_str().ok()?; + if values.next().is_some() { + return None; + } + let (scheme, token) = value.split_once(' ')?; + if !scheme.eq_ignore_ascii_case("Bearer") + || token.is_empty() + || token.bytes().any(|byte| byte.is_ascii_whitespace()) + { + return None; + } + Some(token) +} + +fn rejected(admission: &Admission, status: StatusCode) -> Response { + let mut response = status.into_response(); + if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { + response.headers_mut().insert( + header::WWW_AUTHENTICATE, + if status == StatusCode::FORBIDDEN { + admission.insufficient_scope_challenge.clone() + } else { + admission.challenge.clone() + }, + ); + } + response +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs b/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs index f0ede8ea3..b15b0a473 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/lib.rs @@ -1,12 +1,12 @@ -//! Governed PostgreSQL queries over MCP for a trusted local stdio launcher. +//! Governed PostgreSQL queries over authenticated Streamable HTTP MCP. //! -//! Named connections and user/client identities come exclusively from runtime configuration. +//! Named connections come from runtime configuration; identities come from verified request tokens. //! The bounded SELECT grammar rejects unsupported syntax and catalog features. //! Every referenced column needs an explicit grant for both identities in one SpiceDB snapshot. //! Transactional relation locks and cursor metadata bind authorization to the executed relations. //! Production access and deployment are outside this crate's local implementation contract. //! -//! The launcher and database administrators are trusted. Callers must not control the configuration +//! The operator and database administrators are trusted. Callers must not control the configuration //! or credentials, and the query role must have only the required SELECT privileges. //! The service installs no database objects and requires no migration coordination. A read-only //! transaction locks supported base tables, verifies cursor relation identities, and authorizes @@ -16,33 +16,45 @@ //! Replacing or restoring a database target requires a new connection ID when catalog IDs can repeat. //! Permissions are checked at admission; revocation does not retroactively cancel an admitted query. +mod auth_error; +mod authentication_config; +mod authenticator; mod config; mod connection_id; mod constants; mod database_options; mod environment_variable; mod error; +mod http_config; +mod http_transport; mod identifier; mod object_id; mod parameter; mod parser; mod postgres; +mod request_identity; mod row_limit; mod server; mod spicedb; mod spicedb_endpoint; mod telemetry; +pub use auth_error::AuthError; +pub use authentication_config::{AuthenticationConfigInput, DelegationPolicy, IssuerInput}; +pub use authenticator::Authenticator; pub use config::{ConfigInput, RegisteredConnection, RuntimeConfig}; pub use connection_id::ConnectionId; pub use database_options::DatabaseOptions; pub use environment_variable::EnvironmentVariable; pub use error::Error; +pub use http_config::{HttpTransportConfig, HttpTransportConfigInput}; +pub use http_transport::http_router; pub use identifier::Identifier; pub use object_id::ObjectId; pub use parameter::{Parameter, ParameterInput}; pub use parser::Query; pub use postgres::{QueryResult, QueryService}; +pub use request_identity::RequestIdentity; pub use server::PostgresMcp; pub use spicedb::{PermissionCheck, PermissionRequest, Resource, SpiceDb}; pub use spicedb_endpoint::SpiceDbEndpoint; diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/main.rs b/rsworkspace/crates/mcp/mcp-postgres/src/main.rs index e2bdc4549..3e59c77d0 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/main.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/main.rs @@ -1,7 +1,8 @@ #![cfg_attr(coverage_nightly, feature(coverage_attribute))] -use mcp_postgres::{ConfigInput, PostgresMcp, QueryService, RuntimeConfig, SpiceDb}; -use rmcp::ServiceExt; +use mcp_postgres::{ + Authenticator, ConfigInput, HttpTransportConfig, QueryService, RuntimeConfig, SpiceDb, http_router, +}; use std::{process::ExitCode, sync::Arc}; use tracing_subscriber::{Layer, layer::SubscriberExt, util::SubscriberInitExt}; use trogon_std::env::{ReadEnv, SystemEnv}; @@ -29,12 +30,30 @@ async fn run() -> Result<(), Box> { if args.next().is_some() { return Err(mcp_postgres::Error::Configuration.into()); } - let input: ConfigInput = toml::from_str(&std::fs::read_to_string(path)?)?; + let mut input: ConfigInput = toml::from_str(&std::fs::read_to_string(path)?)?; + let http = HttpTransportConfig::new(input.http.take().ok_or(mcp_postgres::Error::Configuration)?)?; + let authentication = + Arc::new(Authenticator::new(input.authentication.take().ok_or(mcp_postgres::Error::Configuration)?).await?); + let address = http.bind(); let config = RuntimeConfig::from_input(input, |key| SystemEnv.var(key).ok())?; let permissions = Arc::new(SpiceDb::connect(config.spicedb_endpoint().into(), config.spicedb_token().into()).await?); let service = QueryService::connect(config, permissions).await?; - let server = PostgresMcp::new(service).serve(rmcp::transport::stdio()).await?; - server.waiting().await?; + let app = http_router(service, authentication, http); + #[cfg(unix)] + let mut terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?; + #[cfg(unix)] + let shutdown = async move { + tokio::select! { + _ = tokio::signal::ctrl_c() => {}, + _ = terminate.recv() => {}, + } + }; + #[cfg(not(unix))] + let shutdown = async { + let _ = tokio::signal::ctrl_c().await; + }; + let listener = tokio::net::TcpListener::bind(address).await?; + axum::serve(listener, app).with_graceful_shutdown(shutdown).await?; Ok(()) } diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs index 240f9f28f..109423412 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/postgres.rs @@ -1,5 +1,5 @@ use crate::{ - Error, ObjectId, PermissionCheck, Query, Resource, RuntimeConfig, + Error, ObjectId, PermissionCheck, Query, RequestIdentity, Resource, RuntimeConfig, parser::{Column, Expr}, }; use futures_util::TryStreamExt; @@ -43,8 +43,6 @@ impl QueryResult { } } pub struct QueryService { - user: ObjectId, - client: ObjectId, connections: BTreeMap, permissions: Arc, permits: Semaphore, @@ -68,27 +66,23 @@ impl QueryService { connections.insert(name, (connection.id.into_object_id(), pool)); } Ok(Self { - user: config.user, - client: config.client, connections, permissions, permits: Semaphore::new(4), }) } - pub async fn query(&self, connection: &str, sql: &str) -> Result { - self.query_with_parameters(connection, sql, &[]).await + pub async fn query(&self, identity: &RequestIdentity, connection: &str, sql: &str) -> Result { + self.query_with_parameters(identity, connection, sql, &[]).await } pub async fn query_with_parameters( &self, + identity: &RequestIdentity, connection: &str, sql: &str, parameters: &[crate::Parameter], ) -> Result { - let mut audit = crate::telemetry::QueryAudit::start( - &self.user, - &self.client, - self.connections.get(connection).map(|(id, _)| id), - ); + let mut audit = + crate::telemetry::QueryAudit::start(identity, self.connections.get(connection).map(|(id, _)| id)); let result = async { if parameters.len() > 32 { return Err(Error::Limit); @@ -96,7 +90,7 @@ impl QueryService { let _permit = self.permits.try_acquire().map_err(|_| Error::Limit)?; tokio::time::timeout( Duration::from_secs(10), - self.execute(connection, Query::parse(sql)?, parameters), + self.execute(identity, connection, Query::parse(sql)?, parameters), ) .await .map_err(|_| Error::Timeout)? @@ -105,21 +99,25 @@ impl QueryService { audit.complete(&result); result } - async fn authorize(&self, resource: &Resource) -> Result<(), Error> { - self.permissions.check(resource, "sql_user", &self.user).await?; - self.permissions.check(resource, "sql_client", &self.client).await + async fn authorize(&self, identity: &RequestIdentity, resource: &Resource) -> Result<(), Error> { + self.permissions.check(resource, "sql_user", identity.user()).await?; + self.permissions.check(resource, "sql_client", identity.client()).await } async fn execute( &self, + identity: &RequestIdentity, connection: &str, query: Query, parameters: &[crate::Parameter], ) -> Result { let (connection_id, pool) = self.connections.get(connection).ok_or(Error::UnknownConnection)?; - self.authorize(&Resource { - kind: "sql_connection", - id: connection_id.clone(), - }) + self.authorize( + identity, + &Resource { + kind: "sql_connection", + id: connection_id.clone(), + }, + ) .await?; let mut connection = pool.acquire().await?; connection.close_on_drop(); @@ -284,7 +282,7 @@ impl QueryService { .await?; let visible_columns = query.select.len(); verify_bindings(metadata.columns(), visible_columns, &relations)?; - crate::telemetry::references(&resources, &self.user, &self.client, connection_id); + crate::telemetry::references(&resources, identity.user(), identity.client(), connection_id); let requests = resources .into_iter() .flat_map(|resource| { @@ -292,12 +290,12 @@ impl QueryService { crate::PermissionRequest { resource: resource.clone(), subject_kind: "sql_user", - subject: self.user.clone(), + subject: identity.user().clone(), }, crate::PermissionRequest { resource, subject_kind: "sql_client", - subject: self.client.clone(), + subject: identity.client().clone(), }, ] }) diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/request_identity.rs b/rsworkspace/crates/mcp/mcp-postgres/src/request_identity.rs new file mode 100644 index 000000000..f0fbfff08 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/src/request_identity.rs @@ -0,0 +1,30 @@ +use crate::ObjectId; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RequestIdentity { + user: ObjectId, + client: ObjectId, + delegation: Vec, +} + +impl RequestIdentity { + pub fn new(user: ObjectId, client: ObjectId, delegation: Vec) -> Self { + Self { + user, + client, + delegation, + } + } + + pub fn user(&self) -> &ObjectId { + &self.user + } + + pub fn client(&self) -> &ObjectId { + &self.client + } + + pub fn delegation(&self) -> &[ObjectId] { + &self.delegation + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/server.rs b/rsworkspace/crates/mcp/mcp-postgres/src/server.rs index 925de8510..8b63b77ac 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/server.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/server.rs @@ -1,4 +1,4 @@ -use crate::QueryService; +use crate::{QueryService, RequestIdentity}; use rmcp::schemars; use rmcp::{ ErrorData, ServerHandler, @@ -31,13 +31,14 @@ struct QueryRequest { } impl ServerHandler for PostgresMcp { fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()).with_instructions("Queries use registered connections and the launcher's trusted identity. Only the documented bounded SELECT subset is supported.") + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()).with_instructions("Queries use registered connections and the authenticated request identity. Only the documented bounded SELECT subset is supported.") } async fn list_tools( &self, _request: Option, - _context: RequestContext, + context: RequestContext, ) -> Result { + request_identity(&context)?; let input_schema = rmcp::handler::server::common::schema_for_type::(); Ok(ListToolsResult { tools: vec![Tool::new( @@ -53,6 +54,7 @@ impl ServerHandler for PostgresMcp { request: CallToolRequestParams, context: RequestContext, ) -> Result { + let identity = request_identity(&context)?; if request.name != "query" { return Err(ErrorData::invalid_params("unknown tool", None)); } @@ -68,7 +70,7 @@ impl ServerHandler for PostgresMcp { let outcome = tokio::select! { biased; _ = context.ct.cancelled() => return Err(ErrorData::internal_error("request cancelled",None)), - outcome = self.service.query_with_parameters(&args.connection,&args.sql,¶meters) => outcome, + outcome = self.service.query_with_parameters(&identity,&args.connection,&args.sql,¶meters) => outcome, }; let result = match outcome { Ok(result) => CallToolResult::success(vec![ContentBlock::text(result.into_json().to_string())]), @@ -89,3 +91,12 @@ impl ServerHandler for PostgresMcp { Ok(CallToolResponse::Complete(result)) } } + +fn request_identity(context: &RequestContext) -> Result { + context + .extensions + .get::() + .and_then(|parts| parts.extensions.get::()) + .cloned() + .ok_or_else(|| ErrorData::internal_error("authenticated request context unavailable", None)) +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs b/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs index 4fb2f29e2..a1664cbc3 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/src/telemetry.rs @@ -1,4 +1,4 @@ -use crate::{Error, ObjectId, Resource}; +use crate::{Error, ObjectId, RequestIdentity, Resource}; use std::{collections::BTreeSet, time::Instant}; #[derive(Debug)] @@ -20,12 +20,15 @@ pub(crate) struct QueryAudit<'a> { outcome: QueryOutcome, } impl<'a> QueryAudit<'a> { - pub(crate) fn start(user: &'a ObjectId, client: &'a ObjectId, connection: Option<&'a ObjectId>) -> Self { + pub(crate) fn start(identity: &'a RequestIdentity, connection: Option<&'a ObjectId>) -> Self { + let user = identity.user(); + let client = identity.client(); { let user = user.as_str(); let client = client.as_str(); let connection = connection.map(ObjectId::as_str); - tracing::info!(event = "sql_query_admission", user, client, connection); + let delegation: Vec<_> = identity.delegation().iter().map(ObjectId::as_str).collect(); + tracing::info!(event = "sql_query_admission", user, client, connection, delegation = ?delegation); } Self { user, diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/authentication.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/authentication.rs new file mode 100644 index 000000000..b944739c7 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/authentication.rs @@ -0,0 +1,402 @@ +use axum::{Json, Router, extract::State, routing::get}; +use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, encode, jwk::Jwk}; +use mcp_postgres::{AuthError, AuthenticationConfigInput, Authenticator}; +use rsa::{RsaPrivateKey, pkcs8::EncodePrivateKey, rand_core::OsRng}; +use serde_json::{Value, json}; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; +use std::time::{SystemTime, UNIX_EPOCH}; +use tokio::sync::RwLock; + +type TestResult = Result<(), Box>; + +#[derive(Clone)] +struct StateData { + document: Arc>, + requests: Arc, +} + +struct Fixture { + issuer: String, + key: EncodingKey, + state: StateData, + task: tokio::task::JoinHandle<()>, +} + +impl Drop for Fixture { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl Fixture { + async fn new() -> Result> { + let pair = rcgen::KeyPair::generate()?; + let key = EncodingKey::from_ec_der(&pair.serialize_der()); + let mut jwk = Jwk::from_encoding_key(&key, Algorithm::ES256)?; + jwk.common.key_id = Some("key-one".to_owned()); + let state = StateData { + document: Arc::new(RwLock::new(json!({"keys":[jwk]}))), + requests: Arc::new(AtomicUsize::new(0)), + }; + let app = Router::new() + .route( + "/jwks", + get(|State(state): State| async move { + state.requests.fetch_add(1, Ordering::SeqCst); + Json(state.document.read().await.clone()) + }), + ) + .with_state(state.clone()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let issuer = format!("http://{}", listener.local_addr()?); + let task = tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + Ok(Self { + issuer, + key, + state, + task, + }) + } + + fn config(&self) -> Value { + json!({"issuers":[{ + "issuer":self.issuer,"jwks_uri":format!("{}/jwks",self.issuer),"audience":"sql-mcp", + "required_scopes":["sql:read"], + "users":{"external-user":"user-one"}, + "clients":{"external-client":"client-one","agent":"agent-one","portal":"portal-one"} + }]}) + } + + fn claims(&self) -> Value { + json!({"iss":self.issuer,"aud":"sql-mcp","exp":now()+300,"sub":"external-user", + "client_id":"external-client","scope":"sql:read"}) + } + + fn token(&self, claims: &Value) -> Result { + let mut header = Header::new(Algorithm::ES256); + header.kid = Some("key-one".to_owned()); + encode(&header, claims, &self.key) + } + + async fn auth(&self) -> Result> { + let input: AuthenticationConfigInput = serde_json::from_value(self.config())?; + Ok(Authenticator::new(input).await?) + } +} + +fn now() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |time| time.as_secs()) +} + +#[tokio::test] +async fn maps_verified_issuer_identity_and_delegated_actor() -> TestResult { + let fixture = Fixture::new().await?; + let auth = fixture.auth().await?; + let direct = auth.authenticate(&fixture.token(&fixture.claims())?).await?; + assert_eq!(direct.user().as_str(), "user-one"); + assert_eq!(direct.client().as_str(), "client-one"); + assert!(direct.delegation().is_empty()); + let mut delegated = fixture.claims(); + delegated["act"] = json!({"sub":"agent","act":{"sub":"portal"}}); + let principal = auth.authenticate(&fixture.token(&delegated)?).await?; + assert_eq!(principal.client().as_str(), "agent-one"); + assert_eq!( + principal + .delegation() + .iter() + .map(|actor| actor.as_str()) + .collect::>(), + ["agent-one", "portal-one"] + ); + assert_eq!(auth.issuer_urls(), std::slice::from_ref(&fixture.issuer)); + assert_eq!(fixture.state.requests.load(Ordering::SeqCst), 1); + Ok(()) +} + +#[tokio::test] +async fn maps_configured_okta_claims_without_using_email_subject_as_local_identity() -> TestResult { + let fixture = Fixture::new().await?; + let mut config = fixture.config(); + config["issuers"][0]["user_claim"] = json!("uid"); + config["issuers"][0]["client_claim"] = json!("cid"); + config["issuers"][0]["scope_claim"] = json!("scp"); + config["issuers"][0]["users"] = json!({"00u-test-user":"user-one"}); + config["issuers"][0]["clients"] = json!({"0oa-test-client":"client-one"}); + let auth = Authenticator::new(serde_json::from_value(config)?).await?; + let mut claims = json!({ + "iss":fixture.issuer,"aud":"sql-mcp","exp":now()+300, + "sub":"reader@example.test","uid":"00u-test-user","cid":"0oa-test-client", + "scp":["openid","sql:read"] + }); + let identity = auth.authenticate(&fixture.token(&claims)?).await?; + assert_eq!(identity.user().as_str(), "user-one"); + assert_eq!(identity.client().as_str(), "client-one"); + claims["uid"] = json!("unmapped-user"); + assert!(matches!( + auth.authenticate(&fixture.token(&claims)?).await, + Err(AuthError::Unauthenticated) + )); + claims["uid"] = json!("00u-test-user"); + claims["scp"] = json!(["openid"]); + assert!(matches!( + auth.authenticate(&fixture.token(&claims)?).await, + Err(AuthError::InsufficientScope) + )); + Ok(()) +} + +#[tokio::test] +async fn verifies_ephemeral_rsa_signed_access_token_against_published_jwks() -> TestResult { + let fixture = Fixture::new().await?; + let private = tokio::task::spawn_blocking(|| RsaPrivateKey::new(&mut OsRng, 2048)).await??; + let pem = private.to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)?; + let key = EncodingKey::from_rsa_pem(pem.as_bytes())?; + let mut jwk = Jwk::from_encoding_key(&key, Algorithm::RS256)?; + jwk.common.key_id = Some("rsa-key".to_owned()); + *fixture.state.document.write().await = json!({"keys":[jwk]}); + let auth = fixture.auth().await?; + let mut header = Header::new(Algorithm::RS256); + header.kid = Some("rsa-key".to_owned()); + let token = encode(&header, &fixture.claims(), &key)?; + let identity = auth.authenticate(&token).await?; + assert_eq!(identity.user().as_str(), "user-one"); + assert_eq!(identity.client().as_str(), "client-one"); + assert_eq!(fixture.state.requests.load(Ordering::SeqCst), 1); + Ok(()) +} + +#[tokio::test] +async fn verifies_ephemeral_ed25519_signed_access_token_against_published_jwks() -> TestResult { + let fixture = Fixture::new().await?; + let pair = rcgen::KeyPair::generate_for(&rcgen::PKCS_ED25519)?; + let key = EncodingKey::from_ed_der(&pair.serialize_der()); + let public = DecodingKey::from_ed_der(pair.public_key_raw()); + let mut jwk = Jwk::from_decoding_key(&public, Some(Algorithm::EdDSA))?; + jwk.common.key_id = Some("ed25519-key".to_owned()); + *fixture.state.document.write().await = json!({"keys":[jwk]}); + let auth = fixture.auth().await?; + let mut header = Header::new(Algorithm::EdDSA); + header.kid = Some("ed25519-key".to_owned()); + let token = encode(&header, &fixture.claims(), &key)?; + let identity = auth.authenticate(&token).await?; + assert_eq!(identity.user().as_str(), "user-one"); + assert_eq!(identity.client().as_str(), "client-one"); + Ok(()) +} + +#[tokio::test] +async fn rejects_invalid_claims_signatures_and_scopes() -> TestResult { + let fixture = Fixture::new().await?; + let auth = fixture.auth().await?; + for (claim, value) in [ + ("iss", json!("https://untrusted.example.test")), + ("aud", json!("other")), + ("exp", json!(now() - 1)), + ("nbf", json!(now() + 300)), + ("sub", json!("")), + ("sub", json!("unknown-user")), + ("client_id", json!("unknown-client")), + ("chain_act", json!([])), + ("act", json!({})), + ("act", json!({"sub":"unknown"})), + ("act", json!({"sub":"agent","act":{"sub":"agent"}})), + ("act", json!({"agent":"agent"})), + ] { + let mut claims = fixture.claims(); + claims[claim] = value; + assert!(matches!( + auth.authenticate(&fixture.token(&claims)?).await, + Err(AuthError::Unauthenticated) + )); + } + for claim in ["exp", "iss", "aud", "sub", "client_id"] { + let mut claims = fixture.claims(); + claims.as_object_mut().ok_or("claims object")?.remove(claim); + assert!(matches!( + auth.authenticate(&fixture.token(&claims)?).await, + Err(AuthError::Unauthenticated) + )); + } + let mut claims = fixture.claims(); + claims["scope"] = json!("other:scope"); + assert!(matches!( + auth.authenticate(&fixture.token(&claims)?).await, + Err(AuthError::InsufficientScope) + )); + let other_key = rcgen::KeyPair::generate()?; + let mut header = Header::new(Algorithm::ES256); + header.kid = Some("key-one".to_owned()); + let wrong_signature = encode( + &header, + &fixture.claims(), + &EncodingKey::from_ec_der(&other_key.serialize_der()), + )?; + assert!(matches!( + auth.authenticate(&wrong_signature).await, + Err(AuthError::Unauthenticated) + )); + assert!(matches!( + auth.authenticate(&"a".repeat(16385)).await, + Err(AuthError::Unauthenticated) + )); + assert_eq!(fixture.state.requests.load(Ordering::SeqCst), 1); + Ok(()) +} + +#[tokio::test] +async fn refreshes_rotated_keys_once_and_limits_unknown_key_fetches() -> TestResult { + let fixture = Fixture::new().await?; + let auth = fixture.auth().await?; + let rotated_pair = rcgen::KeyPair::generate()?; + let rotated_key = EncodingKey::from_ec_der(&rotated_pair.serialize_der()); + let mut jwk = Jwk::from_encoding_key(&rotated_key, Algorithm::ES256)?; + jwk.common.key_id = Some("key-two".to_owned()); + *fixture.state.document.write().await = json!({"keys":[jwk]}); + let mut header = Header::new(Algorithm::ES256); + header.kid = Some("key-two".to_owned()); + let token = encode(&header, &fixture.claims(), &rotated_key)?; + assert!(auth.authenticate(&token).await.is_ok()); + for kid in ["unknown-one", "unknown-two"] { + header.kid = Some(kid.to_owned()); + let token = encode(&header, &fixture.claims(), &rotated_key)?; + assert!(matches!( + auth.authenticate(&token).await, + Err(AuthError::Unauthenticated) + )); + } + assert_eq!(fixture.state.requests.load(Ordering::SeqCst), 2); + Ok(()) +} + +#[tokio::test] +async fn rejects_disallowed_delegation_and_unsafe_configuration() -> TestResult { + let fixture = Fixture::new().await?; + for (policy, act) in [("required", None), ("forbidden", Some(json!({"sub":"agent"})))] { + let mut config = fixture.config(); + config["issuers"][0]["delegation"] = json!(policy); + let auth = Authenticator::new(serde_json::from_value(config)?).await?; + let mut claims = fixture.claims(); + if let Some(act) = act { + claims["act"] = act; + } + assert!(matches!( + auth.authenticate(&fixture.token(&claims)?).await, + Err(AuthError::Unauthenticated) + )); + } + for (field, value) in [ + ("jwks_uri", "http://remote.example.test/jwks"), + ("issuer", "http://remote.example.test"), + ("audience", ""), + ] { + let mut config = fixture.config(); + config["issuers"][0][field] = json!(value); + assert!(matches!( + Authenticator::new(serde_json::from_value(config)?).await, + Err(AuthError::Configuration) + )); + } + Ok(()) +} + +#[tokio::test] +async fn rejects_incompatible_and_ambiguous_jwk_metadata() -> TestResult { + let fixture = Fixture::new().await?; + for (field, value) in [ + ("use", json!("enc")), + ("key_ops", json!(["encrypt"])), + ("alg", json!("RS256")), + ] { + let mut jwk = serde_json::to_value(Jwk::from_encoding_key(&fixture.key, Algorithm::ES256)?)?; + jwk["kid"] = json!("key-one"); + jwk[field] = value; + *fixture.state.document.write().await = json!({"keys":[jwk]}); + let auth = fixture.auth().await?; + assert!(matches!( + auth.authenticate(&fixture.token(&fixture.claims())?).await, + Err(AuthError::Unauthenticated) + )); + } + let mut jwk = serde_json::to_value(Jwk::from_encoding_key(&fixture.key, Algorithm::ES256)?)?; + jwk["kid"] = json!("key-one"); + *fixture.state.document.write().await = json!({"keys":[jwk.clone(),jwk]}); + let auth = fixture.auth().await?; + assert!(matches!( + auth.authenticate(&fixture.token(&fixture.claims())?).await, + Err(AuthError::Unauthenticated) + )); + Ok(()) +} + +#[tokio::test] +async fn rejects_delegation_from_a_different_actor_issuer() -> TestResult { + let fixture = Fixture::new().await?; + let auth = fixture.auth().await?; + for actor in [ + json!({"sub":"agent","iss":"https://other.example.test"}), + json!({"sub":"agent","iss":1}), + json!({"sub":"agent","act":{"sub":"portal","iss":"https://other.example.test"}}), + ] { + let mut claims = fixture.claims(); + claims["act"] = actor; + assert!(matches!( + auth.authenticate(&fixture.token(&claims)?).await, + Err(AuthError::Unauthenticated) + )); + } + let mut claims = fixture.claims(); + claims["act"] = json!({"sub":"agent","iss":fixture.issuer}); + assert!(auth.authenticate(&fixture.token(&claims)?).await.is_ok()); + Ok(()) +} + +#[tokio::test] +async fn refuses_unsupported_token_headers_and_malformed_tokens() -> TestResult { + let fixture = Fixture::new().await?; + let auth = fixture.auth().await?; + let mut header = Header::new(Algorithm::ES256); + header.kid = Some("key-one".to_owned()); + header.crit = Some(vec!["unrecognized".to_owned()]); + let token = encode(&header, &fixture.claims(), &fixture.key)?; + assert!(matches!( + auth.authenticate(&token).await, + Err(AuthError::Unauthenticated) + )); + header.crit = None; + header.kid = None; + let token = encode(&header, &fixture.claims(), &fixture.key)?; + assert!(matches!( + auth.authenticate(&token).await, + Err(AuthError::Unauthenticated) + )); + let mut header = Header::new(Algorithm::HS256); + header.kid = Some("key-one".to_owned()); + let token = encode( + &header, + &fixture.claims(), + &EncodingKey::from_secret(b"test-only-signing-material"), + )?; + assert!(matches!( + auth.authenticate(&token).await, + Err(AuthError::Unauthenticated) + )); + assert!(matches!( + auth.authenticate("not-a-jwt").await, + Err(AuthError::Unauthenticated) + )); + let mut config = fixture.config(); + let duplicate = config["issuers"][0].clone(); + config["issuers"].as_array_mut().ok_or("issuer array")?.push(duplicate); + assert!(matches!( + Authenticator::new(serde_json::from_value(config)?).await, + Err(AuthError::Configuration) + )); + Ok(()) +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/common/mod.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/common/mod.rs new file mode 100644 index 000000000..7085fd8e1 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/common/mod.rs @@ -0,0 +1,131 @@ +#![allow(dead_code)] +use axum::{Json, Router, routing::get}; +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use jsonwebtoken::{Algorithm, EncodingKey, Header}; +use mcp_postgres::{ + AuthenticationConfigInput, Authenticator, HttpTransportConfig, HttpTransportConfigInput, QueryService, http_router, +}; +use rmcp::{ + ServiceExt, + service::{RoleClient, RunningService}, + transport::{StreamableHttpClientTransport, streamable_http_client::StreamableHttpClientTransportConfig}, +}; +use serde_json::{Value, json}; +use std::{ + sync::Arc, + time::{SystemTime, UNIX_EPOCH}, +}; +use tokio::{net::TcpListener, task::JoinHandle}; + +pub struct TestIssuer { + pub issuer: String, + pub jwks_uri: String, + key: EncodingKey, + task: JoinHandle<()>, +} +impl TestIssuer { + pub async fn new() -> Self { + let pair = rcgen::KeyPair::generate_for(&rcgen::PKCS_ECDSA_P256_SHA256).unwrap(); + let public = pair.public_key_raw(); + let jwks = json!({"keys":[{"kty":"EC","crv":"P-256","alg":"ES256","use":"sig","kid":"test-key","x":URL_SAFE_NO_PAD.encode(&public[1..33]),"y":URL_SAFE_NO_PAD.encode(&public[33..65])}]}); + let key = EncodingKey::from_ec_pem(pair.serialize_pem().as_bytes()).unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let issuer = format!("http://{}", listener.local_addr().unwrap()); + let jwks_uri = format!("{issuer}/jwks"); + let task = tokio::spawn(async move { + axum::serve( + listener, + Router::new().route("/jwks", get(move || async move { Json(jwks) })), + ) + .await + .unwrap(); + }); + Self { + issuer, + jwks_uri, + key, + task, + } + } + pub fn authentication(&self) -> AuthenticationConfigInput { + serde_json::from_value(json!({"issuers":[{ + "issuer":self.issuer,"jwks_uri":self.jwks_uri,"audience":"test-mcp", + "required_scopes":["sql:query"], + "users":{"test-user":"example-user"},"clients":{"test-client":"example-client"} + }]})) + .unwrap() + } + pub fn claims(&self) -> Value { + let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(); + json!({"iss":self.issuer,"aud":"test-mcp","sub":"test-user","client_id":"test-client","scope":"sql:query","iat":now,"exp":now+300}) + } + pub fn token(&self) -> String { + self.token_with_claims(self.claims()) + } + pub fn token_with_claims(&self, claims: Value) -> String { + self.token_with_key_id(claims, "test-key") + } + pub fn token_with_key_id(&self, claims: Value, key_id: &str) -> String { + let mut header = Header::new(Algorithm::ES256); + header.kid = Some(key_id.into()); + jsonwebtoken::encode(&header, &claims, &self.key).unwrap() + } + pub async fn stop(&self) { + self.task.abort(); + while !self.task.is_finished() { + tokio::task::yield_now().await; + } + } +} +impl Drop for TestIssuer { + fn drop(&mut self) { + self.task.abort(); + } +} + +pub struct HttpFixture { + pub issuer: TestIssuer, + pub url: String, + pub token: String, + task: JoinHandle<()>, +} +impl HttpFixture { + pub async fn new(service: QueryService) -> Self { + Self::with_origins(service, &[]).await + } + pub async fn with_origins(service: QueryService, origins: &[&str]) -> Self { + let issuer = TestIssuer::new().await; + let authenticator = Arc::new(Authenticator::new(issuer.authentication()).await.unwrap()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let bind = listener.local_addr().unwrap(); + let url = format!("http://{bind}/mcp"); + let config = HttpTransportConfig::new(HttpTransportConfigInput { + bind, + resource_url: url.clone(), + allowed_origins: origins.iter().map(|origin| (*origin).into()).collect(), + }) + .unwrap(); + let router = http_router(service, authenticator, config); + let task = tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + let token = issuer.token(); + Self { + issuer, + url, + token, + task, + } + } + pub async fn client(&self) -> RunningService { + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(self.url.clone()).auth_header(self.token.clone()), + ); + ().serve(transport).await.unwrap() + } +} +impl Drop for HttpFixture { + fn drop(&mut self) { + self.task.abort(); + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs index ac8d973dc..9e7a5fa10 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/config.rs @@ -9,7 +9,7 @@ fn rejects_duplicate_connection_names_and_ids() -> Result<(), Box Result<(), Box(valid).is_ok()); + for field in ["user", "client"] { + assert!(toml::from_str::(&format!("{field}=\"legacy-identity\"\n{valid}")).is_err()); + } assert!(toml::from_str::(&format!("permissions_disabled=true\n{valid}")).is_err()); } @@ -93,15 +96,9 @@ fn secrets_and_remote_transport_settings_are_validated_at_the_boundary() -> Resu #[test] fn authorization_ids_reject_dots_at_configuration_boundary() -> Result<(), Box> { assert!(ObjectId::new("invalid.id".into()).is_err()); - for field in ["user", "client", "connection"] { - let mut input: ConfigInput = toml::from_str(include_str!("../examples/config.toml"))?; - match field { - "user" => input.user = "invalid.id".into(), - "client" => input.client = "invalid.id".into(), - _ => input.connections[0].id = "invalid.id".into(), - } - assert!(RuntimeConfig::from_input(input, config_environment).is_err(), "{field}"); - } + let mut input: ConfigInput = toml::from_str(include_str!("../examples/config.toml"))?; + input.connections[0].id = "invalid.id".into(); + assert!(RuntimeConfig::from_input(input, config_environment).is_err()); Ok(()) } diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/http_config.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/http_config.rs new file mode 100644 index 000000000..5f2d1a976 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/http_config.rs @@ -0,0 +1,58 @@ +#![allow(clippy::unwrap_used)] +use mcp_postgres::{HttpTransportConfig, HttpTransportConfigInput}; +fn input(resource_url: &str, origins: &[&str]) -> HttpTransportConfigInput { + HttpTransportConfigInput { + bind: "127.0.0.1:8080".parse().unwrap(), + resource_url: resource_url.into(), + allowed_origins: origins.iter().map(|origin| (*origin).into()).collect(), + } +} +#[test] +fn public_resources_require_canonical_https_and_local_resources_allow_http() { + for resource in [ + "https://sql.example.test/mcp", + "http://127.0.0.1:8080/mcp", + "http://localhost:8080/mcp", + "http://[::1]:8080/mcp", + ] { + let config = HttpTransportConfig::new(input(resource, &[])).unwrap(); + assert_eq!(config.resource_url(), resource); + assert_eq!(config.bind(), "127.0.0.1:8080".parse().unwrap()); + } + for resource in [ + "http://sql.example.test/mcp", + "https://sql.example.test/other", + "https://sql.example.test/mcp/", + "https://user@sql.example.test/mcp", + "https://user:password@sql.example.test/mcp", + "https://sql.example.test/mcp?query", + "https://sql.example.test/mcp#fragment", + "https://SQL.example.test/mcp", + "invalid", + "file:///mcp", + ] { + assert!(HttpTransportConfig::new(input(resource, &[])).is_err()); + } +} +#[test] +fn browser_origins_are_explicit_exact_origins() { + for origin in [ + "https://app.example.test", + "https://app.example.test:8443", + "http://localhost:8080", + ] { + assert!(HttpTransportConfig::new(input("https://sql.example.test/mcp", &[origin])).is_ok()); + } + for origin in [ + "*", + "null", + "https://app.example.test/", + "https://app.example.test/path", + "http://app.example.test", + "https://app.example.test?x", + "https://app.example.test#x", + "https://user@app.example.test", + ] { + assert!(HttpTransportConfig::new(input("https://sql.example.test/mcp", &[origin])).is_err()); + } +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/http_identity.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/http_identity.rs new file mode 100644 index 000000000..26abb9451 --- /dev/null +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/http_identity.rs @@ -0,0 +1,217 @@ +#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +mod common; + +use async_trait::async_trait; +use common::TestIssuer; +use mcp_postgres::{ + Authenticator, ConfigInput, Error, HttpTransportConfig, HttpTransportConfigInput, PermissionCheck, + PermissionRequest, QueryService, RuntimeConfig, http_router, +}; +use reqwest::StatusCode; +use serde_json::{Value, json}; +use std::{ + collections::BTreeSet, + sync::{Arc, Mutex}, +}; +use testcontainers_modules::{ + postgres::Postgres, + testcontainers::{ImageExt, runners::AsyncRunner}, +}; + +#[derive(Default)] +struct Permissions { + checked: Mutex>, +} + +#[async_trait] +impl PermissionCheck for Permissions { + async fn check_all(&self, requests: &[PermissionRequest]) -> Result<(), Error> { + for request in requests { + self.checked + .lock() + .unwrap() + .insert((request.subject_kind.into(), request.subject.as_str().into())); + if !matches!( + (request.subject_kind, request.subject.as_str()), + ("sql_user", "allowed-user") | ("sql_client", "allowed-client") + ) { + return Err(Error::Denied); + } + } + Ok(()) + } +} + +struct Server(tokio::task::JoinHandle<()>); +impl Drop for Server { + fn drop(&mut self) { + self.0.abort(); + } +} + +async fn call( + client: &reqwest::Client, + url: &str, + token: &str, + arguments: Value, + metadata: Value, +) -> reqwest::Response { + client + .post(url) + .bearer_auth(token) + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-11-25") + .json(&json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{ + "name":"query","arguments":arguments,"_meta":metadata + }})) + .send() + .await + .unwrap() +} + +async fn success(response: reqwest::Response) { + assert_eq!(response.status(), StatusCode::OK); + let body: Value = response.json().await.unwrap(); + assert_ne!(body["result"]["isError"], true); + let result: Value = serde_json::from_str(body["result"]["content"][0]["text"].as_str().unwrap()).unwrap(); + assert_eq!(result["rows"], json!([[1]])); +} + +async fn denied(response: reqwest::Response) { + assert_eq!(response.status(), StatusCode::OK); + let body: Value = response.json().await.unwrap(); + assert_eq!(body["result"]["isError"], true); + assert_eq!(body["result"]["content"][0]["text"], "access denied"); +} + +#[tokio::test] +async fn http_requests_keep_issuer_namespaces_and_current_actor_permissions_separate() { + let postgres = Postgres::default().with_tag("18.3-alpine").start().await.unwrap(); + let host = postgres.get_host().await.unwrap(); + let port = postgres.get_host_port_ipv4(5432).await.unwrap(); + let admin = sqlx::PgPool::connect(&format!("postgres://postgres:postgres@{host}:{port}/postgres")) + .await + .unwrap(); + sqlx::raw_sql("CREATE ROLE query_role LOGIN PASSWORD 'query-only'; CREATE TABLE public.customers (id integer); INSERT INTO public.customers VALUES (1); GRANT USAGE ON SCHEMA public TO query_role; GRANT SELECT ON public.customers TO query_role;") + .execute(&admin).await.unwrap(); + let dsn = format!("postgres://query_role:query-only@{host}:{port}/postgres"); + let config = RuntimeConfig::from_input( + toml::from_str::(include_str!("../examples/config.toml")).unwrap(), + |key| { + Some(if key.ends_with("DSN") { + dsn.clone() + } else { + "test-token".into() + }) + }, + ) + .unwrap(); + let permissions = Arc::new(Permissions::default()); + let service = QueryService::connect(config, permissions.clone()).await.unwrap(); + let first = TestIssuer::new().await; + let second = TestIssuer::new().await; + let mut authentication = first.authentication(); + authentication.issuers[0] + .users + .insert("test-user".into(), "allowed-user".into()); + authentication.issuers[0].clients = [ + ("test-client", "allowed-client"), + ("agent", "allowed-client"), + ("denied-agent", "denied-client"), + ("history", "history-client"), + ] + .into_iter() + .map(|(external, local)| (external.into(), local.into())) + .collect(); + let mut other = second.authentication(); + other.issuers[0].users.insert("test-user".into(), "other-user".into()); + other.issuers[0] + .clients + .insert("test-client".into(), "other-client".into()); + authentication.issuers.extend(other.issuers); + let authenticator = Arc::new(Authenticator::new(authentication).await.unwrap()); + let other_identity = authenticator.authenticate(&second.token()).await.unwrap(); + assert_eq!(other_identity.user().as_str(), "other-user"); + assert_eq!(other_identity.client().as_str(), "other-client"); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let bind = listener.local_addr().unwrap(); + let url = format!("http://{bind}/mcp"); + let http = HttpTransportConfig::new(HttpTransportConfigInput { + bind, + resource_url: url.clone(), + allowed_origins: vec![], + }) + .unwrap(); + let router = http_router(service, authenticator, http); + let _server = Server(tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + })); + let client = reqwest::Client::new(); + let arguments = json!({"connection":"analytics-readonly","sql":"SELECT id FROM public.customers"}); + let allowed_token = first.token(); + let denied_token = second.token(); + let (allowed_response, denied_response) = tokio::join!( + call(&client, &url, &allowed_token, arguments.clone(), json!({})), + call( + &client, + &url, + &denied_token, + arguments.clone(), + json!({"user":"allowed-user","client":"allowed-client"}) + ), + ); + success(allowed_response).await; + denied(denied_response).await; + let mut claims = first.claims(); + claims["act"] = json!({"sub":"agent","iss":first.issuer,"act":{"sub":"history"}}); + success( + call( + &client, + &url, + &first.token_with_claims(claims), + arguments.clone(), + json!({}), + ) + .await, + ) + .await; + let mut claims = first.claims(); + claims["act"] = json!({"sub":"denied-agent","act":{"sub":"agent"}}); + denied( + call( + &client, + &url, + &first.token_with_claims(claims), + arguments.clone(), + json!({}), + ) + .await, + ) + .await; + for actor in [json!({"sub":"unknown"}), json!({"sub":"agent","iss":second.issuer})] { + let mut claims = first.claims(); + claims["act"] = actor; + let response = call( + &client, + &url, + &first.token_with_claims(claims), + arguments.clone(), + json!({}), + ) + .await; + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + } + let mut spoofed_arguments = arguments.clone(); + spoofed_arguments["user"] = json!("allowed-user"); + spoofed_arguments["client"] = json!("allowed-client"); + let rejected = call(&client, &url, &denied_token, spoofed_arguments, json!({})).await; + assert_eq!(rejected.status(), StatusCode::OK); + let rejected: Value = rejected.json().await.unwrap(); + assert_eq!(rejected["error"]["code"], -32602); + assert_eq!(rejected["error"]["message"], "invalid query arguments"); + success(call(&client, &url, &allowed_token, arguments, json!({})).await).await; + let checked = permissions.checked.lock().unwrap(); + assert!(checked.contains(&("sql_user".into(), "other-user".into()))); + assert!(checked.contains(&("sql_client".into(), "denied-client".into()))); + assert!(!checked.iter().any(|(_, subject)| subject == "history-client")); +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs index 235fe343e..f30cec2ac 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres.rs @@ -1,4 +1,5 @@ #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +mod common; use std::collections::BTreeSet; use std::io; @@ -12,11 +13,9 @@ use authzed::v1::{ schema_service_client::SchemaServiceClient, }; use mcp_postgres::{ - ConfigInput, Error, Parameter, ParameterInput, PermissionCheck, PermissionRequest, PostgresMcp, QueryService, - RuntimeConfig, + ConfigInput, Error, Parameter, ParameterInput, PermissionCheck, PermissionRequest, QueryService, RuntimeConfig, }; use rmcp::{ - ServiceExt, model::{CallToolRequest, CallToolRequestParams, ClientRequest}, service::PeerRequestOptions, }; @@ -122,6 +121,7 @@ async fn ordinary_readonly_database_requires_no_installed_objects() { let service = fixture.service(Arc::new(Permissions::default())).await; let result = service .query( + &identity(), "analytics-readonly", "SELECT id FROM public.customers ORDER BY id LIMIT 1", ) @@ -153,6 +153,7 @@ async fn actual_postgres_returns_authorized_results_and_preserves_limits() { let service = fixture.service(permissions.clone()).await; let result = service .query( + &identity(), "analytics-readonly", "SELECT c.id, c.active FROM public.customers c ORDER BY c.id LIMIT 1", ) @@ -161,29 +162,39 @@ async fn actual_postgres_returns_authorized_results_and_preserves_limits() { assert_eq!(result.columns, ["id", "active"]); assert_eq!(result.rows, vec![vec![json!(1), json!(true)]]); let empty = service - .query("analytics-readonly", "SELECT id FROM public.customers LIMIT 0") + .query( + &identity(), + "analytics-readonly", + "SELECT id FROM public.customers LIMIT 0", + ) .await .unwrap(); assert_eq!(empty.columns, ["id"]); assert!(empty.rows.is_empty()); let absent = service - .query("analytics-readonly", "SELECT id FROM public.customers WHERE id = 999") + .query( + &identity(), + "analytics-readonly", + "SELECT id FROM public.customers WHERE id = 999", + ) .await .unwrap(); assert_eq!(absent.columns, ["id"]); assert!(absent.rows.is_empty()); let joined = service .query( + &identity(), "analytics-readonly", "SELECT c.id, o.id AS order_id FROM public.customers c INNER JOIN public.orders o ON c.id = o.customer_id", ) .await .unwrap(); assert_eq!(joined.rows, vec![vec![json!(1), json!(9)]]); - let self_join = service.query("analytics-readonly", "SELECT c.id, other.id FROM public.customers c INNER JOIN public.customers other ON c.id = other.id ORDER BY c.id").await.unwrap(); + let self_join = service.query(&identity(), "analytics-readonly", "SELECT c.id, other.id FROM public.customers c INNER JOIN public.customers other ON c.id = other.id ORDER BY c.id").await.unwrap(); assert_eq!(self_join.rows, vec![vec![json!(1), json!(1)], vec![json!(2), json!(2)]]); let unused_join = service .query( + &identity(), "analytics-readonly", "SELECT c.id FROM public.customers c INNER JOIN public.orders o ON true ORDER BY c.id", ) @@ -218,13 +229,16 @@ async fn forbidden_columns_block_projections_predicates_joins_and_ordering_for_e "SELECT id FROM public.customers ORDER BY email", ] { assert!( - matches!(service.query("analytics-readonly", sql).await, Err(Error::Denied)), + matches!( + service.query(&identity(), "analytics-readonly", sql).await, + Err(Error::Denied) + ), "{sql}" ); } assert!( service - .query("analytics-readonly", "SELECT id FROM public.customers") + .query(&identity(), "analytics-readonly", "SELECT id FROM public.customers") .await .is_ok() ); @@ -257,7 +271,11 @@ async fn catalog_features_and_unsupported_sql_cannot_expose_hidden_dependencies( ] { assert!( service - .query("analytics-readonly", &format!("SELECT id FROM public.{table}")) + .query( + &identity(), + "analytics-readonly", + &format!("SELECT id FROM public.{table}") + ) .await .is_err(), "{table}" @@ -278,7 +296,10 @@ async fn catalog_features_and_unsupported_sql_cannot_expose_hidden_dependencies( "SELECT id FROM public.customers; SELECT email FROM public.customers", "SELECT id FROM public.customers WHERE id IN (SELECT id FROM public.orders)", ] { - assert!(service.query("analytics-readonly", sql).await.is_err(), "{sql}"); + assert!( + service.query(&identity(), "analytics-readonly", sql).await.is_err(), + "{sql}" + ); } } @@ -312,7 +333,7 @@ async fn namespace_replacement_during_authorization_cannot_redirect_the_bound_qu let service = service.clone(); async move { service - .query("analytics-readonly", "SELECT id FROM public.customers") + .query(&identity(), "analytics-readonly", "SELECT id FROM public.customers") .await } }); @@ -376,7 +397,11 @@ async fn namespace_replacement_before_binding_is_rejected_before_column_authoriz let service = service.clone(); async move { service - .query("analytics-readonly", "SELECT id FROM public.customers LIMIT 1") + .query( + &identity(), + "analytics-readonly", + "SELECT id FROM public.customers LIMIT 1", + ) .await } }); @@ -467,7 +492,7 @@ async fn column_ddl_waits_for_the_bound_query_without_a_database_guard() { let service = service.clone(); async move { service - .query("analytics-readonly", "SELECT id FROM public.customers") + .query(&identity(), "analytics-readonly", "SELECT id FROM public.customers") .await } }); @@ -509,19 +534,27 @@ async fn resource_limits_abort_without_returning_partial_results() { let service = fixture.service(Arc::new(Permissions::default())).await; assert!(matches!( service - .query("analytics-readonly", "SELECT id FROM public.many_rows") + .query(&identity(), "analytics-readonly", "SELECT id FROM public.many_rows") .await, Err(Error::Limit) )); assert!(matches!( service - .query("analytics-readonly", "SELECT content FROM public.large_rows") + .query( + &identity(), + "analytics-readonly", + "SELECT content FROM public.large_rows" + ) .await, Err(Error::Limit) )); assert_eq!( service - .query("analytics-readonly", "SELECT id FROM public.customers LIMIT 1") + .query( + &identity(), + "analytics-readonly", + "SELECT id FROM public.customers LIMIT 1" + ) .await .unwrap() .rows @@ -537,14 +570,19 @@ async fn admission_rejects_privileged_roles_and_excess_parameters() { let parameters = vec![Parameter::try_from(ParameterInput::Boolean(true)).unwrap(); 33]; assert!(matches!( service - .query_with_parameters("analytics-readonly", "SELECT id FROM public.customers", ¶meters) + .query_with_parameters( + &identity(), + "analytics-readonly", + "SELECT id FROM public.customers", + ¶meters + ) .await, Err(Error::Limit) )); fixture.admin.execute("ALTER ROLE query_role SUPERUSER").await.unwrap(); assert!(matches!( service - .query("analytics-readonly", "SELECT id FROM public.customers") + .query(&identity(), "analytics-readonly", "SELECT id FROM public.customers") .await, Err(Error::UnsupportedObject) )); @@ -556,6 +594,7 @@ async fn boolean_predicates_alias_rejection_and_aggregate_byte_budget() { let service = fixture.service(Arc::new(Permissions::default())).await; let result = service .query( + &identity(), "analytics-readonly", "SELECT id FROM public.customers WHERE (id = 1 OR id = 2) AND active = true ORDER BY id", ) @@ -565,6 +604,7 @@ async fn boolean_predicates_alias_rejection_and_aggregate_byte_budget() { assert!(matches!( service .query( + &identity(), "analytics-readonly", "SELECT id AS renamed FROM public.customers ORDER BY renamed" ) @@ -574,7 +614,11 @@ async fn boolean_predicates_alias_rejection_and_aggregate_byte_budget() { sqlx::raw_sql("CREATE TABLE public.aggregate_large (content text); INSERT INTO public.aggregate_large SELECT repeat('x',100000) FROM generate_series(1,11); GRANT SELECT ON public.aggregate_large TO query_role;").execute(&fixture.admin).await.unwrap(); assert!(matches!( service - .query("analytics-readonly", "SELECT content FROM public.aggregate_large") + .query( + &identity(), + "analytics-readonly", + "SELECT content FROM public.aggregate_large" + ) .await, Err(Error::Limit) )); @@ -598,6 +642,7 @@ async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_width let parameters = [Parameter::try_from(ParameterInput::Integer(7)).unwrap()]; let result = service .query_with_parameters( + &identity(), "analytics-readonly", &format!("SELECT {column} FROM public.widths WHERE {column} = $1"), ¶meters, @@ -609,6 +654,7 @@ async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_width let parameters = [Parameter::try_from(ParameterInput::Text("' OR true --".into())).unwrap()]; let result = service .query_with_parameters( + &identity(), "analytics-readonly", "SELECT id FROM public.customers WHERE email = $1", ¶meters, @@ -619,6 +665,7 @@ async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_width let parameters = [Parameter::try_from(ParameterInput::Boolean(true)).unwrap()]; let result = service .query_with_parameters( + &identity(), "analytics-readonly", "SELECT id FROM public.customers WHERE active = $1", ¶meters, @@ -628,7 +675,11 @@ async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_width assert_eq!(result.rows, vec![vec![json!(1)]]); assert!( service - .query("analytics-readonly", "SELECT id FROM public.customers WHERE id = $1") + .query( + &identity(), + "analytics-readonly", + "SELECT id FROM public.customers WHERE id = $1" + ) .await .is_err() ); @@ -638,17 +689,8 @@ async fn bound_parameters_keep_sql_looking_text_as_data_and_handle_integer_width async fn mcp_client_initializes_discovers_queries_and_cannot_spoof_identity() { let fixture = Fixture::new().await; let service = fixture.service(Arc::new(Permissions::default())).await; - let (client_io, server_io) = tokio::io::duplex(65536); - let server_task = tokio::spawn(async move { - PostgresMcp::new(service) - .serve(server_io) - .await - .unwrap() - .waiting() - .await - .unwrap(); - }); - let client = ().serve(client_io).await.unwrap(); + let http = common::HttpFixture::new(service).await; + let client = http.client().await; let tools = client.list_tools(None).await.unwrap(); assert_eq!(tools.tools.len(), 1); assert_eq!(tools.tools[0].name, "query"); @@ -680,10 +722,6 @@ async fn mcp_client_initializes_discovers_queries_and_cannot_spoof_identity() { assert_eq!(invalid.is_error, Some(true)); assert!(!serde_json::to_string(&invalid).unwrap().contains("secret-marker-839")); client.cancel().await.unwrap(); - tokio::time::timeout(Duration::from_secs(3), server_task) - .await - .unwrap() - .unwrap(); } #[tokio::test] @@ -694,17 +732,8 @@ async fn cancelling_an_mcp_request_releases_its_relation_locks() { release: Notify::new(), }); let service = fixture.service(gate.clone()).await; - let (client_io, server_io) = tokio::io::duplex(65536); - let server = tokio::spawn(async move { - PostgresMcp::new(service) - .serve(server_io) - .await - .unwrap() - .waiting() - .await - .unwrap(); - }); - let client = ().serve(client_io).await.unwrap(); + let http = common::HttpFixture::new(service).await; + let client = http.client().await; let params = CallToolRequestParams::new("query").with_arguments( json!({"connection":"analytics-readonly","sql":"SELECT id FROM public.customers"}) .as_object() @@ -749,10 +778,6 @@ async fn cancelling_an_mcp_request_releases_its_relation_locks() { .await .unwrap(); client.cancel().await.unwrap(); - tokio::time::timeout(Duration::from_secs(3), server) - .await - .unwrap() - .unwrap(); } fn spice_request(value: T) -> tonic::Request { @@ -841,17 +866,8 @@ async fn mcp_spicedb_and_postgres_enforce_hard_column_denial_together() { .unwrap(), ); let service = fixture.service(adapter).await; - let (client_io, server_io) = tokio::io::duplex(65536); - let server = tokio::spawn(async move { - PostgresMcp::new(service) - .serve(server_io) - .await - .unwrap() - .waiting() - .await - .unwrap(); - }); - let client = ().serve(client_io).await.unwrap(); + let http = common::HttpFixture::new(service).await; + let client = http.client().await; let query = |sql: &str| { CallToolRequestParams::new("query").with_arguments( json!({"connection":"analytics-readonly","sql":sql}) @@ -891,8 +907,100 @@ async fn mcp_spicedb_and_postgres_enforce_hard_column_denial_together() { .unwrap(); assert_ne!(still_allowed.is_error, Some(true)); client.cancel().await.unwrap(); - tokio::time::timeout(Duration::from_secs(3), server) - .await - .unwrap() - .unwrap(); +} + +fn identity() -> mcp_postgres::RequestIdentity { + mcp_postgres::RequestIdentity::new( + mcp_postgres::ObjectId::new("example-user".into()).unwrap(), + mcp_postgres::ObjectId::new("example-client".into()).unwrap(), + Vec::new(), + ) +} + +struct IsolatedPermissions { + concurrent_admissions: tokio::sync::Barrier, + admitted: Mutex>, +} + +#[async_trait] +impl PermissionCheck for IsolatedPermissions { + async fn check_all(&self, requests: &[PermissionRequest]) -> Result<(), Error> { + for request in requests { + let allowed = match request.subject_kind { + "sql_user" => matches!(request.subject.as_str(), "alice" | "bob"), + "sql_client" => matches!(request.subject.as_str(), "desktop-a" | "desktop-b"), + _ => false, + }; + if !allowed { + return Err(Error::Denied); + } + } + if requests.len() > 1 { + let subjects = |kind| { + requests + .iter() + .filter(|request| request.subject_kind == kind) + .map(|request| request.subject.as_str()) + .collect::>() + }; + let users = subjects("sql_user"); + let clients = subjects("sql_client"); + assert_eq!(users.len(), 1); + assert_eq!(clients.len(), 1); + let pair = ( + (*users.first().unwrap()).to_owned(), + (*clients.first().unwrap()).to_owned(), + ); + if !matches!( + (pair.0.as_str(), pair.1.as_str()), + ("alice", "desktop-a") | ("bob", "desktop-b") + ) { + return Err(Error::Denied); + } + self.concurrent_admissions.wait().await; + self.admitted.lock().unwrap().insert(pair); + } + Ok(()) + } +} + +#[tokio::test] +async fn shared_service_isolates_concurrent_request_identities_and_denies_ungranted_subjects() { + let fixture = Fixture::new().await; + let permissions = Arc::new(IsolatedPermissions { + concurrent_admissions: tokio::sync::Barrier::new(2), + admitted: Mutex::new(BTreeSet::new()), + }); + let service = fixture.service(permissions.clone()).await; + let principal = |user: &str, client: &str| { + mcp_postgres::RequestIdentity::new( + mcp_postgres::ObjectId::new(user.into()).unwrap(), + mcp_postgres::ObjectId::new(client.into()).unwrap(), + vec![mcp_postgres::ObjectId::new("prior-actor".into()).unwrap()], + ) + }; + let alice = principal("alice", "desktop-a"); + let bob = principal("bob", "desktop-b"); + assert_eq!(alice.delegation()[0].as_str(), "prior-actor"); + let sql = "SELECT id FROM public.customers ORDER BY id LIMIT 1"; + let (alice_result, bob_result) = tokio::join!( + service.query(&alice, "analytics-readonly", sql), + service.query(&bob, "analytics-readonly", sql), + ); + assert_eq!(alice_result.unwrap().rows, vec![vec![json!(1)]]); + assert_eq!(bob_result.unwrap().rows, vec![vec![json!(1)]]); + assert_eq!( + *permissions.admitted.lock().unwrap(), + BTreeSet::from([("alice".into(), "desktop-a".into()), ("bob".into(), "desktop-b".into())]) + ); + for denied in [ + principal("eve", "desktop-a"), + principal("alice", "ungranted-client"), + principal("alice", "desktop-b"), + ] { + assert!(matches!( + service.query(&denied, "analytics-readonly", sql).await, + Err(Error::Denied) + )); + } } diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres_tls.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres_tls.rs index 547c94435..01de31482 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/postgres_tls.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/postgres_tls.rs @@ -145,7 +145,7 @@ async fn trusted_ca_and_matching_hostname_protect_the_query_connection() { .unwrap(); assert!(encrypted); let result = service - .query("analytics-readonly", "SELECT id FROM public.customers") + .query(&identity(), "analytics-readonly", "SELECT id FROM public.customers") .await .unwrap(); assert_eq!(result.rows, vec![vec![json!(1)]]); @@ -181,3 +181,11 @@ async fn verify_full_refuses_a_server_without_tls() { Err(Error::Database(sqlx::Error::Tls(_))) )); } + +fn identity() -> mcp_postgres::RequestIdentity { + mcp_postgres::RequestIdentity::new( + mcp_postgres::ObjectId::new("example-user".into()).unwrap(), + mcp_postgres::ObjectId::new("example-client".into()).unwrap(), + Vec::new(), + ) +} diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs index c2b5af397..6ae10302a 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/server.rs @@ -1,7 +1,8 @@ #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] +mod common; use async_trait::async_trait; -use mcp_postgres::{ConfigInput, Error, PermissionCheck, PermissionRequest, PostgresMcp, QueryService, RuntimeConfig}; -use rmcp::{ServiceExt, model::CallToolRequestParams}; +use mcp_postgres::{ConfigInput, Error, PermissionCheck, PermissionRequest, QueryService, RuntimeConfig}; +use rmcp::model::CallToolRequestParams; use serde_json::json; use std::sync::{Arc, Mutex}; use testcontainers_modules::{ @@ -76,17 +77,134 @@ async fn mcp_errors_are_safe_and_audits_exclude_sql_parameters_and_rows() { .with_filter(tracing_subscriber::filter::Targets::new().with_target("mcp_postgres", tracing::Level::INFO)), ); let _guard = tracing::subscriber::set_default(subscriber); - let (client_io, server_io) = tokio::io::duplex(65536); - let server = tokio::spawn(async move { - PostgresMcp::new(service) - .serve(server_io) - .await + let http = common::HttpFixture::with_origins(service, &["https://portal.example.test"]).await; + let http_client = reqwest::Client::new(); + let preflight = http_client + .request(reqwest::Method::OPTIONS, &http.url) + .header("Origin", "https://portal.example.test") + .header("Access-Control-Request-Method", "POST") + .header( + "Access-Control-Request-Headers", + "authorization,content-type,mcp-protocol-version", + ) + .send() + .await + .unwrap(); + assert_eq!(preflight.status(), reqwest::StatusCode::OK); + assert_eq!( + preflight.headers()["access-control-allow-origin"], + "https://portal.example.test" + ); + assert!( + preflight.headers()["access-control-allow-headers"] + .to_str() + .unwrap() + .contains("authorization") + ); + let allowed_origin = http_client + .post(&http.url) + .header("Origin", "https://portal.example.test") + .send() + .await + .unwrap(); + assert_eq!(allowed_origin.status(), reqwest::StatusCode::UNAUTHORIZED); + assert_eq!( + allowed_origin.headers()["access-control-allow-origin"], + "https://portal.example.test" + ); + assert!( + allowed_origin.headers()["access-control-expose-headers"] + .to_str() .unwrap() - .waiting() + .contains("www-authenticate") + ); + let missing = http_client.post(&http.url).send().await.unwrap(); + assert_eq!(missing.status(), reqwest::StatusCode::UNAUTHORIZED); + assert!( + missing.headers()["www-authenticate"] + .to_str() + .unwrap() + .contains("/.well-known/oauth-protected-resource/mcp") + ); + for authorization in ["Basic example", "Bearer", "Bearer bad token", "Bearer invalid"] { + let denied = http_client + .post(&http.url) + .header("Authorization", authorization) + .send() .await .unwrap(); - }); - let client = ().serve(client_io).await.unwrap(); + assert_eq!(denied.status(), reqwest::StatusCode::UNAUTHORIZED); + } + let duplicate = http_client + .post(&http.url) + .header("Authorization", format!("Bearer {}", http.issuer.token())) + .header("Authorization", format!("Bearer {}", http.issuer.token())) + .send() + .await + .unwrap(); + assert_eq!(duplicate.status(), reqwest::StatusCode::UNAUTHORIZED); + let mut claims = http.issuer.claims(); + claims["scope"] = json!("unrelated"); + let insufficient = http_client + .post(&http.url) + .bearer_auth(http.issuer.token_with_claims(claims)) + .send() + .await + .unwrap(); + assert_eq!(insufficient.status(), reqwest::StatusCode::FORBIDDEN); + assert!( + insufficient.headers()["www-authenticate"] + .to_str() + .unwrap() + .contains("error=\"insufficient_scope\"") + ); + let blocked_origin = http_client + .post(&http.url) + .bearer_auth(http.issuer.token()) + .header("Origin", "https://untrusted.example.test") + .send() + .await + .unwrap(); + assert_eq!(blocked_origin.status(), reqwest::StatusCode::FORBIDDEN); + let blocked_port = http_client + .post(&http.url) + .bearer_auth(http.issuer.token()) + .header("Origin", "https://portal.example.test:8443") + .send() + .await + .unwrap(); + assert_eq!(blocked_port.status(), reqwest::StatusCode::FORBIDDEN); + let blocked_host = http_client + .post(&http.url) + .bearer_auth(http.issuer.token()) + .header("Host", "untrusted.example.test") + .send() + .await + .unwrap(); + assert_eq!(blocked_host.status(), reqwest::StatusCode::FORBIDDEN); + let oversized = http_client + .post(&http.url) + .bearer_auth(http.issuer.token()) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .body("x".repeat(131_073)) + .send() + .await + .unwrap(); + assert_eq!(oversized.status(), reqwest::StatusCode::PAYLOAD_TOO_LARGE); + let metadata_url = http.url.replace("/mcp", "/.well-known/oauth-protected-resource/mcp"); + let metadata: serde_json::Value = http_client + .get(metadata_url) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(metadata["resource"], http.url); + assert_eq!(metadata["authorization_servers"], json!([http.issuer.issuer])); + assert_eq!(metadata["bearer_methods_supported"], json!(["header"])); + let client = http.client().await; assert!(client.call_tool(CallToolRequestParams::new("missing")).await.is_err()); for (connection, failure, expected) in [ ("missing", None, "unknown connection"), @@ -115,9 +233,28 @@ async fn mcp_errors_are_safe_and_audits_exclude_sql_parameters_and_rows() { assert_eq!(result["columns"], json!(["email"])); assert_eq!(result["rows"], json!([["private-row@example.test"]])); client.cancel().await.unwrap(); - server.await.unwrap(); + let bearer_token = http.token.clone(); + http.issuer.stop().await; + let unavailable_token = http.issuer.token_with_key_id(http.issuer.claims(), "rotated-key"); + let unavailable = http_client + .post(&http.url) + .bearer_auth(unavailable_token) + .send() + .await + .unwrap(); + assert_eq!(unavailable.status(), reqwest::StatusCode::SERVICE_UNAVAILABLE); + assert!(unavailable.headers().get("www-authenticate").is_none()); let output = String::from_utf8(logs.0.lock().unwrap().clone()).unwrap(); - assert_audit_output(&output, &["SELECT", "private-row@example.test", "private-token", &dsn]); + assert_audit_output( + &output, + &[ + "SELECT", + "private-row@example.test", + "private-token", + &dsn, + &bearer_token, + ], + ); } fn assert_audit_output(output: &str, forbidden: &[&str]) { diff --git a/rsworkspace/crates/mcp/mcp-postgres/tests/startup.rs b/rsworkspace/crates/mcp/mcp-postgres/tests/startup.rs index ae3560bef..662ccc6d5 100644 --- a/rsworkspace/crates/mcp/mcp-postgres/tests/startup.rs +++ b/rsworkspace/crates/mcp/mcp-postgres/tests/startup.rs @@ -1,5 +1,5 @@ #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)] -use rmcp::ServiceExt; +mod common; use std::{ process::{Command, Stdio}, time::Duration, @@ -34,7 +34,7 @@ fn invalid_startup_reports_only_safe_failure() { } #[tokio::test] -async fn configured_binary_initializes_stdio_and_exits_when_client_disconnects() { +async fn configured_binary_authenticates_http_and_stops_on_termination() { let postgres = Postgres::default().with_tag("18.3-alpine").start().await.unwrap(); let spicedb = GenericImage::new("authzed/spicedb", "v1.53.0") .with_exposed_port(50051.tcp()) @@ -59,32 +59,108 @@ async fn configured_binary_initializes_stdio_and_exits_when_client_disconnects() postgres.get_host().await.unwrap(), postgres.get_host_port_ipv4(5432).await.unwrap() ); - let config = include_str!("../examples/config.toml").replace("http://127.0.0.1:50051", &endpoint); + let issuer = common::TestIssuer::new().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let bind = listener.local_addr().unwrap(); + drop(listener); + let resource = format!("http://{bind}/mcp"); + let config = include_str!("../examples/config.toml") + .replace("http://127.0.0.1:50051", &endpoint) + .replace("https://identity.example.com/.well-known/jwks.json", &issuer.jwks_uri) + .replace("https://identity.example.com", &issuer.issuer) + .replace("https://sql.example.com/mcp", &resource) + .replace("0.0.0.0:8080", &bind.to_string()) + .replace("external-user-id", "test-user") + .replace("external-client-id", "test-client"); let path = std::env::temp_dir().join(format!("mcp-startup-{}.toml", std::process::id())); std::fs::write(&path, config).unwrap(); - let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_mcp-postgres")) - .arg(&path) - .env("MCP_POSTGRES_SPICEDB_TOKEN", "startup-test-token") - .env("MCP_POSTGRES_ANALYTICS_DSN", &dsn) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .kill_on_drop(true) - .spawn() - .unwrap(); - let client = ().serve((child.stdout.take().unwrap(), child.stdin.take().unwrap())).await.unwrap(); - let tools = client.list_tools(None).await.unwrap(); - assert_eq!(tools.tools.len(), 1); - assert_eq!(tools.tools[0].name, "query"); - client.cancel().await.unwrap(); - let status = tokio::time::timeout(Duration::from_secs(5), child.wait()) + for signal in ["-TERM", "-INT"] { + let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_mcp-postgres")) + .arg(&path) + .env("MCP_POSTGRES_SPICEDB_TOKEN", "startup-test-token") + .env("MCP_POSTGRES_ANALYTICS_DSN", &dsn) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .unwrap(); + let http = reqwest::Client::new(); + let metadata = format!("http://{bind}/.well-known/oauth-protected-resource/mcp"); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + if http + .get(&metadata) + .send() + .await + .is_ok_and(|response| response.status().is_success()) + { + break; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + }) .await - .unwrap() .unwrap(); - assert!(status.success()); - let mut stderr = String::new(); - child.stderr.take().unwrap().read_to_string(&mut stderr).await.unwrap(); - assert!(!stderr.contains("startup-test-token")); - assert!(!stderr.contains(&dsn)); + assert_eq!( + http.post(&resource).send().await.unwrap().status(), + reqwest::StatusCode::UNAUTHORIZED + ); + let mut claims = issuer.claims(); + claims["aud"] = serde_json::json!(resource); + let token = issuer.token_with_claims(claims); + let response = http + .post(&resource) + .bearer_auth(&token) + .header("Accept", "application/json, text/event-stream") + .header("MCP-Protocol-Version", "2025-11-25") + .json(&serde_json::json!({"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}})) + .send() + .await + .unwrap(); + assert!(response.status().is_success()); + let tools: serde_json::Value = response.json().await.unwrap(); + assert_eq!(tools["result"]["tools"][0]["name"], "query"); + assert!( + Command::new("kill") + .args([signal, &child.id().unwrap().to_string()]) + .status() + .unwrap() + .success() + ); + let status = tokio::time::timeout(Duration::from_secs(5), child.wait()) + .await + .unwrap() + .unwrap(); + assert!(status.success()); + let mut stderr = String::new(); + child.stderr.take().unwrap().read_to_string(&mut stderr).await.unwrap(); + assert!(!stderr.contains("startup-test-token")); + assert!(!stderr.contains(&dsn)); + assert!(!stderr.contains(&token)); + let mut stdout = Vec::new(); + child.stdout.take().unwrap().read_to_end(&mut stdout).await.unwrap(); + assert!(stdout.is_empty()); + } std::fs::remove_file(path).unwrap(); } + +#[test] +fn startup_requires_http_and_authentication_configuration() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + for missing in ["http", "authentication"] { + let mut config: toml::Value = toml::from_str(include_str!("../examples/config.toml")).unwrap(); + config.as_table_mut().unwrap().remove(missing); + std::fs::write(&path, toml::to_string(&config).unwrap()).unwrap(); + let output = Command::new(env!("CARGO_BIN_EXE_mcp-postgres")) + .arg(&path) + .output() + .unwrap(); + assert!(!output.status.success()); + assert!(output.stdout.is_empty()); + let stderr = String::from_utf8(output.stderr).unwrap(); + assert!(stderr.contains("startup or transport failed")); + assert!(!stderr.contains("identity.example.com")); + } +}