Currently, cancel request key is a 32-bit token, which isn't very much entropy. If you want to cancel another session's query, you can brute-force it. In most environments, an unauthorized cancellation of a query isn't very serious, but it nevertheless would be nice to have more protection from it. The attached patch makes it longer. It is an optional protocol feature, so it's fully backwards-compatible with clients that don't support longer keys.

If the client requests the "_pq_.extended_query_cancel" protocol feature, the server will generate a longer 256-bit cancellation key. However, the new longer key length is no longer hardcoded in the protocol. The client is expected to deal with variable length keys, up to some reasonable upper limit (TODO: document the maximum). This flexibility allows e.g. a connection pooler to add more information to the cancel key, which could be useful. If the client doesn't request the protocol feature, the server generates a 32-bit key like before.

One complication with this was that because we no longer know how long the key should be, 4-bytes or something longer, until the backend has performed the protocol negotiation, we cannot generate the key in the postmaster before forking the process anymore. The first patch here changes things so that the cancellation key is generated later, in the backend, and the backend advertises the key in the PMSignalState array. This is similar to how this has always worked in EXEC_BACKEND mode with the ShmemBackendArray, but instead of having a separate array, I added fields to the PMSignalState slots. This removes a bunch of EXEC_BACKEND-specific code, which is nice.

Any thoughts on this? Documentation is still missing, and there's one TODO on adding a portable time-constant memcmp() function; I'll add those if there's agreement on this otherwise.

--
Heikki Linnakangas
Neon (https://neon.tech)
From f871e915fee08b3cc27e732646f3a62ec1c3024b Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakan...@iki.fi>
Date: Thu, 29 Feb 2024 23:20:52 +0200
Subject: [PATCH 1/2] Move cancel key generation to after forking the backend

Move responsibility of generating the cancel key to the backend
process. The cancel key is now generated after forking, and the
backend advertises it in the PMSignalState array. When a cancel
request arrives, the backend handling it scans the PMSignalState array
to find the target pid and cancel key. This is similar to how this
previously worked in the EXEC_BACKEND case with the ShmemBackendArray,
but instead of having the separate array, we store the cancel keys in
PMSignalState array, and use it also in the !EXEC_BACKEND case.

While we're at it, switch to using atomics in pmsignal.c for the state
field. That feels easier to reason about than volatile
pointers. (TODO: also switch to pg_atomic_flag for PMSignalFlags
array. Weirdly we don't have "pg_atomic_test_clear_flag" nor
"pg_atomic_set_flag". Those would be the natural opertions for
PMSignalFlags. Maybe we should add those, but I didn't want to go that
deep down the rabbit hole in this patch.)

This is needed by the next patch, so that we can generate the cancel
key later, after negotiating the protocol version with the client.
---
 src/backend/postmaster/postmaster.c | 185 +---------------------------
 src/backend/storage/ipc/ipci.c      |  11 --
 src/backend/storage/ipc/pmsignal.c  | 124 ++++++++++++++-----
 src/backend/storage/lmgr/proc.c     |  16 ++-
 src/include/postmaster/postmaster.h |   3 -
 src/include/storage/pmsignal.h      |   3 +-
 src/tools/pgindent/typedefs.list    |   1 +
 7 files changed, 116 insertions(+), 227 deletions(-)

diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
index da0c627107e..408f8fb9b6d 100644
--- a/src/backend/postmaster/postmaster.c
+++ b/src/backend/postmaster/postmaster.c
@@ -177,7 +177,6 @@
 typedef struct bkend
 {
 	pid_t		pid;			/* process id of backend */
-	int32		cancel_key;		/* cancel key for cancels for this backend */
 	int			child_slot;		/* PMChildSlot for this backend, if any */
 	int			bkend_type;		/* child process flavor, see above */
 	bool		dead_end;		/* is it going to send an error and quit? */
@@ -187,10 +186,6 @@ typedef struct bkend
 
 static dlist_head BackendList = DLIST_STATIC_INIT(BackendList);
 
-#ifdef EXEC_BACKEND
-static Backend *ShmemBackendArray;
-#endif
-
 BackgroundWorker *MyBgworkerEntry = NULL;
 
 
@@ -430,7 +425,6 @@ static void SendNegotiateProtocolVersion(List *unrecognized_protocol_options);
 static void processCancelRequest(Port *port, void *pkt);
 static void report_fork_failure_to_client(Port *port, int errnum);
 static CAC_state canAcceptConnections(int backend_type);
-static bool RandomCancelKey(int32 *cancel_key);
 static void signal_child(pid_t pid, int signal);
 static void sigquit_child(pid_t pid);
 static bool SignalSomeChildren(int signal, int target);
@@ -517,7 +511,6 @@ typedef struct
 #endif
 	void	   *UsedShmemSegAddr;
 	slock_t    *ShmemLock;
-	Backend    *ShmemBackendArray;
 #ifndef HAVE_SPINLOCKS
 	PGSemaphore *SpinlockSemaArray;
 #endif
@@ -559,9 +552,6 @@ static bool save_backend_variables(BackendParameters *param, Port *port, Backgro
 static bool save_backend_variables(BackendParameters *param, Port *port, BackgroundWorker *worker,
 								   HANDLE childProcess, pid_t childPid);
 #endif
-
-static void ShmemBackendArrayAdd(Backend *bn);
-static void ShmemBackendArrayRemove(Backend *bn);
 #endif							/* EXEC_BACKEND */
 
 /* Macros to check exit status of a child process */
@@ -2334,58 +2324,11 @@ processCancelRequest(Port *port, void *pkt)
 	CancelRequestPacket *canc = (CancelRequestPacket *) pkt;
 	int			backendPID;
 	int32		cancelAuthCode;
-	Backend    *bp;
-
-#ifndef EXEC_BACKEND
-	dlist_iter	iter;
-#else
-	int			i;
-#endif
 
 	backendPID = (int) pg_ntoh32(canc->backendPID);
 	cancelAuthCode = (int32) pg_ntoh32(canc->cancelAuthCode);
 
-	/*
-	 * See if we have a matching backend.  In the EXEC_BACKEND case, we can no
-	 * longer access the postmaster's own backend list, and must rely on the
-	 * duplicate array in shared memory.
-	 */
-#ifndef EXEC_BACKEND
-	dlist_foreach(iter, &BackendList)
-	{
-		bp = dlist_container(Backend, elem, iter.cur);
-#else
-	for (i = MaxLivePostmasterChildren() - 1; i >= 0; i--)
-	{
-		bp = (Backend *) &ShmemBackendArray[i];
-#endif
-		if (bp->pid == backendPID)
-		{
-			if (bp->cancel_key == cancelAuthCode)
-			{
-				/* Found a match; signal that backend to cancel current op */
-				ereport(DEBUG2,
-						(errmsg_internal("processing cancel request: sending SIGINT to process %d",
-										 backendPID)));
-				signal_child(bp->pid, SIGINT);
-			}
-			else
-				/* Right PID, wrong key: no way, Jose */
-				ereport(LOG,
-						(errmsg("wrong key in cancel request for process %d",
-								backendPID)));
-			return;
-		}
-#ifndef EXEC_BACKEND			/* make GNU Emacs 26.1 see brace balance */
-	}
-#else
-	}
-#endif
-
-	/* No matching backend */
-	ereport(LOG,
-			(errmsg("PID %d in cancel request did not match any process",
-					backendPID)));
+	SendCancelRequest(backendPID, cancelAuthCode);
 }
 
 /*
@@ -3296,9 +3239,6 @@ CleanupBackgroundWorker(int pid,
 
 		/* Get it out of the BackendList and clear out remaining data */
 		dlist_delete(&rw->rw_backend->elem);
-#ifdef EXEC_BACKEND
-		ShmemBackendArrayRemove(rw->rw_backend);
-#endif
 
 		/*
 		 * It's possible that this background worker started some OTHER
@@ -3384,9 +3324,6 @@ CleanupBackend(int pid,
 					HandleChildCrash(pid, exitstatus, _("server process"));
 					return;
 				}
-#ifdef EXEC_BACKEND
-				ShmemBackendArrayRemove(bp);
-#endif
 			}
 			if (bp->bgworker_notify)
 			{
@@ -3454,9 +3391,6 @@ HandleChildCrash(int pid, int exitstatus, const char *procname)
 			 */
 			(void) ReleasePostmasterChildSlot(rw->rw_child_slot);
 			dlist_delete(&rw->rw_backend->elem);
-#ifdef EXEC_BACKEND
-			ShmemBackendArrayRemove(rw->rw_backend);
-#endif
 			pfree(rw->rw_backend);
 			rw->rw_backend = NULL;
 			rw->rw_pid = 0;
@@ -3489,9 +3423,6 @@ HandleChildCrash(int pid, int exitstatus, const char *procname)
 			if (!bp->dead_end)
 			{
 				(void) ReleasePostmasterChildSlot(bp->child_slot);
-#ifdef EXEC_BACKEND
-				ShmemBackendArrayRemove(bp);
-#endif
 			}
 			dlist_delete(iter.cur);
 			pfree(bp);
@@ -4103,22 +4034,6 @@ BackendStartup(Port *port)
 		return STATUS_ERROR;
 	}
 
-	/*
-	 * Compute the cancel key that will be assigned to this backend. The
-	 * backend will have its own copy in the forked-off process' value of
-	 * MyCancelKey, so that it can transmit the key to the frontend.
-	 */
-	if (!RandomCancelKey(&MyCancelKey))
-	{
-		pfree(bn);
-		ereport(LOG,
-				(errcode(ERRCODE_INTERNAL_ERROR),
-				 errmsg("could not generate random cancel key")));
-		return STATUS_ERROR;
-	}
-
-	bn->cancel_key = MyCancelKey;
-
 	/* Pass down canAcceptConnections state */
 	port->canAcceptConnections = canAcceptConnections(BACKEND_TYPE_NORMAL);
 	bn->dead_end = (port->canAcceptConnections != CAC_OK);
@@ -4182,11 +4097,6 @@ BackendStartup(Port *port)
 	bn->bkend_type = BACKEND_TYPE_NORMAL;	/* Can change later to WALSND */
 	dlist_push_head(&BackendList, &bn->elem);
 
-#ifdef EXEC_BACKEND
-	if (!bn->dead_end)
-		ShmemBackendArrayAdd(bn);
-#endif
-
 	return STATUS_OK;
 }
 
@@ -5252,16 +5162,6 @@ StartupPacketTimeoutHandler(void)
 	_exit(1);
 }
 
-
-/*
- * Generate a random cancel key.
- */
-static bool
-RandomCancelKey(int32 *cancel_key)
-{
-	return pg_strong_random(cancel_key, sizeof(int32));
-}
-
 /*
  * Count up number of child processes of specified types (dead_end children
  * are always excluded).
@@ -5437,25 +5337,9 @@ StartAutovacuumWorker(void)
 	 */
 	if (canAcceptConnections(BACKEND_TYPE_AUTOVAC) == CAC_OK)
 	{
-		/*
-		 * Compute the cancel key that will be assigned to this session. We
-		 * probably don't need cancel keys for autovac workers, but we'd
-		 * better have something random in the field to prevent unfriendly
-		 * people from sending cancels to them.
-		 */
-		if (!RandomCancelKey(&MyCancelKey))
-		{
-			ereport(LOG,
-					(errcode(ERRCODE_INTERNAL_ERROR),
-					 errmsg("could not generate random cancel key")));
-			return;
-		}
-
 		bn = (Backend *) palloc_extended(sizeof(Backend), MCXT_ALLOC_NO_OOM);
 		if (bn)
 		{
-			bn->cancel_key = MyCancelKey;
-
 			/* Autovac workers are not dead_end and need a child slot */
 			bn->dead_end = false;
 			bn->child_slot = MyPMChildSlot = AssignPostmasterChildSlot();
@@ -5466,9 +5350,6 @@ StartAutovacuumWorker(void)
 			{
 				bn->bkend_type = BACKEND_TYPE_AUTOVAC;
 				dlist_push_head(&BackendList, &bn->elem);
-#ifdef EXEC_BACKEND
-				ShmemBackendArrayAdd(bn);
-#endif
 				/* all OK */
 				return;
 			}
@@ -5600,11 +5481,10 @@ CreateOptsFile(int argc, char *argv[], char *fullprogname)
 /*
  * MaxLivePostmasterChildren
  *
- * This reports the number of entries needed in per-child-process arrays
- * (the PMChildFlags array, and if EXEC_BACKEND the ShmemBackendArray).
- * These arrays include regular backends, autovac workers, walsenders
+ * This reports the number of entries needed in the per-child-process array.
+ * The array includes regular backends, autovac workers, walsenders
  * and background workers, but not special children nor dead_end children.
- * This allows the arrays to have a fixed maximum size, to wit the same
+ * This allows the array to have a fixed maximum size, to wit the same
  * too-many-children limit enforced by canAcceptConnections().  The exact value
  * isn't too critical as long as it's more than MaxBackends.
  */
@@ -5804,9 +5684,6 @@ do_start_bgworker(RegisteredBgWorker *rw)
 			ReportBackgroundWorkerPID(rw);
 			/* add new worker to lists of backends */
 			dlist_push_head(&BackendList, &rw->rw_backend->elem);
-#ifdef EXEC_BACKEND
-			ShmemBackendArrayAdd(rw->rw_backend);
-#endif
 			return true;
 	}
 
@@ -5877,20 +5754,6 @@ assign_backendlist_entry(RegisteredBgWorker *rw)
 		return false;
 	}
 
-	/*
-	 * Compute the cancel key that will be assigned to this session. We
-	 * probably don't need cancel keys for background workers, but we'd better
-	 * have something random in the field to prevent unfriendly people from
-	 * sending cancels to them.
-	 */
-	if (!RandomCancelKey(&MyCancelKey))
-	{
-		ereport(LOG,
-				(errcode(ERRCODE_INTERNAL_ERROR),
-				 errmsg("could not generate random cancel key")));
-		return false;
-	}
-
 	bn = palloc_extended(sizeof(Backend), MCXT_ALLOC_NO_OOM);
 	if (bn == NULL)
 	{
@@ -5900,7 +5763,6 @@ assign_backendlist_entry(RegisteredBgWorker *rw)
 		return false;
 	}
 
-	bn->cancel_key = MyCancelKey;
 	bn->child_slot = MyPMChildSlot = AssignPostmasterChildSlot();
 	bn->bkend_type = BACKEND_TYPE_BGWORKER;
 	bn->dead_end = false;
@@ -6129,7 +5991,6 @@ save_backend_variables(BackendParameters *param, Port *port, BackgroundWorker *w
 	param->UsedShmemSegAddr = UsedShmemSegAddr;
 
 	param->ShmemLock = ShmemLock;
-	param->ShmemBackendArray = ShmemBackendArray;
 
 #ifndef HAVE_SPINLOCKS
 	param->SpinlockSemaArray = SpinlockSemaArray;
@@ -6374,7 +6235,6 @@ restore_backend_variables(BackendParameters *param, Port **port, BackgroundWorke
 	UsedShmemSegAddr = param->UsedShmemSegAddr;
 
 	ShmemLock = param->ShmemLock;
-	ShmemBackendArray = param->ShmemBackendArray;
 
 #ifndef HAVE_SPINLOCKS
 	SpinlockSemaArray = param->SpinlockSemaArray;
@@ -6426,43 +6286,6 @@ restore_backend_variables(BackendParameters *param, Port **port, BackgroundWorke
 		ReserveExternalFD();
 #endif
 }
-
-
-Size
-ShmemBackendArraySize(void)
-{
-	return mul_size(MaxLivePostmasterChildren(), sizeof(Backend));
-}
-
-void
-ShmemBackendArrayAllocation(void)
-{
-	Size		size = ShmemBackendArraySize();
-
-	ShmemBackendArray = (Backend *) ShmemAlloc(size);
-	/* Mark all slots as empty */
-	memset(ShmemBackendArray, 0, size);
-}
-
-static void
-ShmemBackendArrayAdd(Backend *bn)
-{
-	/* The array slot corresponding to my PMChildSlot should be free */
-	int			i = bn->child_slot - 1;
-
-	Assert(ShmemBackendArray[i].pid == 0);
-	ShmemBackendArray[i] = *bn;
-}
-
-static void
-ShmemBackendArrayRemove(Backend *bn)
-{
-	int			i = bn->child_slot - 1;
-
-	Assert(ShmemBackendArray[i].pid == bn->pid);
-	/* Mark the slot as empty */
-	ShmemBackendArray[i].pid = 0;
-}
 #endif							/* EXEC_BACKEND */
 
 
diff --git a/src/backend/storage/ipc/ipci.c b/src/backend/storage/ipc/ipci.c
index 7e7941d6259..c2267b7b080 100644
--- a/src/backend/storage/ipc/ipci.c
+++ b/src/backend/storage/ipc/ipci.c
@@ -155,9 +155,6 @@ CalculateShmemSize(int *num_semaphores)
 	size = add_size(size, WaitEventExtensionShmemSize());
 	size = add_size(size, InjectionPointShmemSize());
 	size = add_size(size, SlotSyncShmemSize());
-#ifdef EXEC_BACKEND
-	size = add_size(size, ShmemBackendArraySize());
-#endif
 
 	/* include additional requested shmem from preload libraries */
 	size = add_size(size, total_addin_request);
@@ -247,14 +244,6 @@ CreateSharedMemoryAndSemaphores(void)
 	/* Initialize subsystems */
 	CreateOrAttachShmemStructs();
 
-#ifdef EXEC_BACKEND
-
-	/*
-	 * Alloc the win32 shared backend array
-	 */
-	ShmemBackendArrayAllocation();
-#endif
-
 	/* Initialize dynamic shared memory facilities. */
 	dsm_postmaster_startup(shim);
 
diff --git a/src/backend/storage/ipc/pmsignal.c b/src/backend/storage/ipc/pmsignal.c
index 27844b46a2b..9516b869cb2 100644
--- a/src/backend/storage/ipc/pmsignal.c
+++ b/src/backend/storage/ipc/pmsignal.c
@@ -22,6 +22,7 @@
 #endif
 
 #include "miscadmin.h"
+#include "port/atomics.h"
 #include "postmaster/postmaster.h"
 #include "replication/walsender.h"
 #include "storage/pmsignal.h"
@@ -37,9 +38,8 @@
  * if the same reason is signaled more than once simultaneously, the
  * postmaster will observe it only once.)
  *
- * The flags are actually declared as "volatile sig_atomic_t" for maximum
- * portability.  This should ensure that loads and stores of the flag
- * values are atomic, allowing us to dispense with any explicit locking.
+ * The flag fields use atomics, allowing us to dispense with any explicit
+ * locking.
  *
  * In addition to the per-reason flags, we store a set of per-child-process
  * flags that are currently used only for detecting whether a backend has
@@ -67,20 +67,29 @@
 #define PM_CHILD_ACTIVE		2
 #define PM_CHILD_WALSENDER	3
 
+typedef struct ChildSlotData
+{
+	pg_atomic_uint32 state;
+
+	int			pid;
+	int32		cancel_key;
+} ChildSlotData;
+
 /* "typedef struct PMSignalData PMSignalData" appears in pmsignal.h */
 struct PMSignalData
 {
 	/* per-reason flags for signaling the postmaster */
-	sig_atomic_t PMSignalFlags[NUM_PMSIGNALS];
+	volatile sig_atomic_t PMSignalFlags[NUM_PMSIGNALS];
 	/* global flags for signals from postmaster to children */
 	QuitSignalReason sigquit_reason;	/* why SIGQUIT was sent */
-	/* per-child-process flags */
-	int			num_child_flags;	/* # of entries in PMChildFlags[] */
-	sig_atomic_t PMChildFlags[FLEXIBLE_ARRAY_MEMBER];
+
+	/* per-child-process slots */
+	int			num_child_slots;	/* # of entries in child_slots[] */
+	ChildSlotData child_slots[FLEXIBLE_ARRAY_MEMBER];
 };
 
 /* PMSignalState pointer is valid in both postmaster and child processes */
-NON_EXEC_STATIC volatile PMSignalData *PMSignalState = NULL;
+NON_EXEC_STATIC PMSignalData *PMSignalState = NULL;
 
 /*
  * These static variables are valid only in the postmaster.  We keep a
@@ -130,9 +139,9 @@ PMSignalShmemSize(void)
 {
 	Size		size;
 
-	size = offsetof(PMSignalData, PMChildFlags);
+	size = offsetof(PMSignalData, child_slots);
 	size = add_size(size, mul_size(MaxLivePostmasterChildren(),
-								   sizeof(sig_atomic_t)));
+								   sizeof(ChildSlotData)));
 
 	return size;
 }
@@ -151,9 +160,9 @@ PMSignalShmemInit(void)
 	if (!found)
 	{
 		/* initialize all flags to zeroes */
-		MemSet(unvolatize(PMSignalData *, PMSignalState), 0, PMSignalShmemSize());
+		MemSet(PMSignalState, 0, PMSignalShmemSize());
 		num_child_inuse = MaxLivePostmasterChildren();
-		PMSignalState->num_child_flags = num_child_inuse;
+		PMSignalState->num_child_slots = num_child_inuse;
 
 		/*
 		 * Also allocate postmaster's private PMChildInUse[] array.  We
@@ -262,14 +271,14 @@ AssignPostmasterChildSlot(void)
 		if (!PMChildInUse[slot])
 		{
 			PMChildInUse[slot] = true;
-			PMSignalState->PMChildFlags[slot] = PM_CHILD_ASSIGNED;
+			pg_atomic_write_u32(&PMSignalState->child_slots[slot].state, PM_CHILD_ASSIGNED);
 			next_child_inuse = slot;
 			return slot + 1;
 		}
 	}
 
 	/* Out of slots ... should never happen, else postmaster.c messed up */
-	elog(FATAL, "no free slots in PMChildFlags array");
+	elog(FATAL, "no free slots in postmaster child array");
 	return 0;					/* keep compiler quiet */
 }
 
@@ -283,7 +292,7 @@ AssignPostmasterChildSlot(void)
 bool
 ReleasePostmasterChildSlot(int slot)
 {
-	bool		result;
+	uint32		oldstate;
 
 	Assert(slot > 0 && slot <= num_child_inuse);
 	slot--;
@@ -293,10 +302,9 @@ ReleasePostmasterChildSlot(int slot)
 	 * postmaster.c is such that this might get called twice when a child
 	 * crashes.  So we don't try to Assert anything about the state.
 	 */
-	result = (PMSignalState->PMChildFlags[slot] == PM_CHILD_ASSIGNED);
-	PMSignalState->PMChildFlags[slot] = PM_CHILD_UNUSED;
+	oldstate = pg_atomic_exchange_u32(&PMSignalState->child_slots[slot].state, PM_CHILD_UNUSED);
 	PMChildInUse[slot] = false;
-	return result;
+	return oldstate == PM_CHILD_ASSIGNED;
 }
 
 /*
@@ -309,7 +317,7 @@ IsPostmasterChildWalSender(int slot)
 	Assert(slot > 0 && slot <= num_child_inuse);
 	slot--;
 
-	if (PMSignalState->PMChildFlags[slot] == PM_CHILD_WALSENDER)
+	if (pg_atomic_read_u32(&PMSignalState->child_slots[slot].state) == PM_CHILD_WALSENDER)
 		return true;
 	else
 		return false;
@@ -320,14 +328,17 @@ IsPostmasterChildWalSender(int slot)
  * actively using shared memory.  This is called in the child process.
  */
 void
-MarkPostmasterChildActive(void)
+MarkPostmasterChildActive(int pid, int32 cancelAuthCode)
 {
 	int			slot = MyPMChildSlot;
 
-	Assert(slot > 0 && slot <= PMSignalState->num_child_flags);
+	Assert(slot > 0 && slot <= PMSignalState->num_child_slots);
 	slot--;
-	Assert(PMSignalState->PMChildFlags[slot] == PM_CHILD_ASSIGNED);
-	PMSignalState->PMChildFlags[slot] = PM_CHILD_ACTIVE;
+	Assert(pg_atomic_read_u32(&PMSignalState->child_slots[slot].state) == PM_CHILD_ASSIGNED);
+	PMSignalState->child_slots[slot].pid = pid;
+	PMSignalState->child_slots[slot].cancel_key = cancelAuthCode;
+	pg_memory_barrier();
+	pg_atomic_write_u32(&PMSignalState->child_slots[slot].state, PM_CHILD_ACTIVE);
 }
 
 /*
@@ -342,10 +353,10 @@ MarkPostmasterChildWalSender(void)
 
 	Assert(am_walsender);
 
-	Assert(slot > 0 && slot <= PMSignalState->num_child_flags);
+	Assert(slot > 0 && slot <= PMSignalState->num_child_slots);
 	slot--;
-	Assert(PMSignalState->PMChildFlags[slot] == PM_CHILD_ACTIVE);
-	PMSignalState->PMChildFlags[slot] = PM_CHILD_WALSENDER;
+	Assert(pg_atomic_read_u32(&PMSignalState->child_slots[slot].state) == PM_CHILD_ACTIVE);
+	pg_atomic_write_u32(&PMSignalState->child_slots[slot].state, PM_CHILD_WALSENDER);
 }
 
 /*
@@ -357,11 +368,17 @@ MarkPostmasterChildInactive(void)
 {
 	int			slot = MyPMChildSlot;
 
-	Assert(slot > 0 && slot <= PMSignalState->num_child_flags);
+	Assert(slot > 0 && slot <= PMSignalState->num_child_slots);
 	slot--;
-	Assert(PMSignalState->PMChildFlags[slot] == PM_CHILD_ACTIVE ||
-		   PMSignalState->PMChildFlags[slot] == PM_CHILD_WALSENDER);
-	PMSignalState->PMChildFlags[slot] = PM_CHILD_ASSIGNED;
+#ifdef USE_ASSERT_CHECKING
+	{
+		uint32		oldstate;
+
+		oldstate = pg_atomic_read_u32(&PMSignalState->child_slots[slot].state);
+		Assert(oldstate == PM_CHILD_ACTIVE || oldstate == PM_CHILD_WALSENDER);
+	}
+#endif
+	pg_atomic_write_u32(&PMSignalState->child_slots[slot].state, PM_CHILD_ASSIGNED);
 }
 
 
@@ -460,3 +477,50 @@ PostmasterDeathSignalInit(void)
 	postmaster_possibly_dead = true;
 #endif							/* USE_POSTMASTER_DEATH_SIGNAL */
 }
+
+void
+SendCancelRequest(int backendPID, int32 cancelAuthCode)
+{
+	/*
+	 * See if we have a matching backend.  In the EXEC_BACKEND case, we can no
+	 * longer access the postmaster's own backend list, and must rely on the
+	 * duplicate array in shared memory. XXX
+	 */
+	for (int i = MaxLivePostmasterChildren() - 1; i >= 0; i--)
+	{
+		ChildSlotData *slot = &PMSignalState->child_slots[i];
+		uint32		state = pg_atomic_read_u32(&slot->state);
+
+		if (state != PM_CHILD_ACTIVE)
+			continue;
+
+		pg_read_barrier();
+		if (slot->pid == backendPID)
+		{
+			if (slot->cancel_key == cancelAuthCode)
+			{
+				/* Found a match; signal that backend to cancel current op */
+				ereport(DEBUG2,
+						(errmsg_internal("processing cancel request: sending SIGINT to process %d",
+										 backendPID)));
+
+				/*
+				 * FIXME: we used to use signal_child. I believe kill() is
+				 * maybe even more correct, but verify that.
+				 */
+				kill(backendPID, SIGINT);
+			}
+			else
+				/* Right PID, wrong key: no way, Jose */
+				ereport(LOG,
+						(errmsg("wrong key in cancel request for process %d",
+								backendPID)));
+			return;
+		}
+	}
+
+	/* No matching backend */
+	ereport(LOG,
+			(errmsg("PID %d in cancel request did not match any process",
+					backendPID)));
+}
diff --git a/src/backend/storage/lmgr/proc.c b/src/backend/storage/lmgr/proc.c
index 6e334971dc9..5984fe7d44c 100644
--- a/src/backend/storage/lmgr/proc.c
+++ b/src/backend/storage/lmgr/proc.c
@@ -310,6 +310,20 @@ InitProcess(void)
 	if (MyProc != NULL)
 		elog(ERROR, "you already exist");
 
+	/*
+	 * Generate a random cancel key.
+	 *
+	 * We probably don't need cancel keys for non-backend processes, but we'd
+	 * better have something random in the field to prevent unfriendly people
+	 * from sending cancels to them.
+	 */
+	if (!pg_strong_random(&MyCancelKey, sizeof(int32)))
+	{
+		ereport(ERROR,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("could not generate random cancel key")));
+	}
+
 	/* Decide which list should supply our PGPROC. */
 	if (IsAnyAutoVacuumProcess())
 		procgloballist = &ProcGlobal->autovacFreeProcs;
@@ -373,7 +387,7 @@ InitProcess(void)
 	 */
 	if (IsUnderPostmaster && !IsAutoVacuumLauncherProcess() &&
 		!IsLogicalSlotSyncWorker())
-		MarkPostmasterChildActive();
+		MarkPostmasterChildActive(MyProcPid, MyCancelKey);
 
 	/*
 	 * Initialize all fields of MyProc, except for those previously
diff --git a/src/include/postmaster/postmaster.h b/src/include/postmaster/postmaster.h
index 03f78b5b25f..7cfd07a2bda 100644
--- a/src/include/postmaster/postmaster.h
+++ b/src/include/postmaster/postmaster.h
@@ -61,9 +61,6 @@ extern bool PostmasterMarkPIDForWorkerNotify(int);
 #ifdef EXEC_BACKEND
 extern pid_t postmaster_forkexec(int argc, char *argv[]);
 extern void SubPostmasterMain(int argc, char *argv[]) pg_attribute_noreturn();
-
-extern Size ShmemBackendArraySize(void);
-extern void ShmemBackendArrayAllocation(void);
 #endif
 
 /*
diff --git a/src/include/storage/pmsignal.h b/src/include/storage/pmsignal.h
index 029b7201093..8dff47e86e4 100644
--- a/src/include/storage/pmsignal.h
+++ b/src/include/storage/pmsignal.h
@@ -69,11 +69,12 @@ extern QuitSignalReason GetQuitSignalReason(void);
 extern int	AssignPostmasterChildSlot(void);
 extern bool ReleasePostmasterChildSlot(int slot);
 extern bool IsPostmasterChildWalSender(int slot);
-extern void MarkPostmasterChildActive(void);
+extern void MarkPostmasterChildActive(int pid, int32 cancelAuthCode);
 extern void MarkPostmasterChildInactive(void);
 extern void MarkPostmasterChildWalSender(void);
 extern bool PostmasterIsAliveInternal(void);
 extern void PostmasterDeathSignalInit(void);
+extern void SendCancelRequest(int backendPID, int32 cancelAuthCode);
 
 
 /*
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index fc8b15d0cf2..ac0085a226f 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -375,6 +375,7 @@ CatalogId
 CatalogIdMapEntry
 CatalogIndexState
 ChangeVarNodes_context
+ChildSlotData
 ReplaceVarnoContext
 CheckPoint
 CheckPointStmt
-- 
2.39.2

From 32513dd0c5b48e4d6ac56f8a53085480c30a94bf Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakan...@iki.fi>
Date: Thu, 29 Feb 2024 23:22:08 +0200
Subject: [PATCH 2/2] Make cancel request keys longer, as an optional protocol
 feature

Currently, cancel request key is 32-bit token, which isn't very much
entropy. If you want to cancel another session's query, you can
brute-force it. In most environments, an unauthorized cancellation of
a query isn't very serious, but it nevertheless would be nice to have
more protection from it.

To make it harder to guess the key, make it larger. But because the
length is coded in the FE-BE protocol, make it an optional protocol
feature. If the client requests the "_pq_.extended_query_cancel"
protocol feature, the server will generate a longer 256-bit
cancellation key. However, the new longer key length is not hardcoded
in the protocol anymore, the client is expected to deal with variable
length keys, up to some reasonable upper limit (TODO: document the
maximum).  This flexibility allows e.g. a connection pooler to add
more information to the cancel key, which might be useful for finding
the connection. If the client doesn't request the optional protocol
feature, the server generates a 32-bit key like before.
---
 src/backend/access/transam/parallel.c |  6 +---
 src/backend/postmaster/postmaster.c   | 52 +++++++++++++++------------
 src/backend/storage/ipc/pmsignal.c    | 29 +++++++++++----
 src/backend/storage/lmgr/proc.c       |  9 +++--
 src/backend/tcop/postgres.c           |  2 +-
 src/backend/utils/init/globals.c      |  1 -
 src/include/libpq/libpq-be.h          |  5 +++
 src/include/libpq/pqcomm.h            |  5 +--
 src/include/miscadmin.h               |  3 +-
 src/include/storage/pmsignal.h        |  6 ++--
 src/interfaces/libpq/fe-cancel.c      | 34 +++++++++---------
 src/interfaces/libpq/fe-connect.c     | 26 ++++++++++----
 src/interfaces/libpq/fe-protocol3.c   | 32 ++++++++++++++---
 src/interfaces/libpq/libpq-int.h      | 10 ++++--
 14 files changed, 148 insertions(+), 72 deletions(-)

diff --git a/src/backend/access/transam/parallel.c b/src/backend/access/transam/parallel.c
index 849a03e4b65..652b96225a4 100644
--- a/src/backend/access/transam/parallel.c
+++ b/src/backend/access/transam/parallel.c
@@ -1135,7 +1135,6 @@ HandleParallelMessage(ParallelContext *pcxt, int i, StringInfo msg)
 			{
 				int32		pid = pq_getmsgint(msg, 4);
 
-				(void) pq_getmsgint(msg, 4);	/* discard cancel key */
 				(void) pq_getmsgend(msg);
 				pcxt->worker[i].pid = pid;
 				break;
@@ -1372,13 +1371,10 @@ ParallelWorkerMain(Datum main_arg)
 	/*
 	 * Send a BackendKeyData message to the process that initiated parallelism
 	 * so that it has access to our PID before it receives any other messages
-	 * from us.  Our cancel key is sent, too, since that's the way the
-	 * protocol message is defined, but it won't actually be used for anything
-	 * in this case.
+	 * from us.
 	 */
 	pq_beginmessage(&msgbuf, PqMsg_BackendKeyData);
 	pq_sendint32(&msgbuf, (int32) MyProcPid);
-	pq_sendint32(&msgbuf, (int32) MyCancelKey);
 	pq_endmessage(&msgbuf);
 
 	/*
diff --git a/src/backend/postmaster/postmaster.c b/src/backend/postmaster/postmaster.c
index 408f8fb9b6d..47c29e8c10e 100644
--- a/src/backend/postmaster/postmaster.c
+++ b/src/backend/postmaster/postmaster.c
@@ -422,7 +422,7 @@ static int	ServerLoop(void);
 static int	BackendStartup(Port *port);
 static int	ProcessStartupPacket(Port *port, bool ssl_done, bool gss_done);
 static void SendNegotiateProtocolVersion(List *unrecognized_protocol_options);
-static void processCancelRequest(Port *port, void *pkt);
+static void processCancelRequest(Port *port, void *pkt, int pktlen);
 static void report_fork_failure_to_client(Port *port, int errnum);
 static CAC_state canAcceptConnections(int backend_type);
 static void signal_child(pid_t pid, int signal);
@@ -501,7 +501,7 @@ typedef struct
 	BackgroundWorker bgworker;
 
 	char		DataDir[MAXPGPATH];
-	int32		MyCancelKey;
+	char		MyCancelKey[CANCEL_KEY_LENGTH];
 	int			MyPMChildSlot;
 #ifndef WIN32
 	unsigned long UsedShmemSegID;
@@ -2008,16 +2008,9 @@ ProcessStartupPacket(Port *port, bool ssl_done, bool gss_done)
 	 */
 	port->proto = proto = pg_ntoh32(*((ProtocolVersion *) buf));
 
-	if (proto == CANCEL_REQUEST_CODE)
+	if (proto == EXTENDED_CANCEL_REQUEST_CODE || proto == CANCEL_REQUEST_CODE)
 	{
-		if (len != sizeof(CancelRequestPacket))
-		{
-			ereport(COMMERROR,
-					(errcode(ERRCODE_PROTOCOL_VIOLATION),
-					 errmsg("invalid length of startup packet")));
-			return STATUS_ERROR;
-		}
-		processCancelRequest(port, buf);
+		processCancelRequest(port, buf, len);
 		/* Not really an error, but we don't want to proceed further */
 		return STATUS_ERROR;
 	}
@@ -2202,11 +2195,16 @@ retry1:
 			{
 				/*
 				 * Any option beginning with _pq_. is reserved for use as a
-				 * protocol-level option, but at present no such options are
-				 * defined.
+				 * protocol-level option. At present, there is only this one.
 				 */
-				unrecognized_protocol_options =
-					lappend(unrecognized_protocol_options, pstrdup(nameptr));
+				if (strcmp(nameptr + 5, "extended_query_cancel") == 0)
+				{
+					/* the value is ignored */
+					port->extended_query_cancel = true;
+				}
+				else
+					unrecognized_protocol_options =
+						lappend(unrecognized_protocol_options, pstrdup(nameptr));
 			}
 			else
 			{
@@ -2319,16 +2317,27 @@ SendNegotiateProtocolVersion(List *unrecognized_protocol_options)
  * Nothing is sent back to the client.
  */
 static void
-processCancelRequest(Port *port, void *pkt)
+processCancelRequest(Port *port, void *pkt, int pktlen)
 {
 	CancelRequestPacket *canc = (CancelRequestPacket *) pkt;
-	int			backendPID;
-	int32		cancelAuthCode;
+	int			len;
 
-	backendPID = (int) pg_ntoh32(canc->backendPID);
-	cancelAuthCode = (int32) pg_ntoh32(canc->cancelAuthCode);
+	if (pktlen < offsetof(CancelRequestPacket, cancelAuthCode))
+	{
+		ereport(COMMERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("invalid length of extended quer cancel packet")));
+		return;
+	}
+	canc = (CancelRequestPacket *) pkt;
+	len = pktlen - offsetof(CancelRequestPacket, cancelAuthCode);
+
+	if (len != 4 && canc->cancelRequestCode == CANCEL_REQUEST_CODE)
+		ereport(COMMERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("invalid length of cancel request packet")));
 
-	SendCancelRequest(backendPID, cancelAuthCode);
+	SendCancelRequest(pg_ntoh32(canc->backendPID), canc->cancelAuthCode, len);
 }
 
 /*
@@ -5981,7 +5990,6 @@ save_backend_variables(BackendParameters *param, Port *port, BackgroundWorker *w
 
 	strlcpy(param->DataDir, DataDir, MAXPGPATH);
 
-	param->MyCancelKey = MyCancelKey;
 	param->MyPMChildSlot = MyPMChildSlot;
 
 #ifdef WIN32
diff --git a/src/backend/storage/ipc/pmsignal.c b/src/backend/storage/ipc/pmsignal.c
index 9516b869cb2..79751b513a0 100644
--- a/src/backend/storage/ipc/pmsignal.c
+++ b/src/backend/storage/ipc/pmsignal.c
@@ -72,7 +72,8 @@ typedef struct ChildSlotData
 	pg_atomic_uint32 state;
 
 	int			pid;
-	int32		cancel_key;
+	int			cancel_key_len;
+	char		cancel_key[MAX_CANCEL_KEY_LENGTH];
 } ChildSlotData;
 
 /* "typedef struct PMSignalData PMSignalData" appears in pmsignal.h */
@@ -328,16 +329,18 @@ IsPostmasterChildWalSender(int slot)
  * actively using shared memory.  This is called in the child process.
  */
 void
-MarkPostmasterChildActive(int pid, int32 cancelAuthCode)
+MarkPostmasterChildActive(int pid, char *cancelKey, int len)
 {
 	int			slot = MyPMChildSlot;
 
+	Assert(len <= MAX_CANCEL_KEY_LENGTH);
 	Assert(slot > 0 && slot <= PMSignalState->num_child_slots);
 	slot--;
 	Assert(pg_atomic_read_u32(&PMSignalState->child_slots[slot].state) == PM_CHILD_ASSIGNED);
 	PMSignalState->child_slots[slot].pid = pid;
-	PMSignalState->child_slots[slot].cancel_key = cancelAuthCode;
-	pg_memory_barrier();
+	memcpy(PMSignalState->child_slots[slot].cancel_key, cancelKey, len);
+	PMSignalState->child_slots[slot].cancel_key_len = len;
+	pg_write_barrier();
 	pg_atomic_write_u32(&PMSignalState->child_slots[slot].state, PM_CHILD_ACTIVE);
 }
 
@@ -478,8 +481,18 @@ PostmasterDeathSignalInit(void)
 #endif							/* USE_POSTMASTER_DEATH_SIGNAL */
 }
 
+static int
+pg_const_time_memcmp(const void *a, const void *b, size_t len)
+{
+	/*
+	 * FIXME: need a constant time implementation. Implement one somewhere in
+	 * src/port.
+	 */
+	return memcmp(a, b, len);
+}
+
 void
-SendCancelRequest(int backendPID, int32 cancelAuthCode)
+SendCancelRequest(int backendPID, char *cancelKey, int len)
 {
 	/*
 	 * See if we have a matching backend.  In the EXEC_BACKEND case, we can no
@@ -497,7 +510,11 @@ SendCancelRequest(int backendPID, int32 cancelAuthCode)
 		pg_read_barrier();
 		if (slot->pid == backendPID)
 		{
-			if (slot->cancel_key == cancelAuthCode)
+			/*
+			 * Use pg_const_time_memcmp() to prevent an attacker from using
+			 * timing to reveal the cancel key.
+			 */
+			if (len == slot->cancel_key_len && pg_const_time_memcmp(slot->cancel_key, cancelKey, len) == 0)
 			{
 				/* Found a match; signal that backend to cancel current op */
 				ereport(DEBUG2,
diff --git a/src/backend/storage/lmgr/proc.c b/src/backend/storage/lmgr/proc.c
index 5984fe7d44c..7e4cce4c91f 100644
--- a/src/backend/storage/lmgr/proc.c
+++ b/src/backend/storage/lmgr/proc.c
@@ -36,6 +36,7 @@
 #include "access/transam.h"
 #include "access/twophase.h"
 #include "access/xlogutils.h"
+#include "libpq/libpq-be.h"
 #include "miscadmin.h"
 #include "pgstat.h"
 #include "postmaster/autovacuum.h"
@@ -68,6 +69,9 @@ bool		log_lock_waits = false;
 PGPROC	   *MyProc = NULL;
 int			MyProcNumber = INVALID_PGPROCNO;
 
+char		MyCancelKey[MAX_CANCEL_KEY_LENGTH];
+int			MyCancelKeyLength;
+
 /*
  * This spinlock protects the freelist of recycled PGPROC structures.
  * We cannot use an LWLock because the LWLock manager depends on already
@@ -317,7 +321,8 @@ InitProcess(void)
 	 * better have something random in the field to prevent unfriendly people
 	 * from sending cancels to them.
 	 */
-	if (!pg_strong_random(&MyCancelKey, sizeof(int32)))
+	MyCancelKeyLength = (MyProcPort != NULL && MyProcPort->extended_query_cancel) ? MAX_CANCEL_KEY_LENGTH : 4;
+	if (!pg_strong_random(&MyCancelKey, MyCancelKeyLength))
 	{
 		ereport(ERROR,
 				(errcode(ERRCODE_INTERNAL_ERROR),
@@ -387,7 +392,7 @@ InitProcess(void)
 	 */
 	if (IsUnderPostmaster && !IsAutoVacuumLauncherProcess() &&
 		!IsLogicalSlotSyncWorker())
-		MarkPostmasterChildActive(MyProcPid, MyCancelKey);
+		MarkPostmasterChildActive(MyProcPid, MyCancelKey, MyCancelKeyLength);
 
 	/*
 	 * Initialize all fields of MyProc, except for those previously
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 59ab812d2e8..7164aa31978 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -4274,7 +4274,7 @@ PostgresMain(const char *dbname, const char *username)
 
 		pq_beginmessage(&buf, PqMsg_BackendKeyData);
 		pq_sendint32(&buf, (int32) MyProcPid);
-		pq_sendint32(&buf, (int32) MyCancelKey);
+		pq_sendbytes(&buf, MyCancelKey, MyCancelKeyLength);
 		pq_endmessage(&buf);
 		/* Need not flush since ReadyForQuery will do it. */
 	}
diff --git a/src/backend/utils/init/globals.c b/src/backend/utils/init/globals.c
index 5eaee88d969..7efcdae9657 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -46,7 +46,6 @@ int			MyProcPid;
 pg_time_t	MyStartTime;
 TimestampTz MyStartTimestamp;
 struct Port *MyProcPort;
-int32		MyCancelKey;
 int			MyPMChildSlot;
 
 /*
diff --git a/src/include/libpq/libpq-be.h b/src/include/libpq/libpq-be.h
index 47d66d55241..3aebf92a754 100644
--- a/src/include/libpq/libpq-be.h
+++ b/src/include/libpq/libpq-be.h
@@ -175,6 +175,11 @@ typedef struct Port
 	 */
 	char	   *application_name;
 
+	/*
+	 * Protocol options supported by the client
+	 */
+	bool		extended_query_cancel;
+
 	/*
 	 * Information that needs to be held during the authentication cycle.
 	 */
diff --git a/src/include/libpq/pqcomm.h b/src/include/libpq/pqcomm.h
index 9ae469c86c4..9b5579a5386 100644
--- a/src/include/libpq/pqcomm.h
+++ b/src/include/libpq/pqcomm.h
@@ -129,6 +129,7 @@ typedef uint32 AuthRequest;
  * The cancel request code must not match any protocol version number
  * we're ever likely to use.  This random choice should do.
  */
+#define EXTENDED_CANCEL_REQUEST_CODE PG_PROTOCOL(1234,5677)
 #define CANCEL_REQUEST_CODE PG_PROTOCOL(1234,5678)
 
 typedef struct CancelRequestPacket
@@ -136,10 +137,10 @@ typedef struct CancelRequestPacket
 	/* Note that each field is stored in network byte order! */
 	MsgType		cancelRequestCode;	/* code to identify a cancel request */
 	uint32		backendPID;		/* PID of client's backend */
-	uint32		cancelAuthCode; /* secret key to authorize cancel */
+	char		cancelAuthCode[FLEXIBLE_ARRAY_MEMBER];	/* secret key to
+														 * authorize cancel */
 } CancelRequestPacket;
 
-
 /*
  * A client can also start by sending a SSL or GSSAPI negotiation request to
  * get a secure channel.
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index 756d144c323..b70b412ea90 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -192,7 +192,8 @@ extern PGDLLIMPORT pg_time_t MyStartTime;
 extern PGDLLIMPORT TimestampTz MyStartTimestamp;
 extern PGDLLIMPORT struct Port *MyProcPort;
 extern PGDLLIMPORT struct Latch *MyLatch;
-extern PGDLLIMPORT int32 MyCancelKey;
+extern PGDLLIMPORT char MyCancelKey[];
+extern PGDLLIMPORT int MyCancelKeyLength;
 extern PGDLLIMPORT int MyPMChildSlot;
 
 extern PGDLLIMPORT char OutputFileName[];
diff --git a/src/include/storage/pmsignal.h b/src/include/storage/pmsignal.h
index 8dff47e86e4..94420ecb0b5 100644
--- a/src/include/storage/pmsignal.h
+++ b/src/include/storage/pmsignal.h
@@ -57,6 +57,8 @@ typedef enum
 /* PMSignalData is an opaque struct, details known only within pmsignal.c */
 typedef struct PMSignalData PMSignalData;
 
+#define MAX_CANCEL_KEY_LENGTH	32
+
 /*
  * prototypes for functions in pmsignal.c
  */
@@ -69,12 +71,12 @@ extern QuitSignalReason GetQuitSignalReason(void);
 extern int	AssignPostmasterChildSlot(void);
 extern bool ReleasePostmasterChildSlot(int slot);
 extern bool IsPostmasterChildWalSender(int slot);
-extern void MarkPostmasterChildActive(int pid, int32 cancelAuthCode);
+extern void MarkPostmasterChildActive(int pid, char *cancelKey, int len);
 extern void MarkPostmasterChildInactive(void);
 extern void MarkPostmasterChildWalSender(void);
 extern bool PostmasterIsAliveInternal(void);
 extern void PostmasterDeathSignalInit(void);
-extern void SendCancelRequest(int backendPID, int32 cancelAuthCode);
+extern void SendCancelRequest(int backendPID, char *cancelKey, int len);
 
 
 /*
diff --git a/src/interfaces/libpq/fe-cancel.c b/src/interfaces/libpq/fe-cancel.c
index 51f8d8a78c4..6c9bfeaa33f 100644
--- a/src/interfaces/libpq/fe-cancel.c
+++ b/src/interfaces/libpq/fe-cancel.c
@@ -34,6 +34,8 @@ PGcancel *
 PQgetCancel(PGconn *conn)
 {
 	PGcancel   *cancel;
+	int			cancel_req_len;
+	CancelRequestPacket *req;
 
 	if (!conn)
 		return NULL;
@@ -41,13 +43,13 @@ PQgetCancel(PGconn *conn)
 	if (conn->sock == PGINVALID_SOCKET)
 		return NULL;
 
-	cancel = malloc(sizeof(PGcancel));
+	cancel_req_len = offsetof(CancelRequestPacket, cancelAuthCode) + conn->be_cancel_key_len;
+	cancel = malloc(offsetof(PGcancel, cancel_req) + cancel_req_len);
 	if (cancel == NULL)
 		return NULL;
 
 	memcpy(&cancel->raddr, &conn->raddr, sizeof(SockAddr));
-	cancel->be_pid = conn->be_pid;
-	cancel->be_key = conn->be_key;
+
 	/* We use -1 to indicate an unset connection option */
 	cancel->pgtcp_user_timeout = -1;
 	cancel->keepalives = -1;
@@ -90,6 +92,13 @@ PQgetCancel(PGconn *conn)
 			goto fail;
 	}
 
+	req = (CancelRequestPacket *) &cancel->cancel_req;
+	req->cancelRequestCode = (MsgType) pg_hton32(conn->be_cancel_key_len == 4 ? CANCEL_REQUEST_CODE : EXTENDED_CANCEL_REQUEST_CODE);
+	req->backendPID = pg_hton32(conn->be_pid);
+	memcpy(req->cancelAuthCode, conn->be_cancel_key, conn->be_cancel_key_len);
+	/* include the length field itself in the length */
+	cancel->cancel_pkt_len = pg_hton32(cancel_req_len + 4);
+
 	return cancel;
 
 fail:
@@ -150,11 +159,8 @@ PQcancel(PGcancel *cancel, char *errbuf, int errbufsize)
 	int			save_errno = SOCK_ERRNO;
 	pgsocket	tmpsock = PGINVALID_SOCKET;
 	int			maxlen;
-	struct
-	{
-		uint32		packetlen;
-		CancelRequestPacket cp;
-	}			crp;
+	char		recvbuf;
+	int			cancel_pkt_len;
 
 	if (!cancel)
 	{
@@ -256,15 +262,11 @@ retry3:
 		goto cancel_errReturn;
 	}
 
-	/* Create and send the cancel request packet. */
-
-	crp.packetlen = pg_hton32((uint32) sizeof(crp));
-	crp.cp.cancelRequestCode = (MsgType) pg_hton32(CANCEL_REQUEST_CODE);
-	crp.cp.backendPID = pg_hton32(cancel->be_pid);
-	crp.cp.cancelAuthCode = pg_hton32(cancel->be_key);
+	/* Send the cancel request packet. */
+	cancel_pkt_len = pg_ntoh32(cancel->cancel_pkt_len);
 
 retry4:
-	if (send(tmpsock, (char *) &crp, sizeof(crp), 0) != (int) sizeof(crp))
+	if (send(tmpsock, &cancel->cancel_pkt_len, cancel_pkt_len, 0) != cancel_pkt_len)
 	{
 		if (SOCK_ERRNO == EINTR)
 			/* Interrupted system call - we'll just try again */
@@ -281,7 +283,7 @@ retry4:
 	 * read to obtain any data, we are just waiting for EOF to be signaled.
 	 */
 retry5:
-	if (recv(tmpsock, (char *) &crp, 1, 0) < 0)
+	if (recv(tmpsock, &recvbuf, 1, 0) < 0)
 	{
 		if (SOCK_ERRNO == EINTR)
 			/* Interrupted system call - we'll just try again */
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index d4e10a0c4f3..5148b24ba34 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -617,7 +617,12 @@ pqDropServerData(PGconn *conn)
 	free(conn->write_err_msg);
 	conn->write_err_msg = NULL;
 	conn->be_pid = 0;
-	conn->be_key = 0;
+	if (conn->be_cancel_key != NULL)
+	{
+		free(conn->be_cancel_key);
+		conn->be_cancel_key = NULL;
+	}
+	conn->be_cancel_key_len = 0;
 }
 
 
@@ -3724,14 +3729,21 @@ keep_going:						/* We will come back to here until there is
 				}
 				else if (beresp == PqMsg_NegotiateProtocolVersion)
 				{
-					if (pqGetNegotiateProtocolVersion3(conn))
+					switch (pqGetNegotiateProtocolVersion3(conn))
 					{
-						libpq_append_conn_error(conn, "received invalid protocol negotiation message");
-						goto error_return;
+						case 0:
+							/* OK, we read the message; mark data consumed */
+							conn->inStart = conn->inCursor;
+							/* Stay in the CONNECTION_AWAITING_RESPONSE state */
+							goto keep_going;
+						case 1:
+							/* OK, we read the message; mark data consumed */
+							conn->inStart = conn->inCursor;
+							goto error_return;
+						case EOF:
+							/* We'll come back when there is more data */
+							return PGRES_POLLING_READING;
 					}
-					/* OK, we read the message; mark data consumed */
-					conn->inStart = conn->inCursor;
-					goto error_return;
 				}
 
 				/* It is an authentication request. */
diff --git a/src/interfaces/libpq/fe-protocol3.c b/src/interfaces/libpq/fe-protocol3.c
index 701d58e1087..c079297f736 100644
--- a/src/interfaces/libpq/fe-protocol3.c
+++ b/src/interfaces/libpq/fe-protocol3.c
@@ -310,8 +310,22 @@ pqParseInput3(PGconn *conn)
 					 */
 					if (pqGetInt(&(conn->be_pid), 4, conn))
 						return;
-					if (pqGetInt(&(conn->be_key), 4, conn))
-						return;
+
+					{
+						int			cancel_key_len = 5 + msgLength - (conn->inCursor - conn->inStart);
+
+						conn->be_cancel_key = malloc(cancel_key_len);
+						if (conn->be_cancel_key == NULL)
+						{
+							libpq_append_conn_error(conn, "out of memory");
+							/* discard the message */
+							conn->inCursor = conn->inStart + 5 + msgLength;
+							break;
+						}
+						if (pqGetnchar(conn->be_cancel_key, cancel_key_len, conn))
+							return;
+						conn->be_cancel_key_len = cancel_key_len;
+					}
 					break;
 				case PqMsg_RowDescription:
 					if (conn->error_result ||
@@ -1404,7 +1418,8 @@ reportErrorPosition(PQExpBuffer msg, const char *query, int loc, int encoding)
 /*
  * Attempt to read a NegotiateProtocolVersion message.
  * Entry: 'v' message type and length have already been consumed.
- * Exit: returns 0 if successfully consumed message.
+ * Exit: returns 0 if successfully consumed message and the negotiation succeeded.
+ *		 returns 1 if successfully consumed message and the negotiation failed.
  *		 returns EOF if not enough data.
  */
 int
@@ -1413,6 +1428,7 @@ pqGetNegotiateProtocolVersion3(PGconn *conn)
 	int			tmp;
 	ProtocolVersion their_version;
 	int			num;
+	int			num_required_missing = 0;
 	PQExpBufferData buf;
 
 	if (pqGetInt(&tmp, 4, conn) != 0)
@@ -1430,16 +1446,23 @@ pqGetNegotiateProtocolVersion3(PGconn *conn)
 			termPQExpBuffer(&buf);
 			return EOF;
 		}
+		if (strcmp(conn->workBuffer.data, "_pq_.extended_query_cancel") == 0)
+		{
+			/* that's ok */
+			continue;
+		}
+
 		if (buf.len > 0)
 			appendPQExpBufferChar(&buf, ' ');
 		appendPQExpBufferStr(&buf, conn->workBuffer.data);
+		num_required_missing++;
 	}
 
 	if (their_version < conn->pversion)
 		libpq_append_conn_error(conn, "protocol version not supported by server: client uses %u.%u, server supports up to %u.%u",
 								PG_PROTOCOL_MAJOR(conn->pversion), PG_PROTOCOL_MINOR(conn->pversion),
 								PG_PROTOCOL_MAJOR(their_version), PG_PROTOCOL_MINOR(their_version));
-	if (num > 0)
+	if (num_required_missing > 0)
 	{
 		appendPQExpBuffer(&conn->errorMessage,
 						  libpq_ngettext("protocol extension not supported by server: %s",
@@ -2312,6 +2335,7 @@ build_startup_packet(const PGconn *conn, char *packet,
 				ADD_STARTUP_OPTION(next_eo->pgName, val);
 		}
 	}
+	ADD_STARTUP_OPTION("_pq_.extended_query_cancel", "");
 
 	/* Add trailing terminator */
 	if (packet)
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 82c18f870d2..6187f282fb0 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -491,8 +491,9 @@ struct pg_conn
 	bool		send_appname;	/* okay to send application_name? */
 
 	/* Miscellaneous stuff */
-	int			be_pid;			/* PID of backend --- needed for cancels */
-	int			be_key;			/* key of backend --- needed for cancels */
+	int			be_pid;			/* PID of backend --- needed for XX cancels */
+	char	   *be_cancel_key;
+	int			be_cancel_key_len;
 	pgParameterStatus *pstatus; /* ParameterStatus data */
 	int			client_encoding;	/* encoding id */
 	bool		std_strings;	/* standard_conforming_strings */
@@ -629,7 +630,6 @@ struct pg_cancel
 {
 	SockAddr	raddr;			/* Remote address */
 	int			be_pid;			/* PID of backend --- needed for cancels */
-	int			be_key;			/* key of backend --- needed for cancels */
 	int			pgtcp_user_timeout; /* tcp user timeout */
 	int			keepalives;		/* use TCP keepalives? */
 	int			keepalives_idle;	/* time between TCP keepalives */
@@ -637,6 +637,10 @@ struct pg_cancel
 										 * retransmits */
 	int			keepalives_count;	/* maximum number of TCP keepalive
 									 * retransmits */
+
+	/* Pre-constructed cancel request packet starts here */
+	int32		cancel_pkt_len; /* in network-byte-order */
+	char		cancel_req[FLEXIBLE_ARRAY_MEMBER];	/* CancelRequestPacket */
 };
 
 
-- 
2.39.2

Reply via email to