From 46d665bc26ce57b5afecbc218c8fc3c6848211d8 Mon Sep 17 00:00:00 2001
From: Tom Lane <tgl@sss.pgh.pa.us>
Date: Fri, 19 Nov 2021 12:11:38 -0500
Subject: [PATCH] Allow psql's other uses of simple_prompt() to be interrupted
 by ^C.

This fills in the work left un-done by 5f1148224.  \prompt can
be canceled out of now, and so can password prompts issued during
\connect.  (We don't need to do anything for password prompts
issued during startup, because we aren't yet trapping SIGINT
at that point.)

Nathan Bossart

Discussion: https://postgr.es/m/747443.1635536754@sss.pgh.pa.us
---
 src/bin/psql/command.c | 45 ++++++++++++++++++++++++++++++++----------
 src/bin/psql/startup.c |  3 ++-
 src/common/sprompt.c   |  6 ++++++
 3 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 102bc5956b7..f5b2cd12c0f 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -2132,6 +2132,12 @@ exec_command_prompt(PsqlScanState scan_state, bool active_branch,
 		else
 		{
 			char	   *result;
+			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 (arg2)
 			{
@@ -2143,7 +2149,7 @@ exec_command_prompt(PsqlScanState scan_state, bool active_branch,
 
 			if (!pset.inputfile)
 			{
-				result = simple_prompt(prompt_text, true);
+				result = simple_prompt_extended(prompt_text, true, &prompt_ctx);
 			}
 			else
 			{
@@ -2161,8 +2167,8 @@ exec_command_prompt(PsqlScanState scan_state, bool active_branch,
 				}
 			}
 
-			if (result &&
-				!SetVariable(pset.vars, opt, result))
+			if (prompt_ctx.canceled ||
+				(result && !SetVariable(pset.vars, opt, result)))
 				success = false;
 
 			if (result)
@@ -3058,24 +3064,36 @@ copy_previous_query(PQExpBuffer query_buf, PQExpBuffer previous_buf)
 
 /*
  * Ask the user for a password; 'username' is the username the
- * password is for, if one has been explicitly specified. Returns a
- * malloc'd string.
+ * password is for, if one has been explicitly specified.
+ * Returns a malloc'd string.
+ * If 'canceled' is provided, *canceled will be set to true if the prompt
+ * is canceled via SIGINT, and to false otherwise.
  */
 static char *
-prompt_for_password(const char *username)
+prompt_for_password(const char *username, bool *canceled)
 {
 	char	   *result;
+	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 (username == NULL || username[0] == '\0')
-		result = simple_prompt("Password: ", false);
+		result = simple_prompt_extended("Password: ", false, &prompt_ctx);
 	else
 	{
 		char	   *prompt_text;
 
 		prompt_text = psprintf(_("Password for user %s: "), username);
-		result = simple_prompt(prompt_text, false);
+		result = simple_prompt_extended(prompt_text, false, &prompt_ctx);
 		free(prompt_text);
 	}
+
+	if (canceled)
+		*canceled = prompt_ctx.canceled;
+
 	return result;
 }
 
@@ -3331,6 +3349,8 @@ do_connect(enum trivalue reuse_previous_specification,
 	 */
 	if (pset.getPassword == TRI_YES && success)
 	{
+		bool		canceled = false;
+
 		/*
 		 * If a connstring or URI is provided, we don't know which username
 		 * will be used, since we haven't dug that out of the connstring.
@@ -3338,7 +3358,9 @@ do_connect(enum trivalue reuse_previous_specification,
 		 * not seem worth working harder, since this getPassword setting is
 		 * normally only used in noninteractive cases.
 		 */
-		password = prompt_for_password(has_connection_string ? NULL : user);
+		password = prompt_for_password(has_connection_string ? NULL : user,
+									   &canceled);
+		success = !canceled;
 	}
 
 	/*
@@ -3417,13 +3439,16 @@ do_connect(enum trivalue reuse_previous_specification,
 		 */
 		if (!password && PQconnectionNeedsPassword(n_conn) && pset.getPassword != TRI_NO)
 		{
+			bool		canceled = false;
+
 			/*
 			 * Prompt for password using the username we actually connected
 			 * with --- it might've come out of "dbname" rather than "user".
 			 */
-			password = prompt_for_password(PQuser(n_conn));
+			password = prompt_for_password(PQuser(n_conn), &canceled);
 			PQfinish(n_conn);
 			n_conn = NULL;
+			success = !canceled;
 			continue;
 		}
 
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
index 2931530f338..f7ea4ce3d46 100644
--- a/src/bin/psql/startup.c
+++ b/src/bin/psql/startup.c
@@ -239,7 +239,8 @@ main(int argc, char *argv[])
 		/*
 		 * We can't be sure yet of the username that will be used, so don't
 		 * offer a potentially wrong one.  Typical uses of this option are
-		 * noninteractive anyway.
+		 * noninteractive anyway.  (Note: since we've not yet set up our
+		 * cancel handler, there's no need to use simple_prompt_extended.)
 		 */
 		password = simple_prompt("Password: ", false);
 	}
diff --git a/src/common/sprompt.c b/src/common/sprompt.c
index 917676b58ce..8f5aeeccbc0 100644
--- a/src/common/sprompt.c
+++ b/src/common/sprompt.c
@@ -164,6 +164,12 @@ simple_prompt_extended(const char *prompt, bool echo,
 		fflush(termout);
 #endif
 	}
+	else if (prompt_ctx && prompt_ctx->canceled)
+	{
+		/* also echo \n if prompt was canceled */
+		fputs("\n", termout);
+		fflush(termout);
+	}
 
 	if (termin != stdin)
 	{