Commit dd5fc264 authored by Pieter Noordhuis's avatar Pieter Noordhuis
Browse files

Make command formatters gracefully abort when out of memory

parent d4ebb60d
...@@ -652,59 +652,74 @@ static int intlen(int i) { ...@@ -652,59 +652,74 @@ static int intlen(int i) {
return len; return len;
} }
/* Helper function for redisvFormatCommand(). */ /* Helper that calculates the bulk length given a certain string length. */
static void addArgument(sds a, char ***argv, int *argc, int *totlen) { static size_t bulklen(size_t len) {
(*argc)++; return 1+intlen(len)+2+len+2;
if ((*argv = realloc(*argv, sizeof(char*)*(*argc))) == NULL) redisOOM();
if (totlen) *totlen = *totlen+1+intlen(sdslen(a))+2+sdslen(a)+2;
(*argv)[(*argc)-1] = a;
} }
int redisvFormatCommand(char **target, const char *format, va_list ap) { int redisvFormatCommand(char **target, const char *format, va_list ap) {
size_t size; const char *c = format;
const char *arg, *c = format;
char *cmd = NULL; /* final command */ char *cmd = NULL; /* final command */
int pos; /* position in final command */ int pos; /* position in final command */
sds current; /* current argument */ sds curarg, newarg; /* current argument */
int touched = 0; /* was the current argument touched? */ int touched = 0; /* was the current argument touched? */
char **argv = NULL; char **curargv = NULL, **newargv = NULL;
int argc = 0, j; int argc = 0;
int totlen = 0; int totlen = 0;
int j;
/* Abort if there is not target to set */ /* Abort if there is not target to set */
if (target == NULL) if (target == NULL)
return -1; return -1;
/* Build the command string accordingly to protocol */ /* Build the command string accordingly to protocol */
current = sdsempty(); curarg = sdsempty();
if (curarg == NULL)
return -1;
while(*c != '\0') { while(*c != '\0') {
if (*c != '%' || c[1] == '\0') { if (*c != '%' || c[1] == '\0') {
if (*c == ' ') { if (*c == ' ') {
if (touched) { if (touched) {
addArgument(current, &argv, &argc, &totlen); newargv = realloc(curargv,sizeof(char*)*(argc+1));
current = sdsempty(); if (newargv == NULL) goto err;
curargv = newargv;
curargv[argc++] = curarg;
totlen += bulklen(sdslen(curarg));
/* curarg is put in argv so it can be overwritten. */
curarg = sdsempty();
if (curarg == NULL) goto err;
touched = 0; touched = 0;
} }
} else { } else {
current = sdscatlen(current,c,1); newarg = sdscatlen(curarg,c,1);
if (newarg == NULL) goto err;
curarg = newarg;
touched = 1; touched = 1;
} }
} else { } else {
char *arg;
size_t size;
/* Set newarg so it can be checked even if it is not touched. */
newarg = curarg;
switch(c[1]) { switch(c[1]) {
case 's': case 's':
arg = va_arg(ap,char*); arg = va_arg(ap,char*);
size = strlen(arg); size = strlen(arg);
if (size > 0) if (size > 0)
current = sdscatlen(current,arg,size); newarg = sdscatlen(curarg,arg,size);
break; break;
case 'b': case 'b':
arg = va_arg(ap,char*); arg = va_arg(ap,char*);
size = va_arg(ap,size_t); size = va_arg(ap,size_t);
if (size > 0) if (size > 0)
current = sdscatlen(current,arg,size); newarg = sdscatlen(curarg,arg,size);
break; break;
case '%': case '%':
current = sdscat(current,"%"); newarg = sdscat(curarg,"%");
break; break;
default: default:
/* Try to detect printf format */ /* Try to detect printf format */
...@@ -746,7 +761,7 @@ int redisvFormatCommand(char **target, const char *format, va_list ap) { ...@@ -746,7 +761,7 @@ int redisvFormatCommand(char **target, const char *format, va_list ap) {
memcpy(_format,c,_l); memcpy(_format,c,_l);
_format[_l] = '\0'; _format[_l] = '\0';
va_copy(_cpy,ap); va_copy(_cpy,ap);
current = sdscatvprintf(current,_format,_cpy); newarg = sdscatvprintf(curarg,_format,_cpy);
va_end(_cpy); va_end(_cpy);
/* Update current position (note: outer blocks /* Update current position (note: outer blocks
...@@ -759,6 +774,10 @@ int redisvFormatCommand(char **target, const char *format, va_list ap) { ...@@ -759,6 +774,10 @@ int redisvFormatCommand(char **target, const char *format, va_list ap) {
va_arg(ap,void); va_arg(ap,void);
} }
} }
if (newarg == NULL) goto err;
curarg = newarg;
touched = 1; touched = 1;
c++; c++;
} }
...@@ -767,31 +786,55 @@ int redisvFormatCommand(char **target, const char *format, va_list ap) { ...@@ -767,31 +786,55 @@ int redisvFormatCommand(char **target, const char *format, va_list ap) {
/* Add the last argument if needed */ /* Add the last argument if needed */
if (touched) { if (touched) {
addArgument(current, &argv, &argc, &totlen); newargv = realloc(curargv,sizeof(char*)*(argc+1));
if (newargv == NULL) goto err;
curargv = newargv;
curargv[argc++] = curarg;
totlen += bulklen(sdslen(curarg));
} else { } else {
sdsfree(current); sdsfree(curarg);
} }
/* Clear curarg because it was put in curargv or was free'd. */
curarg = NULL;
/* Add bytes needed to hold multi bulk count */ /* Add bytes needed to hold multi bulk count */
totlen += 1+intlen(argc)+2; totlen += 1+intlen(argc)+2;
/* Build the command at protocol level */ /* Build the command at protocol level */
cmd = malloc(totlen+1); cmd = malloc(totlen+1);
if (!cmd) redisOOM(); if (cmd == NULL) goto err;
pos = sprintf(cmd,"*%d\r\n",argc); pos = sprintf(cmd,"*%d\r\n",argc);
for (j = 0; j < argc; j++) { for (j = 0; j < argc; j++) {
pos += sprintf(cmd+pos,"$%zu\r\n",sdslen(argv[j])); pos += sprintf(cmd+pos,"$%zu\r\n",sdslen(curargv[j]));
memcpy(cmd+pos,argv[j],sdslen(argv[j])); memcpy(cmd+pos,curargv[j],sdslen(curargv[j]));
pos += sdslen(argv[j]); pos += sdslen(curargv[j]);
sdsfree(argv[j]); sdsfree(curargv[j]);
cmd[pos++] = '\r'; cmd[pos++] = '\r';
cmd[pos++] = '\n'; cmd[pos++] = '\n';
} }
assert(pos == totlen); assert(pos == totlen);
free(argv); cmd[pos] = '\0';
cmd[totlen] = '\0';
free(curargv);
*target = cmd; *target = cmd;
return totlen; return totlen;
err:
while(argc--)
sdsfree(curargv[argc]);
free(curargv);
if (curarg != NULL)
sdsfree(curarg);
/* No need to check cmd since it is the last statement that can fail,
* but do it anyway to be as defensive as possible. */
if (cmd != NULL)
free(cmd);
return -1;
} }
/* Format a command according to the Redis protocol. This function /* Format a command according to the Redis protocol. This function
...@@ -830,12 +873,14 @@ int redisFormatCommandArgv(char **target, int argc, const char **argv, const siz ...@@ -830,12 +873,14 @@ int redisFormatCommandArgv(char **target, int argc, const char **argv, const siz
totlen = 1+intlen(argc)+2; totlen = 1+intlen(argc)+2;
for (j = 0; j < argc; j++) { for (j = 0; j < argc; j++) {
len = argvlen ? argvlen[j] : strlen(argv[j]); len = argvlen ? argvlen[j] : strlen(argv[j]);
totlen += 1+intlen(len)+2+len+2; totlen += bulklen(len);
} }
/* Build the command at protocol level */ /* Build the command at protocol level */
cmd = malloc(totlen+1); cmd = malloc(totlen+1);
if (!cmd) redisOOM(); if (cmd == NULL)
return -1;
pos = sprintf(cmd,"*%d\r\n",argc); pos = sprintf(cmd,"*%d\r\n",argc);
for (j = 0; j < argc; j++) { for (j = 0; j < argc; j++) {
len = argvlen ? argvlen[j] : strlen(argv[j]); len = argvlen ? argvlen[j] : strlen(argv[j]);
...@@ -846,7 +891,8 @@ int redisFormatCommandArgv(char **target, int argc, const char **argv, const siz ...@@ -846,7 +891,8 @@ int redisFormatCommandArgv(char **target, int argc, const char **argv, const siz
cmd[pos++] = '\n'; cmd[pos++] = '\n';
} }
assert(pos == totlen); assert(pos == totlen);
cmd[totlen] = '\0'; cmd[pos] = '\0';
*target = cmd; *target = cmd;
return totlen; return totlen;
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment