src/scripting.c

   1 #include "redis.h"
   2 #include "sha1.h"
   3
   4 #include <lua.h>
   5 #include <lauxlib.h>
   6 #include <lualib.h>
   7
   8 char *redisProtocolToLuaType_Int(lua_State *lua, char *reply);
   9 char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply);
  10 char *redisProtocolToLuaType_Status(lua_State *lua, char *reply);
  11
  12 /* Take a Redis reply in the Redis protocol format and convert it into a
  13  * Lua type. Thanks to this function, and the introduction of not connected
  14  * clients, it is trvial to implement the redis() lua function.
  15  *
  16  * Basically we take the arguments, execute the Redis command in the context
  17  * of a non connected client, then take the generated reply and convert it
  18  * into a suitable Lua type. With this trick the scripting feature does not
  19  * need the introduction of a full Redis internals API. Basically the script
  20  * is like a normal client that bypasses all the slow I/O paths.
  21  *
  22  * Note: in this function we do not do any sanity check as the reply is
  23  * generated by Redis directly. This allows use to go faster.
  24  * The reply string can be altered during the parsing as it is discared
  25  * after the conversion is completed.
  26  *
  27  * Errors are returned as a table with a single 'err' field set to the
  28  * error string.
  29  */
  30
  31 char *redisProtocolToLuaType(lua_State *lua, char* reply) {
  32     char *p = reply;
  33
  34     switch(*p) {
  35     case ':':
  36         p = redisProtocolToLuaType_Int(lua,reply);
  37         break;
  38     case '$':
  39         p = redisProtocolToLuaType_Bulk(lua,reply);
  40         break;
  41     case '+':
  42         p = redisProtocolToLuaType_Status(lua,reply);
  43         break;
  44     }
  45     return p;
  46 }
  47
  48 char *redisProtocolToLuaType_Int(lua_State *lua, char *reply) {
  49     char *p = strchr(reply+1,'\r');
  50     long long value;
  51
  52     string2ll(reply+1,p-reply-1,&value);
  53     lua_pushnumber(lua,(lua_Number)value);
  54     return p+2;
  55 }
  56
  57 char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply) {
  58     char *p = strchr(reply+1,'\r');
  59     long long bulklen;
  60
  61     string2ll(reply+1,p-reply-1,&bulklen);
  62     if (bulklen == 0) {
  63         lua_pushnil(lua);
  64         return p+2;
  65     } else {
  66         lua_pushlstring(lua,p+2,bulklen);
  67         return p+2+bulklen+2;
  68     }
  69 }
  70
  71 char *redisProtocolToLuaType_Status(lua_State *lua, char *reply) {
  72     char *p = strchr(reply+1,'\r');
  73
  74     lua_pushlstring(lua,reply+1,p-reply-1);
  75     return p+2;
  76 }
  77
  78 int luaRedisCommand(lua_State *lua) {
  79     int j, argc = lua_gettop(lua);
  80     struct redisCommand *cmd;
  81     robj **argv;
  82     redisClient *c = server.lua_client;
  83     sds reply;
  84
  85     /* Build the arguments vector */
  86     argv = zmalloc(sizeof(robj*)*argc);
  87     for (j = 0; j < argc; j++)
  88         argv[j] = createStringObject((char*)lua_tostring(lua,j+1),
  89                                      lua_strlen(lua,j+1));
  90
  91     /* Command lookup */
  92     cmd = lookupCommand(argv[0]->ptr);
  93     if (!cmd) {
  94         for (j = 0; j < argc; j++) decrRefCount(argv[j]);
  95         zfree(argv);
  96         lua_newtable(lua);
  97         lua_pushstring(lua,"err");
  98         lua_pushstring(lua,"Unknown Redis command called from Lua script");
  99         lua_settable(lua,-3);
 100         return 1;
 101     }
 102
 103     /* Run the command in the context of a fake client */
 104     c->argv = argv;
 105     c->argc = argc;
 106     cmd->proc(c);
 107
 108     /* Convert the result of the Redis command into a suitable Lua type.
 109      * The first thing we need is to create a single string from the client
 110      * output buffers. */
 111     reply = sdsempty();
 112     if (c->bufpos) {
 113         reply = sdscatlen(reply,c->buf,c->bufpos);
 114         c->bufpos = 0;
 115     }
 116     while(listLength(c->reply)) {
 117         robj *o = listNodeValue(listFirst(c->reply));
 118
 119         sdscatlen(reply,o->ptr,sdslen(o->ptr));
 120         listDelNode(c->reply,listFirst(c->reply));
 121     }
 122     redisProtocolToLuaType(lua,reply);
 123     sdsfree(reply);
 124
 125     /* Clean up. Command code may have changed argv/argc so we use the
 126      * argv/argc of the client instead of the local variables. */
 127     for (j = 0; j < c->argc; j++)
 128         decrRefCount(c->argv[j]);
 129     zfree(c->argv);
 130
 131     return 1;
 132 }
 133
 134 void scriptingInit(void) {
 135     lua_State *lua = lua_open();
 136     luaL_openlibs(lua);
 137
 138     /* Register the 'r' command */
 139     lua_pushcfunction(lua,luaRedisCommand);
 140     lua_setglobal(lua,"redis");
 141
 142     /* Create the (non connected) client that we use to execute Redis commands
 143      * inside the Lua interpreter */
 144     server.lua_client = createClient(-1);
 145     server.lua_client->flags |= REDIS_LUA_CLIENT;
 146
 147     server.lua = lua;
 148 }
 149
 150 /* Hash the scripit into a SHA1 digest. We use this as Lua function name.
 151  * Digest should point to a 41 bytes buffer: 40 for SHA1 converted into an
 152  * hexadecimal number, plus 1 byte for null term. */
 153 void hashScript(char *digest, char *script, size_t len) {
 154     SHA1_CTX ctx;
 155     unsigned char hash[20];
 156     char *cset = "0123456789abcdef";
 157     int j;
 158
 159     SHA1Init(&ctx);
 160     SHA1Update(&ctx,(unsigned char*)script,len);
 161     SHA1Final(hash,&ctx);
 162
 163     for (j = 0; j < 20; j++) {
 164         digest[j*2] = cset[((hash[j]&0xF0)>>4)];
 165         digest[j*2+1] = cset[(hash[j]&0xF)];
 166     }
 167     digest[40] = '\0';
 168 }
 169
 170 void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
 171     int t = lua_type(lua,1);
 172
 173     switch(t) {
 174     case LUA_TSTRING:
 175         addReplyBulkCBuffer(c,(char*)lua_tostring(lua,1),lua_strlen(lua,1));
 176         break;
 177     case LUA_TBOOLEAN:
 178         addReply(c,lua_toboolean(lua,1) ? shared.cone : shared.czero);
 179         break;
 180     case LUA_TNUMBER:
 181         addReplyLongLong(c,(long long)lua_tonumber(lua,1));
 182         break;
 183     case LUA_TTABLE:
 184         /* We need to check if it is an array or an error.
 185          * Error are returned as a single element table with 'err' field. */
 186         lua_pushstring(lua,"err");
 187         lua_gettable(lua,-2);
 188         t = lua_type(lua,-1);
 189         if (t == LUA_TSTRING) {
 190             addReplyError(c,(char*)lua_tostring(lua,-1));
 191             lua_pop(lua,1);
 192         } else {
 193             void *replylen = addDeferredMultiBulkLength(c);
 194             int j = 1, mbulklen = 0;
 195
 196             lua_pop(lua,1); /* Discard the 'err' field value we popped */
 197             while(1) {
 198                 lua_pushnumber(lua,j++);
 199                 lua_gettable(lua,-2);
 200                 t = lua_type(lua,-1);
 201                 if (t == LUA_TNIL) {
 202                     lua_pop(lua,1);
 203                     break;
 204                 } else if (t == LUA_TSTRING) {
 205                     size_t len;
 206                     char *s = (char*) lua_tolstring(lua,-1,&len);
 207
 208                     addReplyBulkCBuffer(c,s,len);
 209                     mbulklen++;
 210                 } else if (t == LUA_TNUMBER) {
 211                     addReplyLongLong(c,(long long)lua_tonumber(lua,-1));
 212                     mbulklen++;
 213                 }
 214                 lua_pop(lua,1);
 215             }
 216             setDeferredMultiBulkLength(c,replylen,mbulklen);
 217         }
 218         break;
 219     default:
 220         addReply(c,shared.nullbulk);
 221     }
 222     lua_pop(lua,1);
 223 }
 224
 225 /* Set an array of Redis String Objects as a Lua array (table) stored into a
 226  * global variable. */
 227 void luaSetGlobalArray(lua_State *lua, char *var, robj **elev, int elec) {
 228     int j;
 229
 230     lua_newtable(lua);
 231     for (j = 0; j < elec; j++) {
 232         lua_pushlstring(lua,(char*)elev[j]->ptr,sdslen(elev[j]->ptr));
 233         lua_rawseti(lua,-2,j+1);
 234     }
 235     lua_setglobal(lua,var);
 236 }
 237
 238 void evalCommand(redisClient *c) {
 239     lua_State *lua = server.lua;
 240     char funcname[43];
 241     long long numkeys;
 242
 243     /* Get the number of arguments that are keys */
 244     if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK)
 245         return;
 246     if (numkeys > (c->argc - 3)) {
 247         addReplyError(c,"Number of keys can't be greater than number of args");
 248         return;
 249     }
 250
 251     /* We obtain the script SHA1, then check if this function is already
 252      * defined into the Lua state */
 253     funcname[0] = 'f';
 254     funcname[1] = '_';
 255     hashScript(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
 256     lua_getglobal(lua, funcname);
 257     if (lua_isnil(lua,1)) {
 258         /* Function not defined... let's define it. */
 259         sds funcdef = sdsempty();
 260
 261         lua_pop(lua,1); /* remove the nil from the stack */
 262         funcdef = sdscat(funcdef,"function ");
 263         funcdef = sdscatlen(funcdef,funcname,42);
 264         funcdef = sdscatlen(funcdef," ()\n",4);
 265         funcdef = sdscatlen(funcdef,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
 266         funcdef = sdscatlen(funcdef,"\nend\n",5);
 267         printf("Defining:\n%s\n",funcdef);
 268
 269         if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) {
 270             addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
 271                 lua_tostring(lua,-1));
 272             lua_pop(lua,1);
 273             sdsfree(funcdef);
 274             return;
 275         }
 276         sdsfree(funcdef);
 277         if (lua_pcall(lua,0,0,0)) {
 278             addReplyErrorFormat(c,"Error running script (new function): %s\n",
 279                 lua_tostring(lua,-1));
 280             lua_pop(lua,1);
 281             return;
 282         }
 283         lua_getglobal(lua, funcname);
 284     }
 285
 286     /* Populate the argv and keys table accordingly to the arguments that
 287      * EVAL received. */
 288     luaSetGlobalArray(lua,"KEYS",c->argv+3,numkeys);
 289     luaSetGlobalArray(lua,"ARGV",c->argv+3+numkeys,c->argc-3-numkeys);
 290
 291     /* At this point whatever this script was never seen before or if it was
 292      * already defined, we can call it. We have zero arguments and expect
 293      * a single return value. */
 294     if (lua_pcall(lua,0,1,0)) {
 295         addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
 296             funcname, lua_tostring(lua,-1));
 297         lua_pop(lua,1);
 298         return;
 299     }
 300     luaReplyToRedisReply(c,lua);
 301 }