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(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->bufpos,c->buf);
  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_pushnumber(lua,1);
  47
  48     /* Clean up. Command code may have changed argv/argc so we use the
  49      * argv/argc of the client instead of the local variables. */
  50     for (j = 0; j < c->argc; j++)
  51         decrRefCount(c->argv[j]);
  52     zfree(c->argv);
  53
  54     return 1;
  55 }
  56
  57 void scriptingInit(void) {
  58     lua_State *lua = lua_open();
  59     luaL_openlibs(lua);
  60
  61     /* Register the 'r' command */
  62     lua_pushcfunction(lua,luaRedisCommand);
  63     lua_setglobal(lua,"r");
  64
  65     /* Create the (non connected) client that we use to execute Redis commands
  66      * inside the Lua interpreter */
  67     server.lua_client = createClient(-1);
  68
  69     server.lua = lua;
  70 }
  71
  72 /* Hash the scripit into a SHA1 digest. We use this as Lua function name.
  73  * Digest should point to a 41 bytes buffer: 40 for SHA1 converted into an
  74  * hexadecimal number, plus 1 byte for null term. */
  75 void hashScript(char *digest, char *script, size_t len) {
  76     SHA1_CTX ctx;
  77     unsigned char hash[20];
  78     char *cset = "0123456789abcdef";
  79     int j;
  80
  81     SHA1Init(&ctx);
  82     SHA1Update(&ctx,(unsigned char*)script,len);
  83     SHA1Final(hash,&ctx);
  84
  85     for (j = 0; j < 20; j++) {
  86         digest[j*2] = cset[((hash[j]&0xF0)>>4)];
  87         digest[j*2+1] = cset[(hash[j]&0xF)];
  88     }
  89     digest[40] = '\0';
  90 }
  91
  92 void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
  93     int t = lua_type(lua,1);
  94
  95     switch(t) {
  96     case LUA_TSTRING:
  97         addReplyBulkCBuffer(c,(char*)lua_tostring(lua,1),lua_strlen(lua,1));
  98         break;
  99     case LUA_TBOOLEAN:
 100         addReply(c,lua_toboolean(lua,1) ? shared.cone : shared.czero);
 101         break;
 102     case LUA_TNUMBER:
 103         addReplyLongLong(c,(long long)lua_tonumber(lua,1));
 104         break;
 105     default:
 106         addReply(c,shared.nullbulk);
 107     }
 108     lua_pop(lua,1);
 109 }
 110
 111 void evalCommand(redisClient *c) {
 112     lua_State *lua = server.lua;
 113     char funcname[43];
 114
 115     /* We obtain the script SHA1, then check if this function is already
 116      * defined into the Lua state */
 117     funcname[0] = 'f';
 118     funcname[1] = '_';
 119     hashScript(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
 120     lua_getglobal(lua, funcname);
 121     if (lua_isnil(lua,1)) {
 122         /* Function not defined... let's define it. */
 123         sds funcdef = sdsempty();
 124
 125         lua_pop(lua,1); /* remove the nil from the stack */
 126         funcdef = sdscat(funcdef,"function ");
 127         funcdef = sdscatlen(funcdef,funcname,42);
 128         funcdef = sdscatlen(funcdef," ()\n",4);
 129         funcdef = sdscatlen(funcdef,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
 130         funcdef = sdscatlen(funcdef,"\nend\n",5);
 131         printf("Defining:\n%s\n",funcdef);
 132
 133         if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) {
 134             addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
 135                 lua_tostring(lua,-1));
 136             lua_pop(lua,1);
 137             return;
 138         }
 139         if (lua_pcall(lua,0,0,0)) {
 140             addReplyErrorFormat(c,"Error running script (new function): %s\n",
 141                 lua_tostring(lua,-1));
 142             lua_pop(lua,1);
 143             return;
 144         }
 145         lua_getglobal(lua, funcname);
 146     }
 147
 148     /* At this point whatever this script was never seen before or if it was
 149      * already defined, we can call it. We have zero arguments and expect
 150      * a single return value. */
 151     if (lua_pcall(lua,0,1,0)) {
 152         addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
 153             funcname, lua_tostring(lua,-1));
 154         lua_pop(lua,1);
 155         return;
 156     }
 157     luaReplyToRedisReply(c,lua);
 158 }