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 void evalCommand(redisClient *c) {
 114     lua_State *lua = server.lua;
 115     char funcname[43];
 116
 117     /* We obtain the script SHA1, then check if this function is already
 118      * defined into the Lua state */
 119     funcname[0] = 'f';
 120     funcname[1] = '_';
 121     hashScript(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
 122     lua_getglobal(lua, funcname);
 123     if (lua_isnil(lua,1)) {
 124         /* Function not defined... let's define it. */
 125         sds funcdef = sdsempty();
 126
 127         lua_pop(lua,1); /* remove the nil from the stack */
 128         funcdef = sdscat(funcdef,"function ");
 129         funcdef = sdscatlen(funcdef,funcname,42);
 130         funcdef = sdscatlen(funcdef," ()\n",4);
 131         funcdef = sdscatlen(funcdef,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
 132         funcdef = sdscatlen(funcdef,"\nend\n",5);
 133         printf("Defining:\n%s\n",funcdef);
 134
 135         if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) {
 136             addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
 137                 lua_tostring(lua,-1));
 138             lua_pop(lua,1);
 139             sdsfree(funcdef);
 140             return;
 141         }
 142         sdsfree(funcdef);
 143         if (lua_pcall(lua,0,0,0)) {
 144             addReplyErrorFormat(c,"Error running script (new function): %s\n",
 145                 lua_tostring(lua,-1));
 146             lua_pop(lua,1);
 147             return;
 148         }
 149         lua_getglobal(lua, funcname);
 150     }
 151
 152     /* At this point whatever this script was never seen before or if it was
 153      * already defined, we can call it. We have zero arguments and expect
 154      * a single return value. */
 155     if (lua_pcall(lua,0,1,0)) {
 156         addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
 157             funcname, lua_tostring(lua,-1));
 158         lua_pop(lua,1);
 159         return;
 160     }
 161     luaReplyToRedisReply(c,lua);
 162 }