From 475142e8f0555bb0d5e194e85cfb9bf59e093590 Mon Sep 17 00:00:00 2001 From: Charles Sibbald Date: Mon, 2 Nov 2026 20:59:39 +0200 Subject: [PATCH 1/2] ci, docs: the fork's CI trigger, pre-commit hooks and CONTRIBUTING From the integration branch before this rebase (tag integration-2026-11-02). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y3TSRmEZAYU6RBUXhkn7WF --- .github/workflows/ci.yml | 5 ++++- .gitignore | 3 +++ .pre-commit-config.yaml | 18 ++++++++++++++++++ CONTRIBUTING.md | 33 +++++++++++++++++++++++++++++++++ 4 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 .pre-commit-config.yaml create mode 100644 CONTRIBUTING.md diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index caea076..bbf59e6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,7 +2,7 @@ name: CI on: push: - branches: [master] + branches: [master, integration/microscaler-fork] pull_request: env: @@ -74,6 +74,9 @@ jobs: - uses: Swatinem/rust-cache@v2 - name: Test + env: + # severed_connection proxies to this port (defaults to 5432). + TEST_PG_PORT: "5433" run: cargo test --all - name: PostgreSQL logs diff --git a/.gitignore b/.gitignore index 8e4846f..1dfa938 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,6 @@ Cargo.lock .idea *.iml .vscode + +# NFS silly-rename leftovers +.nfs* diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..6aaf1bc --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,18 @@ +# Guards against pushing what would fail a Rust build or review: install once per checkout with +# pre-commit install +# Run on everything: pre-commit run --all-files +repos: + - repo: local + hooks: + - id: cargo-fmt + name: cargo fmt --check + entry: cargo fmt --all -- --check + language: system + files: \.rs$ + pass_filenames: false + - id: cargo-clippy + name: cargo clippy -D warnings + entry: cargo clippy --all-targets -- -D warnings + language: system + files: (\.rs|Cargo\.toml|Cargo\.lock)$ + pass_filenames: false diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..33727c0 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,33 @@ +# Contributing to this fork + +This repository is a fork with two long-lived branches: + +- **`master` tracks upstream.** It mirrors upstream's default branch; nothing of ours is committed to it. +- **`integration/microscaler-fork` carries every change made in this fork** on top of `master`. It is the + default branch and the only branch that projects depending on the fork should consume (a dependent can + follow just one branch). + +## Making a change + +- Every fix or feature goes to `integration/microscaler-fork` first, through a pull request merged with a + merge commit. +- The pull request for the upstream project is cut separately, fresh from upstream's default branch, as a + single-purpose change (one concern, one commit, following upstream's conventions). It is not merged + from the integration branch. +- Before an upstream pull request merges or closes, check that its content is on the integration branch + (by content, not commit hashes: upstream branches are re-cut). + +## After an upstream update + +1. Update `master` to upstream's default branch. +2. **Before rebasing, tag the current integration head and push the tag**, e.g. + `git tag integration-2026-10-28 origin/integration/microscaler-fork && git push origin integration-2026-10-28`. + Lockfiles pin integration commits; the tag keeps those commits reachable after the rebase, so older + releases still rebuild. +3. Rebase `integration/microscaler-fork` onto `master`, dropping changes upstream now contains, run the + tests, and force-push the integration branch (the only branch that is ever force-pushed). +4. Re-lock every consumer to the new integration head and check each has a single copy of this crate in + its dependency graph. + +Never reset or rewrite `master` while consumers still pin commits that exist only on it; move them to the +integration branch (or tag those commits) first. From 0f9df0e6ad7d8aa1ade256c2c4e85b6132027a59 Mon Sep 17 00:00:00 2001 From: Charles Sibbald Date: Mon, 2 Nov 2026 20:59:39 +0200 Subject: [PATCH 2/2] fork: TLS, connection lifecycle and client additions on upstream master The integration branch rebased onto upstream master 0f93c04, which now contains #2 (notification delivery), #3 (COPY IN deadlock) and #4 (CI). What upstream does not have, as one change; the tree is identical to the old integration head 3a24363, so consumers get exactly the code they ran: - TLS (rustls) with libpq sslmode, root CA and client certificates; explicit plaintext sslmode accepted - graceful shutdown on drop (close the socket), pending responses failed when the I/O loop exits, a 1 MiB stack for the connection coroutine - Client::is_closed, Client::next_notification (wakes LISTEN consumers; never waits less than 1 ms), RowStream affected-row count - Error::user_column_decode_failed - CopyInReceiver drops the unused recv method upstream kept - tests: severed connections, TLS, notifications; examples/leak_probe Per-commit history: tag integration-2026-11-02 (old head 3a24363). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y3TSRmEZAYU6RBUXhkn7WF --- Cargo.toml | 7 + examples/leak_probe.rs | 75 +++++++ src/client.rs | 115 ++++++++--- src/codec.rs | 10 +- src/config.rs | 155 +++++++++++++- src/connect.rs | 3 +- src/connect_raw.rs | 12 +- src/connection.rs | 304 +++++++++++++++++++++------ src/copy_in.rs | 24 --- src/error/mod.rs | 19 +- src/lib.rs | 1 + src/query.rs | 32 ++- src/tls.rs | 377 ++++++++++++++++++++++++++++++++++ tests/severed_connection.rs | 183 +++++++++++++++++ tests/test/main.rs | 194 +++++++++++++++++ tests/test/parse.rs | 27 ++- tests/test/types/chrono_04.rs | 2 +- tests/tls.rs | 271 ++++++++++++++++++++++++ tests/tls/run.sh | 55 +++++ 19 files changed, 1723 insertions(+), 143 deletions(-) create mode 100644 examples/leak_probe.rs create mode 100644 src/tls.rs create mode 100644 tests/severed_connection.rs create mode 100644 tests/tls.rs create mode 100755 tests/tls/run.sh diff --git a/Cargo.toml b/Cargo.toml index 2e9f641..fd0e69b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,6 +48,9 @@ fallible-iterator = "0.2" postgres-protocol = "0.6" may = { version = "0.3", default-features = false } +# TLS (sslmode, sslrootcert, sslcert/sslkey). ring: no C toolchain beyond cc, no cmake. +rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12", "logging"] } +rustls-pemfile = "2" [dev-dependencies] env_logger = "0.11" @@ -70,3 +73,7 @@ lto = true # [patch.crates-io] # may = { path = "../may" } + +[lints.rust] +# Test code that upstream disabled on purpose (no unix-socket support in this port; a geo-types version never added). +unexpected_cfgs = { level = "warn", check-cfg = ["cfg(unix_nosupport)", "cfg(feature, values(\"with-geo-0_10\"))"] } diff --git a/examples/leak_probe.rs b/examples/leak_probe.rs new file mode 100644 index 0000000..717f60f --- /dev/null +++ b/examples/leak_probe.rs @@ -0,0 +1,75 @@ +//! Diagnostic: does dropping a Client actually release its resources? +//! +//! Two independent measurements per iteration: +//! - server side: backends in pg_stat_activity for our database +//! - client side: open socket fds in /proc/self/fd +//! +//! If BOTH climb one-per-iteration, the process itself is holding the socket +//! after drop - a genuine client-side leak. If fds return to baseline but +//! backends linger, the client is fine and the server is just slow to reap. + +fn socket_fds() -> usize { + std::fs::read_dir("/proc/self/fd") + .unwrap() + .filter_map(|e| e.ok()) + .filter(|e| { + std::fs::read_link(e.path()) + .map(|t| t.to_string_lossy().starts_with("socket:")) + .unwrap_or(false) + }) + .count() +} + +fn main() { + let url = std::env::var("TEST_URL").expect("set TEST_URL"); + + let observer = may_postgres::connect(&url).expect("observer connect"); + let backends = |o: &may_postgres::Client| -> i64 { + o.query_one( + "SELECT count(*) FROM pg_stat_activity WHERE datname = current_database()", + &[], + ) + .unwrap() + .get(0) + }; + + let base_be = backends(&observer); + let base_fd = socket_fds(); + println!("BASELINE backends={base_be} socket_fds={base_fd}"); + + let iters: u32 = std::env::var("ITERS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(10); + let verbose = iters <= 20; + for i in 1..=iters { + { + let client = may_postgres::connect(&url).expect("connect"); + let rows = client.query("SELECT 1", &[]).expect("query"); + assert_eq!(rows.len(), 1); + drop(client); + } + std::thread::sleep(std::time::Duration::from_millis(if verbose { + 300 + } else { + 5 + })); + if !verbose && i % 100 != 0 { + continue; + } + println!( + "ITER {i} backends={} socket_fds={}", + backends(&observer), + socket_fds() + ); + } + + // Give the server every chance to reap on its own. + println!("WAITING 15s..."); + std::thread::sleep(std::time::Duration::from_secs(15)); + println!( + "FINAL backends={} socket_fds={} (baseline {base_be}/{base_fd})", + backends(&observer), + socket_fds() + ); +} diff --git a/src/client.rs b/src/client.rs index bf90732..1b8b189 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,7 +1,7 @@ use crate::cancel_query; use crate::codec::BackendMessages; use crate::config::Host; -use crate::connection::{Connection, RefOrValue, Request, RequestMessages}; +use crate::connection::{Connection, Request, RequestMessages}; use crate::copy_out::CopyOutStream; use crate::query::RowStream; use crate::simple_query::SimpleQueryStream; @@ -22,10 +22,12 @@ use std::rc::Rc; use std::sync::Arc; use std::time::Duration; +use crate::connection::ResponseResult; + pub struct Responses { tag: usize, cur: BackendMessages, - rx: Rc>, + rx: Rc>, } impl Responses { @@ -36,12 +38,16 @@ impl Responses { Some((_, Message::ErrorResponse(body))) => return Err(Error::db(body)), Some((_, message)) => return Ok(message), None => match self.rx.recv() { - Ok(messages) => { + Ok(Ok(messages)) => { if messages.tag != self.tag { continue; } self.cur = messages } + // The connection I/O loop exited and failed this pending + // request. The error is untagged — the connection is dead, + // every outstanding request on it fails the same way. + Ok(Err(e)) => return Err(e), Err(_) => return Err(Error::closed()), }, } @@ -63,17 +69,17 @@ pub struct InnerClient { struct CoChannel { tag: Cell, - rx: Rc>, - tx: spsc::Sender, + rx: Rc>, + tx: Arc>, } impl CoChannel { - fn sender(&self) -> RefOrValue<'static, spsc::Sender> { - // Safety: - // 1. there is only one sender - // 2. we will wait until all response come back - let tx: &'static spsc::Sender = unsafe { std::mem::transmute(&self.tx) }; - RefOrValue::Ref(tx) + fn sender(&self) -> Arc> { + // Shared ownership, not a transmuted borrow. The old 'static borrow + // relied on the client outliving every queued request, which drop + // breaks - the io coroutine would then push responses through a + // dangling pointer into freed Client memory. + self.tx.clone() } fn tag(&self) -> usize { @@ -82,7 +88,7 @@ impl CoChannel { tag } - fn receiver(&self) -> Rc> { + fn receiver(&self) -> Rc> { self.rx.clone() } } @@ -91,7 +97,7 @@ impl InnerClient { /// ignore the result pub fn raw_send(&self, messages: RequestMessages) -> Result<(), Error> { let (sender, _rx) = spsc::channel(); - let request = Request::new(0, messages, RefOrValue::Value(sender)); + let request = Request::new(0, messages, Arc::new(sender)); self.sender.send(request); Ok(()) } @@ -160,7 +166,11 @@ impl Clone for Client { let (tx, rx) = spsc::channel(); let rx = Rc::new(rx); let tag = Cell::new(0); - CoChannel { tag, rx, tx } + CoChannel { + tag, + rx, + tx: Arc::new(tx), + } }; Client { inner: self.inner.clone(), @@ -182,26 +192,64 @@ impl Client { /// so `LISTEN` could be issued but nothing could ever be observed. /// /// The queue is unbounded and is only drained by the caller, so a client - /// that issues `LISTEN` and never reads will grow it without limit. Poll it - /// as part of the same loop that does the work: + /// that issues `LISTEN` and never reads will grow it without limit. The + /// connection fills it as notifications arrive, without any request from + /// the client; to wait for the next one rather than poll, use + /// [`Client::next_notification`]. + #[inline] + pub fn notifications(&self) -> &may::queue::mpsc::Queue { + self.inner.sender.notifications().queue() + } + + /// Wait up to `timeout` for the next notification. + /// + /// Returns as soon as one is queued (or immediately if one already is), + /// parking the coroutine (or thread) in the meantime, so a `LISTEN` + /// consumer needs no polling interval. Returns `None` when `timeout` + /// passes with nothing delivered, or when the connection has closed and + /// nothing is left in the queue. /// /// ```no_run /// # fn main() -> Result<(), may_postgres::Error> { + /// # use std::time::Duration; /// let client = may_postgres::connect("host=localhost user=postgres")?; /// client.batch_execute("LISTEN my_channel")?; /// - /// loop { - /// // Any round trip lets the connection surface pending notifications. - /// client.batch_execute("SELECT 1")?; - /// while let Some(notification) = client.notifications().pop() { - /// println!("{}: {}", notification.channel(), notification.payload()); + /// while !client.is_closed() { + /// if let Some(n) = client.next_notification(Duration::from_secs(1)) { + /// println!("{}: {}", n.channel(), n.payload()); /// } /// } + /// # Ok(()) /// # } /// ``` - #[inline] - pub fn notifications(&self) -> &may::queue::mpsc::Queue { - self.inner.sender.notifications() + pub fn next_notification(&self, timeout: std::time::Duration) -> Option { + let notifications = self.inner.sender.notifications(); + let deadline = std::time::Instant::now() + timeout; + loop { + if let Some(n) = notifications.queue().pop() { + return Some(n); + } + if self.is_closed() { + // Pass the wake-up on, so every waiter on this connection + // returns rather than one of them. + notifications.ready().post(); + return notifications.queue().pop(); + } + let left = deadline.saturating_duration_since(std::time::Instant::now()); + if left.is_zero() { + return notifications.queue().pop(); + } + // may keeps park timeouts in whole milliseconds, and releases up to + // 0.3.51 truncate anything shorter to "no timeout" (the wait would + // never return): never ask it for less than 1 ms. + let left = left.max(std::time::Duration::from_millis(1)); + if !notifications.ready().wait_timeout(left) { + return notifications.queue().pop(); + } + // A permit can outnumber queued notifications when the queue was + // also drained with `notifications().pop()`; loop and re-check. + } } pub(crate) fn new(sender: Connection, process_id: i32, secret_key: i32) -> Client { @@ -209,7 +257,11 @@ impl Client { let (tx, rx) = spsc::channel(); let rx = Rc::new(rx); let tag = Cell::new(0); - CoChannel { tag, rx, tx } + CoChannel { + tag, + rx, + tx: Arc::new(tx), + } }; Client { inner: Arc::new(InnerClient { @@ -513,10 +565,17 @@ impl Client { /// Determines if the connection to the server has already closed. /// - /// In that case, all future queries will fail. + /// In that case, all future queries will fail — fast, with a + /// connection-closed error, rather than blocking. Pool implementations can + /// use this to retire a dead client before dispatching work to it. + pub fn is_closed(&self) -> bool { + self.inner.sender.is_dead() + } + + #[doc(hidden)] + #[deprecated(note = "use `is_closed`")] pub fn _is_closed(&self) -> bool { - // self.inner.sender.is_closed() - unimplemented!() + self.is_closed() } #[inline] diff --git a/src/codec.rs b/src/codec.rs index 0dc1667..6c21133 100644 --- a/src/codec.rs +++ b/src/codec.rs @@ -1,6 +1,6 @@ +use crate::tls::PgStream; use bytes::{Buf, BufMut, BytesMut}; use fallible_iterator::FallibleIterator; -use may::net::TcpStream; use postgres_protocol::message::backend; use postgres_protocol::message::frontend::CopyData; use std::io::{self, Read}; @@ -111,13 +111,13 @@ mod frame_codec { use super::*; pub struct Framed { - r_stream: TcpStream, + r_stream: PgStream, read_buf: BytesMut, codec: PostgresCodec, } impl Framed { - pub fn new(s: TcpStream) -> Self { + pub fn new(s: PgStream) -> Self { Framed { r_stream: s, read_buf: BytesMut::with_capacity(4096 * 8), @@ -125,11 +125,11 @@ mod frame_codec { } } - pub fn inner_mut(&mut self) -> &mut TcpStream { + pub fn inner_mut(&mut self) -> &mut PgStream { &mut self.r_stream } - pub fn into_stream(self) -> TcpStream { + pub fn into_stream(self) -> PgStream { self.r_stream } diff --git a/src/config.rs b/src/config.rs index bc3143d..f70e43f 100644 --- a/src/config.rs +++ b/src/config.rs @@ -48,8 +48,11 @@ pub(crate) enum Host { /// * `dbname` - The name of the database to connect to. Defaults to the username. /// * `options` - Command line options used to configure the server. /// * `application_name` - Sets the `application_name` parameter on the server. -/// * `sslmode` - Controls usage of TLS. If set to `disable`, TLS will not be used. If set to `prefer`, TLS will be used -/// if available, but not used otherwise. If set to `require`, TLS will be forced to be used. Defaults to `prefer`. +/// * `sslmode` - `disable`, `prefer` (default), `require`, `verify-ca` or `verify-full`, as libpq. `require` with +/// `sslrootcert` behaves as `verify-ca`. +/// * `sslrootcert` - PEM file of the certificate authority the server certificate must chain to (`verify-ca`, +/// `verify-full`). +/// * `sslcert`, `sslkey` - PEM client certificate and private key, presented when the server asks for one. /// * `host` - The host to connect to. On Unix platforms, if the host starts with a `/` character it is treated as the /// path to the directory containing Unix domain sockets. Otherwise, it is treated as a hostname. Multiple hosts /// can be specified, separated by commas. Each host will be tried in turn when connecting. Required if connecting @@ -121,6 +124,10 @@ pub struct Config { pub(crate) keepalives: bool, pub(crate) keepalives_idle: Duration, pub(crate) target_session_attrs: TargetSessionAttrs, + pub(crate) ssl_mode: SslMode, + pub(crate) ssl_root_cert: Option, + pub(crate) ssl_cert: Option, + pub(crate) ssl_key: Option, } impl Default for Config { @@ -144,6 +151,10 @@ impl Config { keepalives: true, keepalives_idle: Duration::from_secs(2 * 60 * 60), target_session_attrs: TargetSessionAttrs::Any, + ssl_mode: SslMode::Prefer, + ssl_root_cert: None, + ssl_cert: None, + ssl_key: None, } } @@ -260,6 +271,50 @@ impl Config { self } + /// Sets the TLS mode. Defaults to `Prefer`. + pub fn ssl_mode(&mut self, ssl_mode: SslMode) -> &mut Config { + self.ssl_mode = ssl_mode; + self + } + + /// The TLS mode. + pub fn get_ssl_mode(&self) -> SslMode { + self.ssl_mode + } + + /// Sets the PEM file of the certificate authority the server certificate must chain to. + pub fn ssl_root_cert(&mut self, path: impl Into) -> &mut Config { + self.ssl_root_cert = Some(path.into()); + self + } + + /// The root certificate file, if set. + pub fn get_ssl_root_cert(&self) -> Option<&std::path::Path> { + self.ssl_root_cert.as_deref() + } + + /// Sets the PEM client certificate file (with [`Config::ssl_key`]). + pub fn ssl_cert(&mut self, path: impl Into) -> &mut Config { + self.ssl_cert = Some(path.into()); + self + } + + /// The client certificate file, if set. + pub fn get_ssl_cert(&self) -> Option<&std::path::Path> { + self.ssl_cert.as_deref() + } + + /// Sets the PEM private key file of the client certificate. + pub fn ssl_key(&mut self, path: impl Into) -> &mut Config { + self.ssl_key = Some(path.into()); + self + } + + /// The client key file, if set. + pub fn get_ssl_key(&self) -> Option<&std::path::Path> { + self.ssl_key.as_deref() + } + fn param(&mut self, key: &str, value: &str) -> Result<(), Error> { match key { "user" => { @@ -277,6 +332,20 @@ impl Config { "application_name" => { self.application_name(value); } + "sslmode" => { + let mode = SslMode::parse(value) + .ok_or_else(|| Error::config_parse(Box::new(InvalidValue("sslmode"))))?; + self.ssl_mode(mode); + } + "sslrootcert" => { + self.ssl_root_cert(value); + } + "sslcert" => { + self.ssl_cert(value); + } + "sslkey" => { + self.ssl_key(value); + } "host" => { for host in value.split(',') { self.host(host); @@ -348,7 +417,8 @@ impl Config { /// /// All of the settings other than `user`, `password`, `dbname`, `options`, and `application` name are ignored. pub fn connect_raw(&self, stream: TcpStream) -> Result { - connect_raw(stream, self) + let host = self.host.first().map(|Host::Tcp(h)| h.as_str()); + connect_raw(stream, self, host) } } @@ -379,7 +449,10 @@ impl fmt::Debug for Config { .field("dbname", &self.dbname) .field("options", &self.options) .field("application_name", &self.application_name) - // .field("ssl_mode", &self.ssl_mode) + .field("ssl_mode", &self.ssl_mode) + .field("ssl_root_cert", &self.ssl_root_cert) + .field("ssl_cert", &self.ssl_cert) + .field("ssl_key", &self.ssl_key) .field("host", &self.host) .field("port", &self.port) .field("connect_timeout", &self.connect_timeout) @@ -760,3 +833,77 @@ impl<'a> UrlParser<'a> { .map_err(|e| Error::config_parse(e.into())) } } + +/// TLS mode, as libpq's `sslmode`. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum SslMode { + /// Never TLS. + Disable, + /// TLS when the server offers it; no certificate checks. + Prefer, + /// TLS or fail; no certificate checks (unless a root certificate is set: then as `VerifyCa`). + Require, + /// TLS; the server certificate must chain to the root certificate. + VerifyCa, + /// As `VerifyCa`, and the certificate must name the host. + VerifyFull, +} + +impl SslMode { + fn parse(s: &str) -> Option { + Some(match s { + "disable" => SslMode::Disable, + "prefer" => SslMode::Prefer, + "require" => SslMode::Require, + "verify-ca" => SslMode::VerifyCa, + "verify-full" => SslMode::VerifyFull, + _ => return None, + }) + } + + pub(crate) fn as_str(&self) -> &'static str { + match self { + SslMode::Disable => "disable", + SslMode::Prefer => "prefer", + SslMode::Require => "require", + SslMode::VerifyCa => "verify-ca", + SslMode::VerifyFull => "verify-full", + } + } +} + +#[cfg(test)] +mod ssl_config_tests { + use super::*; + + #[test] + fn default_is_prefer() { + assert_eq!(Config::new().get_ssl_mode(), SslMode::Prefer); + } + + #[test] + fn key_value_ssl_params() { + let c: Config = + "host=db user=u sslmode=verify-full sslrootcert=/ca.crt sslcert=/c.crt sslkey=/c.key" + .parse() + .unwrap(); + assert_eq!(c.get_ssl_mode(), SslMode::VerifyFull); + assert_eq!(c.get_ssl_root_cert(), Some(std::path::Path::new("/ca.crt"))); + assert_eq!(c.get_ssl_cert(), Some(std::path::Path::new("/c.crt"))); + assert_eq!(c.get_ssl_key(), Some(std::path::Path::new("/c.key"))); + } + + #[test] + fn url_ssl_params() { + let c: Config = "postgres://u:p@db:5432/x?sslmode=require".parse().unwrap(); + assert_eq!(c.get_ssl_mode(), SslMode::Require); + let c: Config = "postgres://u@db/x?sslmode=disable".parse().unwrap(); + assert_eq!(c.get_ssl_mode(), SslMode::Disable); + } + + #[test] + fn bad_sslmode_rejected() { + assert!("host=db sslmode=allow-ish".parse::().is_err()); + } +} diff --git a/src/connect.rs b/src/connect.rs index c366901..bdcf650 100644 --- a/src/connect.rs +++ b/src/connect.rs @@ -44,7 +44,8 @@ fn connect_once(host: &Host, port: u16, config: &Config) -> Result io::Result { @@ -40,7 +43,12 @@ impl StartupStream { } } -pub fn connect_raw(stream: TcpStream, config: &Config) -> Result { +pub fn connect_raw( + stream: TcpStream, + config: &Config, + host: Option<&str>, +) -> Result { + let stream = tls::negotiate(stream, config, host)?; let mut stream = StartupStream { inner: Framed::new(stream), buf: BackendMessages::empty(0), diff --git a/src/connection.rs b/src/connection.rs index 730be21..4e7bff1 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -1,11 +1,11 @@ +use crate::tls::{tls_read, tls_write, PgStream}; use bytes::{Buf, BufMut, BytesMut}; use fallible_iterator::FallibleIterator; use may::coroutine::JoinHandle; use may::go; use may::io::{WaitIo, WaitIoWaker}; -use may::net::TcpStream; use may::queue::mpsc::Queue; -use may::sync::spsc; +use may::sync::{spsc, Semphore}; use postgres_protocol::message::backend::Message; use postgres_protocol::message::frontend; @@ -15,36 +15,31 @@ use crate::{Error, Notification}; use std::collections::{HashMap, VecDeque}; use std::io::{self, Read, Write}; -use std::ops::Deref; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; const IO_BUF_SIZE: usize = 4096 * 2; -pub enum RefOrValue<'a, T> { - Ref(&'a T), - Value(T), -} - -impl Deref for RefOrValue<'_, T> { - type Target = T; - #[inline] - fn deref(&self) -> &Self::Target { - match self { - RefOrValue::Ref(r) => r, - RefOrValue::Value(ref v) => v, - } - } -} - pub enum RequestMessages { Single(FrontendMessage), CopyIn(CopyInReceiver), } +/// What the connection coroutine sends back to a waiting `Responses::next`. +/// +/// `Err` is sent exactly once per pending request when the I/O loop exits, so +/// a caller blocked on `rx.recv()` wakes with a connection-closed error +/// instead of waiting forever on a channel whose only producer is gone. +pub type ResponseResult = Result; + pub struct Request { tag: usize, messages: RequestMessages, - sender: RefOrValue<'static, spsc::Sender>, + // OWNED (shared) sender, never a transmuted borrow: requests can sit in + // the queues after the originating Client is dropped, so a borrow here + // would dangle into freed Client memory the moment drop stopped being + // an instant kill. + sender: Arc>, } impl Request { @@ -52,7 +47,7 @@ impl Request { pub fn new( tag: usize, messages: RequestMessages, - sender: RefOrValue<'static, spsc::Sender>, + sender: Arc>, ) -> Request { Request { tag, @@ -64,24 +59,101 @@ impl Request { pub struct Response { tag: usize, - tx: RefOrValue<'static, spsc::Sender>, + tx: Arc>, +} + +/// Asynchronous notifications delivered by the server: a queue the connection +/// coroutine fills and the client drains, plus a count of pushes a consumer can +/// park on instead of polling the queue. +pub(crate) struct Notifications { + queue: Queue, + ready: Semphore, +} + +impl Notifications { + fn new() -> Self { + Notifications { + queue: Queue::new(), + ready: Semphore::new(0), + } + } + + fn push(&self, notification: Notification) { + self.queue.push(notification); + self.ready.post(); + } + + /// Wake a consumer without a notification (the connection closed). + fn wake(&self) { + self.ready.post(); + } + + pub(crate) fn queue(&self) -> &Queue { + &self.queue + } + + pub(crate) fn ready(&self) -> &Semphore { + &self.ready + } } /// A connection to a PostgreSQL database. pub(crate) struct Connection { - io_handle: JoinHandle<()>, + /// Held only to keep the io coroutine owned; dropping it DETACHES. The + /// coroutine is never cancelled (cancel skips destructors, leaking the + /// socket) and never joined (drop can run inside a coroutine that is + /// itself being cancelled, where parking to join aborts the process). + _io_handle: JoinHandle<()>, req_queue: Arc>, /// Asynchronous notifications delivered by the server, filled by the /// connection coroutine and drained by the client. - notifications: Arc>, + notifications: Arc, waker: Arc, + /// Cooperative shutdown: `drop` sets this and wakes the coroutine; the + /// connection loop polls it and exits on its own stack, closing its own + /// socket. Reaching in from outside (cancelling, or shutting the fd down + /// from `drop`) corrupts the coroutine mid-flight. + shutdown: Arc, + /// Set (only) by the I/O coroutine when its loop has exited, on error, EOF + /// or cooperative shutdown. Once set, the connection is permanently dead: + /// nothing will ever drain `req_queue` or complete a response again, so + /// `send` fails new requests immediately instead of queueing them forever. + dead: Arc, + /// Serializes draining `req_queue` on the death path. The queue is MPSC — + /// while the loop runs it is the only consumer; after `dead` is set both + /// the exiting coroutine and any sender that lost the race may need to + /// drain, and this lock keeps "single consumer" true. + drain_lock: Arc>, id: usize, } +/// Pop every queued request and fail its response channel with +/// [`Error::closed`]. Callers must hold `drain_lock` and only call this after +/// `dead` is set (the I/O loop must no longer be consuming). +fn fail_queued_requests(req_queue: &Queue) { + while let Some(req) = req_queue.pop() { + req.sender.send(Err(Error::closed())).ok(); + } +} + +/// Fail every response the server still owed us. Runs once, on the I/O +/// coroutine's own exit path. +fn fail_pending_responses(rsp_queue: &mut VecDeque) { + for rsp in rsp_queue.drain(..) { + rsp.tx.send(Err(Error::closed())).ok(); + } +} + impl Drop for Connection { fn drop(&mut self) { - let rx = self.io_handle.coroutine(); - unsafe { rx.cancel() }; + // Signal the loop and wake it out of wait_io(); it exits on its own + // stack, sends the protocol Terminate and closes its own socket. This + // must NOT block: a Client can be dropped inside a coroutine that is + // itself being cancelled, and parking there (to join) aborts. The + // waker holds an Arc to the reactor slot, so waking after the loop + // has already exited is a harmless no-op. + self.shutdown.store(true, Ordering::Release); + self.waker.wakeup(); } } @@ -139,7 +211,7 @@ fn decode_messages( read_buf: &mut BytesMut, rsp_queue: &mut VecDeque, parameters: &mut HashMap, - notifications: &Queue, + notifications: &Notifications, ) -> Result<(), Error> { use crate::codec::PostgresCodec; @@ -159,7 +231,7 @@ fn decode_messages( }; messages.tag = response.tag; - response.tx.send(messages).ok(); + response.tx.send(Ok(messages)).ok(); if request_complete { rsp_queue.pop_front(); @@ -266,83 +338,183 @@ fn process_req( } #[cold] -fn terminate_connection(stream: &mut TcpStream) { +fn terminate_connection(stream: &mut PgStream) { let mut request = BytesMut::new(); frontend::terminate(&mut request); stream.write_all(&request.freeze()).ok(); - stream.shutdown(std::net::Shutdown::Both).ok(); + if let Some(conn) = stream.tls.as_deref_mut() { + conn.send_close_notify(); + } + stream.flush().ok(); + stream.tcp.shutdown(std::net::Shutdown::Both).ok(); } #[inline] fn connection_loop( - stream: &mut TcpStream, - req_queue: Arc>, + stream: &mut PgStream, + req_queue: &Queue, + rsp_queue: &mut VecDeque, mut params: HashMap, - notifications: Arc>, + notifications: Arc, + shutdown: &AtomicBool, ) -> Result<(), Error> { let mut read_buf = BytesMut::with_capacity(IO_BUF_SIZE); let mut write_buf = BytesMut::with_capacity(IO_BUF_SIZE); - let mut rsp_queue = VecDeque::with_capacity(512); let mut pending_copy_in: Option = None; let mut io_events = 1; // allow read loop { - let inner_stream = stream.inner_mut(); - - process_req( - &req_queue, - &mut rsp_queue, - &mut write_buf, - &mut pending_copy_in, - ) - .map_err(Error::io)?; + // Cooperative shutdown, set by Connection::drop. Exiting HERE - + // rather than being cancelled from outside - is what guarantees the + // stream is dropped on this stack with destructors run. + if shutdown.load(Ordering::Acquire) { + return Ok(()); + } + let inner_stream = stream.tcp.inner_mut(); + + process_req(req_queue, rsp_queue, &mut write_buf, &mut pending_copy_in) + .map_err(Error::io)?; // An in-flight copy may have had more data queued since the last pass. drain_copy_in(&mut pending_copy_in, &mut write_buf); - nonblock_write(inner_stream, &mut write_buf).map_err(Error::io)?; + match stream.tls.as_deref_mut() { + None => nonblock_write(inner_stream, &mut write_buf).map(|_| ()), + Some(conn) => tls_write(conn, inner_stream, &mut write_buf), + } + .map_err(Error::io)?; let mut read_blocked = true; if io_events & 1 != 0 { - read_blocked = nonblock_read(inner_stream, &mut read_buf).map_err(Error::io)?; - decode_messages(&mut read_buf, &mut rsp_queue, &mut params, ¬ifications)?; + read_blocked = match stream.tls.as_deref_mut() { + None => nonblock_read(inner_stream, &mut read_buf), + Some(conn) => tls_read(conn, inner_stream, &mut read_buf), + } + .map_err(Error::io)?; + decode_messages(&mut read_buf, rsp_queue, &mut params, ¬ifications)?; } - io_events = if read_blocked { stream.wait_io() } else { 1 } + io_events = if read_blocked { + stream.tcp.wait_io() + } else { + 1 + } } } impl Connection { - pub(crate) fn new(mut stream: TcpStream, parameters: HashMap) -> Connection { + pub(crate) fn new(mut stream: PgStream, parameters: HashMap) -> Connection { use std::os::fd::AsRawFd; - let id = stream.as_raw_fd() as usize; - let waker = Arc::new(stream.waker()); + let id = stream.tcp.as_raw_fd() as usize; + let waker = Arc::new(stream.tcp.waker()); let req_queue = Arc::new(Queue::new()); let req_queue_dup = req_queue.clone(); - let notifications = Arc::new(Queue::new()); + let notifications = Arc::new(Notifications::new()); let notifications_dup = notifications.clone(); - let io_handle = go!(move || { - if let Err(e) = - connection_loop(&mut stream, req_queue_dup, parameters, notifications_dup) - { - log::error!("connection error = {:?}", e); + let notifications_exit = notifications.clone(); + let shutdown = Arc::new(AtomicBool::new(false)); + let shutdown_dup = shutdown.clone(); + let dead = Arc::new(AtomicBool::new(false)); + let dead_dup = dead.clone(); + let drain_lock = Arc::new(spin::Mutex::new(())); + let drain_lock_dup = drain_lock.clone(); + // EXPLICIT stack for the I/O coroutine. Bare go! inherits the + // embedder ambient coroutine stack config, and inside BRRTRouter + // that proved fatal: a connection spawned at RUNTIME (lifeguard + // max-lifetime rotation, 2026-08-30 market wedge) overflowed its + // stack and may killed it WITHOUT unwinding - so the exit path + // below (mark dead, fail pending) never ran and every request + // into that slot hung forever instead of failing fast. Same class + // as the may-redis overflow recorded in PriceWhisperer hot_tier.rs. + // 1 MiB: 256 KiB still overflowed once in ~8 rotation cycles under the + // compressed-lifetime repro, and stack here is cheap at pool scale. + let io_stack: usize = std::env::var("MAY_POSTGRES_IO_STACK") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0x100000); + let io_handle = go!( + may::coroutine::Builder::new() + .name("pg-io".to_owned()) + .stack_size(io_stack), + move || { + let mut rsp_queue: VecDeque = VecDeque::with_capacity(512); + if let Err(e) = connection_loop( + &mut stream, + &req_queue_dup, + &mut rsp_queue, + parameters, + notifications_dup, + &shutdown_dup, + ) { + log::error!("connection error = {:?}", e); + } + // On EVERY exit (error, EOF, cooperative shutdown): the loop will + // never write another request or decode another response, so fail + // everything that is still waiting. Without this, a caller blocked + // in `Responses::next` waits forever — its channel's only producer + // is this coroutine — and that wedge is what starved lifeguard's + // pool workers (2026-08-29 tickers outage). + // + // Order matters: mark dead FIRST, so a sender that enqueues after + // our drain sees `dead` on its post-push check and drains its own + // request; then drain under the lock. SeqCst pairs with the SeqCst + // load in `send` (Dekker pattern: push / load-dead on one side, + // store-dead / drain on the other) so no push can slip between the + // final drain and a stale `dead == false` read. + dead_dup.store(true, Ordering::SeqCst); + fail_pending_responses(&mut rsp_queue); + { + let _g = drain_lock_dup.lock(); + fail_queued_requests(&req_queue_dup); + } + // A consumer parked waiting for a notification would otherwise + // wait out its timeout on a connection that will never deliver. + notifications_exit.wake(); + // Send the protocol Terminate and close the socket, on this + // coroutine's own stack. This is what releases the server backend + // the moment the client goes away, instead of at process exit. terminate_connection(&mut stream); } - }); + ); Connection { - io_handle, + // Builder spawn is fallible (stack allocation); a connection we + // cannot give a safe stack is a connection we must not hand out. + _io_handle: io_handle.expect("spawn pg-io coroutine"), req_queue, notifications, waker, + shutdown, + dead, + drain_lock, id, } } /// send a request to the connection + /// + /// If the I/O loop has already exited, the request's response channel is + /// completed with [`Error::closed`] immediately (fail-fast) instead of + /// leaving the request on a queue no coroutine will ever drain. #[inline] pub fn send(&self, req: Request) { self.req_queue.push(req); self.waker.wakeup(); + // Post-push check: either the loop was alive after our push (it will + // consume the request normally, or fail it on its own exit drain), or + // it was already dead — then nothing consumes the queue anymore and + // WE drain it. Both racers may drain; the lock keeps it single-file. + // SeqCst pairs with the SeqCst store on the exit path (see above). + if self.dead.load(Ordering::SeqCst) { + let _g = self.drain_lock.lock(); + fail_queued_requests(&self.req_queue); + } + } + + /// `true` once the connection I/O loop has exited; the connection can + /// never carry another request and every call on it will fail fast. + #[inline] + pub fn is_dead(&self) -> bool { + self.dead.load(Ordering::Acquire) } #[inline] @@ -350,6 +522,14 @@ impl Connection { self.id } + /// Asynchronous notifications received from the server. + /// + /// Populated by `LISTEN`; see [`crate::Client::notifications`]. + #[inline] + pub fn notifications(&self) -> &Notifications { + &self.notifications + } + /// Waker for the connection coroutine. /// /// `send` wakes the loop when a request is queued. A copy-in streams its @@ -359,12 +539,4 @@ impl Connection { pub fn waker(&self) -> Arc { self.waker.clone() } - - /// Asynchronous notifications received from the server. - /// - /// Populated by `LISTEN`; see [`crate::Client::notifications`]. - #[inline] - pub fn notifications(&self) -> &Queue { - &self.notifications - } } diff --git a/src/copy_in.rs b/src/copy_in.rs index ddcface..168b876 100644 --- a/src/copy_in.rs +++ b/src/copy_in.rs @@ -53,30 +53,6 @@ impl CopyInReceiver { } } } - - pub fn recv(&mut self) -> Result, ()> { - if self.done { - return Err(()); - } - - match self.receiver.recv() { - Ok(CopyInMessage::Message(message)) => Ok(Some(message)), - Ok(CopyInMessage::Done) => { - self.done = true; - let mut buf = BytesMut::new(); - frontend::copy_done(&mut buf); - frontend::sync(&mut buf); - Ok(Some(FrontendMessage::Raw(buf))) - } - Err(_) => { - self.done = true; - let mut buf = BytesMut::new(); - frontend::copy_fail("", &mut buf).unwrap(); - frontend::sync(&mut buf); - Ok(Some(FrontendMessage::Raw(buf))) - } - } - } } /// A sink for `COPY ... FROM STDIN` query data. diff --git a/src/error/mod.rs b/src/error/mod.rs index f0b258e..ef7cdb1 100644 --- a/src/error/mod.rs +++ b/src/error/mod.rs @@ -341,7 +341,7 @@ pub enum ErrorPosition { enum Kind { Io, UnexpectedMessage, - // Tls, + Tls, ToSql(usize), FromSql(usize), Column(String), @@ -380,7 +380,7 @@ impl fmt::Display for Error { match &self.0.kind { Kind::Io => fmt.write_str("error communicating with the server")?, Kind::UnexpectedMessage => fmt.write_str("unexpected message from server")?, - // Kind::Tls => fmt.write_str("error performing TLS handshake")?, + Kind::Tls => fmt.write_str("error performing TLS handshake")?, Kind::ToSql(idx) => write!(fmt, "error serializing parameter {}", idx)?, Kind::FromSql(idx) => write!(fmt, "error deserializing column {}", idx)?, Kind::Column(column) => write!(fmt, "invalid column `{}`", column)?, @@ -481,9 +481,9 @@ impl Error { // Error::new(Kind::Parameters(real, expected), None) // } - // pub(crate) fn tls(e: Box) -> Error { - // Error::new(Kind::Tls, Some(e)) - // } + pub(crate) fn tls(e: Box) -> Error { + Error::new(Kind::Tls, Some(e)) + } pub(crate) fn io(e: io::Error) -> Error { Error::new(Kind::Io, Some(Box::new(e))) @@ -513,4 +513,13 @@ impl Error { pub fn __private_api_timeout() -> Error { Error::new(Kind::Timeout, None) } + + /// Application-level failure converting a successfully-read column value (e.g. malformed JSON + /// in a `TEXT` / `JSON` column). Surfaced as [`Kind::FromSql`] for `column_index`. + pub fn user_column_decode_failed(column_index: usize, message: String) -> Error { + Error::from_sql( + Box::new(io::Error::new(io::ErrorKind::InvalidData, message)), + column_index, + ) + } } diff --git a/src/lib.rs b/src/lib.rs index 5cfefb5..d435c8a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -108,6 +108,7 @@ mod query; pub mod row; mod simple_query; mod statement; +mod tls; mod to_statement; mod transaction; mod transaction_builder; diff --git a/src/query.rs b/src/query.rs index e84b5c6..cebfe8c 100644 --- a/src/query.rs +++ b/src/query.rs @@ -17,6 +17,7 @@ pub fn query( Ok(RowStream { statement, responses, + rows_affected: None, }) } @@ -32,6 +33,7 @@ pub fn query_portal(client: &Client, portal: &Portal, max_rows: i32) -> Result, +} + +impl RowStream { + /// The number of rows the statement affected, from its `CommandComplete` + /// tag (`INSERT 0 3`, `UPDATE 5`, `SELECT 2`, ...). + /// + /// `None` until the stream has been read to its end, and for a portal + /// that was suspended before completing. This lets a statement sent with + /// `query_raw` (for example to pipeline it with others) report its count + /// the way `execute` does, without a second round trip. + pub fn rows_affected(&self) -> Option { + self.rows_affected + } } impl Iterator for RowStream { @@ -135,9 +151,19 @@ impl Iterator for RowStream { Message::DataRow(body) => { return Some(Ok(o_try!(Row::new(self.statement.clone(), body)))) } - Message::EmptyQueryResponse - | Message::CommandComplete(_) - | Message::PortalSuspended => return None, + Message::CommandComplete(body) => { + self.rows_affected = body + .tag() + .ok() + .and_then(|tag| tag.rsplit(' ').next()) + .and_then(|n| n.parse().ok()); + return None; + } + Message::EmptyQueryResponse => { + self.rows_affected = Some(0); + return None; + } + Message::PortalSuspended => return None, Message::ErrorResponse(body) => return Some(Err(Error::db(body))), Message::BindComplete => {} _ => return Some(Err(Error::unexpected_message())), diff --git a/src/tls.rs b/src/tls.rs new file mode 100644 index 0000000..3ec560c --- /dev/null +++ b/src/tls.rs @@ -0,0 +1,377 @@ +//! TLS for the connection (rustls, ring provider), negotiated the way libpq does it. +//! +//! `sslmode`: +//! * `disable` - never TLS. +//! * `prefer` (default) - TLS when the server offers it, otherwise plaintext. No certificate checks. +//! * `require` - TLS or fail. No certificate checks, unless `sslrootcert` is set, which makes it `verify-ca` (libpq). +//! * `verify-ca` - TLS, server certificate chain checked against `sslrootcert`. +//! * `verify-full` - as `verify-ca`, and the certificate must name the host connected to. +//! +//! `sslcert` + `sslkey` (PEM files) present a client certificate when the server asks for one. +//! +//! The startup phase (TLS handshake, startup, authentication) runs on the blocking coroutine stream ([`PgStream`] +//! implements `Read`/`Write`); the connection's I/O coroutine drives rustls non-blocking against the raw socket +//! ([`tls_write`], [`tls_read`]). + +use std::fs::File; +use std::io::{self, BufReader, Read, Write}; +use std::path::Path; +use std::sync::Arc; + +use bytes::{Buf, BufMut, BytesMut}; +use may::net::TcpStream; +use postgres_protocol::message::frontend; +use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier}; +use rustls::client::WebPkiServerVerifier; +use rustls::crypto::{verify_tls12_signature, verify_tls13_signature, CryptoProvider}; +use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime}; +use rustls::{ + CertificateError, ClientConfig, ClientConnection, DigitallySignedStruct, RootCertStore, + SignatureScheme, +}; + +use crate::config::{Config, SslMode}; +use crate::Error; + +/// The socket of one connection, plain or TLS. +pub(crate) struct PgStream { + pub(crate) tcp: TcpStream, + pub(crate) tls: Option>, +} + +impl PgStream { + pub(crate) fn plain(tcp: TcpStream) -> PgStream { + PgStream { tcp, tls: None } + } +} + +// Blocking (coroutine-suspending) I/O, used during startup and for the final Terminate. +impl Read for PgStream { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + match &mut self.tls { + None => self.tcp.read(buf), + Some(conn) => rustls::Stream::new(conn.as_mut(), &mut self.tcp).read(buf), + } + } +} + +impl Write for PgStream { + fn write(&mut self, buf: &[u8]) -> io::Result { + match &mut self.tls { + None => self.tcp.write(buf), + Some(conn) => rustls::Stream::new(conn.as_mut(), &mut self.tcp).write(buf), + } + } + + fn flush(&mut self) -> io::Result<()> { + match &mut self.tls { + None => self.tcp.flush(), + Some(conn) => rustls::Stream::new(conn.as_mut(), &mut self.tcp).flush(), + } + } +} + +/// Sends SSLRequest when the configuration asks for TLS and completes the handshake. `host` names the server for +/// SNI and `verify-full`. +pub(crate) fn negotiate( + mut tcp: TcpStream, + config: &Config, + host: Option<&str>, +) -> Result { + let mode = config.get_ssl_mode(); + if mode == SslMode::Disable { + return Ok(PgStream::plain(tcp)); + } + let mut buf = BytesMut::new(); + frontend::ssl_request(&mut buf); + tcp.write_all(&buf).map_err(Error::io)?; + tcp.flush().map_err(Error::io)?; + let mut answer = [0u8; 1]; + tcp.read_exact(&mut answer).map_err(Error::io)?; + if answer[0] != b'S' { + if mode == SslMode::Prefer { + return Ok(PgStream::plain(tcp)); + } + return Err(Error::tls("the server does not support TLS".into())); + } + + let tls_config = client_config(config)?; + let name = server_name(host, mode)?; + let mut conn = + ClientConnection::new(Arc::new(tls_config), name).map_err(|e| Error::tls(Box::new(e)))?; + while conn.is_handshaking() { + conn.complete_io(&mut tcp) + .map_err(|e| Error::tls(Box::new(e)))?; + } + Ok(PgStream { + tcp, + tls: Some(Box::new(conn)), + }) +} + +fn server_name(host: Option<&str>, mode: SslMode) -> Result, Error> { + match host { + Some(h) => ServerName::try_from(h.to_owned()) + .map_err(|e| Error::tls(format!("invalid server name {h:?}: {e}").into())), + None if mode == SslMode::VerifyFull => { + Err(Error::tls("verify-full needs a host name".into())) + } + // Not checked in the other modes; SNI only. + None => Ok(ServerName::try_from("localhost").expect("static name")), + } +} + +fn provider() -> Arc { + Arc::new(rustls::crypto::ring::default_provider()) +} + +fn client_config(config: &Config) -> Result { + let provider = provider(); + let mut mode = config.get_ssl_mode(); + let roots = match config.get_ssl_root_cert() { + Some(path) => Some(root_store(path)?), + None => None, + }; + // libpq: require with a root certificate behaves as verify-ca. + if mode == SslMode::Require && roots.is_some() { + mode = SslMode::VerifyCa; + } + let builder = ClientConfig::builder_with_provider(provider.clone()) + .with_safe_default_protocol_versions() + .map_err(|e| Error::tls(Box::new(e)))?; + let verifier: Arc = match mode { + SslMode::Prefer | SslMode::Require => Arc::new(NoServerVerification(provider.clone())), + SslMode::VerifyCa | SslMode::VerifyFull => { + let roots = roots.ok_or_else(|| { + Error::config(format!("sslmode={} needs sslrootcert", mode.as_str()).into()) + })?; + let webpki = + WebPkiServerVerifier::builder_with_provider(Arc::new(roots), provider.clone()) + .build() + .map_err(|e| Error::tls(Box::new(e)))?; + if mode == SslMode::VerifyFull { + webpki + } else { + Arc::new(ChainOnly(webpki)) + } + } + SslMode::Disable => unreachable!("no TLS configuration for sslmode=disable"), + }; + let builder = builder + .dangerous() + .with_custom_certificate_verifier(verifier); + let tls = match (config.get_ssl_cert(), config.get_ssl_key()) { + (Some(cert), Some(key)) => builder + .with_client_auth_cert(certs(cert)?, private_key(key)?) + .map_err(|e| Error::tls(Box::new(e)))?, + (None, None) => builder.with_no_client_auth(), + _ => return Err(Error::config("sslcert and sslkey go together".into())), + }; + Ok(tls) +} + +fn pem(path: &Path) -> Result, Error> { + File::open(path) + .map(BufReader::new) + .map_err(|e| Error::config(format!("{}: {e}", path.display()).into())) +} + +fn certs(path: &Path) -> Result>, Error> { + let certs = rustls_pemfile::certs(&mut pem(path)?) + .collect::, _>>() + .map_err(|e| Error::config(format!("{}: {e}", path.display()).into()))?; + if certs.is_empty() { + return Err(Error::config( + format!("{}: no certificate", path.display()).into(), + )); + } + Ok(certs) +} + +fn private_key(path: &Path) -> Result, Error> { + rustls_pemfile::private_key(&mut pem(path)?) + .map_err(|e| Error::config(format!("{}: {e}", path.display()).into()))? + .ok_or_else(|| Error::config(format!("{}: no private key", path.display()).into())) +} + +fn root_store(path: &Path) -> Result { + let mut roots = RootCertStore::empty(); + for cert in certs(path)? { + roots + .add(cert) + .map_err(|e| Error::config(format!("{}: {e}", path.display()).into()))?; + } + Ok(roots) +} + +/// `prefer` / `require`: encrypt, but accept any server certificate (handshake signatures are still checked). +#[derive(Debug)] +struct NoServerVerification(Arc); + +impl ServerCertVerifier for NoServerVerification { + fn verify_server_cert( + &self, + _end_entity: &CertificateDer<'_>, + _intermediates: &[CertificateDer<'_>], + _server_name: &ServerName<'_>, + _ocsp_response: &[u8], + _now: UnixTime, + ) -> Result { + Ok(ServerCertVerified::assertion()) + } + + fn verify_tls12_signature( + &self, + message: &[u8], + cert: &CertificateDer<'_>, + dss: &DigitallySignedStruct, + ) -> Result { + verify_tls12_signature( + message, + cert, + dss, + &self.0.signature_verification_algorithms, + ) + } + + fn verify_tls13_signature( + &self, + message: &[u8], + cert: &CertificateDer<'_>, + dss: &DigitallySignedStruct, + ) -> Result { + verify_tls13_signature( + message, + cert, + dss, + &self.0.signature_verification_algorithms, + ) + } + + fn supported_verify_schemes(&self) -> Vec { + self.0.signature_verification_algorithms.supported_schemes() + } +} + +/// `verify-ca`: the chain must lead to the root certificate; the host name is not checked. +#[derive(Debug)] +struct ChainOnly(Arc); + +impl ServerCertVerifier for ChainOnly { + fn verify_server_cert( + &self, + end_entity: &CertificateDer<'_>, + intermediates: &[CertificateDer<'_>], + server_name: &ServerName<'_>, + ocsp_response: &[u8], + now: UnixTime, + ) -> Result { + match self + .0 + .verify_server_cert(end_entity, intermediates, server_name, ocsp_response, now) + { + Err(rustls::Error::InvalidCertificate(CertificateError::NotValidForName)) + | Err(rustls::Error::InvalidCertificate(CertificateError::NotValidForNameContext { + .. + })) => Ok(ServerCertVerified::assertion()), + other => other, + } + } + + fn verify_tls12_signature( + &self, + message: &[u8], + cert: &CertificateDer<'_>, + dss: &DigitallySignedStruct, + ) -> Result { + self.0.verify_tls12_signature(message, cert, dss) + } + + fn verify_tls13_signature( + &self, + message: &[u8], + cert: &CertificateDer<'_>, + dss: &DigitallySignedStruct, + ) -> Result { + self.0.verify_tls13_signature(message, cert, dss) + } + + fn supported_verify_schemes(&self) -> Vec { + self.0.supported_verify_schemes() + } +} + +fn closed() -> io::Error { + io::Error::new(io::ErrorKind::BrokenPipe, "closed") +} + +/// Non-blocking: moves `write_buf` into rustls and as much TLS as the socket takes onto it. +pub(crate) fn tls_write( + conn: &mut ClientConnection, + sock: &mut impl Write, + write_buf: &mut BytesMut, +) -> io::Result<()> { + loop { + if write_buf.has_remaining() { + let n = conn.writer().write(write_buf.chunk())?; + write_buf.advance(n); + } + if !conn.wants_write() { + if write_buf.has_remaining() { + continue; + } + return Ok(()); + } + match conn.write_tls(sock) { + Ok(0) => return Err(closed()), + Ok(_) => {} + Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(()), + Err(e) => return Err(e), + } + } +} + +/// Non-blocking: reads TLS records from the socket until it would block and appends the plaintext to `read_buf`. +/// Returns true when the socket would block (the caller waits for I/O). +pub(crate) fn tls_read( + conn: &mut ClientConnection, + sock: &mut (impl Read + Write), + read_buf: &mut BytesMut, +) -> io::Result { + loop { + let eof = match conn.read_tls(sock) { + Ok(0) => true, + Ok(_) => false, + Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(true), + Err(e) => return Err(e), + }; + let state = conn + .process_new_packets() + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + let mut plain = state.plaintext_bytes_to_read(); + while plain > 0 { + if read_buf.capacity() - read_buf.len() < plain { + read_buf.reserve(plain); + } + let dst: &mut [u8] = unsafe { std::mem::transmute(read_buf.chunk_mut()) }; + let want = plain.min(dst.len()); + let n = conn.reader().read(&mut dst[..want])?; + if n == 0 { + break; + } + unsafe { read_buf.advance_mut(n) }; + plain -= n; + } + // Alerts, key updates: let rustls answer now rather than on the next request. + while conn.wants_write() { + match conn.write_tls(sock) { + Ok(_) => {} + Err(e) if e.kind() == io::ErrorKind::WouldBlock => break, + Err(e) => return Err(e), + } + } + if eof || state.peer_has_closed() { + return Err(closed()); + } + } +} diff --git a/tests/severed_connection.rs b/tests/severed_connection.rs new file mode 100644 index 0000000..c6b46bf --- /dev/null +++ b/tests/severed_connection.rs @@ -0,0 +1,183 @@ +//! Regression tests for the pending-response starvation defect (2026-08-29 +//! lifeguard pool worker wedge, PriceWhisperer /tickers outage). +//! +//! Mechanism under test: when a connection's I/O loop exits (error / EOF / +//! shutdown) every pending in-flight request must be FAILED — the caller's +//! `Responses::next` must return `Err`, never block forever — and subsequent +//! sends on the dead client must fail fast instead of queueing onto a channel +//! no coroutine will ever drain. +//! +//! The tests sever the TCP path abruptly through a local one-shot proxy, so +//! the server never gets a chance to send a clean ErrorResponse first (that +//! cleaner path already worked; the abrupt cut is the wedge). +//! +//! Requires a PostgreSQL at 127.0.0.1:$TEST_PG_PORT (default 5432) accepting +//! user=postgres password=postgres. + +use std::io::{Read, Write}; +use std::net::{Shutdown, TcpListener, TcpStream}; +use std::sync::mpsc; +use std::thread; +use std::time::{Duration, Instant}; + +use may_postgres::{Client, Config}; + +/// How long a failed pending query may take to surface as `Err`. Generous for +/// CI; the point is "bounded", versus the pre-fix behavior of forever. +const FAIL_BOUND: Duration = Duration::from_secs(5); + +fn pg_port() -> u16 { + std::env::var("TEST_PG_PORT") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(5432) +} + +/// One-shot TCP proxy: accepts a single client connection, pipes it to the +/// real server, and severs both directions abruptly on command. +struct SeverableProxy { + addr: std::net::SocketAddr, + conns: mpsc::Receiver<(TcpStream, TcpStream)>, +} + +impl SeverableProxy { + fn start(backend_port: u16) -> SeverableProxy { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let (tx, rx) = mpsc::channel(); + thread::spawn(move || { + let (client_sock, _) = listener.accept().unwrap(); + let server_sock = + TcpStream::connect(("127.0.0.1", backend_port)).expect("backend postgres"); + client_sock.set_nodelay(true).ok(); + server_sock.set_nodelay(true).ok(); + // Hand clones to the test so it can sever; pump bytes with the originals. + tx.send(( + client_sock.try_clone().unwrap(), + server_sock.try_clone().unwrap(), + )) + .unwrap(); + let c2 = client_sock.try_clone().unwrap(); + let s2 = server_sock.try_clone().unwrap(); + thread::spawn(move || pump(client_sock, server_sock)); + thread::spawn(move || pump(s2, c2)); + }); + SeverableProxy { addr, conns: rx } + } + + /// Abruptly closes both sides of the proxied connection. + fn sever(&self) { + let (c, s) = self + .conns + .recv_timeout(Duration::from_secs(10)) + .expect("proxy accepted a connection"); + c.shutdown(Shutdown::Both).ok(); + s.shutdown(Shutdown::Both).ok(); + } +} + +fn pump(mut from: TcpStream, mut to: TcpStream) { + let mut buf = [0u8; 8192]; + loop { + match from.read(&mut buf) { + Ok(0) | Err(_) => { + to.shutdown(Shutdown::Both).ok(); + return; + } + Ok(n) => { + if to.write_all(&buf[..n]).is_err() { + from.shutdown(Shutdown::Both).ok(); + return; + } + } + } + } +} + +fn connect_via(addr: std::net::SocketAddr) -> Client { + let socket = may::net::TcpStream::connect(addr).unwrap(); + let config = "user=postgres password=postgres dbname=postgres" + .parse::() + .unwrap(); + config.connect_raw(socket).expect("connect through proxy") +} + +/// A query that is in flight when the transport dies must return `Err` within +/// a bound — not hang forever (the lifeguard pool-worker wedge). +#[test] +fn pending_query_fails_when_connection_severed_mid_flight() { + let proxy = SeverableProxy::start(pg_port()); + let client = connect_via(proxy.addr); + + let (done_tx, done_rx) = mpsc::channel(); + let worker = thread::spawn(move || { + // In flight for 30s unless the severance fails it early. + let result = client.query_one("SELECT pg_sleep(30)", &[]); + done_tx.send(result.map(|_| ())).ok(); + client // keep the client alive for the second phase + }); + + // Let the query hit the wire, then cut the transport abruptly. + thread::sleep(Duration::from_millis(500)); + let start = Instant::now(); + proxy.sever(); + + let result = done_rx + .recv_timeout(FAIL_BOUND) + .expect("pending query must fail within a bound after severance, not hang"); + let waited = start.elapsed(); + let err = result.expect_err("severed mid-flight query cannot succeed"); + assert!( + waited < FAIL_BOUND, + "query failed but only after {waited:?}" + ); + // The error must read as a dead-connection error (lifeguard's heal + // classifier keys on is_closed / io kinds / SQLSTATE 08***). + assert!( + err.is_closed() || err.to_string().contains("io error"), + "expected a connection-death error, got: {err:?}" + ); + + // Phase 2: the client is now dead — new queries fail fast, not block. + let client = worker.join().unwrap(); + assert!( + client.is_closed(), + "client must report closed after severance" + ); + let start = Instant::now(); + let again = client.query_one("SELECT 1", &[]); + assert!(again.is_err(), "queries on a dead client must error"); + assert!( + start.elapsed() < Duration::from_secs(2), + "queries on a dead client must fail fast, took {:?}", + start.elapsed() + ); +} + +/// Same wedge, but for a request QUEUED (not yet written) when the loop dies: +/// the exit drain must fail queued requests too. +#[test] +fn queued_request_fails_after_connection_death() { + let proxy = SeverableProxy::start(pg_port()); + let client = connect_via(proxy.addr); + + // Prove liveness, then kill the transport with nothing in flight. + client.query_one("SELECT 1", &[]).expect("warmup"); + proxy.sever(); + + // Give the I/O loop a moment to observe the death. + let deadline = Instant::now() + Duration::from_secs(5); + while !client.is_closed() && Instant::now() < deadline { + thread::sleep(Duration::from_millis(50)); + } + assert!(client.is_closed(), "loop must mark the connection dead"); + + let start = Instant::now(); + let result = client.query_one("SELECT 1", &[]); + assert!(result.is_err(), "query on dead connection must error"); + assert!( + start.elapsed() < Duration::from_secs(2), + "must fail fast, took {:?}", + start.elapsed() + ); +} diff --git a/tests/test/main.rs b/tests/test/main.rs index 541c8a4..9b82f4e 100644 --- a/tests/test/main.rs +++ b/tests/test/main.rs @@ -649,6 +649,119 @@ fn a_client_that_never_listens_receives_nothing() { ); } +#[test] +fn next_notification_wakes_a_waiting_client_without_polling() { + let client = connect("user=postgres"); + client.batch_execute("LISTEN test_wakeup").unwrap(); + + // Raise the notification from another coroutine after the waiter is parked; + // the waiter issues no request in the meantime, so only the connection's + // own read can deliver it. + let notify = may::go!(|| { + may::coroutine::sleep(std::time::Duration::from_millis(200)); + let notifier = connect("user=postgres"); + notifier + .batch_execute("NOTIFY test_wakeup, 'woken'") + .unwrap(); + }); + + let started = std::time::Instant::now(); + let n = client + .next_notification(std::time::Duration::from_secs(10)) + .expect("the notification should wake the waiter"); + let waited = started.elapsed(); + notify.join().unwrap(); + + assert_eq!(n.payload(), "woken"); + assert!( + waited >= std::time::Duration::from_millis(150), + "returned before the notification was raised: {waited:?}" + ); + assert!( + waited < std::time::Duration::from_secs(5), + "the waiter must wake on arrival, not at its timeout: {waited:?}" + ); +} + +#[test] +fn next_notification_returns_queued_notifications_first_and_times_out_empty() { + let client = connect("user=postgres"); + client.batch_execute("LISTEN test_queued").unwrap(); + let notifier = connect("user=postgres"); + notifier + .batch_execute("NOTIFY test_queued, 'a'; NOTIFY test_queued, 'b';") + .unwrap(); + + let a = client + .next_notification(std::time::Duration::from_secs(5)) + .expect("first"); + let b = client + .next_notification(std::time::Duration::from_secs(5)) + .expect("second"); + assert_eq!((a.payload(), b.payload()), ("a", "b")); + + let started = std::time::Instant::now(); + assert!(client + .next_notification(std::time::Duration::from_millis(100)) + .is_none()); + let waited = started.elapsed(); + assert!( + waited >= std::time::Duration::from_millis(90) + && waited < std::time::Duration::from_secs(2), + "an empty wait lasts about its timeout: {waited:?}" + ); +} + +#[test] +fn next_notification_returns_for_sub_millisecond_timeouts() { + // In a coroutine, so the wait parks on may's timer; a remaining time under + // 1 ms must not turn into "no timeout". + let (tx, rx) = std::sync::mpsc::channel(); + let _waiter = may::go!(move || { + let client = connect("user=postgres"); + client.batch_execute("LISTEN test_short_wait").unwrap(); + for us in [1u64, 100, 500, 999, 1500] { + assert!(client + .next_notification(std::time::Duration::from_micros(us)) + .is_none()); + } + tx.send(()).unwrap(); + }); + assert!( + rx.recv_timeout(std::time::Duration::from_secs(10)).is_ok(), + "a sub-millisecond wait did not return" + ); +} + +#[test] +fn next_notification_returns_when_the_connection_closes() { + let client = connect("user=postgres"); + client.batch_execute("LISTEN test_closed").unwrap(); + let pid: i32 = client + .query_one("SELECT pg_backend_pid()", &[]) + .unwrap() + .get(0); + + let killer = may::go!(move || { + may::coroutine::sleep(std::time::Duration::from_millis(200)); + let admin = connect("user=postgres"); + admin + .query_one("SELECT pg_terminate_backend($1)", &[&pid]) + .unwrap(); + }); + + let started = std::time::Instant::now(); + assert!(client + .next_notification(std::time::Duration::from_secs(10)) + .is_none()); + killer.join().unwrap(); + assert!( + started.elapsed() < std::time::Duration::from_secs(5), + "a closed connection must release the waiter, not hold it to its timeout" + ); + assert!(client.is_closed()); +} + /// Poll for a notification, giving the connection round trips to surface it. /// /// Notifications arrive asynchronously: the server sends them when it chooses, @@ -665,6 +778,54 @@ fn wait_for_notification(client: &Client) -> Option None } +#[test] +fn row_stream_reports_rows_affected() { + let client = connect("user=postgres"); + client + .batch_execute( + "CREATE TEMPORARY TABLE rows_affected_test (id INT PRIMARY KEY, n INT); + INSERT INTO rows_affected_test VALUES (1, 0), (2, 0), (3, 0);", + ) + .unwrap(); + + // A write without RETURNING: no rows, but a count, as `execute` reports. + let stmt = client + .prepare("UPDATE rows_affected_test SET n = n + 1 WHERE id <= $1") + .unwrap(); + let mut stream = client.query_raw(&stmt, &[&2i32]).unwrap(); + assert_eq!( + stream.rows_affected(), + None, + "unknown until the stream ends" + ); + assert!(stream.next().is_none()); + assert_eq!(stream.rows_affected(), Some(2)); + + // With RETURNING: the rows and the count. + let stmt = client + .prepare("UPDATE rows_affected_test SET n = n + 1 RETURNING id") + .unwrap(); + let mut stream = client.query_raw(&stmt, &[]).unwrap(); + let rows = stream.by_ref().collect::, _>>().unwrap(); + assert_eq!(rows.len(), 3); + assert_eq!(stream.rows_affected(), Some(3)); + + // Pipelined: both statements are sent before either is read, and each + // stream reports its own count. + let insert = client + .prepare("INSERT INTO rows_affected_test VALUES ($1, 0)") + .unwrap(); + let delete = client + .prepare("DELETE FROM rows_affected_test WHERE id <= $1") + .unwrap(); + let mut a = client.query_raw(&insert, &[&10i32]).unwrap(); + let mut b = client.query_raw(&delete, &[&3i32]).unwrap(); + assert!(a.next().is_none()); + assert!(b.next().is_none()); + assert_eq!(a.rows_affected(), Some(1)); + assert_eq!(b.rows_affected(), Some(3)); +} + #[test] fn query_portal() { let mut client = connect("user=postgres"); @@ -787,3 +948,36 @@ fn query_opt() { .unwrap(); client.query_one("SELECT * FROM foo", &[]).err().unwrap(); } + +#[test] +fn dropping_a_client_closes_the_server_backend() { + // The probe connections carry their own application_name so this count is + // immune to every other test opening and dropping connections in parallel. + const TAG: &str = "drop_leak_probe"; + let observer = connect("user=postgres dbname=postgres"); + let backends = || -> i64 { + observer + .query_one( + "SELECT count(*) FROM pg_stat_activity WHERE application_name = $1", + &[&TAG], + ) + .unwrap() + .get(0) + }; + + for _ in 0..5 { + let client = connect(&format!( + "user=postgres dbname=postgres application_name={TAG}" + )); + client.query("SELECT 1", &[]).unwrap(); + drop(client); + // The server reaps the backend once the socket closes; give it a moment. + std::thread::sleep(std::time::Duration::from_millis(250)); + } + + let after = backends(); + assert_eq!( + after, 0, + "{after} dropped client(s) left their backend open" + ); +} diff --git a/tests/test/parse.rs b/tests/test/parse.rs index 507ec56..2356903 100644 --- a/tests/test/parse.rs +++ b/tests/test/parse.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use may_postgres::config::{Config, TargetSessionAttrs}; +use may_postgres::config::{Config, SslMode, TargetSessionAttrs}; fn check(s: &str, config: &Config) { assert_eq!(s.parse::().expect(s), *config, "`{}`", s); @@ -28,11 +28,12 @@ fn pairs_ws() { #[test] fn settings() { check( - "connect_timeout=3 keepalives=0 keepalives_idle=30 target_session_attrs=read-write", + "connect_timeout=3 keepalives=0 keepalives_idle=30 sslmode=disable target_session_attrs=read-write", Config::new() .connect_timeout(Duration::from_secs(3)) .keepalives(false) .keepalives_idle(Duration::from_secs(30)) + .ssl_mode(SslMode::Disable) .target_session_attrs(TargetSessionAttrs::ReadWrite), ); } @@ -65,14 +66,15 @@ fn url() { .port(5432), ); check( - "postgresql://other@localhost/otherdb?connect_timeout=10&application_name=myapp", + "postgresql://other@localhost/otherdb?connect_timeout=10&application_name=myapp&sslmode=disable", Config::new() .user("other") .host("localhost") .port(5432) .dbname("otherdb") .connect_timeout(Duration::from_secs(10)) - .application_name("myapp"), + .application_name("myapp") + .ssl_mode(SslMode::Disable), ); check( "postgresql://host1:123,host2:456/somedb?target_session_attrs=any&application_name=myapp", @@ -119,3 +121,20 @@ fn url() { .dbname("dbname"), ) } + +#[test] +fn ssl_modes() { + for (mode, want) in [ + ("disable", SslMode::Disable), + ("prefer", SslMode::Prefer), + ("require", SslMode::Require), + ("verify-ca", SslMode::VerifyCa), + ("verify-full", SslMode::VerifyFull), + ] { + let url = format!("postgresql://localhost/postgres?sslmode={mode}"); + assert_eq!(url.parse::().unwrap().get_ssl_mode(), want, "{url}"); + } + assert!("postgresql://localhost/postgres?sslmode=invalid" + .parse::() + .is_err()); +} diff --git a/tests/test/types/chrono_04.rs b/tests/test/types/chrono_04.rs index 1e9f308..97fd5a0 100644 --- a/tests/test/types/chrono_04.rs +++ b/tests/test/types/chrono_04.rs @@ -1,4 +1,4 @@ -use chrono_04::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc}; +use ::chrono_04::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc}; use may_postgres::types::{Date, Timestamp}; use crate::types::test_type; diff --git a/tests/tls.rs b/tests/tls.rs new file mode 100644 index 0000000..44c997e --- /dev/null +++ b/tests/tls.rs @@ -0,0 +1,271 @@ +//! TLS against real servers, set up by tests/tls/run.sh (two containers and a throwaway CA): +//! +//! * TLS_PG_PORT: ssl=on, certificate for DNS:localhost only (so 127.0.0.1 fails verify-full), pg_hba +//! `hostssl ... certuser ... clientcert=verify-full`, `hostssl all all scram-sha-256`, plaintext rejected. +//! * PLAIN_PG_PORT: ssl=off. +//! * TLS_TEST_DIR: ca.crt, client.crt + client.key (CN=certuser), other.crt + other.key (CN=certuser, from +//! other-ca.crt, another CA). +//! +//! Skipped (passes) when TLS_PG_PORT is not set. + +use may_postgres::{config::SslMode, Client, Config}; + +struct Env { + tls: u16, + plain: u16, + dir: String, +} + +fn env() -> Option { + let tls = std::env::var("TLS_PG_PORT").ok()?.parse().ok()?; + let plain = std::env::var("PLAIN_PG_PORT").ok()?.parse().ok()?; + let dir = std::env::var("TLS_TEST_DIR").ok()?; + Some(Env { tls, plain, dir }) +} + +fn connect(s: &str) -> Result { + may_postgres::connect(s) +} + +fn one(c: &Client) -> i32 { + c.query_one("SELECT 1", &[]).unwrap().get(0) +} + +#[test] +fn require_encrypts_and_queries() { + let Some(e) = env() else { return }; + let c = connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=require", + e.tls + )) + .unwrap(); + assert_eq!(one(&c), 1); + let ssl: bool = c + .query_one( + "SELECT ssl FROM pg_stat_ssl WHERE pid = pg_backend_pid()", + &[], + ) + .unwrap() + .get(0); + assert!(ssl, "pg_stat_ssl says the connection is not encrypted"); +} + +#[test] +fn large_results_and_many_round_trips_over_tls() { + let Some(e) = env() else { return }; + let c = connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=require", + e.tls + )) + .unwrap(); + // Many TLS records per response: 200k rows, and one 4 MB value. + let rows = c + .query( + "SELECT g, md5(g::text) FROM generate_series(1, 200000) g", + &[], + ) + .unwrap(); + assert_eq!(rows.len(), 200000); + let big: String = c + .query_one("SELECT repeat('x', 4 * 1024 * 1024)", &[]) + .unwrap() + .get(0); + assert_eq!(big.len(), 4 * 1024 * 1024); + // Large request: a 2 MB parameter. + let payload = "y".repeat(2 * 1024 * 1024); + let n: i32 = c + .query_one("SELECT length($1::text)", &[&payload]) + .unwrap() + .get(0); + assert_eq!(n as usize, payload.len()); + for i in 0..500 { + let v: i32 = c.query_one("SELECT $1::int4", &[&i]).unwrap().get(0); + assert_eq!(v, i); + } +} + +#[test] +fn concurrent_coroutines_share_a_tls_connection() { + let Some(e) = env() else { return }; + let c = connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=require", + e.tls + )) + .unwrap(); + let handles: Vec<_> = (0..32) + .map(|i| { + let c = c.clone(); + may::go!(move || { + for j in 0..50 { + let v: i32 = c + .query_one("SELECT $1::int4", &[&(i * 1000 + j)]) + .unwrap() + .get(0); + assert_eq!(v, i * 1000 + j); + } + }) + }) + .collect(); + for h in handles { + h.join().unwrap(); + } +} + +#[test] +fn verify_full_checks_the_name() { + let Some(e) = env() else { return }; + let ca = format!("{}/ca.crt", e.dir); + let ok = connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=verify-full sslrootcert={ca}", + e.tls + )) + .unwrap(); + assert_eq!(one(&ok), 1); + // The certificate names localhost only. + let err = connect(&format!( + "host=127.0.0.1 port={} user=pw password=pw dbname=postgres sslmode=verify-full sslrootcert={ca}", + e.tls + )) + .err() + .expect("verify-full must reject a certificate for another name"); + assert!(err.to_string().contains("TLS"), "{err}"); + // verify-ca does not check the name. + let c = connect(&format!( + "host=127.0.0.1 port={} user=pw password=pw dbname=postgres sslmode=verify-ca sslrootcert={ca}", + e.tls + )) + .unwrap(); + assert_eq!(one(&c), 1); +} + +#[test] +fn verify_ca_rejects_another_ca() { + let Some(e) = env() else { return }; + let other = format!("{}/other-ca.crt", e.dir); + assert!(connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=verify-ca sslrootcert={other}", + e.tls + )) + .is_err()); + // require + sslrootcert behaves as verify-ca (libpq). + assert!(connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=require sslrootcert={other}", + e.tls + )) + .is_err()); +} + +#[test] +fn verify_modes_need_a_root_certificate() { + let Some(e) = env() else { return }; + let err = connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=verify-full", + e.tls + )) + .err() + .unwrap(); + assert!( + err.to_string().contains("sslrootcert") || format!("{err:?}").contains("sslrootcert"), + "{err:?}" + ); +} + +#[test] +fn client_certificate() { + let Some(e) = env() else { return }; + let d = &e.dir; + let c = connect(&format!( + "host=localhost port={} user=certuser password=cu dbname=postgres sslmode=verify-full \ + sslrootcert={d}/ca.crt sslcert={d}/client.crt sslkey={d}/client.key", + e.tls + )) + .unwrap(); + assert_eq!(one(&c), 1); + // Without a certificate the server refuses (clientcert=verify-full) even with the password. + assert!(connect(&format!( + "host=localhost port={} user=certuser password=cu dbname=postgres sslmode=verify-full sslrootcert={d}/ca.crt", + e.tls + )) + .is_err()); + // A certificate from another CA is refused. + assert!(connect(&format!( + "host=localhost port={} user=certuser password=cu dbname=postgres sslmode=verify-full \ + sslrootcert={d}/ca.crt sslcert={d}/other.crt sslkey={d}/other.key", + e.tls + )) + .is_err()); + // Config builder API. + let mut cfg = Config::new(); + cfg.host("localhost") + .port(e.tls) + .user("certuser") + .password("cu") + .dbname("postgres") + .ssl_mode(SslMode::VerifyFull) + .ssl_root_cert(format!("{d}/ca.crt")) + .ssl_cert(format!("{d}/client.crt")) + .ssl_key(format!("{d}/client.key")); + assert_eq!(one(&cfg.connect().unwrap()), 1); +} + +#[test] +fn plaintext_is_refused_by_the_tls_server() { + let Some(e) = env() else { return }; + assert!(connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=disable", + e.tls + )) + .is_err()); +} + +#[test] +fn prefer_falls_back_and_require_fails_on_a_plain_server() { + let Some(e) = env() else { return }; + let c = connect(&format!( + "host=localhost port={} user=postgres password=postgres dbname=postgres", + e.plain + )) + .unwrap(); + assert_eq!(one(&c), 1); + let c = connect(&format!( + "host=localhost port={} user=postgres password=postgres dbname=postgres sslmode=disable", + e.plain + )) + .unwrap(); + assert_eq!(one(&c), 1); + let err = connect(&format!( + "host=localhost port={} user=postgres password=postgres dbname=postgres sslmode=require", + e.plain + )) + .err() + .expect("require against a server without TLS must fail"); + assert!(err.to_string().contains("TLS"), "{err}"); +} + +#[test] +fn drop_closes_the_tls_connection_cleanly() { + let Some(e) = env() else { return }; + let admin = connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=require application_name=admin", + e.tls + )) + .unwrap(); + for _ in 0..20 { + let c = connect(&format!( + "host=localhost port={} user=pw password=pw dbname=postgres sslmode=require application_name=churn", + e.tls + )) + .unwrap(); + assert_eq!(one(&c), 1); + drop(c); + } + std::thread::sleep(std::time::Duration::from_millis(500)); + let left: i64 = admin + .query_one( + "SELECT count(*) FROM pg_stat_activity WHERE application_name = 'churn'", + &[], + ) + .unwrap() + .get(0); + assert_eq!(left, 0, "backends left behind after drop"); +} diff --git a/tests/tls/run.sh b/tests/tls/run.sh new file mode 100755 index 0000000..d342fd4 --- /dev/null +++ b/tests/tls/run.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +# TLS integration tests (tests/tls.rs) against two throwaway Postgres containers: one TLS-only with client +# certificates for `certuser`, one plaintext. Generates its own CAs. Needs docker and openssl. +# +# tests/tls/run.sh # PG_IMAGE=postgres:17 TLS_PG_PORT=55440 PLAIN_PG_PORT=55441 to override +set -euo pipefail + +here="$(cd "$(dirname "$0")/../.." && pwd)" +IMAGE="${PG_IMAGE:-postgres:17}" +TLS_PORT="${TLS_PG_PORT:-55440}" +PLAIN_PORT="${PLAIN_PG_PORT:-55441}" +dir="$(mktemp -d)" +tls_name="may-pg-tls-$$" plain_name="may-pg-plain-$$" +cleanup() { docker rm -f "$tls_name" "$plain_name" >/dev/null 2>&1 || true; rm -rf "$dir"; } +trap cleanup EXIT + +cd "$dir" +ec() { openssl ecparam -name prime256v1 -genkey -noout -out "$1" 2>/dev/null; } +ec ca.key +openssl req -x509 -new -key ca.key -subj /CN=test-ca -days 2 -out ca.crt 2>/dev/null +ec other-ca.key +openssl req -x509 -new -key other-ca.key -subj /CN=other-ca -days 2 -out other-ca.crt 2>/dev/null +sign() { # $1 name $2 CN $3 CA prefix $4 extensions + ec "$1.key" + openssl req -new -key "$1.key" -subj "/CN=$2" -out "$1.csr" 2>/dev/null + printf '%s\n' "$4" > "$1.ext" + openssl x509 -req -in "$1.csr" -CA "$3.crt" -CAkey "$3.key" -CAcreateserial -days 2 -extfile "$1.ext" -out "$1.crt" 2>/dev/null +} +sign server localhost ca "subjectAltName=DNS:localhost" +sign client certuser ca "extendedKeyUsage=clientAuth" +sign other certuser other-ca "extendedKeyUsage=clientAuth" +cat > pg_hba.conf <<'EOF' +local all all trust +hostssl all certuser all scram-sha-256 clientcert=verify-full +hostssl all all all scram-sha-256 +host all all all reject +EOF +mkdir srv +cp ca.crt server.crt server.key pg_hba.conf srv/ +chmod 644 srv/* + +docker run -d --name "$tls_name" -p "127.0.0.1:${TLS_PORT}:5432" -e POSTGRES_PASSWORD=postgres -v "$dir/srv:/src:ro" \ + --entrypoint bash "$IMAGE" -c 'mkdir -p /certs; cp /src/* /certs/; chown -R postgres /certs; chmod 600 /certs/server.key; + exec docker-entrypoint.sh postgres -c ssl=on -c ssl_cert_file=/certs/server.crt -c ssl_key_file=/certs/server.key \ + -c ssl_ca_file=/certs/ca.crt -c hba_file=/certs/pg_hba.conf -c listen_addresses=*' >/dev/null +docker run -d --name "$plain_name" -p "127.0.0.1:${PLAIN_PORT}:5432" -e POSTGRES_PASSWORD=postgres "$IMAGE" >/dev/null + +ready() { for _ in $(seq 1 60); do docker exec "$1" pg_isready -h 127.0.0.1 -U postgres >/dev/null 2>&1 && return 0; sleep 1; done; docker logs "$1"; return 1; } +ready "$tls_name"; ready "$plain_name" +# The entrypoint restarts the server after init; wait for the final one. +sleep 2; ready "$tls_name"; ready "$plain_name" +docker exec "$tls_name" psql -U postgres -qc "CREATE ROLE pw LOGIN PASSWORD 'pw'; CREATE ROLE certuser LOGIN PASSWORD 'cu'; GRANT pg_read_all_stats TO pw;" >/dev/null + +cd "$here" +TLS_PG_PORT="$TLS_PORT" PLAIN_PG_PORT="$PLAIN_PORT" TLS_TEST_DIR="$dir" cargo test --test tls -- --test-threads=1 "$@"