Be more careful about newline-chomping in pgbench.

process_backslash_command would drop the last character of the input
command on the assumption that it was a newline.  Given a non newline
terminated input file, this could result in dropping the last character
of the command.  Fix that by doing an actual test that we're removing
a newline.

While at it, allow for Windows newlines (\r\n), and suppress multiple
newlines if any.  I do not think either of those cases really occur,
since (a) we read script files in text mode and (b) the lexer stops
when it hits a newline.  But it's cheap enough and it provides a
stronger guarantee about what the result string looks like.

This is just cosmetic, I think, since the possibly-overly-chomped
line was only used for display not for further processing.  So
it doesn't seem necessary to back-patch.

Fabien Coelho, reviewed by Nikolay Shaplov, whacked around a bit by me

Discussion: https://postgr.es/m/alpine.DEB.2.20.1704171422500.4025@lancre

3 files changed, 26 insertions, 22 deletions

diff --git a/src/bin/pgbench/exprscan.l b/src/bin/pgbench/exprscan.l
index 1862fe4030..9bf6d237f5 100644
--- a/src/bin/pgbench/exprscan.l
+++ b/src/bin/pgbench/exprscan.l
@@ -197,7 +197,6 @@ expr_yyerror_more(yyscan_t yyscanner, const char *message, const char *more)
 	int			error_detection_offset = expr_scanner_offset(state) - 1;
 	YYSTYPE		lval;
 	char	   *full_line;
-	size_t		l;
 
 	/*
 	 * While parsing an expression, we may not have collected the whole line
@@ -210,13 +209,11 @@ expr_yyerror_more(yyscan_t yyscanner, const char *message, const char *more)
 			 /* skip */ ;
 	}
 
+	/* Extract the line, trimming trailing newline if any */
 	full_line = expr_scanner_get_substring(state,
 										   expr_start_offset,
-										   expr_scanner_offset(state));
-	/* Trim trailing newline if any */
-	l = strlen(full_line);
-	while (l > 0 && full_line[l - 1] == '\n')
-		full_line[--l] = '\0';
+										   expr_scanner_offset(state),
+										   true);
 
 	syntax_error(expr_source, expr_lineno, full_line, expr_command,
 				 message, more, error_detection_offset - expr_start_offset);
@@ -341,19 +338,30 @@ expr_scanner_offset(PsqlScanState state)
 
 /*
  * Get a malloc'd copy of the lexer input string from start_offset
- * to just before end_offset.
+ * to just before end_offset.  If chomp is true, drop any trailing
+ * newline(s).
  */
 char *
 expr_scanner_get_substring(PsqlScanState state,
-						   int start_offset, int end_offset)
+						   int start_offset, int end_offset,
+						   bool chomp)
 {
 	char	   *result;
+	const char *scanptr = state->scanbuf + start_offset;
 	int			slen = end_offset - start_offset;
 
 	Assert(slen >= 0);
 	Assert(end_offset <= strlen(state->scanbuf));
+
+	if (chomp)
+	{
+		while (slen > 0 &&
+			   (scanptr[slen - 1] == '\n' || scanptr[slen - 1] == '\r'))
+			slen--;
+	}
+
 	result = (char *) pg_malloc(slen + 1);
-	memcpy(result, state->scanbuf + start_offset, slen);
+	memcpy(result, scanptr, slen);
 	result[slen] = '\0';
 
 	return result;
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 364e25447e..e37496c971 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -3065,8 +3065,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 	PQExpBufferData word_buf;
 	int			word_offset;
 	int			offsets[MAX_ARGS];	/* offsets of argument words */
-	int			start_offset,
-				end_offset;
+	int			start_offset;
 	int			lineno;
 	int			j;
 
@@ -3120,13 +3119,11 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 
 		my_command->expr = expr_parse_result;
 
-		/* Get location of the ending newline */
-		end_offset = expr_scanner_offset(sstate) - 1;
-
-		/* Save line */
+		/* Save line, trimming any trailing newline */
 		my_command->line = expr_scanner_get_substring(sstate,
 													  start_offset,
-													  end_offset);
+													  expr_scanner_offset(sstate),
+													  true);
 
 		expr_scanner_finish(yyscanner);
 
@@ -3147,13 +3144,11 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 		my_command->argc++;
 	}
 
-	/* Get location of the ending newline */
-	end_offset = expr_scanner_offset(sstate) - 1;
-
-	/* Save line */
+	/* Save line, trimming any trailing newline */
 	my_command->line = expr_scanner_get_substring(sstate,
 												  start_offset,
-												  end_offset);
+												  expr_scanner_offset(sstate),
+												  true);
 
 	if (pg_strcasecmp(my_command->argv[0], "sleep") == 0)
 	{
diff --git a/src/bin/pgbench/pgbench.h b/src/bin/pgbench/pgbench.h
index abc13e9463..fd428af274 100644
--- a/src/bin/pgbench/pgbench.h
+++ b/src/bin/pgbench/pgbench.h
@@ -128,7 +128,8 @@ extern yyscan_t expr_scanner_init(PsqlScanState state,
 extern void expr_scanner_finish(yyscan_t yyscanner);
 extern int	expr_scanner_offset(PsqlScanState state);
 extern char *expr_scanner_get_substring(PsqlScanState state,
-						   int start_offset, int end_offset);
+						   int start_offset, int end_offset,
+						   bool chomp);
 extern int	expr_scanner_get_lineno(PsqlScanState state, int offset);
 
 extern void syntax_error(const char *source, int lineno, const char *line,