From 457d72df44e177774c72049345649e6a63544bb9 Mon Sep 17 00:00:00 2001 From: Nathan Bossart Date: Sat, 13 Nov 2021 22:56:08 +0000 Subject: [PATCH v1 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 | 47 +++++++++++++++++++++++++++++++++++++-------- src/common/sprompt.c | 17 +++++++++++++++- src/include/common/string.h | 15 +++++++++++++-- 6 files changed, 86 insertions(+), 21 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..bd2f813d82 100644 --- a/src/common/pg_get_line.c +++ b/src/common/pg_get_line.c @@ -47,25 +47,34 @@ * 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. 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; + StringInfo buf; + char *ret; - initStringInfo(&buf); + /* make sure buf is palloc'd so we don't lose changes after a longjmp */ + buf = makeStringInfo(); - 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; - pfree(buf.data); + pfree(buf->data); + pfree(buf); errno = save_errno; return NULL; } - return buf.data; + ret = buf->data; + pfree(buf); + return ret;; } /* @@ -89,7 +98,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 +116,34 @@ 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. 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(*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 */ @@ -126,6 +154,9 @@ pg_get_line_append(FILE *stream, StringInfo buf) enlargeStringInfo(buf, 128); } + 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..6fb133114a 100644 --- a/src/common/sprompt.c +++ b/src/common/sprompt.c @@ -36,6 +36,21 @@ */ 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. 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 +141,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..331f23ef33 100644 --- a/src/include/common/string.h +++ b/src/include/common/string.h @@ -10,8 +10,16 @@ #ifndef COMMON_STRING_H #define COMMON_STRING_H +#include + struct StringInfoData; /* avoid including stringinfo.h here */ +typedef struct SigintInterruptContext { + sigjmp_buf *jmpbuf; /* preexisting longjmp buffer */ + 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 +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, 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