From edd0839967129f7ac8965ebccd9bfe18a7c8967a Mon Sep 17 00:00:00 2001
From: Nathan Bossart <bossartn@amazon.com>
Date: Fri, 19 Nov 2021 22:56:13 +0000
Subject: [PATCH v1 1/1] Add control-C handling for password prompts in
 pg_receivewal and pg_recvlogical.

---
 src/bin/pg_basebackup/nls.mk           |  2 +-
 src/bin/pg_basebackup/pg_receivewal.c  | 11 ++++++++++-
 src/bin/pg_basebackup/pg_recvlogical.c | 11 ++++++++++-
 src/bin/pg_basebackup/streamutil.c     | 36 ++++++++++++++++++++++++++++------
 src/bin/pg_basebackup/streamutil.h     |  5 +++++
 5 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/src/bin/pg_basebackup/nls.mk b/src/bin/pg_basebackup/nls.mk
index 2d521f0683..338048583c 100644
--- a/src/bin/pg_basebackup/nls.mk
+++ b/src/bin/pg_basebackup/nls.mk
@@ -2,5 +2,5 @@
 CATALOG_NAME     = pg_basebackup
 AVAIL_LANGUAGES  = cs de es fr he it ja ko pl pt_BR ru sv tr uk vi zh_CN
 GETTEXT_FILES    = $(FRONTEND_COMMON_GETTEXT_FILES) pg_basebackup.c pg_receivewal.c pg_recvlogical.c receivelog.c streamutil.c walmethods.c ../../common/fe_memutils.c ../../common/file_utils.c ../../fe_utils/recovery_gen.c
-GETTEXT_TRIGGERS = $(FRONTEND_COMMON_GETTEXT_TRIGGERS) simple_prompt tar_set_error
+GETTEXT_TRIGGERS = $(FRONTEND_COMMON_GETTEXT_TRIGGERS) simple_prompt simple_prompt_extended tar_set_error
 GETTEXT_FLAGS    = $(FRONTEND_COMMON_GETTEXT_FLAGS)
diff --git a/src/bin/pg_basebackup/pg_receivewal.c b/src/bin/pg_basebackup/pg_receivewal.c
index 23cf5f8ec7..ecc068fd9b 100644
--- a/src/bin/pg_basebackup/pg_receivewal.c
+++ b/src/bin/pg_basebackup/pg_receivewal.c
@@ -649,7 +649,9 @@ StreamLog(void)
 
 /*
  * When sigint is called, just tell the system to exit at the next possible
- * moment.
+ * moment.  Since this won't work when blocked on user input (e.g., fgets()),
+ * we also provide a mechanism to longjmp out of user prompts.  While blocked,
+ * sigint_interrupt_enabled should be set to true.
  */
 #ifndef WIN32
 
@@ -657,6 +659,13 @@ static void
 sigint_handler(int signum)
 {
 	time_to_stop = true;
+
+	/* if we are waiting for input, longjmp out of it */
+	if (sigint_interrupt_enabled)
+	{
+		sigint_interrupt_enabled = false;
+		siglongjmp(sigint_interrupt_jmp, 1);
+	}
 }
 #endif
 
diff --git a/src/bin/pg_basebackup/pg_recvlogical.c b/src/bin/pg_basebackup/pg_recvlogical.c
index f235d6fecf..8c23abf1f4 100644
--- a/src/bin/pg_basebackup/pg_recvlogical.c
+++ b/src/bin/pg_basebackup/pg_recvlogical.c
@@ -657,12 +657,21 @@ error:
 
 /*
  * When sigint is called, just tell the system to exit at the next possible
- * moment.
+ * moment.  Since this won't work when blocked on user input (e.g., fgets()),
+ * we also provide a mechanism to longjmp out of user prompts.  While blocked,
+ * sigint_interrupt_enabled should be set to true.
  */
 static void
 sigint_handler(int signum)
 {
 	time_to_abort = true;
+
+	/* if we are waiting for input, longjmp out of it */
+	if (sigint_interrupt_enabled)
+	{
+		sigint_interrupt_enabled = false;
+		siglongjmp(sigint_interrupt_jmp, 1);
+	}
 }
 
 /*
diff --git a/src/bin/pg_basebackup/streamutil.c b/src/bin/pg_basebackup/streamutil.c
index 2a3e0c688f..c0adb256b0 100644
--- a/src/bin/pg_basebackup/streamutil.c
+++ b/src/bin/pg_basebackup/streamutil.c
@@ -52,16 +52,21 @@ char	   *dbname = NULL;
 int			dbgetpassword = 0;	/* 0=auto, -1=never, 1=always */
 static char *password = NULL;
 PGconn	   *conn = NULL;
+volatile bool sigint_interrupt_enabled = false;
+sigjmp_buf  sigint_interrupt_jmp;
 
 /*
  * Connect to the server. Returns a valid PGconn pointer if connected,
  * or NULL on non-permanent error. On permanent error, the function will
- * call exit(1) directly.
+ * call exit(1) directly.  The password prompt can be canceled via an
+ * existing SIGINT handler that will longjmp to the place specified by
+ * sigint_interrupt_jmp only when sigint_interrupt_enabled is true.  When
+ * canceled, this function returns NULL.
  */
 PGconn *
 GetConnection(void)
 {
-	PGconn	   *tmpconn;
+	PGconn	   *tmpconn = NULL;
 	int			argcount = 7;	/* dbname, replication, fallback_app_name,
 								 * host, user, port, password */
 	int			i;
@@ -157,10 +162,20 @@ GetConnection(void)
 		/* Get a new password if appropriate */
 		if (need_password)
 		{
+			PromptInterruptContext prompt_ctx;
+
+			/* Set up to let SIGINT cancel simple_prompt_extended() */
+			prompt_ctx.jmpbuf = sigint_interrupt_jmp;
+			prompt_ctx.enabled = &sigint_interrupt_enabled;
+			prompt_ctx.canceled = false;
+
 			if (password)
 				free(password);
-			password = simple_prompt("Password: ", false);
+			password = simple_prompt_extended("Password: ", false, &prompt_ctx);
 			need_password = false;
+
+			if (prompt_ctx.canceled)
+				break;
 		}
 
 		/* Use (or reuse, on a subsequent connection) password if we have it */
@@ -193,15 +208,24 @@ GetConnection(void)
 			dbgetpassword != -1)
 		{
 			PQfinish(tmpconn);
+			tmpconn = NULL;
 			need_password = true;
 		}
 	}
 	while (need_password);
 
-	if (PQstatus(tmpconn) != CONNECTION_OK)
+	if (tmpconn == NULL || PQstatus(tmpconn) != CONNECTION_OK)
 	{
-		pg_log_error("%s", PQerrorMessage(tmpconn));
-		PQfinish(tmpconn);
+		/*
+		 * If tmpconn is NULL, the password prompt was canceled via SIGINT,
+		 * and we don't bother logging an error message.
+		 */
+		if (tmpconn != NULL)
+		{
+			pg_log_error("%s", PQerrorMessage(tmpconn));
+			PQfinish(tmpconn);
+		}
+
 		free(values);
 		free(keywords);
 		if (conn_opts)
diff --git a/src/bin/pg_basebackup/streamutil.h b/src/bin/pg_basebackup/streamutil.h
index 7918935cb3..2c45fdace4 100644
--- a/src/bin/pg_basebackup/streamutil.h
+++ b/src/bin/pg_basebackup/streamutil.h
@@ -12,6 +12,8 @@
 #ifndef STREAMUTIL_H
 #define STREAMUTIL_H
 
+#include <setjmp.h>
+
 #include "access/xlogdefs.h"
 #include "datatype/timestamp.h"
 #include "libpq-fe.h"
@@ -65,4 +67,7 @@ extern bool feTimestampDifferenceExceeds(TimestampTz start_time, TimestampTz sto
 extern void fe_sendint64(int64 i, char *buf);
 extern int64 fe_recvint64(char *buf);
 
+extern volatile bool sigint_interrupt_enabled;
+extern sigjmp_buf sigint_interrupt_jmp;
+
 #endif							/* STREAMUTIL_H */
-- 
2.16.6

