From 05a8fe52db98330aab54179e2062960e927e32da Mon Sep 17 00:00:00 2001 From: Nathan Bossart 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 + #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