From 605d8bef58ec7bafcda0e710009ef07837ea89b7 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Sun, 16 Aug 2026 11:10:46 +0800 Subject: [PATCH 1/5] feat: client api secret key and cancel support --- src/api/client/auth.rs | 6 ++- src/api/client/mod.rs | 8 ++++ src/tokio/client.rs | 101 +++++++++++++++++++++++++++++++++-------- 3 files changed, 94 insertions(+), 21 deletions(-) diff --git a/src/api/client/auth.rs b/src/api/client/auth.rs index 6327b3f5..bcf91214 100644 --- a/src/api/client/auth.rs +++ b/src/api/client/auth.rs @@ -10,7 +10,7 @@ use crate::error::{ErrorInfo, PgWireClientError, PgWireClientResult, PgWireResul use crate::messages::response::ReadyForQuery; use crate::messages::startup::{ Authentication, BackendKeyData, ParameterStatus, Password, PasswordMessageFamily, - SASLInitialResponse, SASLResponse, Startup, + SASLInitialResponse, SASLResponse, SecretKey, Startup, }; use crate::messages::{PgWireBackendMessage, PgWireFrontendMessage}; @@ -116,6 +116,8 @@ pub struct DefaultStartupHandler { server_parameters: BTreeMap, #[new(default)] process_id: Option, + #[new(default)] + secret_key: Option, } #[async_trait] @@ -237,6 +239,7 @@ impl StartupHandler for DefaultStartupHandler { C: ClientInfo + Sink + Unpin + Send, { self.process_id = Some(message.pid); + self.secret_key = Some(message.secret_key); Ok(()) } @@ -251,6 +254,7 @@ impl StartupHandler for DefaultStartupHandler { Ok(ServerInformation { parameters: self.server_parameters.clone(), process_id: self.process_id.unwrap_or(-1), + secret_key: self.secret_key.clone().unwrap_or_default(), }) } } diff --git a/src/api/client/mod.rs b/src/api/client/mod.rs index 380bd09f..d40f414f 100644 --- a/src/api/client/mod.rs +++ b/src/api/client/mod.rs @@ -9,6 +9,7 @@ use std::collections::BTreeMap; pub use config::Config; use crate::messages::ProtocolVersion; +use crate::messages::startup::SecretKey; /// A trait for fetching necessary information from Client pub trait ClientInfo { @@ -21,6 +22,12 @@ pub trait ClientInfo { /// Returns process id received from server fn process_id(&self) -> i32; + /// Returns the secret key received from the server's `BackendKeyData`. + /// + /// Together with [`ClientInfo::process_id`], this identifies the backend + /// session so a `CancelRequest` can be issued against a running query. + fn secret_key(&self) -> &SecretKey; + /// Returns client protocol version fn protocol_version(&self) -> ProtocolVersion; @@ -32,6 +39,7 @@ pub trait ClientInfo { pub struct ServerInformation { pub parameters: BTreeMap, pub process_id: i32, + pub secret_key: SecretKey, } /// Indicate the result of current request diff --git a/src/tokio/client.rs b/src/tokio/client.rs index b15c4e57..f1793c2d 100644 --- a/src/tokio/client.rs +++ b/src/tokio/client.rs @@ -25,6 +25,8 @@ use crate::api::client::config::Host; use crate::api::client::query::{ExtendedQueryClient, ExtendedQueryHandler, SimpleQueryHandler}; use crate::api::client::{ClientInfo, Config, ReadyState, ServerInformation}; use crate::error::{PgWireClientError, PgWireClientResult, PgWireError}; +use crate::messages::cancel::CancelRequest; +use crate::messages::startup::SecretKey; use crate::messages::{ DecodeContext, PgWireBackendMessage, PgWireFrontendMessage, ProtocolVersion, SslNegotiationMetaMessage, @@ -78,6 +80,10 @@ pub struct PgWireClient { socket: Framed, config: Arc, server_information: ServerInformation, + /// TLS connector retained so [`PgWireClient::cancel`] can open a second + /// secured connection to the same server. + #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))] + tls_connector: Option, } impl ClientInfo for PgWireClient { @@ -93,6 +99,10 @@ impl ClientInfo for PgWireClient { self.server_information.process_id } + fn secret_key(&self) -> &SecretKey { + &self.server_information.secret_key + } + fn protocol_version(&self) -> ProtocolVersion { self.socket.codec().decode_context.protocol_version } @@ -130,31 +140,19 @@ impl PgWireClient { where S: StartupHandler, { - // tcp connect - let mut socket = match get_addr(&config)? { - PgSocketAddr::Ip(socket_addr) => { - ClientSocket::Plain(TcpStream::connect(socket_addr).await?) - } - PgSocketAddr::Host(socket_addr) => { - ClientSocket::Plain(TcpStream::connect(socket_addr).await?) - } - #[cfg(unix)] - PgSocketAddr::Unix(socket_addr) => { - ClientSocket::Unix(UnixStream::connect(socket_addr).await?) - } - }; - if let ClientSocket::Plain(tcp_socket) = socket { - // perform ssl handshake based on postgres configuration - // if tls is not enabled, just return the socket and perform startup - // directly - socket = ssl_handshake(tcp_socket, &config, tls_connector).await?; - }; - let socket = Framed::new(socket, PgWireMessageClientCodec::default()); + // The TLS connector is retained so `cancel` can open a second secured + // connection later. When TLS is disabled there is no field to store. + #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))] + let tls_connector_for_cancel = tls_connector.clone(); + + let socket = connect_socket(&config, tls_connector).await?; let mut client = PgWireClient { socket, config: config.clone(), server_information: ServerInformation::default(), + #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))] + tls_connector: tls_connector_for_cancel, }; startup_handler.startup(&mut client).await?; @@ -173,6 +171,41 @@ impl PgWireClient { Err(PgWireClientError::UnexpectedEOF) } + /// Cancel the currently running query on this connection. + /// + /// Per the PostgreSQL wire protocol, a cancel request must be sent on a + /// **separate** connection to the same server — it carries the `pid` and + /// `secret_key` from the original connection's `BackendKeyData`. This + /// method opens that second connection (reusing this client's [`Config`] + /// and TLS connector), sends the `CancelRequest`, and closes it. + /// + /// The server sends no reply on the cancel connection. Whether the cancel + /// succeeded is observed on the original connection: the interrupted query + /// returns an error (typically `57014` / `query_canceled`). + /// + /// Returns an error only if the second connection itself could not be + /// established or the cancel message could not be written. + pub async fn cancel(&self) -> PgWireClientResult<()> { + // TLS connector is only stored when a TLS backend is enabled; without + // TLS the cancel connection is always plaintext. + #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))] + let tls_connector = self.tls_connector.clone(); + #[cfg(not(any(feature = "_ring", feature = "_aws-lc-rs")))] + let tls_connector: Option = None; + + let mut socket = connect_socket(&self.config, tls_connector).await?; + + socket + .send(PgWireFrontendMessage::CancelRequest(CancelRequest::new( + self.server_information.process_id, + self.server_information.secret_key.clone(), + ))) + .await?; + socket.close().await?; + + Ok(()) + } + /// Start a query with simple query subprotocol pub async fn simple_query( &mut self, @@ -370,6 +403,34 @@ pub(crate) async fn ssl_handshake( Ok(socket) } +/// Establish a framed connection to the server: TCP (optionually upgraded to +/// TLS) or Unix domain socket. Shared by [`PgWireClient::connect`] (which then +/// runs startup) and [`PgWireClient::cancel`] (which sends a `CancelRequest` +/// instead of a `Startup`). +async fn connect_socket( + config: &Config, + tls_connector: Option, +) -> PgWireClientResult> { + let mut socket = match get_addr(config)? { + PgSocketAddr::Ip(socket_addr) => { + ClientSocket::Plain(TcpStream::connect(socket_addr).await?) + } + PgSocketAddr::Host(socket_addr) => { + ClientSocket::Plain(TcpStream::connect(socket_addr).await?) + } + #[cfg(unix)] + PgSocketAddr::Unix(socket_addr) => { + ClientSocket::Unix(UnixStream::connect(socket_addr).await?) + } + }; + if let ClientSocket::Plain(tcp_socket) = socket { + // Perform the ssl handshake based on postgres configuration; when TLS + // is disabled `ssl_handshake` returns the plain socket unchanged. + socket = ssl_handshake(tcp_socket, config, tls_connector).await?; + } + Ok(Framed::new(socket, PgWireMessageClientCodec::default())) +} + enum PgSocketAddr { Ip(SocketAddr), Host((String, u16)), From 7c294d01b22287261c30fd97177cee81f979f5ab Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Sun, 16 Aug 2026 11:46:18 +0800 Subject: [PATCH 2/5] feat: client protocol version negotiation Following libpq, add a 3.9999 test version to request the newest minor version a server supports and exercise the negotiation path: - Config::protocol_version selects the version advertised in the startup message (default 3.0), including PROTOCOL3_9999 - StartupHandler::on_negotiate_protocol_version handles the server's NegotiateProtocolVersion response, understanding both the full 32-bit version form (PostgreSQL 18+, pgwire) and the minor-only form used by older servers, and adopts the negotiated version for the connection - PgWireClient now primes its decode context with the advertised version, so a 4-byte protocol 3.0 cancel key decodes as SecretKey::I32 instead of SecretKey::Bytes --- CHANGELOG.md | 19 ++++++++++ src/api/client/auth.rs | 44 +++++++++++++++++++-- src/api/client/config.rs | 23 +++++++++++ src/api/client/mod.rs | 9 +++++ src/api/mod.rs | 2 +- src/messages/mod.rs | 71 ++++++++++++++++++++++++++++++++++ src/messages/startup.rs | 52 ++++++++++++++++++++++++- src/tokio/client.rs | 82 ++++++++++++++++++++++++++++++++++++++++ 8 files changed, 297 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d28f2b5c..607588f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,25 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- Client API: protocol version negotiation. `Config::protocol_version` + selects the version to advertise in the startup message, defaulting to 3.0. + Following libpq, a test version `3.9999` is available to request the newest + minor version a server supports: the client handles the server's + `NegotiateProtocolVersion` response (both the full 32-bit version form used + by PostgreSQL 18+ and the minor-only form used by older servers) and adopts + the negotiated version for the rest of the connection. + +### Fixed + +- Client API: backend messages are now decoded with the rules of the protocol + version the client actually advertised, instead of always 3.2. Previously a + 4-byte protocol 3.0 cancel key was decoded as `SecretKey::Bytes` instead of + `SecretKey::I32`. + ## [0.40.7] - 2026-08-12 ### Fixed diff --git a/src/api/client/auth.rs b/src/api/client/auth.rs index bcf91214..a78473a5 100644 --- a/src/api/client/auth.rs +++ b/src/api/client/auth.rs @@ -9,8 +9,8 @@ use crate::api::auth::sasl::scram::ScramClientAuth; use crate::error::{ErrorInfo, PgWireClientError, PgWireClientResult, PgWireResult}; use crate::messages::response::ReadyForQuery; use crate::messages::startup::{ - Authentication, BackendKeyData, ParameterStatus, Password, PasswordMessageFamily, - SASLInitialResponse, SASLResponse, SecretKey, Startup, + Authentication, BackendKeyData, NegotiateProtocolVersion, ParameterStatus, Password, + PasswordMessageFamily, SASLInitialResponse, SASLResponse, SecretKey, Startup, }; use crate::messages::{PgWireBackendMessage, PgWireFrontendMessage}; @@ -43,6 +43,10 @@ pub trait StartupHandler: Send { PgWireBackendMessage::Authentication(authentication) => { self.on_authentication(client, authentication).await?; } + PgWireBackendMessage::NegotiateProtocolVersion(negotiation) => { + self.on_negotiate_protocol_version(client, negotiation) + .await?; + } PgWireBackendMessage::ParameterStatus(parameter_status) => { self.on_parameter_status(client, parameter_status).await?; } @@ -78,6 +82,34 @@ pub trait StartupHandler: Send { + Send, PgWireClientError: From<>::Error>; + /// Handle a `NegotiateProtocolVersion` message from the server. + /// + /// The default implementation adopts the negotiated version for the rest + /// of the connection. Both the full 32-bit version number used by + /// PostgreSQL 18+ and pgwire, and the minor-only form used by older + /// servers, are understood. Startup parameters reported as unrecognized + /// by the server are ignored, matching libpq's behavior for non-`_pq_` + /// options. + async fn on_negotiate_protocol_version( + &mut self, + client: &mut C, + message: NegotiateProtocolVersion, + ) -> PgWireClientResult<()> + where + C: ClientInfo + Sink + Unpin + Send, + PgWireClientError: From<>::Error>, + { + match message.negotiated_version() { + Some(version) => { + client.set_protocol_version(version); + Ok(()) + } + None => Err(PgWireClientError::UnexpectedMessage(Box::new( + PgWireBackendMessage::NegotiateProtocolVersion(message), + ))), + } + } + /// Handle a parameter status message from the server. async fn on_parameter_status( &mut self, @@ -127,11 +159,17 @@ impl StartupHandler for DefaultStartupHandler { C: ClientInfo + Sink + Unpin + Send, PgWireClientError: From<>::Error>, { - // TODO: customize protocol version let mut startup = Startup::new(); let config = client.config(); + // Advertise the configured protocol version. The connection decodes + // with the rules of the advertised version until the server lowers it + // via NegotiateProtocolVersion. + let (major, minor) = config.get_protocol_version().version_number(); + startup.protocol_number_major = major; + startup.protocol_number_minor = minor; + if let Some(application_name) = &config.application_name { startup .parameters diff --git a/src/api/client/config.rs b/src/api/client/config.rs index cda78832..2ed6bf79 100644 --- a/src/api/client/config.rs +++ b/src/api/client/config.rs @@ -17,6 +17,7 @@ use std::time::Duration; use std::{fmt, iter, mem, str}; use crate::error::PgWireClientError; +use crate::messages::ProtocolVersion; /// Properties required of a session. #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -222,6 +223,7 @@ pub struct Config { pub(crate) target_session_attrs: TargetSessionAttrs, pub(crate) channel_binding: ChannelBinding, pub(crate) load_balance_hosts: LoadBalanceHosts, + pub(crate) protocol_version: ProtocolVersion, } impl Default for Config { @@ -256,6 +258,7 @@ impl Config { target_session_attrs: TargetSessionAttrs::Any, channel_binding: ChannelBinding::Prefer, load_balance_hosts: LoadBalanceHosts::Disable, + protocol_version: ProtocolVersion::PROTOCOL3_0, } } @@ -551,6 +554,25 @@ impl Config { self.load_balance_hosts } + /// Sets the protocol version to advertise in the startup message. + /// + /// Defaults to `PROTOCOL3_0`. + /// + /// Following libpq, you can set [`ProtocolVersion::PROTOCOL3_9999`] to + /// request the newest minor version the server supports: the server + /// replies with `NegotiateProtocolVersion` and the connection proceeds at + /// the negotiated version. This is useful for exercising the negotiation + /// path against any server. + pub fn protocol_version(&mut self, protocol_version: ProtocolVersion) -> &mut Config { + self.protocol_version = protocol_version; + self + } + + /// Gets the protocol version to advertise in the startup message. + pub fn get_protocol_version(&self) -> ProtocolVersion { + self.protocol_version + } + fn param(&mut self, key: &str, value: &str) -> Result<(), PgWireClientError> { match key { "user" => { @@ -749,6 +771,7 @@ impl fmt::Debug for Config { config_dbg .field("target_session_attrs", &self.target_session_attrs) .field("channel_binding", &self.channel_binding) + .field("protocol_version", &self.protocol_version) .finish() } } diff --git a/src/api/client/mod.rs b/src/api/client/mod.rs index d40f414f..b4a434c6 100644 --- a/src/api/client/mod.rs +++ b/src/api/client/mod.rs @@ -31,6 +31,15 @@ pub trait ClientInfo { /// Returns client protocol version fn protocol_version(&self) -> ProtocolVersion; + /// Sets the protocol version in effect for this connection. + /// + /// Custom [`StartupHandler`](auth::StartupHandler) implementations should + /// call this with the version they advertise in the `Startup` message, so + /// that subsequent backend messages are decoded with the rules of that + /// version. The default startup handler and the negotiation flow handle + /// this automatically. + fn set_protocol_version(&mut self, version: ProtocolVersion); + // TODO: transaction state } diff --git a/src/api/mod.rs b/src/api/mod.rs index 8985c295..781e9afb 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -347,7 +347,7 @@ impl PidSecretKeyGenerator for RandomPidSecretKeyGenerator { let pid = self.next_pid.fetch_add(1, Ordering::Relaxed); let secret_key = match client.protocol_version() { ProtocolVersion::PROTOCOL3_0 => SecretKey::I32(rand::random::()), - ProtocolVersion::PROTOCOL3_2 => { + ProtocolVersion::PROTOCOL3_2 | ProtocolVersion::PROTOCOL3_9999 => { let mut bytes = vec![0u8; 32]; rand::fill(&mut bytes); SecretKey::Bytes(bytes.into()) diff --git a/src/messages/mod.rs b/src/messages/mod.rs index c23a0f4b..0d23a25b 100644 --- a/src/messages/mod.rs +++ b/src/messages/mod.rs @@ -17,6 +17,19 @@ pub enum ProtocolVersion { #[default] /// Protocol version 3.2 PROTOCOL3_2, + /// Protocol version 3.9999 + /// + /// This is not a real protocol version. Following libpq, a client can + /// request this version in its startup message to ask for the newest + /// minor version the server supports: a compliant server replies with + /// `NegotiateProtocolVersion` and the connection proceeds at the + /// negotiated version. This exercises the negotiation path against any + /// server. + /// + /// Note that [`ProtocolVersion::from_version_number`] deliberately does + /// not map to this variant: a server must never accept 3.9999 as-is, it + /// always negotiates it down to a supported version. + PROTOCOL3_9999, } impl ProtocolVersion { @@ -25,6 +38,7 @@ impl ProtocolVersion { match &self { Self::PROTOCOL3_0 => (3, 0), Self::PROTOCOL3_2 => (3, 2), + Self::PROTOCOL3_9999 => (3, 9999), } } @@ -1024,10 +1038,67 @@ mod test { roundtrip!(negotiate_protocol_version, NegotiateProtocolVersion, &ctx); } + #[test] + fn test_negotiated_version() { + // Full 32-bit version number form (PostgreSQL 18+, pgwire) + let full_32 = NegotiateProtocolVersion::new(196610, vec![]); + assert_eq!( + full_32.negotiated_version(), + Some(ProtocolVersion::PROTOCOL3_2) + ); + let full_30 = NegotiateProtocolVersion::new(196608, vec![]); + assert_eq!( + full_30.negotiated_version(), + Some(ProtocolVersion::PROTOCOL3_0) + ); + + // Historical minor-only form + let minor_2 = NegotiateProtocolVersion::new(2, vec![]); + assert_eq!( + minor_2.negotiated_version(), + Some(ProtocolVersion::PROTOCOL3_2) + ); + let minor_0 = NegotiateProtocolVersion::new(0, vec![]); + assert_eq!( + minor_0.negotiated_version(), + Some(ProtocolVersion::PROTOCOL3_0) + ); + + // Unknown minors: below 2 behaves like 3.0, at or above 2 falls back + // to our newest 3.x + let minor_1 = NegotiateProtocolVersion::new(1, vec![]); + assert_eq!( + minor_1.negotiated_version(), + Some(ProtocolVersion::PROTOCOL3_0) + ); + let newer_minor = NegotiateProtocolVersion::new(5, vec![]); + assert_eq!( + newer_minor.negotiated_version(), + Some(ProtocolVersion::PROTOCOL3_2) + ); + + // An unknown major version cannot be mapped + let major_4 = NegotiateProtocolVersion::new((4 << 16) | 2, vec![]); + assert_eq!(major_4.negotiated_version(), None); + // A negative value is invalid + let negative = NegotiateProtocolVersion::new(-1, vec![]); + assert_eq!(negative.negotiated_version(), None); + } + #[test] fn test_protocol_version() { assert_eq!(196608i32, i32::from(ProtocolVersion::PROTOCOL3_0)); assert_eq!(196610i32, i32::from(ProtocolVersion::PROTOCOL3_2)); + assert_eq!(206607i32, i32::from(ProtocolVersion::PROTOCOL3_9999)); + // 3.9999 is encodable (a client requests it to trigger negotiation) + // but is never accepted as a supported version by a server. + assert_eq!( + None, + ProtocolVersion::from_version_number( + ProtocolVersion::PROTOCOL3_9999.version_number().0, + ProtocolVersion::PROTOCOL3_9999.version_number().1 + ) + ); } // A count read as signed (0xffff -> -1 -> usize::MAX) must become a decode diff --git a/src/messages/startup.rs b/src/messages/startup.rs index 5de357a0..8ed445de 100644 --- a/src/messages/startup.rs +++ b/src/messages/startup.rs @@ -505,7 +505,9 @@ impl SecretKey { Self::validate_bytes_len(data_len)?; match ctx.protocol_version { - ProtocolVersion::PROTOCOL3_2 => Ok(SecretKey::Bytes(buf.split_to(data_len).freeze())), + ProtocolVersion::PROTOCOL3_2 | ProtocolVersion::PROTOCOL3_9999 => { + Ok(SecretKey::Bytes(buf.split_to(data_len).freeze())) + } ProtocolVersion::PROTOCOL3_0 => Ok(SecretKey::I32(buf.get_i32())), } } @@ -776,10 +778,58 @@ impl Message for SASLResponse { #[non_exhaustive] #[derive(PartialEq, Eq, Debug, new)] pub struct NegotiateProtocolVersion { + /// Version number reported by the server. + /// + /// PostgreSQL 18+ (protocol 3.2) and pgwire send the **full 32-bit + /// protocol version number** here (e.g. `196610` for 3.2), while older + /// servers send only the **minor version** (e.g. `2`). Use + /// [`NegotiateProtocolVersion::negotiated_version`] to interpret both + /// forms. pub newest_minor_protocol: i32, pub unsupported_options: Vec, } +impl NegotiateProtocolVersion { + /// Interpret the version field of this message as a + /// [`ProtocolVersion`](crate::messages::ProtocolVersion) to use for the + /// rest of the connection. + /// + /// Two wire dialects exist: + /// + /// - PostgreSQL 18+ and pgwire send the full 32-bit version number + /// (major `3` implies a value `> 65535`, e.g. `196610`), + /// - older servers send only the minor version (`<= 65535`, e.g. `2`), + /// leaving the major version unchanged from the client's request. + /// + /// Returns `None` if the reported version cannot be mapped to a version + /// this crate supports (e.g. a newer major version). + pub fn negotiated_version(&self) -> Option { + let value = self.newest_minor_protocol; + if value < 0 { + return None; + } + + let (major, minor) = if value > i32::from(u16::MAX) { + // full 32-bit protocol version number + (((value >> 16) & 0xFFFF) as u16, (value & 0xFFFF) as u16) + } else { + // historical minor-only form, major version is unchanged + (3, value as u16) + }; + + match ProtocolVersion::from_version_number(major, minor) { + Some(version) => Some(version), + // A minor version we don't have a variant for: protocol 3.2 is + // what changed the wire formats we care about (secret keys), so + // an unknown minor below 2 behaves like 3.0, and one at or above + // 2 falls back to our newest 3.x. + None if major == 3 && minor >= 2 => Some(ProtocolVersion::PROTOCOL3_2), + None if major == 3 => Some(ProtocolVersion::PROTOCOL3_0), + None => None, + } + } +} + /// Message type byte for NegotiateProtocolVersion pub const MESSAGE_TYPE_BYTE_NEGOTIATE_PROTOCOL_VERSION: u8 = b'v'; diff --git a/src/tokio/client.rs b/src/tokio/client.rs index f1793c2d..73cc3fae 100644 --- a/src/tokio/client.rs +++ b/src/tokio/client.rs @@ -106,6 +106,10 @@ impl ClientInfo for PgWireClient { fn protocol_version(&self) -> ProtocolVersion { self.socket.codec().decode_context.protocol_version } + + fn set_protocol_version(&mut self, version: ProtocolVersion) { + self.socket.codec_mut().decode_context.protocol_version = version; + } } impl Sink for PgWireClient { @@ -155,6 +159,10 @@ impl PgWireClient { tls_connector: tls_connector_for_cancel, }; + // Decode backend messages with the rules of the protocol version we + // are about to advertise, until the server negotiates it down. + client.set_protocol_version(config.get_protocol_version()); + startup_handler.startup(&mut client).await?; // loop until finished while let Some(message_result) = client.socket.next().await { @@ -453,3 +461,77 @@ fn get_addr(config: &Config) -> Result { Err(PgWireClientError::InvalidConfig("host".to_string())) } + +#[cfg(all(test, feature = "server-api"))] +mod tests { + use std::sync::Arc; + + use tokio::net::TcpListener; + + use super::PgWireClient; + use crate::api::PgWireServerHandlers; + use crate::api::client::ClientInfo; + use crate::api::client::auth::DefaultStartupHandler; + use crate::api::client::config::Config; + use crate::messages::ProtocolVersion; + use crate::messages::startup::SecretKey; + use crate::tokio::server::process_socket; + + struct TestHandlers; + + impl PgWireServerHandlers for TestHandlers {} + + async fn spawn_test_server() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + tokio::spawn(async move { + loop { + let (socket, _) = listener.accept().await.unwrap(); + let handlers = Arc::new(TestHandlers); + tokio::spawn(async move { + let _ = process_socket(socket, None, handlers).await; + }); + } + }); + port + } + + #[tokio::test] + async fn client_negotiates_3_9999_down_to_3_2() { + let port = spawn_test_server().await; + + let mut config = Config::new(); + config.host("127.0.0.1"); + config.port(port); + config.user("pgwire"); + config.protocol_version(ProtocolVersion::PROTOCOL3_9999); + + let client = PgWireClient::connect(Arc::new(config), DefaultStartupHandler::new(), None) + .await + .unwrap(); + + // The server did not accept 3.9999 as-is; it negotiated down to its + // newest supported version. + assert_eq!(client.protocol_version(), ProtocolVersion::PROTOCOL3_2); + // Protocol 3.2 backend keys are 32 bytes long. + assert!(matches!(client.secret_key(), SecretKey::Bytes(key) if key.len() == 32)); + } + + #[tokio::test] + async fn client_default_3_0_keeps_i32_secret_key() { + let port = spawn_test_server().await; + + let mut config = Config::new(); + config.host("127.0.0.1"); + config.port(port); + config.user("pgwire"); + + let client = PgWireClient::connect(Arc::new(config), DefaultStartupHandler::new(), None) + .await + .unwrap(); + + assert_eq!(client.protocol_version(), ProtocolVersion::PROTOCOL3_0); + // A protocol 3.0 cancel key is decoded as a 4-byte i32, not as bytes. + assert!(matches!(client.secret_key(), SecretKey::I32(_))); + } +} From 1f36f282723196d2d42ed6e92449a70ffb5dc5a2 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Sun, 16 Aug 2026 17:30:28 +0800 Subject: [PATCH 3/5] chore: fmt --- src/api/client/auth.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/api/client/auth.rs b/src/api/client/auth.rs index 90cd3e87..a78473a5 100644 --- a/src/api/client/auth.rs +++ b/src/api/client/auth.rs @@ -10,7 +10,7 @@ use crate::error::{ErrorInfo, PgWireClientError, PgWireClientResult, PgWireResul use crate::messages::response::ReadyForQuery; use crate::messages::startup::{ Authentication, BackendKeyData, NegotiateProtocolVersion, ParameterStatus, Password, - PasswordMessageFamily, SASLInitialResponse, SASLResponse, SecretKey, Startup + PasswordMessageFamily, SASLInitialResponse, SASLResponse, SecretKey, Startup, }; use crate::messages::{PgWireBackendMessage, PgWireFrontendMessage}; From c17b1684a791f6e6897b1ed2878fb3aa37049f22 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Sun, 16 Aug 2026 18:05:53 +0800 Subject: [PATCH 4/5] refactor: move protocol negotiation to client specific --- src/api/client/auth.rs | 94 ++++++++++++++++++++++++++++++++++++++++- src/messages/mod.rs | 47 --------------------- src/messages/startup.rs | 45 +------------------- 3 files changed, 94 insertions(+), 92 deletions(-) diff --git a/src/api/client/auth.rs b/src/api/client/auth.rs index a78473a5..cd92fce7 100644 --- a/src/api/client/auth.rs +++ b/src/api/client/auth.rs @@ -12,7 +12,7 @@ use crate::messages::startup::{ Authentication, BackendKeyData, NegotiateProtocolVersion, ParameterStatus, Password, PasswordMessageFamily, SASLInitialResponse, SASLResponse, SecretKey, Startup, }; -use crate::messages::{PgWireBackendMessage, PgWireFrontendMessage}; +use crate::messages::{PgWireBackendMessage, PgWireFrontendMessage, ProtocolVersion}; use super::{ClientInfo, ReadyState, ServerInformation}; @@ -99,7 +99,7 @@ pub trait StartupHandler: Send { C: ClientInfo + Sink + Unpin + Send, PgWireClientError: From<>::Error>, { - match message.negotiated_version() { + match negotiated_version(&message) { Some(version) => { client.set_protocol_version(version); Ok(()) @@ -359,3 +359,93 @@ where }; auth_client.verify_server_final(&message) } + +/// Interpret the version number reported by a `NegotiateProtocolVersion` +/// message and return the protocol version to use for the rest of the +/// connection. +/// +/// This is the client-side counterpart of the server's +/// `api::auth::protocol_negotiation`. Two wire dialects exist: +/// +/// - PostgreSQL 18+ and pgwire report the **full 32-bit protocol version +/// number** (e.g. `196610` for 3.2), which is always greater than `65535`, +/// - older servers report only the **minor version** (e.g. `2`), leaving the +/// major version unchanged from the client's request. +/// +/// Returns `None` if the reported version cannot be mapped to a version this +/// crate supports (e.g. a newer major version). +fn negotiated_version(message: &NegotiateProtocolVersion) -> Option { + let value = message.newest_minor_protocol; + if value < 0 { + return None; + } + + let (major, minor) = if value > i32::from(u16::MAX) { + // full 32-bit protocol version number + (((value >> 16) & 0xFFFF) as u16, (value & 0xFFFF) as u16) + } else { + // minor-only form, the major version is unchanged + (3, value as u16) + }; + + match ProtocolVersion::from_version_number(major, minor) { + Some(version) => Some(version), + // A minor version we don't have a variant for: protocol 3.2 is what + // changed the wire formats we care about (secret keys), so an unknown + // minor below 2 behaves like 3.0, and one at or above 2 falls back to + // our newest 3.x. + None if major == 3 && minor >= 2 => Some(ProtocolVersion::PROTOCOL3_2), + None if major == 3 => Some(ProtocolVersion::PROTOCOL3_0), + None => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_negotiated_version() { + // Full 32-bit version number form (PostgreSQL 18+, pgwire) + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(196610, vec![])), + Some(ProtocolVersion::PROTOCOL3_2) + ); + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(196608, vec![])), + Some(ProtocolVersion::PROTOCOL3_0) + ); + + // Historical minor-only form + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(2, vec![])), + Some(ProtocolVersion::PROTOCOL3_2) + ); + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(0, vec![])), + Some(ProtocolVersion::PROTOCOL3_0) + ); + + // Unknown minors: below 2 behaves like 3.0, at or above 2 falls back + // to our newest 3.x + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(1, vec![])), + Some(ProtocolVersion::PROTOCOL3_0) + ); + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(5, vec![])), + Some(ProtocolVersion::PROTOCOL3_2) + ); + + // An unknown major version cannot be mapped + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new((4 << 16) | 2, vec![])), + None + ); + // A negative value is invalid + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(-1, vec![])), + None + ); + } +} diff --git a/src/messages/mod.rs b/src/messages/mod.rs index 0d23a25b..458b19f5 100644 --- a/src/messages/mod.rs +++ b/src/messages/mod.rs @@ -1038,53 +1038,6 @@ mod test { roundtrip!(negotiate_protocol_version, NegotiateProtocolVersion, &ctx); } - #[test] - fn test_negotiated_version() { - // Full 32-bit version number form (PostgreSQL 18+, pgwire) - let full_32 = NegotiateProtocolVersion::new(196610, vec![]); - assert_eq!( - full_32.negotiated_version(), - Some(ProtocolVersion::PROTOCOL3_2) - ); - let full_30 = NegotiateProtocolVersion::new(196608, vec![]); - assert_eq!( - full_30.negotiated_version(), - Some(ProtocolVersion::PROTOCOL3_0) - ); - - // Historical minor-only form - let minor_2 = NegotiateProtocolVersion::new(2, vec![]); - assert_eq!( - minor_2.negotiated_version(), - Some(ProtocolVersion::PROTOCOL3_2) - ); - let minor_0 = NegotiateProtocolVersion::new(0, vec![]); - assert_eq!( - minor_0.negotiated_version(), - Some(ProtocolVersion::PROTOCOL3_0) - ); - - // Unknown minors: below 2 behaves like 3.0, at or above 2 falls back - // to our newest 3.x - let minor_1 = NegotiateProtocolVersion::new(1, vec![]); - assert_eq!( - minor_1.negotiated_version(), - Some(ProtocolVersion::PROTOCOL3_0) - ); - let newer_minor = NegotiateProtocolVersion::new(5, vec![]); - assert_eq!( - newer_minor.negotiated_version(), - Some(ProtocolVersion::PROTOCOL3_2) - ); - - // An unknown major version cannot be mapped - let major_4 = NegotiateProtocolVersion::new((4 << 16) | 2, vec![]); - assert_eq!(major_4.negotiated_version(), None); - // A negative value is invalid - let negative = NegotiateProtocolVersion::new(-1, vec![]); - assert_eq!(negative.negotiated_version(), None); - } - #[test] fn test_protocol_version() { assert_eq!(196608i32, i32::from(ProtocolVersion::PROTOCOL3_0)); diff --git a/src/messages/startup.rs b/src/messages/startup.rs index 8ed445de..8eaf11a1 100644 --- a/src/messages/startup.rs +++ b/src/messages/startup.rs @@ -782,54 +782,13 @@ pub struct NegotiateProtocolVersion { /// /// PostgreSQL 18+ (protocol 3.2) and pgwire send the **full 32-bit /// protocol version number** here (e.g. `196610` for 3.2), while older - /// servers send only the **minor version** (e.g. `2`). Use - /// [`NegotiateProtocolVersion::negotiated_version`] to interpret both + /// servers send only the **minor version** (e.g. `2`). The client-side + /// negotiation logic in `pgwire::api::client::auth` understands both /// forms. pub newest_minor_protocol: i32, pub unsupported_options: Vec, } -impl NegotiateProtocolVersion { - /// Interpret the version field of this message as a - /// [`ProtocolVersion`](crate::messages::ProtocolVersion) to use for the - /// rest of the connection. - /// - /// Two wire dialects exist: - /// - /// - PostgreSQL 18+ and pgwire send the full 32-bit version number - /// (major `3` implies a value `> 65535`, e.g. `196610`), - /// - older servers send only the minor version (`<= 65535`, e.g. `2`), - /// leaving the major version unchanged from the client's request. - /// - /// Returns `None` if the reported version cannot be mapped to a version - /// this crate supports (e.g. a newer major version). - pub fn negotiated_version(&self) -> Option { - let value = self.newest_minor_protocol; - if value < 0 { - return None; - } - - let (major, minor) = if value > i32::from(u16::MAX) { - // full 32-bit protocol version number - (((value >> 16) & 0xFFFF) as u16, (value & 0xFFFF) as u16) - } else { - // historical minor-only form, major version is unchanged - (3, value as u16) - }; - - match ProtocolVersion::from_version_number(major, minor) { - Some(version) => Some(version), - // A minor version we don't have a variant for: protocol 3.2 is - // what changed the wire formats we care about (secret keys), so - // an unknown minor below 2 behaves like 3.0, and one at or above - // 2 falls back to our newest 3.x. - None if major == 3 && minor >= 2 => Some(ProtocolVersion::PROTOCOL3_2), - None if major == 3 => Some(ProtocolVersion::PROTOCOL3_0), - None => None, - } - } -} - /// Message type byte for NegotiateProtocolVersion pub const MESSAGE_TYPE_BYTE_NEGOTIATE_PROTOCOL_VERSION: u8 = b'v'; From bf40ec4dae6a5aeea315bb2aefa56de197f6e6c7 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Sun, 16 Aug 2026 19:31:07 +0800 Subject: [PATCH 5/5] test: clarify tests --- src/api/client/auth.rs | 52 +++++++++++++++++++++++++++--------------- 1 file changed, 33 insertions(+), 19 deletions(-) diff --git a/src/api/client/auth.rs b/src/api/client/auth.rs index cd92fce7..fbfe21f9 100644 --- a/src/api/client/auth.rs +++ b/src/api/client/auth.rs @@ -406,25 +406,39 @@ mod tests { #[test] fn test_negotiated_version() { - // Full 32-bit version number form (PostgreSQL 18+, pgwire) - assert_eq!( - negotiated_version(&NegotiateProtocolVersion::new(196610, vec![])), - Some(ProtocolVersion::PROTOCOL3_2) - ); - assert_eq!( - negotiated_version(&NegotiateProtocolVersion::new(196608, vec![])), - Some(ProtocolVersion::PROTOCOL3_0) - ); - - // Historical minor-only form - assert_eq!( - negotiated_version(&NegotiateProtocolVersion::new(2, vec![])), - Some(ProtocolVersion::PROTOCOL3_2) - ); - assert_eq!( - negotiated_version(&NegotiateProtocolVersion::new(0, vec![])), - Some(ProtocolVersion::PROTOCOL3_0) - ); + // The same version can arrive in two wire dialects, and both forms + // must negotiate to the same version: + // + // - the full 32-bit protocol version number (`i32::from(version)`), + // used by PostgreSQL 18+ and current pgwire, + // - the minor version alone, used by older servers, which leaves the + // major version unchanged from the client's request. + // + // Realistic on the wire: full 3.2 (PostgreSQL 18+), bare 0 (older + // PostgreSQL, which only supports 3.0) and bare 2 (pgwire before the + // #439 fix sent the minor alone). Full-form 3.0 never occurs, but is + // unambiguous. + for (full_form, minor_form, expected) in [ + ( + i32::from(ProtocolVersion::PROTOCOL3_2), + 2, + ProtocolVersion::PROTOCOL3_2, + ), + ( + i32::from(ProtocolVersion::PROTOCOL3_0), + 0, + ProtocolVersion::PROTOCOL3_0, + ), + ] { + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(full_form, vec![])), + Some(expected) + ); + assert_eq!( + negotiated_version(&NegotiateProtocolVersion::new(minor_form, vec![])), + Some(expected) + ); + } // Unknown minors: below 2 behaves like 3.0, at or above 2 falls back // to our newest 3.x