diff --git a/src/acl.c b/src/acl.c index e781c9211..e0fd3f728 100644 --- a/src/acl.c +++ b/src/acl.c @@ -1115,8 +1115,9 @@ int ACLCheckCommandPerm(client *c, int *keyidxptr) { if (!(c->user->flags & USER_FLAG_ALLKEYS) && (c->cmd->getkeys_proc || c->cmd->firstkey)) { - int numkeys; - int *keyidx = getKeysFromCommand(c->cmd,c->argv,c->argc,&numkeys); + getKeysResult result = GETKEYS_RESULT_INIT; + int numkeys = getKeysFromCommand(c->cmd,c->argv,c->argc,&result); + int *keyidx = result.keys; for (int j = 0; j < numkeys; j++) { listIter li; listNode *ln; @@ -1137,11 +1138,11 @@ int ACLCheckCommandPerm(client *c, int *keyidxptr) { } if (!match) { if (keyidxptr) *keyidxptr = keyidx[j]; - getKeysFreeResult(keyidx); + getKeysFreeResult(&result); return ACL_DENIED_KEY; } } - getKeysFreeResult(keyidx); + getKeysFreeResult(&result); } /* If we survived all the above checks, the user can execute the diff --git a/src/cluster.c b/src/cluster.c index 43bea155b..7d690e863 100644 --- a/src/cluster.c +++ b/src/cluster.c @@ -5640,7 +5640,10 @@ clusterNode *getNodeByQuery(client *c, struct redisCommand *cmd, robj **argv, in margc = ms->commands[i].argc; margv = ms->commands[i].argv; - keyindex = getKeysFromCommand(mcmd,margv,margc,&numkeys); + getKeysResult result = GETKEYS_RESULT_INIT; + numkeys = getKeysFromCommand(mcmd,margv,margc,&result); + keyindex = result.keys; + for (j = 0; j < numkeys; j++) { robj *thiskey = margv[keyindex[j]]; int thisslot = keyHashSlot((char*)thiskey->ptr, @@ -5658,7 +5661,7 @@ clusterNode *getNodeByQuery(client *c, struct redisCommand *cmd, robj **argv, in * not trapped earlier in processCommand(). Report the same * error to the client. */ if (n == NULL) { - getKeysFreeResult(keyindex); + getKeysFreeResult(&result); if (error_code) *error_code = CLUSTER_REDIR_DOWN_UNBOUND; return NULL; @@ -5682,7 +5685,7 @@ clusterNode *getNodeByQuery(client *c, struct redisCommand *cmd, robj **argv, in if (!equalStringObjects(firstkey,thiskey)) { if (slot != thisslot) { /* Error: multiple keys from different slots. */ - getKeysFreeResult(keyindex); + getKeysFreeResult(&result); if (error_code) *error_code = CLUSTER_REDIR_CROSS_SLOT; return NULL; @@ -5701,7 +5704,7 @@ clusterNode *getNodeByQuery(client *c, struct redisCommand *cmd, robj **argv, in missing_keys++; } } - getKeysFreeResult(keyindex); + getKeysFreeResult(&result); } /* No key at all in command? then we can serve the request diff --git a/src/db.c b/src/db.c index 7ed746f9a..eb87cebc5 100644 --- a/src/db.c +++ b/src/db.c @@ -1322,27 +1322,54 @@ int expireIfNeeded(redisDb *db, robj *key) { /* ----------------------------------------------------------------------------- * API to get key arguments from commands * ---------------------------------------------------------------------------*/ -#define MAX_KEYS_BUFFER 256 -static int getKeysTempBuffer[MAX_KEYS_BUFFER]; + +/* Prepare the getKeysResult struct to hold numkeys, either by using the + * pre-allocated keysbuf or by allocating a new array on the heap. + * + * This function must be called at least once before starting to populate + * the result, and can be called repeatedly to enlarge the result array. + */ +int *getKeysPrepareResult(getKeysResult *result, int numkeys) { + /* GETKEYS_RESULT_INIT initializes keys to NULL, point it to the pre-allocated stack + * buffer here. */ + if (!result->keys) { + serverAssert(!result->numkeys); + result->keys = result->keysbuf; + } + + /* Resize if necessary */ + if (numkeys > result->size) { + if (result->keys != result->keysbuf) { + /* We're not using a static buffer, just (re)alloc */ + result->keys = zrealloc(result->keys, numkeys * sizeof(int)); + } else { + /* We are using a static buffer, copy its contents */ + result->keys = zmalloc(numkeys * sizeof(int)); + if (result->numkeys) + memcpy(result->keys, result->keysbuf, result->numkeys * sizeof(int)); + } + result->size = numkeys; + } + + return result->keys; +} /* The base case is to use the keys position as given in the command table * (firstkey, lastkey, step). */ -int *getKeysUsingCommandTable(struct redisCommand *cmd,robj **argv, int argc, int *numkeys) { +int getKeysUsingCommandTable(struct redisCommand *cmd,robj **argv, int argc, getKeysResult *result) { int j, i = 0, last, *keys; UNUSED(argv); if (cmd->firstkey == 0) { - *numkeys = 0; - return NULL; + result->numkeys = 0; + return 0; } last = cmd->lastkey; if (last < 0) last = argc+last; int count = ((last - cmd->firstkey)+1); - keys = getKeysTempBuffer; - if (count > MAX_KEYS_BUFFER) - keys = zmalloc(sizeof(int)*count); + keys = getKeysPrepareResult(result, count); for (j = cmd->firstkey; j <= last; j += cmd->keystep) { if (j >= argc) { @@ -1353,17 +1380,17 @@ int *getKeysUsingCommandTable(struct redisCommand *cmd,robj **argv, int argc, in * return no keys and expect the command implementation to report * an arity or syntax error. */ if (cmd->flags & CMD_MODULE || cmd->arity < 0) { - getKeysFreeResult(keys); - *numkeys = 0; - return NULL; + getKeysFreeResult(result); + result->numkeys = 0; + return 0; } else { serverPanic("Redis built-in command declared keys positions not matching the arity requirements."); } } keys[i++] = j; } - *numkeys = i; - return keys; + result->numkeys = i; + return i; } /* Return all the arguments that are keys in the command passed via argc / argv. @@ -1377,26 +1404,26 @@ int *getKeysUsingCommandTable(struct redisCommand *cmd,robj **argv, int argc, in * * This function uses the command table if a command-specific helper function * is not required, otherwise it calls the command-specific function. */ -int *getKeysFromCommand(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { +int getKeysFromCommand(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) { if (cmd->flags & CMD_MODULE_GETKEYS) { - return moduleGetCommandKeysViaAPI(cmd,argv,argc,numkeys); + return moduleGetCommandKeysViaAPI(cmd,argv,argc,result); } else if (!(cmd->flags & CMD_MODULE) && cmd->getkeys_proc) { - return cmd->getkeys_proc(cmd,argv,argc,numkeys); + return cmd->getkeys_proc(cmd,argv,argc,result); } else { - return getKeysUsingCommandTable(cmd,argv,argc,numkeys); + return getKeysUsingCommandTable(cmd,argv,argc,result); } } /* Free the result of getKeysFromCommand. */ -void getKeysFreeResult(int *result) { - if (result != getKeysTempBuffer) - zfree(result); +void getKeysFreeResult(getKeysResult *result) { + if (result && result->keys != result->keysbuf) + zfree(result->keys); } /* Helper function to extract keys from following commands: * ZUNIONSTORE ... * ZINTERSTORE ... */ -int *zunionInterGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *numkeys) { +int zunionInterGetKeys(struct redisCommand *cmd, robj **argv, int argc, getKeysResult *result) { int i, num, *keys; UNUSED(cmd); @@ -1404,30 +1431,30 @@ int *zunionInterGetKeys(struct redisCommand *cmd, robj **argv, int argc, int *nu /* Sanity check. Don't return any key if the command is going to * reply with syntax error. */ if (num < 1 || num > (argc-3)) { - *numkeys = 0; - return NULL; + result->numkeys = 0; + return 0; } /* Keys in z{union,inter}store come from two places: * argv[1] = storage key, * argv[3...n] = keys to intersect */ - keys = getKeysTempBuffer; - if (num+1>MAX_KEYS_BUFFER) - keys = zmalloc(sizeof(int)*(num+1)); + /* Total keys = {union,inter} keys + storage key */ + keys = getKeysPrepareResult(result, num+1); + result->numkeys = num+1; /* Add all key positions for argv[3...n] to keys[] */ for (i = 0; i < num; i++) keys[i] = 3+i; /* Finally add the argv[1] key position (the storage key target). */ keys[num] = 1; - *numkeys = num+1; /* Total keys = {union,inter} keys + storage key */ - return keys; + + return result->numkeys; } /* Helper function to extract keys from the following commands: * EVAL