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 int luaRedisCommand(lua_State *lua) {
   9     int j, argc = lua_gettop(lua);
  10     struct redisCommand *cmd;
  11     robj **argv;
  12     redisClient *c = server.lua_client;
  13     sds reply;
  14
  15     argv = zmalloc(sizeof(robj*)*argc);
  16     for (j = 0; j < argc; j++)
  17         argv[j] = createStringObject((char*)lua_tostring(lua,j+1),lua_strlen(lua,j+1));
  18
  19     /* Command lookup */
  20     cmd = lookupCommand(argv[0]->ptr);
  21     if (!cmd) {
  22         zfree(argv);
  23         lua_pushnil(lua);
  24         lua_pushstring(lua,"Unknown Redis command called from Lua script");
  25         return 2;
  26     }
  27     /* Run the command in the context of a fake client */
  28     c->argv = argv;
  29     c->argc = argc;
  30     cmd->proc(c);
  31
  32     /* Convert the result of the Redis command into a suitable Lua type.
  33      * The first thing we need is to create a single string from the client
  34      * output buffers. */
  35     reply = sdsempty();
  36     if (c->bufpos) {
  37         reply = sdscatlen(reply,c->buf,c->bufpos);
  38         c->bufpos = 0;
  39     }
  40     while(listLength(c->reply)) {
  41         robj *o = listNodeValue(listFirst(c->reply));
  42
  43         sdscatlen(reply,o->ptr,sdslen(o->ptr));
  44         listDelNode(c->reply,listFirst(c->reply));
  45     }
  46     lua_pushlstring(lua,reply,sdslen(reply));
  47     sdsfree(reply);
  48
  49     /* Clean up. Command code may have changed argv/argc so we use the
  50      * argv/argc of the client instead of the local variables. */
  51     for (j = 0; j < c->argc; j++)
  52         decrRefCount(c->argv[j]);
  53     zfree(c->argv);
  54
  55     return 1;
  56 }
  57
  58 void scriptingInit(void) {
  59     lua_State *lua = lua_open();
  60     luaL_openlibs(lua);
  61
  62     /* Register the 'r' command */
  63     lua_pushcfunction(lua,luaRedisCommand);
  64     lua_setglobal(lua,"redis");
  65
  66     /* Create the (non connected) client that we use to execute Redis commands
  67      * inside the Lua interpreter */
  68     server.lua_client = createClient(-1);
  69     server.lua_client->flags |= REDIS_LUA_CLIENT;
  70
  71     server.lua = lua;
  72 }
  73
  74 /* Hash the scripit into a SHA1 digest. We use this as Lua function name.
  75  * Digest should point to a 41 bytes buffer: 40 for SHA1 converted into an
  76  * hexadecimal number, plus 1 byte for null term. */
  77 void hashScript(char *digest, char *script, size_t len) {
  78     SHA1_CTX ctx;
  79     unsigned char hash[20];
  80     char *cset = "0123456789abcdef";
  81     int j;
  82
  83     SHA1Init(&ctx);
  84     SHA1Update(&ctx,(unsigned char*)script,len);
  85     SHA1Final(hash,&ctx);
  86
  87     for (j = 0; j < 20; j++) {
  88         digest[j*2] = cset[((hash[j]&0xF0)>>4)];
  89         digest[j*2+1] = cset[(hash[j]&0xF)];
  90     }
  91     digest[40] = '\0';
  92 }
  93
  94 void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
  95     int t = lua_type(lua,1);
  96
  97     switch(t) {
  98     case LUA_TSTRING:
  99         addReplyBulkCBuffer(c,(char*)lua_tostring(lua,1),lua_strlen(lua,1));
 100         break;
 101     case LUA_TBOOLEAN:
 102         addReply(c,lua_toboolean(lua,1) ? shared.cone : shared.czero);
 103         break;
 104     case LUA_TNUMBER:
 105         addReplyLongLong(c,(long long)lua_tonumber(lua,1));
 106         break;
 107     default:
 108         addReply(c,shared.nullbulk);
 109     }
 110     lua_pop(lua,1);
 111 }
 112
 113 /* Set an array of Redis String Objects as a Lua array (table) stored into a
 114  * global variable. */
 115 void luaSetGlobalArray(lua_State *lua, char *var, robj **elev, int elec) {
 116     int j;
 117
 118     lua_newtable(lua);
 119     for (j = 0; j < elec; j++) {
 120         lua_pushlstring(lua,(char*)elev[j]->ptr,sdslen(elev[j]->ptr));
 121         lua_rawseti(lua,-2,j+1);
 122     }
 123     lua_setglobal(lua,var);
 124 }
 125
 126 void evalCommand(redisClient *c) {
 127     lua_State *lua = server.lua;
 128     char funcname[43];
 129     long long numkeys;
 130
 131     /* Get the number of arguments that are keys */
 132     if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK)
 133         return;
 134     if (numkeys > (c->argc - 3)) {
 135         addReplyError(c,"Number of keys can't be greater than number of args");
 136         return;
 137     }
 138
 139     /* We obtain the script SHA1, then check if this function is already
 140      * defined into the Lua state */
 141     funcname[0] = 'f';
 142     funcname[1] = '_';
 143     hashScript(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
 144     lua_getglobal(lua, funcname);
 145     if (lua_isnil(lua,1)) {
 146         /* Function not defined... let's define it. */
 147         sds funcdef = sdsempty();
 148
 149         lua_pop(lua,1); /* remove the nil from the stack */
 150         funcdef = sdscat(funcdef,"function ");
 151         funcdef = sdscatlen(funcdef,funcname,42);
 152         funcdef = sdscatlen(funcdef," ()\n",4);
 153         funcdef = sdscatlen(funcdef,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
 154         funcdef = sdscatlen(funcdef,"\nend\n",5);
 155         printf("Defining:\n%s\n",funcdef);
 156
 157         if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) {
 158             addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
 159                 lua_tostring(lua,-1));
 160             lua_pop(lua,1);
 161             sdsfree(funcdef);
 162             return;
 163         }
 164         sdsfree(funcdef);
 165         if (lua_pcall(lua,0,0,0)) {
 166             addReplyErrorFormat(c,"Error running script (new function): %s\n",
 167                 lua_tostring(lua,-1));
 168             lua_pop(lua,1);
 169             return;
 170         }
 171         lua_getglobal(lua, funcname);
 172     }
 173
 174     /* Populate the argv and keys table accordingly to the arguments that
 175      * EVAL received. */
 176     luaSetGlobalArray(lua,"KEYS",c->argv+3,numkeys);
 177     luaSetGlobalArray(lua,"ARGV",c->argv+3+numkeys,c->argc-3-numkeys);
 178
 179     /* At this point whatever this script was never seen before or if it was
 180      * already defined, we can call it. We have zero arguments and expect
 181      * a single return value. */
 182     if (lua_pcall(lua,0,1,0)) {
 183         addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
 184             funcname, lua_tostring(lua,-1));
 185         lua_pop(lua,1);
 186         return;
 187     }
 188     luaReplyToRedisReply(c,lua);
 189 }