Resolve set_config() arguments from the Bind message and let Postgres answer it

integration/pgdog.toml

@@ -59,6 +59,17 @@ host = "127.0.0.1"
 role = "replica"
 read_only = true
 
+# ------------------------------------------------------------------------------
+# ----- Database :: pgdog_leak -------------------------------------------------
+# One server connection, never reaped: the next client has to get the same one.
+
+[[databases]]
+name = "pgdog_leak"
+host = "127.0.0.1"
+database_name = "pgdog"
+pool_size = 1
+min_pool_size = 1
+
 # ------------------------------------------------------------------------------
 # ----- Database :: pgdog_sharded ----------------------------------------------
 
@@ -505,3 +516,9 @@ password = "pgdog"
 database = "pgdog"
 level = "auto"
 engine = "pg_query_raw"
+
+# These tests are about statements the parser would otherwise opt out of.
+[[query_parsers]]
+database = "pgdog_leak"
+level = "on"
+engine = "pg_query_raw"

integration/python/gate/pgdog.toml

@@ -0,0 +1,18 @@
+[general]
+prepared_statements = "extended"
+
+[[databases]]
+name = "pgdog"
+host = "127.0.0.1"
+
+[[databases]]
+name = "pgdog_leak_auto"
+host = "127.0.0.1"
+database_name = "pgdog"
+pool_size = 1
+min_pool_size = 1
+
+[[query_parsers]]
+database = "pgdog_leak_auto"
+level = "auto"
+engine = "pg_query_raw"

integration/python/gate/users.toml

@@ -0,0 +1,9 @@
+[[users]]
+name = "pgdog"
+database = "pgdog"
+password = "pgdog"
+
+[[users]]
+name = "pgdog"
+database = "pgdog_leak_auto"
+password = "pgdog"

integration/python/run.sh

@@ -11,3 +11,13 @@ wait_for_pgdog
 source ${SCRIPT_DIR}/dev.sh
 
 stop_pgdog
+run_pgdog integration/python/gate
+wait_for_pgdog
+
+pushd ${SCRIPT_DIR}
+source venv/bin/activate
+PGDOG_LEAK_DATABASES=pgdog_leak_auto pytest -x test_session_params_leak.py
+deactivate
+popd
+
+stop_pgdog

integration/python/test_session_params_leak.py

@@ -0,0 +1,184 @@
+"""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.
+
+Every test runs at two parser levels: "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. "auto" needs its own config (gate/), where nothing else turns the parser
+on; run.sh runs these tests a second time against it.
+"""
+
+import os
+import uuid
+
+import psycopg
+import pytest
+
+DATABASES = os.environ.get("PGDOG_LEAK_DATABASES", "pgdog_leak").split(",")
+
+
+def connect(dbname):
+    conn = psycopg.connect(
+        user="pgdog",
+        password="pgdog",
+        dbname=dbname,
+        host="127.0.0.1",
+        port=6432,
+    )
+    # 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
+
+
+def read(dbname, setting):
+    conn = connect(dbname)
+    value = conn.execute(f"SELECT current_setting('{setting}')").fetchone()[0]
+    conn.close()
+
+    return value
+
+
+@pytest.fixture(params=DATABASES)
+def dbname(request):
+    return request.param
+
+
+@pytest.fixture
+def gated(request, dbname):
+    if dbname == "pgdog_leak_auto":
+        request.node.add_marker(
+            pytest.mark.xfail(
+                strict=True,
+                reason="the regex gate doesn't match set_config() yet (#1327)",
+            )
+        )
+
+
+def test_reset_rolled_back(dbname):
+    """A ROLLBACK brings back the value the RESET cleared.
+
+    This is the sequence pg_dump -t <table> emits. SET search_path TO '' leaves
+    an empty quoted identifier, hence the two spellings of "empty".
+    """
+    conn = connect(dbname)
+    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(dbname, "search_path") not in ("", '""')
+
+
+def test_set_committed_after_connecting(dbname):
+    """A SET that lands once the connection is already ours."""
+    conn = connect(dbname)
+    with conn.transaction():
+        conn.execute("SELECT 1")
+        conn.execute("SET statement_timeout TO '5s'")
+    conn.close()
+
+    assert read(dbname, "statement_timeout") == "0"
+
+
+def test_set_config_bound_params(dbname, gated):
+    """set_config() arguments arrive in the Bind message, not as constants."""
+    conn = connect(dbname)
+    conn.execute("SELECT pg_catalog.set_config(%s, %s, false)", ("search_path", ""))
+    conn.close()
+
+    assert read(dbname, "search_path") not in ("", '""')
+
+
+def test_reset_committed(dbname):
+    """A committed RESET is permanent and needs no undoing."""
+    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(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, gated):
+    """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()
+    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"
+
+    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]
+    same_server = second.execute("SELECT pg_backend_pid()").fetchone()[0]
+    second.close()
+
+    # 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"

integration/rust/tests/integration/set_config.rs

@@ -28,12 +28,11 @@ async fn test_set_config_behaves_like_set() {
         .unwrap();
     assert_eq!(lock_timeout, "500s");
 
-    let set_config: Option<String> =
-        query_scalar("SELECT set_config('lock_timeout', NULL, false);")
-            .fetch_one(&mut *conn)
-            .await
-            .unwrap();
-    assert_eq!(set_config, None);
+    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)

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"

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},
@@ -443,6 +443,26 @@ impl Binding {
         }
     }
 
+    pub(crate) 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)),
+            _ => (),
+        }
+    }
+
+    pub(crate) 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(crate) fn transaction_params_hook(&mut self, rollback: bool) {
         match self {