*** pgbench.c	Tue Jan  5 11:55:39 2010
--- pgbench.c.new	Tue Jan  5 11:46:02 2010
*************** getVariable(CState *st, char *name)
*** 431,438 ****
  		return NULL;
  }
  
  static int
! putVariable(CState *st, char *name, char *value)
  {
  	Variable	key,
  			   *var;
--- 431,453 ----
  		return NULL;
  }
  
+ /* check whether the name consists of alphabets, numerals and underscores. */
+ static bool
+ isLegalVariableName(const char *name)
+ {
+ 	int		i;
+ 
+ 	for (i = 0; name[i] != '\0'; i++)
+ 	{
+ 		if (!isalnum((unsigned char) name[i]) && name[i] != '_')
+ 			return false;
+ 	}
+ 
+ 	return true;
+ }
+ 
  static int
! putVariable(CState *st, const char *context, char *name, char *value)
  {
  	Variable	key,
  			   *var;
*************** putVariable(CState *st, char *name, char
*** 452,457 ****
--- 467,482 ----
  	{
  		Variable   *newvars;
  
+ 		/*
+ 		 * Check for the name only when declaring a new variable to avoid
+ 		 * overhead.
+ 		 */
+ 		if (!isLegalVariableName(name))
+ 		{
+ 			fprintf(stderr, "%s: invalid variable name '%s'\n", context, name);
+ 			return false;
+ 		}
+ 
  		if (st->variables)
  			newvars = (Variable *) realloc(st->variables,
  									(st->nvariables + 1) * sizeof(Variable));
*************** putVariable(CState *st, char *name, char
*** 459,465 ****
  			newvars = (Variable *) malloc(sizeof(Variable));
  
  		if (newvars == NULL)
! 			return false;
  
  		st->variables = newvars;
  
--- 484,490 ----
  			newvars = (Variable *) malloc(sizeof(Variable));
  
  		if (newvars == NULL)
! 			goto out_of_memory;
  
  		st->variables = newvars;
  
*************** putVariable(CState *st, char *name, char
*** 493,498 ****
--- 518,527 ----
  	}
  
  	return true;
+ 
+ out_of_memory:
+ 	fprintf(stderr, "%s: out of memory for variable '%s'\n", context, name);
+ 	return false;
  }
  
  static char *
*************** runShellCommand(CState *st, char *variab
*** 687,697 ****
  		return false;
  	}
  	snprintf(res, sizeof(res), "%d", retval);
! 	if (!putVariable(st, variable, res))
! 	{
! 		fprintf(stderr, "%s: out of memory\n", argv[0]);
  		return false;
- 	}
  
  #ifdef DEBUG
  	printf("shell parameter name: %s, value: %s\n", argv[1], res);
--- 716,723 ----
  		return false;
  	}
  	snprintf(res, sizeof(res), "%d", retval);
! 	if (!putVariable(st, "setshell", variable, res))
  		return false;
  
  #ifdef DEBUG
  	printf("shell parameter name: %s, value: %s\n", argv[1], res);
*************** top:
*** 987,995 ****
  #endif
  			snprintf(res, sizeof(res), "%d", getrand(min, max));
  
! 			if (putVariable(st, argv[1], res) == false)
  			{
- 				fprintf(stderr, "%s: out of memory\n", argv[0]);
  				st->ecnt++;
  				return true;
  			}
--- 1013,1020 ----
  #endif
  			snprintf(res, sizeof(res), "%d", getrand(min, max));
  
! 			if (!putVariable(st, argv[0], argv[1], res))
  			{
  				st->ecnt++;
  				return true;
  			}
*************** top:
*** 1057,1065 ****
  				}
  			}
  
! 			if (putVariable(st, argv[1], res) == false)
  			{
- 				fprintf(stderr, "%s: out of memory\n", argv[0]);
  				st->ecnt++;
  				return true;
  			}
--- 1082,1089 ----
  				}
  			}
  
! 			if (!putVariable(st, argv[0], argv[1], res))
  			{
  				st->ecnt++;
  				return true;
  			}
*************** main(int argc, char **argv)
*** 1874,1884 ****
  					}
  
  					*p++ = '\0';
! 					if (putVariable(&state[0], optarg, p) == false)
! 					{
! 						fprintf(stderr, "Couldn't allocate memory for variable\n");
  						exit(1);
- 					}
  				}
  				break;
  			case 'F':
--- 1898,1905 ----
  					}
  
  					*p++ = '\0';
! 					if (!putVariable(&state[0], "option", optarg, p))
  						exit(1);
  				}
  				break;
  			case 'F':
*************** main(int argc, char **argv)
*** 1958,1968 ****
  			state[i].id = i;
  			for (j = 0; j < state[0].nvariables; j++)
  			{
! 				if (putVariable(&state[i], state[0].variables[j].name, state[0].variables[j].value) == false)
! 				{
! 					fprintf(stderr, "Couldn't allocate memory for variable\n");
  					exit(1);
- 				}
  			}
  		}
  	}
--- 1979,1986 ----
  			state[i].id = i;
  			for (j = 0; j < state[0].nvariables; j++)
  			{
! 				if (!putVariable(&state[i], "startup", state[0].variables[j].name, state[0].variables[j].value))
  					exit(1);
  			}
  		}
  	}
*************** main(int argc, char **argv)
*** 2039,2049 ****
  		snprintf(val, sizeof(val), "%d", scale);
  		for (i = 0; i < nclients; i++)
  		{
! 			if (putVariable(&state[i], "scale", val) == false)
! 			{
! 				fprintf(stderr, "Couldn't allocate memory for variable\n");
  				exit(1);
- 			}
  		}
  	}
  
--- 2057,2064 ----
  		snprintf(val, sizeof(val), "%d", scale);
  		for (i = 0; i < nclients; i++)
  		{
! 			if (!putVariable(&state[i], "startup", "scale", val))
  				exit(1);
  		}
  	}