v3-0001-Add-control-C-handling-for-psql-s-password-comman.patch

From 05a8fe52db98330aab54179e2062960e927e32da Mon Sep 17 00:00:00 2001
From: Nathan Bossart <bossartn@amazon.com>
Date: Mon, 15 Nov 2021 21:43:11 +0000
Subject: [PATCH v3 1/1] Add control-C handling for psql's \password command.

---
 src/backend/libpq/hba.c     |  2 +-
 src/bin/initdb/initdb.c     |  2 +-
 src/bin/psql/command.c      | 24 ++++++++++++++++--------
 src/common/pg_get_line.c    | 42 ++++++++++++++++++++++++++++++++++++++----
 src/common/sprompt.c        | 18 +++++++++++++++++-
 src/include/common/string.h | 13 +++++++++++--
 6 files changed, 84 insertions(+), 17 deletions(-)

diff --git a/src/backend/libpq/hba.c b/src/backend/libpq/hba.c
index 3be8778d21..4328eb74fe 100644
--- a/src/backend/libpq/hba.c
+++ b/src/backend/libpq/hba.c
@@ -500,7 +500,7 @@ tokenize_file(const char *filename, FILE *file, List **tok_lines, int elevel)
 		/* Collect the next input line, handling backslash continuations */
 		resetStringInfo(&buf);
 
-		while (pg_get_line_append(file, &buf))
+		while (pg_get_line_append(file, &buf, NULL))
 		{
 			/* Strip trailing newline, including \r in case we're on Windows */
 			buf.len = pg_strip_crlf(buf.data);
diff --git a/src/bin/initdb/initdb.c b/src/bin/initdb/initdb.c
index 31839c1a19..3c61c789e4 100644
--- a/src/bin/initdb/initdb.c
+++ b/src/bin/initdb/initdb.c
@@ -1497,7 +1497,7 @@ get_su_pwd(void)
 						 pwfilename);
 			exit(1);
 		}
-		pwd1 = pg_get_line(pwf);
+		pwd1 = pg_get_line(pwf, NULL);
 		if (!pwd1)
 		{
 			if (ferror(pwf))
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 3de9d096fd..142c51aa15 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -2025,9 +2025,10 @@ exec_command_password(PsqlScanState scan_state, bool active_branch)
 	{
 		char	   *user = psql_scan_slash_option(scan_state,
 												  OT_SQLID, NULL, true);
-		char	   *pw1;
-		char	   *pw2;
+		char	   *pw1 = NULL;
+		char	   *pw2 = NULL;
 		PQExpBufferData buf;
+		SigintInterruptContext sigint_ctx;
 
 		if (user == NULL)
 		{
@@ -2042,18 +2043,23 @@ exec_command_password(PsqlScanState scan_state, bool active_branch)
 			PQclear(res);
 		}
 
+		sigint_ctx.jmpbuf = &sigint_interrupt_jmp;
+		sigint_ctx.enabled = &sigint_interrupt_enabled;
+		sigint_ctx.canceled = false;
+
 		initPQExpBuffer(&buf);
 		printfPQExpBuffer(&buf, _("Enter new password for user \"%s\": "), user);
 
-		pw1 = simple_prompt(buf.data, false);
-		pw2 = simple_prompt("Enter it again: ", false);
+		pw1 = simple_prompt_extended(buf.data, false, &sigint_ctx);
+		if (!sigint_ctx.canceled)
+			pw2 = simple_prompt_extended("Enter it again: ", false, &sigint_ctx);
 
-		if (strcmp(pw1, pw2) != 0)
+		if (!sigint_ctx.canceled && strcmp(pw1, pw2) != 0)
 		{
 			pg_log_error("Passwords didn't match.");
 			success = false;
 		}
-		else
+		else if (!sigint_ctx.canceled)
 		{
 			char	   *encrypted_password;
 
@@ -2081,8 +2087,10 @@ exec_command_password(PsqlScanState scan_state, bool active_branch)
 		}
 
 		free(user);
-		free(pw1);
-		free(pw2);
+		if (pw1)
+			free(pw1);
+		if (pw2)
+			free(pw2);
 		termPQExpBuffer(&buf);
 	}
 	else
diff --git a/src/common/pg_get_line.c b/src/common/pg_get_line.c
index a80d196156..49e18eec82 100644
--- a/src/common/pg_get_line.c
+++ b/src/common/pg_get_line.c
@@ -18,6 +18,8 @@
 #include "postgres_fe.h"
 #endif
 
+#include <setjmp.h>
+
 #include "common/string.h"
 #include "lib/stringinfo.h"
 
@@ -47,15 +49,20 @@
  * to collect lots of long-lived data.  A less memory-hungry option
  * is to use pg_get_line_buf() or pg_get_line_append() in a loop,
  * then pstrdup() each line.
+ *
+ * sigint_ctx can optionally be provided to allow this function to be
+ * canceled via an existing SIGINT signal handler that will longjmp to the
+ * specified place only when *(sigint_ctx->enabled) is true.  If canceled,
+ * this function returns NULL, and sigint_ctx->canceled is set to true.
  */
 char *
-pg_get_line(FILE *stream)
+pg_get_line(FILE *stream, SigintInterruptContext *sigint_ctx)
 {
 	StringInfoData buf;
 
 	initStringInfo(&buf);
 
-	if (!pg_get_line_append(stream, &buf))
+	if (!pg_get_line_append(stream, &buf, sigint_ctx))
 	{
 		/* ensure that free() doesn't mess up errno */
 		int			save_errno = errno;
@@ -89,7 +96,7 @@ pg_get_line_buf(FILE *stream, StringInfo buf)
 {
 	/* We just need to drop any data from the previous call */
 	resetStringInfo(buf);
-	return pg_get_line_append(stream, buf);
+	return pg_get_line_append(stream, buf, NULL);
 }
 
 /*
@@ -107,15 +114,35 @@ pg_get_line_buf(FILE *stream, StringInfo buf)
  *
  * In the false-result case, the contents of *buf are logically unmodified,
  * though it's possible that the buffer has been resized.
+ *
+ * sigint_ctx can optionally be provided to allow this function to be
+ * canceled via an existing SIGINT signal handler that will longjmp to the
+ * specified place only when *(sigint_ctx->enabled) is true.  If canceled,
+ * this function returns false, and sigint_ctx->canceled is set to true.
  */
 bool
-pg_get_line_append(FILE *stream, StringInfo buf)
+pg_get_line_append(FILE *stream, StringInfo buf,
+				   SigintInterruptContext *sigint_ctx)
 {
 	int			orig_len = buf->len;
 
+	if (sigint_ctx && sigsetjmp(*((sigjmp_buf *) sigint_ctx->jmpbuf), 1) != 0)
+	{
+		/* got here with longjmp */
+		sigint_ctx->canceled = true;
+		return false;
+	}
+
+	/* enable longjmp while waiting for input */
+	if (sigint_ctx)
+		*(sigint_ctx->enabled) = true;
+
 	/* Read some data, appending it to whatever we already have */
 	while (fgets(buf->data + buf->len, buf->maxlen - buf->len, stream) != NULL)
 	{
+		if (sigint_ctx)
+			*(sigint_ctx->enabled) = false;
+
 		buf->len += strlen(buf->data + buf->len);
 
 		/* Done if we have collected a newline */
@@ -124,8 +151,15 @@ pg_get_line_append(FILE *stream, StringInfo buf)
 
 		/* Make some more room in the buffer, and loop to read more data */
 		enlargeStringInfo(buf, 128);
+
+		/* enable longjmp while waiting for input */
+		if (sigint_ctx)
+			*(sigint_ctx->enabled) = true;
 	}
 
+	if (sigint_ctx)
+		*(sigint_ctx->enabled) = false;
+
 	/* Check for I/O errors and EOF */
 	if (ferror(stream) || buf->len == orig_len)
 	{
diff --git a/src/common/sprompt.c b/src/common/sprompt.c
index f3a891a260..2a86bced38 100644
--- a/src/common/sprompt.c
+++ b/src/common/sprompt.c
@@ -36,6 +36,22 @@
  */
 char *
 simple_prompt(const char *prompt, bool echo)
+{
+	return simple_prompt_extended(prompt, echo, NULL);
+}
+
+/*
+ * simple_prompt_extended
+ *
+ * This is the same as simple_prompt(), except that sigint_ctx can
+ * optionally be provided to allow this function to be canceled via an
+ * existing SIGINT signal handler that will longjmp to the specified place
+ * only when *(sigint_ctx->enabled) is true.  If canceled, this function
+ * returns an empty string, and sigint_ctx->canceled is set to true.
+ */
+char *
+simple_prompt_extended(const char *prompt, bool echo,
+					   SigintInterruptContext *sigint_ctx)
 {
 	char	   *result;
 	FILE	   *termin,
@@ -126,7 +142,7 @@ simple_prompt(const char *prompt, bool echo)
 		fflush(termout);
 	}
 
-	result = pg_get_line(termin);
+	result = pg_get_line(termin, sigint_ctx);
 
 	/* If we failed to read anything, just return an empty string */
 	if (result == NULL)
diff --git a/src/include/common/string.h b/src/include/common/string.h
index 686c158efe..238c3dc895 100644
--- a/src/include/common/string.h
+++ b/src/include/common/string.h
@@ -12,6 +12,12 @@
 
 struct StringInfoData;			/* avoid including stringinfo.h here */
 
+typedef struct SigintInterruptContext {
+	void	   *jmpbuf;			/* existing longjmp buffer (really a sigjmp_buf *) */
+	volatile bool *enabled;		/* used to enable/disable SIGINT handling */
+	bool		canceled;		/* indicates whether cancellation occurred */
+} SigintInterruptContext;
+
 /* functions in src/common/string.c */
 extern bool pg_str_endswith(const char *str, const char *end);
 extern int	strtoint(const char *pg_restrict str, char **pg_restrict endptr,
@@ -21,11 +27,14 @@ extern int	pg_strip_crlf(char *str);
 extern bool pg_is_ascii(const char *str);
 
 /* functions in src/common/pg_get_line.c */
-extern char *pg_get_line(FILE *stream);
+extern char *pg_get_line(FILE *stream, SigintInterruptContext *sigint_ctx);
 extern bool pg_get_line_buf(FILE *stream, struct StringInfoData *buf);
-extern bool pg_get_line_append(FILE *stream, struct StringInfoData *buf);
+extern bool pg_get_line_append(FILE *stream, struct StringInfoData *buf,
+							   SigintInterruptContext *sigint_ctx);
 
 /* functions in src/common/sprompt.c */
 extern char *simple_prompt(const char *prompt, bool echo);
+extern char *simple_prompt_extended(const char *prompt, bool echo,
+									SigintInterruptContext *sigint_ctx);
 
 #endif							/* COMMON_STRING_H */
-- 
2.16.6