From 250fd6fbfc2c8ec03221efaee0c16f582810bf99 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Tue, 4 Aug 2026 17:27:40 +0300 Subject: [PATCH 01/11] Track a client's parameter changes on the server connection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A SET or RESET the client sends once a server is attached changes that server's session, but we never wrote it down. What we did instead was clear the whole parameter cache whenever a CommandComplete said RESET, which trades one problem for another: the cache is what tells us to undo the change for the next client, and a ROLLBACK undoes the RESET anyway. pg_dump -t walks straight into it: SET search_path TO '', then RESET search_path inside its transaction, then ROLLBACK. Postgres brings the empty search_path back, the cache no longer mentions it, and the next client gets the connection with unqualified names silently broken. A SET has the same hole in the other direction: BEGIN, a query, SET statement_timeout, COMMIT — the setting stays on the server and nothing resets it for whoever comes next. So record the change where it happens. RESET now has the transaction handling SET always had (reset vs reset_transaction), so a rollback restores what it cleared and a commit makes it permanent, and the server connection keeps the same record its client does. The existing parameter diff then does the rest: the next client is handed a precise RESET for what it doesn't want, instead of a connection nobody dares reuse. The CommandComplete fallback stays for RESETs we don't see coming — with the query parser off, that is still all we have. --- pgdog/src/backend/pool/connection/binding.rs | 24 +- pgdog/src/backend/server.rs | 206 +++++++++++++++++- pgdog/src/frontend/client/query_engine/set.rs | 17 +- pgdog/src/net/parameter.rs | 68 +++++- 4 files changed, 299 insertions(+), 16 deletions(-) diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index faa5fe28a..57a05cff9 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -2,7 +2,7 @@ use crate::{ frontend::{ - ClientRequest, + ClientRequest, SetParam, client::query_engine::{ TwoPcPhase, two_pc::{TwoPcTransaction, statement::phase_control}, @@ -444,6 +444,28 @@ impl Binding { } } + /// Record a client parameter change on every server we hold. + pub fn record_params(&mut self, params: &[SetParam], in_transaction: bool) { + match self { + Binding::Direct(server, ..) => server.record_params(params, in_transaction), + Binding::MultiShard(servers, _) => servers + .iter_mut() + .for_each(|server| server.record_params(params, in_transaction)), + _ => (), + } + } + + /// Record a client `RESET ALL` on every server we hold. + pub fn record_reset_all(&mut self, in_transaction: bool) { + match self { + Binding::Direct(server, ..) => server.record_reset_all(in_transaction), + Binding::MultiShard(servers, _) => servers + .iter_mut() + .for_each(|server| server.record_reset_all(in_transaction)), + _ => (), + } + } + /// Handle transaction end. pub fn transaction_params_hook(&mut self, rollback: bool) { match self { diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 96d9d0e7f..17822bf8a 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -20,7 +20,7 @@ use crate::{ auth::{md5, scram::Client}, backend::pool::stats::MemoryStats, config::AuthType, - frontend::ClientRequest, + frontend::{ClientRequest, SetParam}, net::{ Close, Liveness, MessageBuffer, Parameter, ProtocolMessage, Sync, messages::{ @@ -140,6 +140,9 @@ pub struct Server { streaming: bool, schema_changed: bool, sync_prepared: bool, + // The client's parameter change for the statement in flight is already + // recorded, so its CommandComplete tells us nothing we don't know. + params_recorded: bool, in_transaction: bool, re_synced: bool, replication_mode: bool, @@ -431,6 +434,7 @@ impl Server { streaming: false, schema_changed: false, sync_prepared: false, + params_recorded: false, in_transaction: false, statement_executed: false, re_synced: false, @@ -597,6 +601,8 @@ impl Server { match message.code() { 'Z' => { + self.params_recorded = false; + let now = Instant::now(); let rfq = ReadyForQuery::from_bytes(message.payload())?; @@ -654,7 +660,10 @@ impl Server { self.prepared_statements.clear(); self.client_params.clear(); } - "RESET" => self.client_params.clear(), // Someone reset params, we're gonna need to re-sync. + // Someone reset params. If we didn't see which ones (the query + // parser is off, or the RESET came from somewhere we don't + // track), the cache is worthless and we re-sync from scratch. + "RESET" if !self.params_recorded => self.client_params.clear(), _ => (), } self.stats.rows_affected(&cmd); @@ -746,6 +755,37 @@ impl Server { Ok(executed) } + /// Record a parameter change the client is making on this connection, so + /// we know what to undo before handing it to somebody else. + pub fn record_params(&mut self, params: &[SetParam], in_transaction: bool) { + for param in params { + match (¶m.value, in_transaction) { + (Some(value), true) => { + self.client_params + .insert_transaction(¶m.name, value.clone(), param.local); + } + (Some(value), false) => { + self.client_params.insert(¶m.name, value.clone()); + } + (None, true) => self.client_params.reset_transaction(¶m.name), + (None, false) => self.client_params.reset(¶m.name), + } + } + + self.params_recorded = true; + } + + /// Record a `RESET ALL` the client is making on this connection. + pub fn record_reset_all(&mut self, in_transaction: bool) { + if in_transaction { + self.client_params.reset_all_transaction(); + } else { + self.client_params.reset_all(); + } + + self.params_recorded = true; + } + // Handle COMMIT/ROLLBACK for in-transaction params tracking. pub fn transaction_params_hook(&mut self, rollback: bool) { if rollback { @@ -1320,8 +1360,10 @@ pub mod test { }; use crate::{ - backend::pool::token_cache::TokenCache, config::Memory, frontend::PreparedStatements, - net::*, + backend::pool::token_cache::TokenCache, + config::Memory, + frontend::PreparedStatements, + net::{parameter::ParameterValue, *}, }; use super::{Error, *}; @@ -1365,6 +1407,7 @@ pub mod test { streaming: false, schema_changed: false, sync_prepared: false, + params_recorded: false, in_transaction: false, re_synced: false, replication_mode: false, @@ -3096,6 +3139,161 @@ pub mod test { ) } + #[tokio::test] + async fn test_recorded_reset_survives_rollback() { + let mut server = test_server().await; + let mut params = Parameters::default(); + params.insert("search_path", ""); + server + .link_client(FrontendPid::new(), ¶ms, None) + .await + .unwrap(); + + server.execute("BEGIN").await.unwrap(); + server.record_params( + &[SetParam { + name: "search_path".into(), + value: None, + local: false, + }], + true, + ); + server.execute("RESET search_path").await.unwrap(); + server.execute("ROLLBACK").await.unwrap(); + server.transaction_params_hook(true); + + // The ROLLBACK brought search_path back, so we still owe the next + // client a RESET for it. + let queries = server + .client_params + .reset_queries(&Parameters::default()) + .into_iter() + .map(|query| query.query().to_string()) + .collect::>(); + assert_eq!(queries, vec![r#"RESET "search_path""#]); + } + + #[tokio::test] + async fn test_recorded_reset_committed_is_permanent() { + let mut server = test_server().await; + let mut params = Parameters::default(); + params.insert("search_path", ""); + server + .link_client(FrontendPid::new(), ¶ms, None) + .await + .unwrap(); + + server.execute("BEGIN").await.unwrap(); + server.record_params( + &[SetParam { + name: "search_path".into(), + value: None, + local: false, + }], + true, + ); + server.execute("RESET search_path").await.unwrap(); + server.execute("COMMIT").await.unwrap(); + server.transaction_params_hook(false); + + // Committed: the server really is back to its default, nothing to undo. + assert!( + server + .client_params + .reset_queries(&Parameters::default()) + .is_empty() + ); + } + + #[tokio::test] + async fn test_recorded_set_is_undone_for_the_next_client() { + let mut server = test_server().await; + server + .link_client(FrontendPid::new(), &Parameters::default(), None) + .await + .unwrap(); + + // A SET that lands after the connection is already ours. + server.record_params( + &[SetParam { + name: "statement_timeout".into(), + value: Some(ParameterValue::String("5s".into())), + local: false, + }], + false, + ); + server + .execute("SET statement_timeout TO '5s'") + .await + .unwrap(); + + let queries = server + .client_params + .reset_queries(&Parameters::default()) + .into_iter() + .map(|query| query.query().to_string()) + .collect::>(); + assert_eq!(queries, vec![r#"RESET "statement_timeout""#]); + } + + #[tokio::test] + async fn test_reset_keeps_other_recorded_params() { + let mut server = test_server().await; + server + .link_client(FrontendPid::new(), &Parameters::default(), None) + .await + .unwrap(); + + server.record_params( + &[SetParam { + name: "statement_timeout".into(), + value: Some(ParameterValue::String("5s".into())), + local: false, + }], + false, + ); + server + .execute("SET statement_timeout TO '5s'") + .await + .unwrap(); + + // Resetting one parameter says nothing about the others. + server.record_params( + &[SetParam { + name: "search_path".into(), + value: None, + local: false, + }], + false, + ); + server.execute("RESET search_path").await.unwrap(); + + let queries = server + .client_params + .reset_queries(&Parameters::default()) + .into_iter() + .map(|query| query.query().to_string()) + .collect::>(); + assert_eq!(queries, vec![r#"RESET "statement_timeout""#]); + } + + #[tokio::test] + async fn test_untracked_reset_still_clears_client_params() { + let mut server = test_server().await; + let mut params = Parameters::default(); + params.insert("search_path", "public"); + server + .link_client(FrontendPid::new(), ¶ms, None) + .await + .unwrap(); + + // Nobody told us what this RESET touched (query parser off), so the + // cache is worthless and we start over. + server.execute("RESET search_path").await.unwrap(); + + assert!(server.client_params.is_empty()); + } + #[tokio::test] async fn test_reset_clears_client_params() { let mut server = test_server().await; diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index dff7cfde8..d97b18c5c 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -44,7 +44,11 @@ impl QueryEngine { } } else { fake_command = "RESET"; - context.params.reset(¶m.name); + if context.in_transaction() { + context.params.reset_transaction(¶m.name); + } else { + context.params.reset(¶m.name); + } if is_pin { self.manual_lock = false; } @@ -56,6 +60,10 @@ impl QueryEngine { } if self.backend.connected() { + // The server is ours right now, so its session changes with the + // client's. Record it, or we won't know to undo it for whoever + // gets this connection next. + self.backend.record_params(params, context.in_transaction()); self.execute(context).await?; } else { let values_to_return = @@ -96,9 +104,14 @@ impl QueryEngine { &mut self, context: &mut QueryEngineContext<'_>, ) -> Result<(), Error> { - context.params.reset_all(); + if context.in_transaction() { + context.params.reset_all_transaction(); + } else { + context.params.reset_all(); + } if self.backend.connected() { + self.backend.record_reset_all(context.in_transaction()); self.execute(context).await?; } else { self.fake_command_response(context, "RESET", None::>) diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index 6dbe03a0f..2b04e96b2 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -267,11 +267,25 @@ impl Parameters { } } - /// Remove parameter from params temporarily. The transaction - /// is comitted, it will be removed permanently. + /// Remove a parameter. pub fn reset(&mut self, name: impl ToString) { let name = name.to_string().to_lowercase(); + if self.params.remove(&name).is_some() { + self.hash = Self::compute_hash(&self.params); + } + + self.transaction_params.remove(&name); + self.transaction_local_params.remove(&name); + // Nothing left to restore: the value is gone for good. + self.reset_params.remove(&name); + } + + /// Remove a parameter, but only for the duration of the transaction: + /// a ROLLBACK brings its value back. + pub fn reset_transaction(&mut self, name: impl ToString) { + let name = name.to_string().to_lowercase(); + if let Some(value) = self.params.remove(&name) { self.reset_params.insert(name.clone(), value); self.hash = Self::compute_hash(&self.params); @@ -283,17 +297,27 @@ impl Parameters { /// Reset all tracked parameters. pub fn reset_all(&mut self) { + for key in self.resettable_keys() { + self.reset(&key); + } + } + + /// Reset all tracked parameters for the duration of the transaction. + pub fn reset_all_transaction(&mut self) { + for key in self.resettable_keys() { + self.reset_transaction(&key); + } + } + + fn resettable_keys(&self) -> Vec { let mut keys: Vec = self.params.keys().cloned().collect(); keys.extend(self.transaction_params.keys().cloned()); keys.extend(self.transaction_local_params.keys().cloned()); keys.sort(); keys.dedup(); + keys.retain(|key| !UNTRACKED_PARAMS.contains(key)); - for key in keys { - if !UNTRACKED_PARAMS.contains(&key) { - self.reset(&key); - } - } + keys } /// Commit params we saved during the transaction. @@ -998,11 +1022,37 @@ mod test { } #[test] - fn test_reset_rollback_restores_param() { + fn test_reset_outside_transaction_is_permanent() { let mut params = Parameters::default(); params.insert("search_path", "public"); params.reset("search_path"); + + // A transaction that comes later has nothing to do with that RESET. + params.rollback(); + + assert_eq!(params.get("search_path"), None); + } + + #[test] + fn test_reset_all_outside_transaction_is_permanent() { + let mut params = Parameters::default(); + params.insert("search_path", "public"); + params.insert("timezone", "UTC"); + + params.reset_all(); + params.rollback(); + + assert_eq!(params.get("search_path"), None); + assert_eq!(params.get("timezone"), None); + } + + #[test] + fn test_reset_rollback_restores_param() { + let mut params = Parameters::default(); + params.insert("search_path", "public"); + + params.reset_transaction("search_path"); assert_eq!(params.get("search_path"), None); params.rollback(); @@ -1104,7 +1154,7 @@ mod test { params.insert("search_path", "public"); params.insert("timezone", "UTC"); - params.reset_all(); + params.reset_all_transaction(); assert_eq!(params.get("search_path"), None); assert_eq!(params.get("timezone"), None); From d452e2eb3bad6f9e957feaaca68d594d570e1ef6 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Tue, 4 Aug 2026 17:27:56 +0300 Subject: [PATCH 02/11] Add integration tests for session parameters leaking between clients Runs against a database with a single server connection, so the next client always gets the connection the previous one used. --- integration/pgdog.toml | 12 ++++ .../python/test_session_params_leak.py | 71 +++++++++++++++++++ integration/users.toml | 5 ++ 3 files changed, 88 insertions(+) create mode 100644 integration/python/test_session_params_leak.py diff --git a/integration/pgdog.toml b/integration/pgdog.toml index d5fd64d75..eed85ee83 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -53,6 +53,18 @@ host = "127.0.0.1" role = "replica" read_only = true +# ------------------------------------------------------------------------------ +# ----- Database :: pgdog_leak ------------------------------------------------- +# Exactly one server connection, kept around: tests for session state leaking +# between clients need the next client to get the same connection back. + +[[databases]] +name = "pgdog_leak" +host = "127.0.0.1" +database_name = "pgdog" +pool_size = 1 +min_pool_size = 1 + # ------------------------------------------------------------------------------ # ----- Database :: pgdog_sharded ---------------------------------------------- diff --git a/integration/python/test_session_params_leak.py b/integration/python/test_session_params_leak.py new file mode 100644 index 000000000..74b98e68a --- /dev/null +++ b/integration/python/test_session_params_leak.py @@ -0,0 +1,71 @@ +"""Session parameters must not survive the client that set them. + +Runs against a database with a single server connection, so the next client +always gets the connection the previous one used. current_setting() is used +instead of SHOW because SHOW can be answered by PgDog itself. +""" + +import psycopg + + +def connect(): + conn = psycopg.connect( + user="pgdog", + password="pgdog", + dbname="pgdog_leak", + host="127.0.0.1", + port=6432, + ) + # Without autocommit every statement runs in a transaction that is rolled + # back on close, which would undo the very state we're testing for. + conn.autocommit = True + return conn + + +def read(setting): + conn = connect() + value = conn.execute(f"SELECT current_setting('{setting}')").fetchone()[0] + conn.close() + + return value + + +def test_reset_rolled_back(): + """A ROLLBACK brings back the value the RESET cleared. + + This is the sequence pg_dump -t
emits. SET search_path TO '' leaves + an empty quoted identifier, hence the two spellings of "empty". + """ + conn = connect() + conn.execute("SET search_path TO ''") + with conn.transaction(): + conn.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY") + conn.execute("RESET search_path") + conn.execute("SELECT 1") + raise psycopg.Rollback() + conn.close() + + assert read("search_path") not in ("", '""') + + +def test_set_committed_after_connecting(): + """A SET that lands once the connection is already ours.""" + conn = connect() + with conn.transaction(): + conn.execute("SELECT 1") + conn.execute("SET statement_timeout TO '5s'") + conn.close() + + assert read("statement_timeout") == "0" + + +def test_reset_committed(): + """A committed RESET is permanent and needs no undoing.""" + conn = connect() + conn.execute("SET search_path TO public") + with conn.transaction(): + conn.execute("SELECT 1") + conn.execute("RESET search_path") + conn.close() + + assert read("search_path") == '"$user", public' diff --git a/integration/users.toml b/integration/users.toml index bba115a85..360246b5b 100644 --- a/integration/users.toml +++ b/integration/users.toml @@ -3,6 +3,11 @@ name = "pgdog" database = "pgdog" password = "pgdog" +[[users]] +name = "pgdog" +database = "pgdog_leak" +password = "pgdog" + [[users]] name = "pgdog_migrator" database = "pgdog" From ca001877ab98e06314f8371d05005872ff8bd226 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Tue, 4 Aug 2026 23:34:55 +0300 Subject: [PATCH 03/11] Collect resettable keys in a single filtered pass The keys still have to be lifted out of the maps before we can reset them, but there is no reason to clone the ones we are about to drop. --- pgdog/src/net/parameter.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index 2b04e96b2..7baece9a6 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -310,12 +310,18 @@ impl Parameters { } fn resettable_keys(&self) -> Vec { - let mut keys: Vec = self.params.keys().cloned().collect(); - keys.extend(self.transaction_params.keys().cloned()); - keys.extend(self.transaction_local_params.keys().cloned()); + // The keys have to be lifted out before we can reset them: resetting + // borrows the maps we'd be iterating. + let mut keys: Vec = self + .params + .keys() + .chain(self.transaction_params.keys()) + .chain(self.transaction_local_params.keys()) + .filter(|key| !UNTRACKED_PARAMS.contains(key)) + .cloned() + .collect(); keys.sort(); keys.dedup(); - keys.retain(|key| !UNTRACKED_PARAMS.contains(key)); keys } From a42cd4226ef7f776e3c89c2dbc8a0434b91ee83e Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 5 Aug 2026 00:32:24 +0300 Subject: [PATCH 04/11] Trim comments to what the code doesn't already say --- integration/pgdog.toml | 3 +-- .../python/test_session_params_leak.py | 4 ++-- pgdog/src/backend/server.rs | 23 +++++++------------ pgdog/src/frontend/client/query_engine/set.rs | 5 ++-- pgdog/src/net/parameter.rs | 14 +++++------ 5 files changed, 19 insertions(+), 30 deletions(-) diff --git a/integration/pgdog.toml b/integration/pgdog.toml index eed85ee83..b52035d6a 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -55,8 +55,7 @@ read_only = true # ------------------------------------------------------------------------------ # ----- Database :: pgdog_leak ------------------------------------------------- -# Exactly one server connection, kept around: tests for session state leaking -# between clients need the next client to get the same connection back. +# One server connection, never reaped: the next client has to get the same one. [[databases]] name = "pgdog_leak" diff --git a/integration/python/test_session_params_leak.py b/integration/python/test_session_params_leak.py index 74b98e68a..9af9c4533 100644 --- a/integration/python/test_session_params_leak.py +++ b/integration/python/test_session_params_leak.py @@ -16,8 +16,8 @@ def connect(): host="127.0.0.1", port=6432, ) - # Without autocommit every statement runs in a transaction that is rolled - # back on close, which would undo the very state we're testing for. + # Otherwise psycopg wraps each statement in a transaction and rolls it + # back on close, undoing the state we're testing for. conn.autocommit = True return conn diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 17822bf8a..30b38b552 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -140,8 +140,8 @@ pub struct Server { streaming: bool, schema_changed: bool, sync_prepared: bool, - // The client's parameter change for the statement in flight is already - // recorded, so its CommandComplete tells us nothing we don't know. + // The change in flight is already recorded, so its CommandComplete + // tells us nothing new. params_recorded: bool, in_transaction: bool, re_synced: bool, @@ -660,9 +660,8 @@ impl Server { self.prepared_statements.clear(); self.client_params.clear(); } - // Someone reset params. If we didn't see which ones (the query - // parser is off, or the RESET came from somewhere we don't - // track), the cache is worthless and we re-sync from scratch. + // A RESET nobody told us the contents of: we can't tell + // what it touched, so drop everything. "RESET" if !self.params_recorded => self.client_params.clear(), _ => (), } @@ -755,8 +754,8 @@ impl Server { Ok(executed) } - /// Record a parameter change the client is making on this connection, so - /// we know what to undo before handing it to somebody else. + /// Record a client's parameter change, so we know what to undo before + /// this connection goes to somebody else. pub fn record_params(&mut self, params: &[SetParam], in_transaction: bool) { for param in params { match (¶m.value, in_transaction) { @@ -775,7 +774,7 @@ impl Server { self.params_recorded = true; } - /// Record a `RESET ALL` the client is making on this connection. + /// Record a client's `RESET ALL`. pub fn record_reset_all(&mut self, in_transaction: bool) { if in_transaction { self.client_params.reset_all_transaction(); @@ -3162,8 +3161,6 @@ pub mod test { server.execute("ROLLBACK").await.unwrap(); server.transaction_params_hook(true); - // The ROLLBACK brought search_path back, so we still owe the next - // client a RESET for it. let queries = server .client_params .reset_queries(&Parameters::default()) @@ -3196,7 +3193,6 @@ pub mod test { server.execute("COMMIT").await.unwrap(); server.transaction_params_hook(false); - // Committed: the server really is back to its default, nothing to undo. assert!( server .client_params @@ -3213,7 +3209,6 @@ pub mod test { .await .unwrap(); - // A SET that lands after the connection is already ours. server.record_params( &[SetParam { name: "statement_timeout".into(), @@ -3257,7 +3252,6 @@ pub mod test { .await .unwrap(); - // Resetting one parameter says nothing about the others. server.record_params( &[SetParam { name: "search_path".into(), @@ -3287,8 +3281,7 @@ pub mod test { .await .unwrap(); - // Nobody told us what this RESET touched (query parser off), so the - // cache is worthless and we start over. + // A RESET we never recorded, as when the query parser is off. server.execute("RESET search_path").await.unwrap(); assert!(server.client_params.is_empty()); diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index d97b18c5c..0bac26a82 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -60,9 +60,8 @@ impl QueryEngine { } if self.backend.connected() { - // The server is ours right now, so its session changes with the - // client's. Record it, or we won't know to undo it for whoever - // gets this connection next. + // The server is ours, so its session changes with the client's: + // record it or we won't know what to undo for the next client. self.backend.record_params(params, context.in_transaction()); self.execute(context).await?; } else { diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index 7baece9a6..eb5654c0e 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -267,7 +267,7 @@ impl Parameters { } } - /// Remove a parameter. + /// Remove a parameter permanently. pub fn reset(&mut self, name: impl ToString) { let name = name.to_string().to_lowercase(); @@ -277,12 +277,11 @@ impl Parameters { self.transaction_params.remove(&name); self.transaction_local_params.remove(&name); - // Nothing left to restore: the value is gone for good. + // Nothing left to restore on a rollback. self.reset_params.remove(&name); } - /// Remove a parameter, but only for the duration of the transaction: - /// a ROLLBACK brings its value back. + /// Remove a parameter until the transaction ends: a ROLLBACK brings it back. pub fn reset_transaction(&mut self, name: impl ToString) { let name = name.to_string().to_lowercase(); @@ -302,7 +301,7 @@ impl Parameters { } } - /// Reset all tracked parameters for the duration of the transaction. + /// Reset all tracked parameters until the transaction ends. pub fn reset_all_transaction(&mut self) { for key in self.resettable_keys() { self.reset_transaction(&key); @@ -310,8 +309,7 @@ impl Parameters { } fn resettable_keys(&self) -> Vec { - // The keys have to be lifted out before we can reset them: resetting - // borrows the maps we'd be iterating. + // Lifted out first: resetting borrows the maps we'd be iterating. let mut keys: Vec = self .params .keys() @@ -1034,7 +1032,7 @@ mod test { params.reset("search_path"); - // A transaction that comes later has nothing to do with that RESET. + // A later transaction has nothing to do with that RESET. params.rollback(); assert_eq!(params.get("search_path"), None); From 022da10a62b60f7fdcb83e1faa4c2b20a9ce9865 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Tue, 4 Aug 2026 17:57:30 +0300 Subject: [PATCH 05/11] Resolve set_config() arguments from the Bind message MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The parser only understood constants, so a parameterized call went through untouched and whatever it changed stayed on the connection: SELECT pg_catalog.set_config($1, $2, false) Read $n from the Bind message instead. The is_local argument is decoded from either wire format. Arguments that still don't resolve — no Bind message, a parameter that isn't text, an expression — leave the statement an ordinary query, same as today: it runs, and we don't know what it changed. --- .../router/parser/query/set_config.rs | 67 +++++-- .../router/parser/query/test/test_set.rs | 164 +++++++++++++++++- 2 files changed, 214 insertions(+), 17 deletions(-) diff --git a/pgdog/src/frontend/router/parser/query/set_config.rs b/pgdog/src/frontend/router/parser/query/set_config.rs index 7f3540f4e..d709bc5a1 100644 --- a/pgdog/src/frontend/router/parser/query/set_config.rs +++ b/pgdog/src/frontend/router/parser/query/set_config.rs @@ -1,20 +1,21 @@ use super::*; +use crate::net::messages::{Bind, Format}; impl QueryParser { /// Handle SELECT set_config('key', 'value', is_local) /// - /// If the function arguments are a form we cannot handle, we warn and - /// pass through + /// Arguments we can't resolve leave the statement an ordinary query: it + /// runs, but we don't know what it changed. pub(super) fn set_config( &mut self, fcall: &nodes::FuncCall, context: &QueryParserContext, ) -> Command { - if let Some(param) = parse_args(fcall) { + if let Some(param) = parse_args(fcall, context.router_context.bind) { Command::Set { params: vec![param], route: Route::write(context.shards_calculator.shard()), - behave_like_select: true, + is_select: true, } } else { Command::Query( @@ -22,28 +23,65 @@ impl QueryParser { ) } } + } /// Returns None if the arguments could not be parsed -fn parse_args(fcall: &nodes::FuncCall) -> Option { - let name = parse_config_name(fcall.args().first()?)?; - let value = parse_config_value(fcall.args().get(1)?)?; - let local = parse_is_local(fcall.args().get(2)?)?; +fn parse_args(fcall: &nodes::FuncCall, bind: Option<&Bind>) -> Option { + let name = parse_config_name(fcall.args().first()?, bind)?; + let value = parse_config_value(fcall.args().get(1)?, bind)?; + let local = parse_is_local(fcall.args().get(2)?, bind)?; Some(SetParam { name, value, local }) } +/// Get the value bound to `$number`. The inner Option is the SQL NULL. +fn bound_text(bind: Option<&Bind>, number: i32) -> Option> { + let index = usize::try_from(number).ok()?.checked_sub(1)?; + let param = bind?.parameter(index).ok()??; + + if param.is_null() { + Some(None) + } else { + Some(Some(param.text()?.to_owned())) + } +} + +fn bound_bool(bind: Option<&Bind>, number: i32) -> Option { + let index = usize::try_from(number).ok()?.checked_sub(1)?; + let param = bind?.parameter(index).ok()??; + + if param.is_null() { + return None; + } + + match param.format() { + Format::Binary => match param.data() { + [0] => Some(false), + [1] => Some(true), + _ => None, + }, + Format::Text => match param.text()?.trim().to_lowercase().as_str() { + "t" | "true" | "y" | "yes" | "on" | "1" => Some(true), + "f" | "false" | "n" | "no" | "off" | "0" => Some(false), + _ => None, + }, + } +} + /// Returns None if the name could not be parsed -fn parse_config_name(arg: Node<'_>) -> Option { +fn parse_config_name(arg: Node<'_>, bind: Option<&Bind>) -> Option { match arg { Node::A_Const(c) => c.val()?.string_value().map(ToOwned::to_owned), - // Only constant strings can be handled for now + Node::ParamRef(nodes::ParamRef { number, .. }) => bound_text(bind, *number)?, _ => None, } } +/// Returns None if the name could not be parsed + /// Returns None if the value could not be parsed, Some(None) if the value /// is NULL, and Some if the value was successfully parsed -fn parse_config_value(arg: Node<'_>) -> Option> { +fn parse_config_value(arg: Node<'_>, bind: Option<&Bind>) -> Option> { match arg { Node::A_Const(c) => match c.val() { Some(value) => Some(Some(ParameterValue::String( @@ -51,14 +89,19 @@ fn parse_config_value(arg: Node<'_>) -> Option> { ))), None => Some(None), }, + Node::ParamRef(nodes::ParamRef { number, .. }) => { + Some(bound_text(bind, *number)?.map(ParameterValue::String)) + } _ => None, } } /// Returns None if the node was not a constant boolean -fn parse_is_local(arg: Node<'_>) -> Option { +fn parse_is_local(arg: Node<'_>, bind: Option<&Bind>) -> Option { match arg { Node::A_Const(c) => c.val()?.bool_value(), + Node::ParamRef(nodes::ParamRef { number, .. }) => bound_bool(bind, *number), _ => None, } } + diff --git a/pgdog/src/frontend/router/parser/query/test/test_set.rs b/pgdog/src/frontend/router/parser/query/test/test_set.rs index d8094e4cc..e035abc62 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_set.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_set.rs @@ -7,7 +7,7 @@ use crate::{ route::{OverrideReason, ShardSource}, }, }, - net::parameter::ParameterValue, + net::{Format, messages::Parameter, parameter::ParameterValue}, }; use super::setup::*; @@ -66,15 +66,13 @@ fn test_set_config_null_value() { match command { Command::Set { - params, - behave_like_select, - .. + params, is_select, .. } => { assert_eq!(params.len(), 1); assert_eq!(params[0].name, "lock_timeout"); assert_eq!(params[0].value, None); assert!(!params[0].local); - assert!(behave_like_select); + assert!(is_select); } _ => panic!("expected Command::Set, got {command:#?}"), } @@ -252,3 +250,159 @@ fn test_single_shard_set() { _ => panic!("not a set"), } } + +#[test] +fn test_set_config_bound_params() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_set_config", + "SELECT pg_catalog.set_config($1, $2, $3)", + ) + .into(), + Bind::new_params( + "__test_set_config", + &[ + Parameter::new(b"search_path"), + Parameter::new(b""), + Parameter::new(b"f"), + ], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + match command { + Command::Set { + ref params, + is_select, + .. + } => { + assert_eq!(params.len(), 1); + assert_eq!(params[0].name, "search_path"); + assert_eq!(params[0].value, Some(ParameterValue::String("".into()))); + assert!(!params[0].local); + assert!(is_select); + } + _ => panic!("expected Command::Set, got {command:#?}"), + } +} + +#[test] +fn test_set_config_bound_null_value() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named("__test_set_config_null", "SELECT set_config($1, $2, false)").into(), + Bind::new_params( + "__test_set_config_null", + &[Parameter::new(b"lock_timeout"), Parameter::new_null()], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + match command { + Command::Set { ref params, .. } => { + assert_eq!(params[0].name, "lock_timeout"); + assert_eq!(params[0].value, None); + } + _ => panic!("expected Command::Set, got {command:#?}"), + } +} + +#[test] +fn test_set_config_unresolvable_args_stay_a_query() { + let mut test = QueryParserTest::new(); + + // No Bind message, so the parameters can't be resolved. + let command = test.execute(vec![ + Query::new("SELECT pg_catalog.set_config($1, $2, false)").into(), + ]); + + assert!( + matches!(command, Command::Query(_)), + "expected Command::Query, got {command:#?}", + ); + assert!(command.route().is_write()); +} + +#[test] +fn test_set_config_expression_stays_a_query() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("SELECT set_config('search_path', current_setting('search_path'), false)") + .into(), + ]); + + assert!( + matches!(command, Command::Query(_)), + "expected Command::Query, got {command:#?}", + ); + assert!(command.route().is_write()); +} + +#[test] +fn test_set_config_bound_binary_is_local() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named("__test_set_config_bin", "SELECT set_config($1, $2, $3)").into(), + Bind::new_params_codes( + "__test_set_config_bin", + &[ + Parameter::new(b"statement_timeout"), + Parameter::new(b"1000"), + Parameter::new(&[1]), + ], + &[Format::Text, Format::Text, Format::Binary], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + match command { + Command::Set { ref params, .. } => { + assert_eq!(params[0].name, "statement_timeout"); + assert_eq!(params[0].value, Some(ParameterValue::String("1000".into()))); + assert!(params[0].local, "binary true must be read as SET LOCAL"); + } + _ => panic!("expected Command::Set, got {command:#?}"), + } +} + +#[test] +fn test_set_config_bound_non_utf8_stays_a_query() { + let mut test = QueryParserTest::new(); + + // set_config() only takes text; a value we can't read as text is one we + // can't track. + let command = test.execute(vec![ + Parse::named( + "__test_set_config_bytes", + "SELECT set_config($1, $2, false)", + ) + .into(), + Bind::new_params( + "__test_set_config_bytes", + &[ + Parameter::new(b"search_path"), + Parameter::new(&[0xff, 0xfe]), + ], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!( + matches!(command, Command::Query(_)), + "expected Command::Query, got {command:#?}", + ); + assert!(command.route().is_write()); +} From ec1bb54ce037c5e5fdc7550f7795fac746290f77 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Tue, 4 Aug 2026 17:57:51 +0300 Subject: [PATCH 06/11] Let Postgres answer set_config() instead of faking it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `SELECT set_config(...)` was intercepted and answered locally, which meant inventing a response for a query the client asked Postgres to run: the tag was SET where Postgres sends SELECT 1, and describing the portal claimed no rows and then sent one, which libpq rejects outright. It is a query, so treat it as one — take a server, record what the statement changes on that connection, and forward it. The client gets Postgres' own answer, and the next client gets the connection with that parameter reset. Renamed the flag that marks these statements: it no longer describes how we imitate a SELECT, it says the statement is one. --- integration/pgdog.toml | 7 +++++++ .../python/test_session_params_leak.py | 9 ++++++++ pgdog/src/frontend/client/query_engine/mod.rs | 6 ++---- pgdog/src/frontend/client/query_engine/set.rs | 21 +++++++++++++++---- pgdog/src/frontend/router/parser/command.rs | 4 +++- pgdog/src/frontend/router/parser/query/set.rs | 5 +++-- 6 files changed, 41 insertions(+), 11 deletions(-) diff --git a/integration/pgdog.toml b/integration/pgdog.toml index b52035d6a..8269593f9 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -506,3 +506,10 @@ password = "pgdog" database = "pgdog" level = "auto" engine = "pg_query_raw" + +# Session state leaks are what these tests are about, so don't let the parser +# opt out of looking at the statements that cause them. +[[query_parsers]] +database = "pgdog_leak" +level = "on" +engine = "pg_query_raw" diff --git a/integration/python/test_session_params_leak.py b/integration/python/test_session_params_leak.py index 9af9c4533..b07059912 100644 --- a/integration/python/test_session_params_leak.py +++ b/integration/python/test_session_params_leak.py @@ -59,6 +59,15 @@ def test_set_committed_after_connecting(): assert read("statement_timeout") == "0" +def test_set_config_bound_params(): + """set_config() arguments arrive in the Bind message, not as constants.""" + conn = connect() + conn.execute("SELECT pg_catalog.set_config(%s, %s, false)", ("search_path", "")) + conn.close() + + assert read("search_path") not in ("", '""') + + def test_reset_committed(): """A committed RESET is permanent and needs no undoing.""" conn = connect() diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index d90994c3a..47366627c 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -236,12 +236,10 @@ impl QueryEngine { } Command::Unlisten(channel) => self.unlisten(context, &channel.clone()).await?, Command::Set { - params, - behave_like_select, - .. + params, is_select, .. } => { let params = params.clone(); - self.set(context, ¶ms, *behave_like_select).await?; + self.set(context, ¶ms, *is_select).await?; } Command::ResetAll => { self.reset_all(context).await?; diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index 0bac26a82..76fbdc242 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -15,13 +15,28 @@ impl QueryEngine { &mut self, context: &mut QueryEngineContext<'_>, params: &[SetParam], - behave_like_select: bool, + is_select: bool, ) -> Result<(), Error> { // Make sure client isn't changing route mid-transaction. if self.route_change_check(context, params).await? { return Ok(()); } + // `SELECT set_config(...)` is a query and Postgres answers it, so take a + // server before touching the parameters: syncing a change the statement + // is about to make itself would just send it twice. + if is_select && !self.backend.connected() { + let connected = if context.in_transaction() { + self.connect_transaction(context).await? + } else { + self.connect(context, None).await? + }; + + if !connected { + return Ok(()); + } + } + let mut fake_command = "SET"; for param in params { let is_pin = param.name == PGDOG_PIN; @@ -65,9 +80,7 @@ impl QueryEngine { self.backend.record_params(params, context.in_transaction()); self.execute(context).await?; } else { - let values_to_return = - behave_like_select.then(|| params.iter().map(|p| p.value.as_ref())); - self.fake_command_response(context, fake_command, values_to_return) + self.fake_command_response(context, fake_command, None::>) .await?; } diff --git a/pgdog/src/frontend/router/parser/command.rs b/pgdog/src/frontend/router/parser/command.rs index 9a50cb9cc..c0a0e14f9 100644 --- a/pgdog/src/frontend/router/parser/command.rs +++ b/pgdog/src/frontend/router/parser/command.rs @@ -32,7 +32,9 @@ pub enum Command { Set { params: Vec, route: Route, - behave_like_select: bool, + /// The statement is `SELECT set_config(...)`, not `SET`: Postgres has + /// to answer it, we only note what it changes. + is_select: bool, }, ResetAll, InternalField { diff --git a/pgdog/src/frontend/router/parser/query/set.rs b/pgdog/src/frontend/router/parser/query/set.rs index 08b652409..fa929d95e 100644 --- a/pgdog/src/frontend/router/parser/query/set.rs +++ b/pgdog/src/frontend/router/parser/query/set.rs @@ -27,7 +27,7 @@ impl QueryParser { Ok(Command::Set { params: vec![param], route: Route::write(context.shards_calculator.shard()), - behave_like_select: false, + is_select: false, }) } } @@ -99,7 +99,7 @@ impl QueryParser { Ok(Some(Command::Set { params, route: Route::write(context.shards_calculator.shard()), - behave_like_select: false, + is_select: false, })) } } @@ -130,4 +130,5 @@ impl QueryParser { Ok(value) } + } From 14ae66a3a10869c6d5eb577425c95f86a2764fdc Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Tue, 4 Aug 2026 23:59:47 +0300 Subject: [PATCH 07/11] Assert the value Postgres returns for a NULL set_config() The test pinned down the value PgDog made up while pretending to run the statement: the SetParam it parsed, so NULL for a reset. Postgres resets the setting and answers with the value it landed on, and that is what the client gets now that the statement reaches it. --- integration/rust/tests/integration/set_config.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/integration/rust/tests/integration/set_config.rs b/integration/rust/tests/integration/set_config.rs index 8d294d1a1..4affe217b 100644 --- a/integration/rust/tests/integration/set_config.rs +++ b/integration/rust/tests/integration/set_config.rs @@ -28,12 +28,13 @@ async fn test_set_config_behaves_like_set() { .unwrap(); assert_eq!(lock_timeout, "500s"); - let set_config: Option = - query_scalar("SELECT set_config('lock_timeout', NULL, false);") - .fetch_one(&mut *conn) - .await - .unwrap(); - assert_eq!(set_config, None); + // A NULL resets the setting, and set_config answers with the value the + // setting landed on, not with the NULL it was handed. + let set_config: String = query_scalar("SELECT set_config('lock_timeout', NULL, false);") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!(set_config, "0"); let lock_timeout: String = query_scalar("SHOW lock_timeout") .fetch_one(&mut *conn) From 282814a3f26fb55b19b628d9fffdd536c3619017 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Wed, 5 Aug 2026 00:40:06 +0300 Subject: [PATCH 08/11] Trim comments to what the code doesn't already say --- integration/pgdog.toml | 3 +-- .../rust/tests/integration/set_config.rs | 2 -- pgdog/src/frontend/client/query_engine/set.rs | 5 ++--- pgdog/src/frontend/router/parser/command.rs | 3 +-- .../frontend/router/parser/query/set_config.rs | 18 +++++++++++++++--- .../router/parser/query/test/test_set.rs | 4 +--- 6 files changed, 20 insertions(+), 15 deletions(-) diff --git a/integration/pgdog.toml b/integration/pgdog.toml index 8269593f9..1e5838b92 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -507,8 +507,7 @@ database = "pgdog" level = "auto" engine = "pg_query_raw" -# Session state leaks are what these tests are about, so don't let the parser -# opt out of looking at the statements that cause them. +# These tests are about statements the parser would otherwise opt out of. [[query_parsers]] database = "pgdog_leak" level = "on" diff --git a/integration/rust/tests/integration/set_config.rs b/integration/rust/tests/integration/set_config.rs index 4affe217b..15e9b03d6 100644 --- a/integration/rust/tests/integration/set_config.rs +++ b/integration/rust/tests/integration/set_config.rs @@ -28,8 +28,6 @@ async fn test_set_config_behaves_like_set() { .unwrap(); assert_eq!(lock_timeout, "500s"); - // A NULL resets the setting, and set_config answers with the value the - // setting landed on, not with the NULL it was handed. let set_config: String = query_scalar("SELECT set_config('lock_timeout', NULL, false);") .fetch_one(&mut *conn) .await diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index 76fbdc242..a9f76eae2 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -22,9 +22,8 @@ impl QueryEngine { return Ok(()); } - // `SELECT set_config(...)` is a query and Postgres answers it, so take a - // server before touching the parameters: syncing a change the statement - // is about to make itself would just send it twice. + // Take a server before touching the parameters: syncing a change the + // statement is about to make itself would send it twice. if is_select && !self.backend.connected() { let connected = if context.in_transaction() { self.connect_transaction(context).await? diff --git a/pgdog/src/frontend/router/parser/command.rs b/pgdog/src/frontend/router/parser/command.rs index c0a0e14f9..2526e6811 100644 --- a/pgdog/src/frontend/router/parser/command.rs +++ b/pgdog/src/frontend/router/parser/command.rs @@ -32,8 +32,7 @@ pub enum Command { Set { params: Vec, route: Route, - /// The statement is `SELECT set_config(...)`, not `SET`: Postgres has - /// to answer it, we only note what it changes. + /// `SELECT set_config(...)`, not `SET`: Postgres has to answer it. is_select: bool, }, ResetAll, diff --git a/pgdog/src/frontend/router/parser/query/set_config.rs b/pgdog/src/frontend/router/parser/query/set_config.rs index d709bc5a1..5cebf906f 100644 --- a/pgdog/src/frontend/router/parser/query/set_config.rs +++ b/pgdog/src/frontend/router/parser/query/set_config.rs @@ -4,8 +4,8 @@ use crate::net::messages::{Bind, Format}; impl QueryParser { /// Handle SELECT set_config('key', 'value', is_local) /// - /// Arguments we can't resolve leave the statement an ordinary query: it - /// runs, but we don't know what it changed. + /// Arguments we can't resolve leave it an ordinary query: it runs, but + /// we don't learn what it changed. pub(super) fn set_config( &mut self, fcall: &nodes::FuncCall, @@ -34,7 +34,19 @@ fn parse_args(fcall: &nodes::FuncCall, bind: Option<&Bind>) -> Option Some(SetParam { name, value, local }) } -/// Get the value bound to `$number`. The inner Option is the SQL NULL. +cfg_select! { + not(feature = "new_parser") => { + fn parse_args(fcall: &FuncCall, bind: Option<&Bind>) -> Option { + let name = parse_config_name(fcall.args.first()?, bind)?; + let value = parse_config_value(fcall.args.get(1)?, bind)?; + let local = parse_is_local(fcall.args.get(2)?, bind)?; + Some(SetParam { name, value, local }) + } + } + _ => {} +} + +/// Value bound to `$number`; the inner Option is the SQL NULL. fn bound_text(bind: Option<&Bind>, number: i32) -> Option> { let index = usize::try_from(number).ok()?.checked_sub(1)?; let param = bind?.parameter(index).ok()??; diff --git a/pgdog/src/frontend/router/parser/query/test/test_set.rs b/pgdog/src/frontend/router/parser/query/test/test_set.rs index e035abc62..5586defba 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_set.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_set.rs @@ -318,7 +318,6 @@ fn test_set_config_bound_null_value() { fn test_set_config_unresolvable_args_stay_a_query() { let mut test = QueryParserTest::new(); - // No Bind message, so the parameters can't be resolved. let command = test.execute(vec![ Query::new("SELECT pg_catalog.set_config($1, $2, false)").into(), ]); @@ -380,8 +379,7 @@ fn test_set_config_bound_binary_is_local() { fn test_set_config_bound_non_utf8_stays_a_query() { let mut test = QueryParserTest::new(); - // set_config() only takes text; a value we can't read as text is one we - // can't track. + // set_config() takes text; what we can't read as text we can't track. let command = test.execute(vec![ Parse::named( "__test_set_config_bytes", From 1bd0195d4ffb1087de98c4eaf6a6f8739c865713 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Fri, 7 Aug 2026 16:48:20 +0300 Subject: [PATCH 09/11] Run the leak tests at both parser levels, and against RLS The fixture pinned the leak database to level "on", which is the one level that always parses. Everything a statement has to get past to reach the parser at the default level went untested, and set_config() is a function call inside a SELECT: it matches no statement-start keyword, so the gate drops it and the value stays on the connection. A second copy of the same single-primary database at "auto" covers that path. Both copies run every test. The new test covers the half of the leak that reports nothing: row-level security keyed on a custom GUC is how multi-tenant applications isolate tenants, set_config() is how that GUC gets set, and a value that outlives its client makes the next one read as the previous tenant. Reads go through a plain role because the pooler's user is a superuser and superusers ignore RLS. --- integration/pgdog.toml | 17 +++ .../python/test_session_params_leak.py | 117 +++++++++++++++--- integration/users.toml | 5 + 3 files changed, 123 insertions(+), 16 deletions(-) diff --git a/integration/pgdog.toml b/integration/pgdog.toml index 1e5838b92..e2cacb8b0 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -56,6 +56,8 @@ read_only = true # ------------------------------------------------------------------------------ # ----- Database :: pgdog_leak ------------------------------------------------- # One server connection, never reaped: the next client has to get the same one. +# Two copies of the same single-primary database, differing only in parser level, +# so the leak tests cover both ways a statement reaches the parser. [[databases]] name = "pgdog_leak" @@ -64,6 +66,13 @@ database_name = "pgdog" pool_size = 1 min_pool_size = 1 +[[databases]] +name = "pgdog_leak_auto" +host = "127.0.0.1" +database_name = "pgdog" +pool_size = 1 +min_pool_size = 1 + # ------------------------------------------------------------------------------ # ----- Database :: pgdog_sharded ---------------------------------------------- @@ -512,3 +521,11 @@ engine = "pg_query_raw" database = "pgdog_leak" level = "on" engine = "pg_query_raw" + +# The same tests at the default level. A single primary with no replicas is the +# topology where "auto" doesn't force the parser on, so the statement has to get +# past the regex gate on its own -- the path "on" skips. +[[query_parsers]] +database = "pgdog_leak_auto" +level = "auto" +engine = "pg_query_raw" diff --git a/integration/python/test_session_params_leak.py b/integration/python/test_session_params_leak.py index b07059912..47f6f4cd1 100644 --- a/integration/python/test_session_params_leak.py +++ b/integration/python/test_session_params_leak.py @@ -3,16 +3,27 @@ Runs against a database with a single server connection, so the next client always gets the connection the previous one used. current_setting() is used instead of SHOW because SHOW can be answered by PgDog itself. + +Every test runs against two copies of that database that differ only in parser +level: "on" always parses, "auto" leaves a single-primary cluster to the regex +gate. A statement that only the gate can let through -- set_config(), a function +call inside a SELECT rather than a statement-start keyword -- reaches the parser +on one and has to earn it on the other. """ +import uuid + import psycopg +import pytest +DATABASES = ["pgdog_leak", "pgdog_leak_auto"] -def connect(): + +def connect(dbname): conn = psycopg.connect( user="pgdog", password="pgdog", - dbname="pgdog_leak", + dbname=dbname, host="127.0.0.1", port=6432, ) @@ -22,21 +33,26 @@ def connect(): return conn -def read(setting): - conn = connect() +def read(dbname, setting): + conn = connect(dbname) value = conn.execute(f"SELECT current_setting('{setting}')").fetchone()[0] conn.close() return value -def test_reset_rolled_back(): +@pytest.fixture(params=DATABASES) +def dbname(request): + return request.param + + +def test_reset_rolled_back(dbname): """A ROLLBACK brings back the value the RESET cleared. This is the sequence pg_dump -t
emits. SET search_path TO '' leaves an empty quoted identifier, hence the two spellings of "empty". """ - conn = connect() + conn = connect(dbname) conn.execute("SET search_path TO ''") with conn.transaction(): conn.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY") @@ -45,36 +61,105 @@ def test_reset_rolled_back(): raise psycopg.Rollback() conn.close() - assert read("search_path") not in ("", '""') + assert read(dbname, "search_path") not in ("", '""') -def test_set_committed_after_connecting(): +def test_set_committed_after_connecting(dbname): """A SET that lands once the connection is already ours.""" - conn = connect() + conn = connect(dbname) with conn.transaction(): conn.execute("SELECT 1") conn.execute("SET statement_timeout TO '5s'") conn.close() - assert read("statement_timeout") == "0" + assert read(dbname, "statement_timeout") == "0" -def test_set_config_bound_params(): +def test_set_config_bound_params(dbname): """set_config() arguments arrive in the Bind message, not as constants.""" - conn = connect() + conn = connect(dbname) conn.execute("SELECT pg_catalog.set_config(%s, %s, false)", ("search_path", "")) conn.close() - assert read("search_path") not in ("", '""') + assert read(dbname, "search_path") not in ("", '""') -def test_reset_committed(): +def test_reset_committed(dbname): """A committed RESET is permanent and needs no undoing.""" - conn = connect() + conn = connect(dbname) conn.execute("SET search_path TO public") with conn.transaction(): conn.execute("SELECT 1") conn.execute("RESET search_path") conn.close() - assert read("search_path") == '"$user", public' + assert read(dbname, "search_path") == '"$user", public' + + +TENANT_A = "11111111-1111-1111-1111-111111111111" +TENANT_B = "22222222-2222-2222-2222-222222222222" + + +@pytest.fixture +def tenants(dbname): + """A table whose rows are visible only to the tenant named in a GUC. + + Reads go through a plain role: the pooler's own user is a superuser, and + superusers ignore row-level security however the table is configured. + + NULLIF() is deliberate: a targeted RESET leaves the placeholder GUC as an + empty string rather than unset, and '' would fail the ::uuid cast. + """ + table = "rls_probe_" + uuid.uuid4().hex[:8] + conn = connect(dbname) + conn.execute( + "DO $$ BEGIN CREATE ROLE rls_tenant NOLOGIN; " + "EXCEPTION WHEN duplicate_object THEN NULL; END $$" + ) + conn.execute(f"CREATE TABLE public.{table} (org_id uuid, note text)") + conn.execute( + f"INSERT INTO public.{table} VALUES ('{TENANT_A}', 'a'), ('{TENANT_B}', 'b')" + ) + conn.execute(f"GRANT SELECT ON public.{table} TO rls_tenant") + conn.execute(f"ALTER TABLE public.{table} ENABLE ROW LEVEL SECURITY") + conn.execute( + f"CREATE POLICY tenant_isolation ON public.{table} USING " + "(org_id = NULLIF(current_setting('app.current_org_id', true), '')::uuid)" + ) + conn.close() + + yield table + + conn = connect(dbname) + conn.execute("RESET ROLE") + conn.execute(f"DROP TABLE public.{table}") + conn.close() + + +def test_tenant_guc_does_not_outlive_its_client(dbname, tenants): + """The silent half of the leak: no error, just another tenant's rows. + + Row-level security keyed on a custom GUC is how multi-tenant applications + isolate tenants, and set_config() with a bound parameter is how that GUC + gets set. A value that survives checkin makes the next client read as the + previous tenant. + """ + first = connect(dbname) + first.execute("SET ROLE rls_tenant") + first.execute( + "SELECT pg_catalog.set_config('app.current_org_id', %s, false)", (TENANT_A,) + ) + mine = first.execute(f"SELECT note FROM public.{tenants}").fetchall() + first.close() + + assert mine == [("a",)], "the tenant that set the GUC sees its own row" + + second = connect(dbname) + second.execute("SET ROLE rls_tenant") + theirs = second.execute(f"SELECT note FROM public.{tenants}").fetchall() + leaked = second.execute( + "SELECT current_setting('app.current_org_id', true)" + ).fetchone()[0] + second.close() + + assert theirs == [], f"next client read as tenant {leaked!r}" diff --git a/integration/users.toml b/integration/users.toml index 360246b5b..25a24d84f 100644 --- a/integration/users.toml +++ b/integration/users.toml @@ -8,6 +8,11 @@ name = "pgdog" database = "pgdog_leak" password = "pgdog" +[[users]] +name = "pgdog" +database = "pgdog_leak_auto" +password = "pgdog" + [[users]] name = "pgdog_migrator" database = "pgdog" From ff3c38736aac1e66e70ba08e3f783ea915dbe5f6 Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Sat, 8 Aug 2026 02:26:08 +0300 Subject: [PATCH 10/11] Finish removing the old parser's copies after the rebase main dropped the second parser in #1324, and this branch still carried the changes for both. Resolving the rebase left one cfg_select! block and a stray blank line behind. --- pgdog/src/frontend/router/parser/query/set.rs | 1 - .../frontend/router/parser/query/set_config.rs | 16 ---------------- 2 files changed, 17 deletions(-) diff --git a/pgdog/src/frontend/router/parser/query/set.rs b/pgdog/src/frontend/router/parser/query/set.rs index fa929d95e..b28d86602 100644 --- a/pgdog/src/frontend/router/parser/query/set.rs +++ b/pgdog/src/frontend/router/parser/query/set.rs @@ -130,5 +130,4 @@ impl QueryParser { Ok(value) } - } diff --git a/pgdog/src/frontend/router/parser/query/set_config.rs b/pgdog/src/frontend/router/parser/query/set_config.rs index 5cebf906f..583965405 100644 --- a/pgdog/src/frontend/router/parser/query/set_config.rs +++ b/pgdog/src/frontend/router/parser/query/set_config.rs @@ -23,7 +23,6 @@ impl QueryParser { ) } } - } /// Returns None if the arguments could not be parsed @@ -34,18 +33,6 @@ fn parse_args(fcall: &nodes::FuncCall, bind: Option<&Bind>) -> Option Some(SetParam { name, value, local }) } -cfg_select! { - not(feature = "new_parser") => { - fn parse_args(fcall: &FuncCall, bind: Option<&Bind>) -> Option { - let name = parse_config_name(fcall.args.first()?, bind)?; - let value = parse_config_value(fcall.args.get(1)?, bind)?; - let local = parse_is_local(fcall.args.get(2)?, bind)?; - Some(SetParam { name, value, local }) - } - } - _ => {} -} - /// Value bound to `$number`; the inner Option is the SQL NULL. fn bound_text(bind: Option<&Bind>, number: i32) -> Option> { let index = usize::try_from(number).ok()?.checked_sub(1)?; @@ -89,8 +76,6 @@ fn parse_config_name(arg: Node<'_>, bind: Option<&Bind>) -> Option { } } -/// Returns None if the name could not be parsed - /// Returns None if the value could not be parsed, Some(None) if the value /// is NULL, and Some if the value was successfully parsed fn parse_config_value(arg: Node<'_>, bind: Option<&Bind>) -> Option> { @@ -116,4 +101,3 @@ fn parse_is_local(arg: Node<'_>, bind: Option<&Bind>) -> Option { _ => None, } } - From 29a631acf4d7930f803515efa417413a619d57ba Mon Sep 17 00:00:00 2001 From: Igor Ohrimenko Date: Sat, 8 Aug 2026 02:49:27 +0300 Subject: [PATCH 11/11] Make the tenant test prove the clients shared a connection It passed in CI and failed locally, which means it was answering a question it never asked: with a different server connection there is nothing for the first client to have left behind, and the assertion holds for the wrong reason. Compare pg_backend_pid() across the two clients, and separate the GUC outliving its client from row-level security failing to filter, so a failure says which of the two happened. --- integration/python/test_session_params_leak.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/integration/python/test_session_params_leak.py b/integration/python/test_session_params_leak.py index 47f6f4cd1..e044146a1 100644 --- a/integration/python/test_session_params_leak.py +++ b/integration/python/test_session_params_leak.py @@ -150,6 +150,7 @@ def test_tenant_guc_does_not_outlive_its_client(dbname, tenants): "SELECT pg_catalog.set_config('app.current_org_id', %s, false)", (TENANT_A,) ) mine = first.execute(f"SELECT note FROM public.{tenants}").fetchall() + served_by = first.execute("SELECT pg_backend_pid()").fetchone()[0] first.close() assert mine == [("a",)], "the tenant that set the GUC sees its own row" @@ -160,6 +161,11 @@ def test_tenant_guc_does_not_outlive_its_client(dbname, tenants): leaked = second.execute( "SELECT current_setting('app.current_org_id', true)" ).fetchone()[0] + same_server = second.execute("SELECT pg_backend_pid()").fetchone()[0] second.close() - assert theirs == [], f"next client read as tenant {leaked!r}" + # Without this the test passes whenever the pool happens to hand out a + # different connection, which proves nothing about what the first one left. + assert same_server == served_by, "clients did not share a server connection" + assert leaked in (None, ""), f"the tenant GUC outlived its client: {leaked!r}" + assert theirs == [], "next client read the previous tenant's rows"