Fix role leakage between connections

pgdog/src/backend/pool/cleanup.rs

@@ -14,6 +14,7 @@ static PREPARED: Lazy<Vec<Query>> = Lazy::new(|| vec![Query::new("DEALLOCATE ALL
 static DIRTY: Lazy<Vec<Query>> = Lazy::new(|| {
     vec![
         Query::new("RESET ALL"),                       // Reset all parameters.
+        Query::new("RESET SESSION AUTHORIZATION"),     // Reset all skips session_authorization.
         Query::new("SELECT pg_advisory_unlock_all()"), // Remove all advisory locks.
         Query::new("DISCARD TEMP"),                    // Drop all temporary tables.
     ]

pgdog/src/backend/pool/guard.rs

@@ -314,6 +314,83 @@ mod test {
         drop(guard);
     }
 
+    /// A client set a role and the connection was returned dirty. The next client
+    /// must not inherit it.
+    ///
+    /// `RESET ALL` does not clear `role`: Postgres flags it `GUC_NO_RESET_ALL` and
+    /// `ResetAllOptions()` skips such settings, so cleanup has to reset it explicitly.
+    ///
+    /// <https://github.com/pgdogdev/pgdog/issues/1341>
+    #[tokio::test]
+    async fn test_cleanup_resets_role() {
+        crate::logger();
+        let pool = pool();
+
+        let mut guard = pool.get(&Request::default()).await.unwrap();
+        let pid_before: Vec<i32> = guard.fetch_all("SELECT pg_backend_pid()").await.unwrap();
+
+        guard.execute_checked("SET ROLE pgdog1").await.unwrap();
+        let role: Vec<String> = guard.fetch_all("SELECT current_user").await.unwrap();
+        assert_eq!(role, vec!["pgdog1".to_string()]);
+
+        guard.mark_dirty(true);
+        drop(guard);
+
+        // Our test pool is only 1 connection, so this is the same backend.
+        let mut guard = pool.get(&Request::default()).await.unwrap();
+        let pid_after: Vec<i32> = guard.fetch_all("SELECT pg_backend_pid()").await.unwrap();
+        assert_eq!(
+            pid_before, pid_after,
+            "1-connection pool should hand back the same backend"
+        );
+
+        let role: Vec<String> = guard.fetch_all("SELECT current_user").await.unwrap();
+        assert_eq!(
+            role,
+            vec!["pgdog".to_string()],
+            "SET ROLE leaked across a dirty check-in"
+        );
+    }
+
+    /// Same as [`test_cleanup_resets_role`], for `SET SESSION AUTHORIZATION`.
+    ///
+    /// This one leaks `session_user` as well, and is not covered by `RESET ROLE`:
+    /// only `RESET SESSION AUTHORIZATION` restores the authenticated user.
+    ///
+    /// <https://github.com/pgdogdev/pgdog/issues/1341>
+    #[tokio::test]
+    async fn test_cleanup_resets_session_authorization() {
+        crate::logger();
+        let pool = pool();
+
+        let mut guard = pool.get(&Request::default()).await.unwrap();
+        let pid_before: Vec<i32> = guard.fetch_all("SELECT pg_backend_pid()").await.unwrap();
+
+        guard
+            .execute_checked("SET SESSION AUTHORIZATION pgdog1")
+            .await
+            .unwrap();
+        let session_user: Vec<String> = guard.fetch_all("SELECT session_user").await.unwrap();
+        assert_eq!(session_user, vec!["pgdog1".to_string()]);
+
+        guard.mark_dirty(true);
+        drop(guard);
+
+        let mut guard = pool.get(&Request::default()).await.unwrap();
+        let pid_after: Vec<i32> = guard.fetch_all("SELECT pg_backend_pid()").await.unwrap();
+        assert_eq!(
+            pid_before, pid_after,
+            "1-connection pool should hand back the same backend"
+        );
+
+        let session_user: Vec<String> = guard.fetch_all("SELECT session_user").await.unwrap();
+        assert_eq!(
+            session_user,
+            vec!["pgdog".to_string()],
+            "SET SESSION AUTHORIZATION leaked across a dirty check-in"
+        );
+    }
+
     #[tokio::test]
     async fn test_cleanup_prepared_statements() {
         crate::logger();

pgdog/src/frontend/client/query_engine/test/mod.rs

@@ -31,6 +31,7 @@ mod rewrite_extended;
 mod rewrite_insert_split;
 mod rewrite_offset;
 mod rewrite_simple_prepared;
+mod role;
 mod schema_changed;
 mod set;
 mod set_schema_sharding;

pgdog/src/frontend/client/query_engine/test/role.rs

@@ -0,0 +1,238 @@
+//! `SET ROLE` must not leak from one client to the next in transaction mode.
+//!
+//! <https://github.com/pgdogdev/pgdog/issues/1341>
+
+use crate::{
+    backend::databases::reload_from_existing,
+    config::{config, load_test, set},
+    expect_message,
+    net::{CommandComplete, DataRow, ReadyForQuery, RowDescription},
+};
+
+use super::prelude::*;
+
+/// Run a statement that returns no rows.
+///
+/// Tolerates `ParameterStatus`: `session_authorization` is a reported (GUC_REPORT)
+/// parameter, so `SET SESSION AUTHORIZATION` emits an extra 'S' message that
+/// `SET ROLE` does not.
+async fn run_simple(client: &mut TestClient, query: &str) -> ReadyForQuery {
+    client.send_simple(Query::new(query)).await;
+
+    let mut command_complete = false;
+    loop {
+        let message = client.read().await;
+        match message.code() {
+            'S' => continue,
+            'C' => {
+                expect_message!(message, CommandComplete);
+                command_complete = true;
+            }
+            _ => {
+                assert!(
+                    command_complete,
+                    "expected CommandComplete before ReadyForQuery for {query:?}"
+                );
+                return expect_message!(message, ReadyForQuery);
+            }
+        }
+    }
+}
+
+/// Read a single-row, single-column text result through the proxy.
+async fn fetch_text(client: &mut TestClient, query: &str) -> String {
+    client.send_simple(Query::new(query)).await;
+    expect_message!(client.read().await, RowDescription);
+    let row = expect_message!(client.read().await, DataRow);
+    let value = row.get_text(0).expect("one text column");
+    client.read_until('Z').await.unwrap();
+    value
+}
+
+fn load_single_connection_test_pool() {
+    load_test();
+
+    let mut config = (*config()).clone();
+    config.config.general.default_pool_size = 1;
+    config.config.general.min_pool_size = 0;
+    set(config).unwrap();
+    reload_from_existing().unwrap();
+}
+
+/// The reported bug. Pinning the backend marks it dirty, so check-in runs the
+/// `DIRTY` cleanup queries — and `RESET ALL` does not clear `role`. The role
+/// survives on the backend while check-in clears `client_params`, so the
+/// check-out path no longer knows to reset it either.
+#[tokio::test]
+async fn test_set_role_does_not_leak_to_next_client() {
+    load_single_connection_test_pool();
+
+    let pinned_pid = {
+        // `leak_pool`: dropping a TestClient otherwise shuts the pools down, and the
+        // second client would get a brand new backend, making the assertion vacuous.
+        let mut client = TestClient::new(Parameters::default()).await.leak_pool();
+
+        assert_eq!(
+            run_simple(&mut client, "SET pgdog.pin TO true")
+                .await
+                .status,
+            'I'
+        );
+
+        // Attaches and locks the backend. `SET ROLE` has to come after this: with no
+        // backend attached the SET is answered locally and only materialises at
+        // check-out, so it would never reach this connection.
+        let pid = client.backend_pid().await;
+        assert!(client.backend_locked());
+
+        assert_eq!(run_simple(&mut client, "SET ROLE pgdog1").await.status, 'I');
+        assert_eq!(
+            fetch_text(&mut client, "SELECT current_user").await,
+            "pgdog1"
+        );
+
+        pid
+    };
+
+    let mut next = TestClient::new(Parameters::default()).await;
+    assert_eq!(
+        next.backend_pid().await,
+        pinned_pid,
+        "single connection test pool should reuse the same backend"
+    );
+    assert_eq!(
+        fetch_text(&mut next, "SELECT current_user").await,
+        "pgdog",
+        "SET ROLE leaked to the next client"
+    );
+}
+
+/// Same shape for `SET SESSION AUTHORIZATION`, which also leaks `session_user`.
+/// `session_authorization` is in `UNTRACKED_PARAMS`, so unlike `role` it is never
+/// synced or reset by the check-out path at all.
+///
+/// Requires the connecting user to be a superuser; `integration/setup.sh` creates
+/// `pgdog` and `pgdog1` as `LOGIN SUPERUSER`.
+#[tokio::test]
+async fn test_set_session_authorization_does_not_leak_to_next_client() {
+    load_single_connection_test_pool();
+
+    let pinned_pid = {
+        let mut client = TestClient::new(Parameters::default()).await.leak_pool();
+
+        assert_eq!(
+            run_simple(&mut client, "SET pgdog.pin TO true")
+                .await
+                .status,
+            'I'
+        );
+
+        let pid = client.backend_pid().await;
+        assert!(client.backend_locked());
+
+        assert_eq!(
+            run_simple(&mut client, "SET SESSION AUTHORIZATION pgdog1")
+                .await
+                .status,
+            'I'
+        );
+        assert_eq!(
+            fetch_text(&mut client, "SELECT session_user").await,
+            "pgdog1"
+        );
+
+        pid
+    };
+
+    let mut next = TestClient::new(Parameters::default()).await;
+    assert_eq!(
+        next.backend_pid().await,
+        pinned_pid,
+        "single connection test pool should reuse the same backend"
+    );
+    assert_eq!(
+        fetch_text(&mut next, "SELECT session_user").await,
+        "pgdog",
+        "SET SESSION AUTHORIZATION leaked to the next client"
+    );
+}
+
+/// `SET ROLE` issued *after* a backend is attached, with no pin anywhere.
+///
+/// A plain (non-`LOCAL`) `SET` survives `COMMIT` in Postgres, so the role stays on
+/// the connection once the transaction ends. But `client_params` is only populated
+/// at check-out (`Server::link_client`), and in-transaction sets are tracked
+/// separately, so nothing records that this connection now carries a role — and the
+/// next check-out has nothing to reset.
+///
+/// This is the scenario cleanup cannot reach: no pin means the connection is never
+/// dirty, so the `DIRTY` queries never run.
+#[tokio::test]
+async fn test_set_role_in_transaction_does_not_leak() {
+    load_single_connection_test_pool();
+
+    let pid = {
+        let mut client = TestClient::new(Parameters::default()).await.leak_pool();
+
+        assert_eq!(run_simple(&mut client, "BEGIN").await.status, 'T');
+
+        // Attaches the backend, so the SET below lands on it directly.
+        let pid = client.backend_pid().await;
+
+        assert_eq!(run_simple(&mut client, "SET ROLE pgdog1").await.status, 'T');
+        assert_eq!(
+            fetch_text(&mut client, "SELECT current_user").await,
+            "pgdog1"
+        );
+
+        assert_eq!(run_simple(&mut client, "COMMIT").await.status, 'I');
+
+        pid
+    };
+
+    let mut next = TestClient::new(Parameters::default()).await;
+    assert_eq!(
+        next.backend_pid().await,
+        pid,
+        "single connection test pool should reuse the same backend"
+    );
+    assert_eq!(
+        fetch_text(&mut next, "SELECT current_user").await,
+        "pgdog",
+        "SET ROLE inside a transaction leaked to the next client"
+    );
+}
+
+/// Without a pin the connection is never dirty, so no cleanup runs, `client_params`
+/// still records `role`, and the check-out path resets it. This passes before the
+/// fix as well as after it — it is here to document that the pin is what breaks the
+/// invariant, and to catch a regression in the check-out reset path.
+#[tokio::test]
+async fn test_set_role_without_pin_does_not_leak() {
+    load_single_connection_test_pool();
+
+    let pid = {
+        let mut client = TestClient::new(Parameters::default()).await.leak_pool();
+
+        let pid = client.backend_pid().await;
+        assert_eq!(run_simple(&mut client, "SET ROLE pgdog1").await.status, 'I');
+        assert_eq!(
+            fetch_text(&mut client, "SELECT current_user").await,
+            "pgdog1"
+        );
+
+        pid
+    };
+
+    let mut next = TestClient::new(Parameters::default()).await;
+    assert_eq!(
+        next.backend_pid().await,
+        pid,
+        "single connection test pool should reuse the same backend"
+    );
+    assert_eq!(
+        fetch_text(&mut next, "SELECT current_user").await,
+        "pgdog",
+        "SET ROLE leaked to the next client without a pin"
+    );
+}