diff --git a/include/compatibility.h b/include/compatibility.h index 35587d5..771402e 100644 --- a/include/compatibility.h +++ b/include/compatibility.h @@ -614,6 +614,37 @@ CastCreate(Oid sourcetypeid, Oid targettypeid, Oid funcid, char castcontext, #define WAIT_EVENT_MESSAGE_QUEUE_RECEIVE WAIT_EVENT_MQ_RECEIVE #endif +/* + * PostgreSQL 13 introduced ConditionVariableTimedSleep(). On PG 12 we + * emulate it with WaitLatch(); the wait-loop protocol (PrepareToSleep + + * CancelSleep) is otherwise identical, and returns true iff the timeout + * elapsed. + */ +#include "storage/condition_variable.h" +#include "storage/ipc.h" +#include "storage/latch.h" +#include "miscadmin.h" +#if PG_VERSION_NUM < 130000 +static inline bool +PGTLE_ConditionVariableTimedSleep(ConditionVariable *cv, long timeout_ms, + uint32 wait_event_info) +{ + int rc; + + (void) cv; + rc = WaitLatch(MyLatch, + WL_LATCH_SET | WL_TIMEOUT | WL_POSTMASTER_DEATH, + timeout_ms, wait_event_info); + ResetLatch(MyLatch); + if (rc & WL_POSTMASTER_DEATH) + proc_exit(1); + return (rc & WL_TIMEOUT) != 0; +} +#else +#define PGTLE_ConditionVariableTimedSleep(cv, t, w) \ + ConditionVariableTimedSleep((cv), (t), (w)) +#endif + /* * PostgreSQL version 18+ * diff --git a/src/clientauth.c b/src/clientauth.c index 1127d1a..6b02b15 100644 --- a/src/clientauth.c +++ b/src/clientauth.c @@ -96,6 +96,8 @@ * user function */ #define CLIENT_AUTH_USER_ERROR_MAX_STRLEN 256 +/* How long clientauth_hook() waits for a worker to attach before falling back. */ +#define CLIENTAUTH_WORKER_READY_TIMEOUT_MS 3000 /* * Fixed-length subset of Port, passed to user function. A corresponding SQL @@ -188,6 +190,21 @@ typedef struct ClientAuthBgwShmemSharedState /* Connection queue state */ ClientAuthStatusEntry requests[CLIENT_AUTH_MAX_PENDING_ENTRIES]; + + /* + * Per-shard liveness. worker_live[i] is set by worker i after its + * BackgroundWorkerInitializeConnection() succeeds and cleared by + * before_shmem_exit on any exit. Requests are sharded (requests[j] is + * owned by worker j % num_parallel_workers), so the hook must verify the + * target shard specifically. The worker publishes liveness because + * clientauth_hook runs before InitPostgres and cannot itself read + * pg_database. Only the first num_parallel_workers entries are used. + * + * worker_live_cv is signalled on the false->true transition; worker exits + * wake queued clients directly via each request's client_cv. + */ + bool worker_live[CLIENT_AUTH_MAX_PENDING_ENTRIES]; + ConditionVariable worker_live_cv; } ClientAuthBgwShmemSharedState; static const char *clientauth_shmem_name = "pgtle_clientauth"; @@ -212,6 +229,9 @@ static void clientauth_shmem_request(void); /* Helper functions */ static Size clientauth_shared_memsize(void); static void clientauth_sighup(SIGNAL_ARGS); +static bool wait_for_clientauth_worker_live(int bgw_idx); +static void clientauth_worker_before_shmem_exit(int code, Datum arg); +static void clientauth_fallback_no_worker(const char *reason); void clientauth_init(void); static bool can_allow_without_executing(void); @@ -377,6 +397,16 @@ clientauth_launcher_main(Datum arg) /* Initialize connection to the database */ BackgroundWorkerInitializeConnection(clientauth_database_name, NULL, 0); + /* + * Register cleanup before publishing, so any later FATAL still clears the + * flag. + */ + before_shmem_exit(clientauth_worker_before_shmem_exit, Int32GetDatum(bgw_idx)); + LWLockAcquire(clientauth_ss->lock, LW_EXCLUSIVE); + clientauth_ss->worker_live[bgw_idx] = true; + LWLockRelease(clientauth_ss->lock); + ConditionVariableBroadcast(&clientauth_ss->worker_live_cv); + /* Main worker loop */ while (true) { @@ -654,6 +684,16 @@ clientauth_hook(Port *port, int status) if (check_string_in_guc_list(port->database_name, clientauth_databases_to_skip, "pgtle.clientauth_databases_to_skip")) return; + /* + * Don't enqueue if the request's shard has no live worker (bad db, + * exhausted slots). + */ + if (!wait_for_clientauth_worker_live(idx % clientauth_num_parallel_workers)) + { + clientauth_fallback_no_worker("did not attach within timeout"); + return; + } + /* * If the queue entry is not available, wait until another client using it * has signalled that they are done @@ -721,17 +761,51 @@ clientauth_hook(Port *port, int status) clientauth_ss->requests[idx].done_processing = false; LWLockRelease(clientauth_ss->lock); - ConditionVariablePrepareToSleep(&clientauth_ss->requests[idx].client_cv); - while (true) + /* + * Wait for the worker to signal done_processing, polling worker liveness + * so a worker that dies mid-request doesn't hang us. On exit the lock is + * still held in LW_SHARED for the trailing "erase" block, unless the + * worker died -- in which case we roll the entry back and fall back. + */ { - LWLockAcquire(clientauth_ss->lock, LW_SHARED); - if (clientauth_ss->requests[idx].done_processing) - break; + int shard = idx % clientauth_num_parallel_workers; + bool worker_died = false; - LWLockRelease(clientauth_ss->lock); - ConditionVariableSleep(&clientauth_ss->requests[idx].client_cv, WAIT_EVENT_MESSAGE_QUEUE_RECEIVE); + ConditionVariablePrepareToSleep(&clientauth_ss->requests[idx].client_cv); + while (true) + { + LWLockAcquire(clientauth_ss->lock, LW_SHARED); + if (clientauth_ss->requests[idx].done_processing) + break; /* keep lock held */ + if (!clientauth_ss->worker_live[shard]) + { + LWLockRelease(clientauth_ss->lock); + worker_died = true; + break; + } + LWLockRelease(clientauth_ss->lock); + + CHECK_FOR_INTERRUPTS(); + (void) PGTLE_ConditionVariableTimedSleep(&clientauth_ss->requests[idx].client_cv, + 1000, + PG_WAIT_EXTENSION); + } + ConditionVariableCancelSleep(); + + if (worker_died) + { + LWLockAcquire(clientauth_ss->lock, LW_EXCLUSIVE); + clientauth_ss->requests[idx].pid = 0; + memset(&clientauth_ss->requests[idx].port_info, 0, sizeof(PortSubset)); + clientauth_ss->requests[idx].status = 0; + clientauth_ss->requests[idx].done_processing = true; + clientauth_ss->requests[idx].available_entry = true; + LWLockRelease(clientauth_ss->lock); + ConditionVariableSignal(clientauth_ss->requests[idx].available_entry_cv_ptr); + clientauth_fallback_no_worker("exited before returning"); + return; + } } - ConditionVariableCancelSleep(); /* Copy results of BGW processing from shared memory */ snprintf(error_msg, CLIENT_AUTH_USER_ERROR_MAX_STRLEN, "%s", clientauth_ss->requests[idx].error_msg); @@ -789,6 +863,9 @@ clientauth_shmem_startup(void) clientauth_ss->requests[i].done_processing = true; clientauth_ss->requests[i].available_entry = true; } + + memset(clientauth_ss->worker_live, 0, sizeof(clientauth_ss->worker_live)); + ConditionVariableInit(&clientauth_ss->worker_live_cv); } LWLockRelease(AddinShmemInitLock); @@ -823,6 +900,79 @@ clientauth_sighup(SIGNAL_ARGS) clientauth_reload_config = true; } +/* Wait (bounded) for worker_live[bgw_idx] to become true. */ +static bool +wait_for_clientauth_worker_live(int bgw_idx) +{ + TimestampTz deadline = TimestampTzPlusMilliseconds(GetCurrentTimestamp(), + CLIENTAUTH_WORKER_READY_TIMEOUT_MS); + bool live = false; + + ConditionVariablePrepareToSleep(&clientauth_ss->worker_live_cv); + for (;;) + { + long remaining_ms; + + CHECK_FOR_INTERRUPTS(); + + LWLockAcquire(clientauth_ss->lock, LW_SHARED); + live = clientauth_ss->worker_live[bgw_idx]; + LWLockRelease(clientauth_ss->lock); + if (live) + break; + + remaining_ms = TimestampDifferenceMilliseconds(GetCurrentTimestamp(), deadline); + if (remaining_ms <= 0) + break; + + (void) PGTLE_ConditionVariableTimedSleep(&clientauth_ss->worker_live_cv, + remaining_ms, PG_WAIT_EXTENSION); + } + ConditionVariableCancelSleep(); + return live; +} + +/* + * Clear this worker's liveness flag on exit and wake any client currently + * waiting for this shard, so they observe the worker death promptly instead + * of paying up to 1s (per client) for their next poll cycle. + */ +static void +clientauth_worker_before_shmem_exit(int code, Datum arg) +{ + int bgw_idx = DatumGetInt32(arg); + + LWLockAcquire(clientauth_ss->lock, LW_EXCLUSIVE); + clientauth_ss->worker_live[bgw_idx] = false; + LWLockRelease(clientauth_ss->lock); + + for (int j = bgw_idx; j < CLIENT_AUTH_MAX_PENDING_ENTRIES; + j += clientauth_num_parallel_workers) + ConditionVariableBroadcast(&clientauth_ss->requests[j].client_cv); +} + +/* + * Fallback when no live worker will service this connection: LOG (with an + * actionable hint) and, under FEATURE_REQUIRE, FATAL. Otherwise return so + * the hook accepts the connection (matches can_allow_without_executing()). + */ +static void +clientauth_fallback_no_worker(const char *reason) +{ + ereport(LOG, + errmsg("\"%s.clientauth\" background worker %s", + PG_TLE_NSPNAME, reason), + errhint("Check that pgtle.clientauth_db_name (\"%s\") names an existing database and that max_worker_processes is large enough.", + clientauth_database_name)); + + if (enable_clientauth_feature == FEATURE_REQUIRE) + ereport(FATAL, + errcode(ERRCODE_CONNECTION_EXCEPTION), + errmsg("pgtle.enable_clientauth is set to require, but no clientauth background worker is running for this connection"), + errhint("Check that pgtle.clientauth_db_name (\"%s\") names an existing database.", + clientauth_database_name)); +} + /* * If one (or more) of the following is true, then the connection can be * accepted without executing user functions. diff --git a/test/t/004_pg_tle_clientauth.pl b/test/t/004_pg_tle_clientauth.pl index fa6b4a5..cb48b8d 100644 --- a/test/t/004_pg_tle_clientauth.pl +++ b/test/t/004_pg_tle_clientauth.pl @@ -31,6 +31,9 @@ ### 17. Malformed strings cannot be used for SQL injection ### 18. pg_tle can be updated from 1.4.0 to 1.5.0 without affecting clientauth functions ### 19. application_name field works +### 20. Nonexistent pgtle.clientauth_db_name does not lock users out (enable_clientauth = 'on') +### 21. Nonexistent pgtle.clientauth_db_name returns a clear error (enable_clientauth = 'require') +### 22. Workers dying after startup (DROP DATABASE) do not hang new connections use strict; use warnings; @@ -361,5 +364,59 @@ like($psql_err, qr/FATAL: 004_pg_tle_clientauth.pl/, "application_name field works on pg_tle 1.5.0"); +### 20-22 use bounded psql timeouts so a regression that reintroduces the +### hang is caught deterministically instead of stalling the suite. + +### 20. Nonexistent clientauth_db_name + 'on' -> accept the connection. +$node->append_conf('postgresql.conf', qq(pgtle.clientauth_users_to_skip = '')); +$node->append_conf('postgresql.conf', qq(pgtle.clientauth_databases_to_skip = '')); +$node->append_conf('postgresql.conf', qq(pgtle.enable_clientauth = 'on')); +$node->append_conf('postgresql.conf', qq(pgtle.clientauth_db_name = 'ghost_clientauth_db_does_not_exist')); +$node->restart; + +my ($missing_out, $missing_err) = ('', ''); +my $missing_rc = $node->psql('postgres', 'SELECT 1', + stdout => \$missing_out, stderr => \$missing_err, timeout => 15); +is($missing_rc, 0, "enable_clientauth=on: nonexistent clientauth_db_name still allows connections"); +like($missing_out, qr/^1$/, "enable_clientauth=on: client backend actually executes the query"); + +### 21. Nonexistent clientauth_db_name + 'require' -> clear FATAL. +$node->append_conf('postgresql.conf', qq(pgtle.enable_clientauth = 'require')); +$node->restart; + +my $require_err = ''; +$node->psql('postgres', 'SELECT 1', stderr => \$require_err, timeout => 15); +like($require_err, + qr/FATAL: pgtle\.enable_clientauth is set to require, but no clientauth background worker is running for this connection/, + "enable_clientauth=require: nonexistent clientauth_db_name rejects with actionable error"); + +### 22. Workers dying after startup (DROP DATABASE clientauth_db_name) must not +### hang new connections. Turn clientauth off first because we're currently in +### 'require' with a ghost DB, then create/point-at a real DB, drop it, retry. +$node->append_conf('postgresql.conf', qq(pgtle.enable_clientauth = 'off')); +$node->restart; +$node->safe_psql('postgres', 'CREATE DATABASE clientauth_db'); +$node->append_conf('postgresql.conf', qq(pgtle.enable_clientauth = 'on')); +$node->append_conf('postgresql.conf', qq(pgtle.clientauth_db_name = 'clientauth_db')); +$node->restart; + +is($node->psql('postgres', 'SELECT 1', timeout => 15), 0, + "die-after-start: connections succeed while clientauth_db exists"); + +# Worker is attached to clientauth_db, so FORCE evicts it. bgw_restart_time=1s +# + attach FATAL is < 3s; 5s of headroom for CI jitter. +$node->safe_psql('postgres', 'DROP DATABASE clientauth_db WITH (FORCE)'); +sleep 5; + +my $post_out = ''; +my $post_rc = $node->psql('postgres', 'SELECT 1', + stdout => \$post_out, timeout => 15); +is($post_rc, 0, "die-after-start: connection does not hang after DROP DATABASE"); +like($post_out, qr/^1$/, "die-after-start: client backend executes the query"); + +# Restore a valid clientauth_db_name so the trailing $node->stop is clean. +$node->append_conf('postgresql.conf', qq(pgtle.clientauth_db_name = 'postgres')); +$node->restart; + $node->stop; done_testing();