variables-escape.diff

*** ./command.c.orig	2010-01-02 17:57:59.000000000 +0100
--- ./command.c	2010-01-22 13:14:28.325976489 +0100
***************
*** 938,944 ****
  	{
  		char	   *opt0 = psql_scan_slash_option(scan_state,
  												  OT_NORMAL, NULL, false);
! 
  		if (!opt0)
  		{
  			/* list all variables */
--- 938,944 ----
  	{
  		char	   *opt0 = psql_scan_slash_option(scan_state,
  												  OT_NORMAL, NULL, false);
! 		
  		if (!opt0)
  		{
  			/* list all variables */
***************
*** 947,984 ****
  		}
  		else
  		{
! 			/*
! 			 * Set variable to the concatenation of the arguments.
! 			 */
! 			char	   *newval;
! 			char	   *opt;
! 
! 			opt = psql_scan_slash_option(scan_state,
! 										 OT_NORMAL, NULL, false);
! 			newval = pg_strdup(opt ? opt : "");
! 			free(opt);
! 
! 			while ((opt = psql_scan_slash_option(scan_state,
! 												 OT_NORMAL, NULL, false)))
  			{
! 				newval = realloc(newval, strlen(newval) + strlen(opt) + 1);
! 				if (!newval)
  				{
! 					psql_error("out of memory\n");
! 					exit(EXIT_FAILURE);
  				}
- 				strcat(newval, opt);
- 				free(opt);
  			}
! 
! 			if (!SetVariable(pset.vars, opt0, newval))
! 			{
! 				psql_error("\\%s: error\n", cmd);
  				success = false;
! 			}
! 			free(newval);
  		}
- 		free(opt0);
  	}
  
  	/* \t -- turn off headers and row count */
--- 947,1009 ----
  		}
  		else
  		{
! 			if (psql_scan_is_valid(scan_state))
  			{
! 				 /*
! 				 * Set variable to the concatenation of the arguments. Don't allow add variable,
! 				 * when scanning isn't valid - these check are more strict, we protect content of
! 				 * variables against to invalid value.
! 				 */
! 				char	   *newval;
! 				char	   *opt;
! 
! 				opt = psql_scan_slash_option(scan_state,
! 											OT_NORMAL, NULL, false);
! 											
! 				if (psql_scan_is_valid(scan_state))
! 				{
! 					newval = pg_strdup(opt ? opt : "");
! 					free(opt);
! 					
! 					while ((opt = psql_scan_slash_option(scan_state,
! 													 OT_NORMAL, NULL, false)))
! 					{
! 						if (psql_scan_is_valid(scan_state))
! 						{
! 							newval = realloc(newval, strlen(newval) + strlen(opt) + 1);
! 							if (!newval)
! 							{
! 								psql_error("out of memory\n");
! 								exit(EXIT_FAILURE);
! 							}
! 							strcat(newval, opt);
! 							free(opt);
! 						}
! 						else
! 						{
! 							success = false;
! 							free(opt);
! 							break;
! 						}
! 					}
! 
! 					if (success && !SetVariable(pset.vars, opt0, newval))
! 					{
! 						psql_error("\\%s: error\n", cmd);
! 						success = false;
! 					}
! 					free(newval);
! 				}
! 				else
  				{
! 					success = false;
! 					free(opt);
  				}
  			}
! 			else
  				success = false;
! 			free(opt0);
  		}
  	}
  
  	/* \t -- turn off headers and row count */
***************
*** 1163,1171 ****
  	else
  		status = PSQL_CMD_UNKNOWN;
  
! 	if (!success)
  		status = PSQL_CMD_ERROR;
! 
  	return status;
  }
  
--- 1188,1196 ----
  	else
  		status = PSQL_CMD_UNKNOWN;
  
! 	if (!success || !psql_scan_is_valid(scan_state))
  		status = PSQL_CMD_ERROR;
! 	
  	return status;
  }
  
*** ./mainloop.c.orig	2010-01-02 17:57:59.000000000 +0100
--- ./mainloop.c	2010-01-22 12:02:49.273610342 +0100
***************
*** 253,260 ****
  					line_saved_in_history = true;
  				}
  
! 				/* execute query */
! 				success = SendQuery(query_buf->data);
  				slashCmdStatus = success ? PSQL_CMD_SEND : PSQL_CMD_ERROR;
  
  				/* transfer query to previous_buf by pointer-swapping */
--- 253,270 ----
  					line_saved_in_history = true;
  				}
  
! 				/* 
! 				 * execute query
! 				 *
! 				 * When scanning isn't valid from out of memory
! 				 * or broken chars reason do nothing. Error message 
! 				 * was displayed. So only don't send invalid commands
! 				 */
! 				if (psql_scan_is_valid(scan_state))
! 					success = SendQuery(query_buf->data);
! 				else
! 					success = false;
! 					
  				slashCmdStatus = success ? PSQL_CMD_SEND : PSQL_CMD_ERROR;
  
  				/* transfer query to previous_buf by pointer-swapping */
***************
*** 267,273 ****
  				resetPQExpBuffer(query_buf);
  
  				added_nl_pos = -1;
! 				/* we need not do psql_scan_reset() here */
  			}
  			else if (scan_result == PSCAN_BACKSLASH)
  			{
--- 277,283 ----
  				resetPQExpBuffer(query_buf);
  
  				added_nl_pos = -1;
! 				psql_scan_reset(scan_state);
  			}
  			else if (scan_result == PSCAN_BACKSLASH)
  			{
***************
*** 342,347 ****
--- 352,360 ----
  				}
  				else if (slashCmdStatus == PSQL_CMD_TERMINATE)
  					break;
+ 				else if (slashCmdStatus != PSQL_CMD_UNKNOWN)
+ 					/* reset parsing state - reset flag valid */
+ 					psql_scan_reset(scan_state);
  			}
  
  			/* fall out of loop if lexer reached EOL */
*** ./psqlscan.h.orig	2010-01-02 17:57:59.000000000 +0100
--- ./psqlscan.h	2010-01-22 11:30:35.236914189 +0100
***************
*** 61,64 ****
--- 61,66 ----
  
  extern void psql_scan_slash_command_end(PsqlScanState state);
  
+ extern bool psql_scan_is_valid(PsqlScanState state);
+ 
  #endif   /* PSQLSCAN_H */
*** ./psqlscan.l.orig	2010-01-02 17:57:59.000000000 +0100
--- ./psqlscan.l	2010-01-22 12:03:31.210739617 +0100
***************
*** 93,98 ****
--- 93,99 ----
  	int			paren_depth;	/* depth of nesting in parentheses */
  	int			xcdepth;		/* depth of nesting in slash-star comments */
  	char	   *dolqstart;		/* current $foo$ quote start string */
+ 	bool	   valid;		/* current scan is valid */
  } PsqlScanStateData;
  
  static PsqlScanState cur_state;	/* current state while active */
***************
*** 119,124 ****
--- 120,128 ----
  static void emit(const char *txt, int len);
  static bool is_utf16_surrogate_first(uint32 c);
  
+ static void appendLiteral(PGconn *conn, PQExpBuffer buf, const char *str);
+ static void appendIdentifier(PGconn *conn, PQExpBuffer buf, const char *str);
+ 
  #define ECHO emit(yytext, yyleng)
  
  %}
***************
*** 706,711 ****
--- 710,778 ----
  						ECHO;
  					}
  				}
+ 				
+ :'[A-Za-z0-9_]+'	{
+ 					/* 
+ 					* Possible psql variable substitution with 
+ 					* literal escaping.
+ 					*/
+ 					const char *value;
+ 				     
+ 					yytext[yyleng - 1] = '\0';
+ 					value = GetVariable(pset.vars, yytext + 2);
+ 					
+ 					if (value)
+ 					{
+ 						/* It is a variable, perform substitution */
+ 						PQExpBufferData   buf;
+ 						
+ 						initPQExpBuffer(&buf);
+ 						appendLiteral(pset.db, &buf, value);
+ 						push_new_buffer(buf.data);
+ 						termPQExpBuffer(&buf);
+ 						
+ 						/* yy_scan_string already made buffer active */
+ 					}
+ 					else
+ 					{
+ 						/*
+ 						 * if the variable doesn't exist we'll copy the
+ 						 * string as is
+ 						 */
+ 						ECHO;
+ 					}
+ 				}
+ 		
+ :\"[A-Za-z0-9_]+\"	{
+ 					/* 
+ 					* Possible psql variable substitution with 
+ 					* identifier escaping.
+ 					*/
+ 					const char *value;
+ 				     
+ 					yytext[yyleng - 1] = '\0';
+ 					value = GetVariable(pset.vars, yytext + 2);
+ 					
+ 					if (value)
+ 					{
+ 						/* It is a variable, perform substitution */
+ 						PQExpBufferData   buf;
+ 						
+ 						initPQExpBuffer(&buf);
+ 						appendIdentifier(pset.db, &buf, value);
+ 						push_new_buffer(buf.data);
+ 						termPQExpBuffer(&buf);
+ 						/* yy_scan_string already made buffer active */
+ 					}
+ 					else
+ 					{
+ 						/*
+ 						 * if the variable doesn't exist we'll copy the
+ 						 * string as is
+ 						 */
+ 						ECHO;
+ 					}
+ 				}
  
  	/*
  	 * Back to backend-compatible rules.
***************
*** 927,932 ****
--- 994,1053 ----
  					return LEXRES_OK;
  				}
  
+ 
+ :'[A-Za-z0-9_]+'	{
+ 					/* Possible psql variable substitution with literal escaping */
+ 					if (option_type == OT_VERBATIM)
+ 						ECHO;
+ 					else
+ 					{
+ 						const char *value;
+ 						
+ 						yytext[yyleng - 1] = '\0';
+ 						value = GetVariable(pset.vars, yytext + 2);
+ 
+ 						/*
+ 						 * The variable value is just emitted without any
+ 						 * further examination.  This is consistent with the
+ 						 * pre-8.0 code behavior, if not with the way that
+ 						 * variables are handled outside backslash commands.
+ 						 */
+ 						if (value)
+ 							appendLiteral(pset.db, output_buf, value);
+ 					}
+ 
+ 					*option_quote = ':';
+ 
+ 					return LEXRES_OK;
+ 				}
+ 		
+ :\"[A-Za-z0-9_]+\"	{
+ 					/* Possible psql variable substitution with identifier escaping */
+ 					if (option_type == OT_VERBATIM)
+ 						ECHO;
+ 					else
+ 					{
+ 						const char *value;
+ 						
+ 						yytext[yyleng - 1] = '\0';
+ 						value = GetVariable(pset.vars, yytext + 2);
+ 
+ 						/*
+ 						 * The variable value is just emitted without any
+ 						 * further examination.  This is consistent with the
+ 						 * pre-8.0 code behavior, if not with the way that
+ 						 * variables are handled outside backslash commands.
+ 						 */
+ 						if (value)
+ 							appendIdentifier(pset.db, output_buf, value);
+ 					}
+ 
+ 					*option_quote = ':';
+ 
+ 					return LEXRES_OK;
+ 				}
+ 
+ 
  "|"				{
  					ECHO;
  					if (option_type == OT_FILEPIPE)
***************
*** 1325,1330 ****
--- 1446,1452 ----
  	if (state->dolqstart)
  		free(state->dolqstart);
  	state->dolqstart = NULL;
+ 	state->valid = true;
  }
  
  /*
***************
*** 1341,1346 ****
--- 1463,1478 ----
  }
  
  /*
+  * Return true if lexer doesn't detect "out of memory" or 
+  * bad multibyte characters problem.
+  */
+ bool 
+ psql_scan_is_valid(PsqlScanState state)
+ {
+ 	return state->valid;
+ }
+ 
+ /*
   * Scan the command name of a psql backslash command.  This should be called
   * after psql_scan() returns PSCAN_BACKSLASH.  It is assumed that the input
   * has been consumed through the leading backslash.
***************
*** 1740,1742 ****
--- 1872,1930 ----
  {
  	return (c >= 0xD800 && c <= 0xDBFF);
  }
+ 
+ /*
+  * Convert a string value to an SQL Literal and append it to
+  * the given buffer.
+  */
+ static void
+ appendLiteral(PGconn *conn, PQExpBuffer buf, const char *str)
+ {
+ 	char	*escaped_str;
+ 	size_t		len;
+ 	
+ 	len = strlen(str);
+ 	escaped_str = PQescapeLiteral(conn, str, len);
+ 	
+ 	if (escaped_str == NULL)
+ 	{
+ 		const char *error_message = PQerrorMessage(pset.db);
+ 
+ 		if (strlen(error_message))
+ 			psql_error("%s", error_message);
+ 	}
+ 	else
+ 	{
+ 		/* only valid value is appended */
+ 		appendPQExpBufferStr(buf, escaped_str);
+ 		free(escaped_str);
+ 	}
+ }
+ 
+ /*
+  * Convert a string value to an SQL identifier and append it to
+  * the given buffer.
+  */
+ static void
+ appendIdentifier(PGconn *conn, PQExpBuffer buf, const char *str)
+ {
+ 	char	*escaped_str;
+ 	size_t		len;
+ 	
+ 	len = strlen(str);
+ 	escaped_str = PQescapeIdentifier(conn, str, len);
+ 	
+ 	if (escaped_str == NULL)
+ 	{
+ 		const char *error_message = PQerrorMessage(pset.db);
+ 
+ 		if (strlen(error_message))
+ 			psql_error("%s", error_message);
+ 	}
+ 	else
+ 	{
+ 		/* only valid value is appended */
+ 		appendPQExpBufferStr(buf, escaped_str);
+ 		free(escaped_str);
+ 	}
+ }