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..102bc5956b 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; + PromptInterruptContext prompt_ctx; if (user == NULL) { @@ -2042,13 +2043,24 @@ exec_command_password(PsqlScanState scan_state, bool active_branch) PQclear(res); } + /* 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; + 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, &prompt_ctx); + if (!prompt_ctx.canceled) + pw2 = simple_prompt_extended("Enter it again: ", false, &prompt_ctx); - if (strcmp(pw1, pw2) != 0) + if (prompt_ctx.canceled) + { + /* fail silently */ + success = false; + } + else if (strcmp(pw1, pw2) != 0) { pg_log_error("Passwords didn't match."); success = false; @@ -2081,8 +2093,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/bin/psql/nls.mk b/src/bin/psql/nls.mk index 5da216f8f6..a1b7328943 100644 --- a/src/bin/psql/nls.mk +++ b/src/bin/psql/nls.mk @@ -10,5 +10,5 @@ GETTEXT_FILES = $(FRONTEND_COMMON_GETTEXT_FILES) \ ../../common/exec.c ../../common/fe_memutils.c ../../common/username.c \ ../../common/wait_error.c GETTEXT_TRIGGERS = $(FRONTEND_COMMON_GETTEXT_TRIGGERS) \ - N_ simple_prompt + N_ simple_prompt simple_prompt_extended GETTEXT_FLAGS = $(FRONTEND_COMMON_GETTEXT_FLAGS) diff --git a/src/common/pg_get_line.c b/src/common/pg_get_line.c index a80d196156..889472b28d 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 + #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. + * + * prompt_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 *(prompt_ctx->enabled) is true. If canceled, + * this function returns NULL, and prompt_ctx->canceled is set to true. */ char * -pg_get_line(FILE *stream) +pg_get_line(FILE *stream, PromptInterruptContext *prompt_ctx) { StringInfoData buf; initStringInfo(&buf); - if (!pg_get_line_append(stream, &buf)) + if (!pg_get_line_append(stream, &buf, prompt_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,48 @@ 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. + * + * prompt_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 *(prompt_ctx->enabled) is true. If canceled, + * this function returns false, and prompt_ctx->canceled is set to true. */ bool -pg_get_line_append(FILE *stream, StringInfo buf) +pg_get_line_append(FILE *stream, StringInfo buf, + PromptInterruptContext *prompt_ctx) { int orig_len = buf->len; - /* Read some data, appending it to whatever we already have */ - while (fgets(buf->data + buf->len, buf->maxlen - buf->len, stream) != NULL) + if (prompt_ctx && sigsetjmp(*((sigjmp_buf *) prompt_ctx->jmpbuf), 1) != 0) { + /* Got here with longjmp */ + prompt_ctx->canceled = true; + /* Discard any data we collected before detecting error */ + buf->len = orig_len; + buf->data[orig_len] = '\0'; + return false; + } + + /* Loop until newline or EOF/error */ + for (;;) + { + char *res; + + /* Enable longjmp while waiting for input */ + if (prompt_ctx) + *(prompt_ctx->enabled) = true; + + /* Read some data, appending it to whatever we already have */ + res = fgets(buf->data + buf->len, buf->maxlen - buf->len, stream); + + /* Disable longjmp again, then break if fgets failed */ + if (prompt_ctx) + *(prompt_ctx->enabled) = false; + + if (res == NULL) + break; + + /* Got data, so update buf->len */ buf->len += strlen(buf->data + buf->len); /* Done if we have collected a newline */ diff --git a/src/common/sprompt.c b/src/common/sprompt.c index f3a891a260..917676b58c 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 prompt_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 *(prompt_ctx->enabled) is true. If canceled, this function + * returns an empty string, and prompt_ctx->canceled is set to true. + */ +char * +simple_prompt_extended(const char *prompt, bool echo, + PromptInterruptContext *prompt_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, prompt_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..8eb5271ec8 100644 --- a/src/include/common/string.h +++ b/src/include/common/string.h @@ -12,6 +12,14 @@ struct StringInfoData; /* avoid including stringinfo.h here */ +typedef struct PromptInterruptContext +{ + /* To avoid including here, jmpbuf is declared "void *" */ + void *jmpbuf; /* existing longjmp buffer */ + volatile bool *enabled; /* flag that enables longjmp-on-interrupt */ + bool canceled; /* indicates whether cancellation occurred */ +} PromptInterruptContext; + /* 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 +29,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, PromptInterruptContext *prompt_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, + PromptInterruptContext *prompt_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, + PromptInterruptContext *prompt_ctx); #endif /* COMMON_STRING_H */