xref: /redis-3.2.3/src/scripting.c (revision 22c9c407)
1 #include "redis.h"
2 #include "sha1.h"
3 #include "rand.h"
4 
5 #include <lua.h>
6 #include <lauxlib.h>
7 #include <lualib.h>
8 #include <ctype.h>
9 #include <math.h>
10 
11 char *redisProtocolToLuaType_Int(lua_State *lua, char *reply);
12 char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply);
13 char *redisProtocolToLuaType_Status(lua_State *lua, char *reply);
14 char *redisProtocolToLuaType_Error(lua_State *lua, char *reply);
15 char *redisProtocolToLuaType_MultiBulk(lua_State *lua, char *reply);
16 int redis_math_random (lua_State *L);
17 int redis_math_randomseed (lua_State *L);
18 void sha1hex(char *digest, char *script, size_t len);
19 
20 /* Take a Redis reply in the Redis protocol format and convert it into a
21  * Lua type. Thanks to this function, and the introduction of not connected
22  * clients, it is trvial to implement the redis() lua function.
23  *
24  * Basically we take the arguments, execute the Redis command in the context
25  * of a non connected client, then take the generated reply and convert it
26  * into a suitable Lua type. With this trick the scripting feature does not
27  * need the introduction of a full Redis internals API. Basically the script
28  * is like a normal client that bypasses all the slow I/O paths.
29  *
30  * Note: in this function we do not do any sanity check as the reply is
31  * generated by Redis directly. This allows us to go faster.
32  * The reply string can be altered during the parsing as it is discared
33  * after the conversion is completed.
34  *
35  * Errors are returned as a table with a single 'err' field set to the
36  * error string.
37  */
38 
39 char *redisProtocolToLuaType(lua_State *lua, char* reply) {
40     char *p = reply;
41 
42     switch(*p) {
43     case ':':
44         p = redisProtocolToLuaType_Int(lua,reply);
45         break;
46     case '$':
47         p = redisProtocolToLuaType_Bulk(lua,reply);
48         break;
49     case '+':
50         p = redisProtocolToLuaType_Status(lua,reply);
51         break;
52     case '-':
53         p = redisProtocolToLuaType_Error(lua,reply);
54         break;
55     case '*':
56         p = redisProtocolToLuaType_MultiBulk(lua,reply);
57         break;
58     }
59     return p;
60 }
61 
62 char *redisProtocolToLuaType_Int(lua_State *lua, char *reply) {
63     char *p = strchr(reply+1,'\r');
64     long long value;
65 
66     string2ll(reply+1,p-reply-1,&value);
67     lua_pushnumber(lua,(lua_Number)value);
68     return p+2;
69 }
70 
71 char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply) {
72     char *p = strchr(reply+1,'\r');
73     long long bulklen;
74 
75     string2ll(reply+1,p-reply-1,&bulklen);
76     if (bulklen == -1) {
77         lua_pushboolean(lua,0);
78         return p+2;
79     } else {
80         lua_pushlstring(lua,p+2,bulklen);
81         return p+2+bulklen+2;
82     }
83 }
84 
85 char *redisProtocolToLuaType_Status(lua_State *lua, char *reply) {
86     char *p = strchr(reply+1,'\r');
87 
88     lua_newtable(lua);
89     lua_pushstring(lua,"ok");
90     lua_pushlstring(lua,reply+1,p-reply-1);
91     lua_settable(lua,-3);
92     return p+2;
93 }
94 
95 char *redisProtocolToLuaType_Error(lua_State *lua, char *reply) {
96     char *p = strchr(reply+1,'\r');
97 
98     lua_newtable(lua);
99     lua_pushstring(lua,"err");
100     lua_pushlstring(lua,reply+1,p-reply-1);
101     lua_settable(lua,-3);
102     return p+2;
103 }
104 
105 char *redisProtocolToLuaType_MultiBulk(lua_State *lua, char *reply) {
106     char *p = strchr(reply+1,'\r');
107     long long mbulklen;
108     int j = 0;
109 
110     string2ll(reply+1,p-reply-1,&mbulklen);
111     p += 2;
112     if (mbulklen == -1) {
113         lua_pushboolean(lua,0);
114         return p;
115     }
116     lua_newtable(lua);
117     for (j = 0; j < mbulklen; j++) {
118         lua_pushnumber(lua,j+1);
119         p = redisProtocolToLuaType(lua,p);
120         lua_settable(lua,-3);
121     }
122     return p;
123 }
124 
125 void luaPushError(lua_State *lua, char *error) {
126     lua_newtable(lua);
127     lua_pushstring(lua,"err");
128     lua_pushstring(lua, error);
129     lua_settable(lua,-3);
130 }
131 
132 /* Sort the array currently in the stack. We do this to make the output
133  * of commands like KEYS or SMEMBERS something deterministic when called
134  * from Lua (to play well with AOf/replication).
135  *
136  * The array is sorted using table.sort itself, and assuming all the
137  * list elements are strings. */
138 void luaSortArray(lua_State *lua) {
139     /* Initial Stack: array */
140     lua_getglobal(lua,"table");
141     lua_pushstring(lua,"sort");
142     lua_gettable(lua,-2);       /* Stack: array, table, table.sort */
143     lua_pushvalue(lua,-3);      /* Stack: array, table, table.sort, array */
144     if (lua_pcall(lua,1,0,0)) {
145         /* Stack: array, table, error */
146 
147         /* We are not interested in the error, we assume that the problem is
148          * that there are 'false' elements inside the array, so we try
149          * again with a slower function but able to handle this case, that
150          * is: table.sort(table, __redis__compare_helper) */
151         lua_pop(lua,1);             /* Stack: array, table */
152         lua_pushstring(lua,"sort"); /* Stack: array, table, sort */
153         lua_gettable(lua,-2);       /* Stack: array, table, table.sort */
154         lua_pushvalue(lua,-3);      /* Stack: array, table, table.sort, array */
155         lua_getglobal(lua,"__redis__compare_helper");
156         /* Stack: array, table, table.sort, array, __redis__compare_helper */
157         lua_call(lua,2,0);
158     }
159     /* Stack: array (sorted), table */
160     lua_pop(lua,1);             /* Stack: array (sorted) */
161 }
162 
163 int luaRedisGenericCommand(lua_State *lua, int raise_error) {
164     int j, argc = lua_gettop(lua);
165     struct redisCommand *cmd;
166     robj **argv;
167     redisClient *c = server.lua_client;
168     sds reply;
169 
170     /* Build the arguments vector */
171     argv = zmalloc(sizeof(robj*)*argc);
172     for (j = 0; j < argc; j++) {
173         if (!lua_isstring(lua,j+1)) break;
174         argv[j] = createStringObject((char*)lua_tostring(lua,j+1),
175                                      lua_strlen(lua,j+1));
176     }
177 
178     /* Check if one of the arguments passed by the Lua script
179      * is not a string or an integer (lua_isstring() return true for
180      * integers as well). */
181     if (j != argc) {
182         j--;
183         while (j >= 0) {
184             decrRefCount(argv[j]);
185             j--;
186         }
187         zfree(argv);
188         luaPushError(lua,
189             "Lua redis() command arguments must be strings or integers");
190         return 1;
191     }
192 
193     /* Setup our fake client for command execution */
194     c->argv = argv;
195     c->argc = argc;
196 
197     /* Command lookup */
198     cmd = lookupCommand(argv[0]->ptr);
199     if (!cmd || ((cmd->arity > 0 && cmd->arity != argc) ||
200                    (argc < -cmd->arity)))
201     {
202         if (cmd)
203             luaPushError(lua,
204                 "Wrong number of args calling Redis command From Lua script");
205         else
206             luaPushError(lua,"Unknown Redis command called from Lua script");
207         goto cleanup;
208     }
209 
210     /* There are commands that are not allowed inside scripts. */
211     if (cmd->flags & REDIS_CMD_NOSCRIPT) {
212         luaPushError(lua, "This Redis command is not allowed from scripts");
213         goto cleanup;
214     }
215 
216     /* Write commands are forbidden against read-only slaves, or if a
217      * command marked as non-deterministic was already called in the context
218      * of this script. */
219     if (cmd->flags & REDIS_CMD_WRITE) {
220         if (server.lua_random_dirty) {
221             luaPushError(lua,
222                 "Write commands not allowed after non deterministic commands");
223             goto cleanup;
224         } else if (server.masterhost && server.repl_slave_ro &&
225                    !(server.lua_caller->flags & REDIS_MASTER))
226         {
227             luaPushError(lua, shared.roslaveerr->ptr);
228             goto cleanup;
229         } else if (server.stop_writes_on_bgsave_err &&
230                    server.saveparamslen > 0 &&
231                    server.lastbgsave_status == REDIS_ERR)
232         {
233             luaPushError(lua, shared.bgsaveerr->ptr);
234             goto cleanup;
235         }
236     }
237 
238     /* If we reached the memory limit configured via maxmemory, commands that
239      * could enlarge the memory usage are not allowed, but only if this is the
240      * first write in the context of this script, otherwise we can't stop
241      * in the middle. */
242     if (server.maxmemory && server.lua_write_dirty == 0 &&
243         (cmd->flags & REDIS_CMD_DENYOOM))
244     {
245         if (freeMemoryIfNeeded() == REDIS_ERR) {
246             luaPushError(lua, shared.oomerr->ptr);
247             goto cleanup;
248         }
249     }
250 
251     if (cmd->flags & REDIS_CMD_RANDOM) server.lua_random_dirty = 1;
252     if (cmd->flags & REDIS_CMD_WRITE) server.lua_write_dirty = 1;
253 
254     /* Run the command */
255     c->cmd = cmd;
256     call(c,REDIS_CALL_SLOWLOG | REDIS_CALL_STATS);
257 
258     /* Convert the result of the Redis command into a suitable Lua type.
259      * The first thing we need is to create a single string from the client
260      * output buffers. */
261     reply = sdsempty();
262     if (c->bufpos) {
263         reply = sdscatlen(reply,c->buf,c->bufpos);
264         c->bufpos = 0;
265     }
266     while(listLength(c->reply)) {
267         robj *o = listNodeValue(listFirst(c->reply));
268 
269         reply = sdscatlen(reply,o->ptr,sdslen(o->ptr));
270         listDelNode(c->reply,listFirst(c->reply));
271     }
272     if (raise_error && reply[0] != '-') raise_error = 0;
273     redisProtocolToLuaType(lua,reply);
274     /* Sort the output array if needed, assuming it is a non-null multi bulk
275      * reply as expected. */
276     if ((cmd->flags & REDIS_CMD_SORT_FOR_SCRIPT) &&
277         (reply[0] == '*' && reply[1] != '-')) {
278         /* Skip this step if command is SORT but output was already sorted */
279         if (cmd->proc != sortCommand || server.sort_dontsort)
280             luaSortArray(lua);
281     }
282     sdsfree(reply);
283 
284 cleanup:
285     /* Clean up. Command code may have changed argv/argc so we use the
286      * argv/argc of the client instead of the local variables. */
287     for (j = 0; j < c->argc; j++)
288         decrRefCount(c->argv[j]);
289     zfree(c->argv);
290 
291     if (raise_error) {
292         /* If we are here we should have an error in the stack, in the
293          * form of a table with an "err" field. Extract the string to
294          * return the plain error. */
295         lua_pushstring(lua,"err");
296         lua_gettable(lua,-2);
297         return lua_error(lua);
298     }
299     return 1;
300 }
301 
302 int luaRedisCallCommand(lua_State *lua) {
303     return luaRedisGenericCommand(lua,1);
304 }
305 
306 int luaRedisPCallCommand(lua_State *lua) {
307     return luaRedisGenericCommand(lua,0);
308 }
309 
310 /* This adds redis.sha1hex(string) to Lua scripts using the same hashing
311  * function used for sha1ing lua scripts. */
312 int luaRedisSha1hexCommand(lua_State *lua) {
313     int argc = lua_gettop(lua);
314     char digest[41];
315     size_t len;
316     char *s;
317 
318     if (argc != 1) {
319         luaPushError(lua, "wrong number of arguments");
320         return 1;
321     }
322 
323     s = (char*)lua_tolstring(lua,1,&len);
324     sha1hex(digest,s,len);
325     lua_pushstring(lua,digest);
326     return 1;
327 }
328 
329 int luaLogCommand(lua_State *lua) {
330     int j, argc = lua_gettop(lua);
331     int level;
332     sds log;
333 
334     if (argc < 2) {
335         luaPushError(lua, "redis.log() requires two arguments or more.");
336         return 1;
337     } else if (!lua_isnumber(lua,-argc)) {
338         luaPushError(lua, "First argument must be a number (log level).");
339         return 1;
340     }
341     level = lua_tonumber(lua,-argc);
342     if (level < REDIS_DEBUG || level > REDIS_WARNING) {
343         luaPushError(lua, "Invalid debug level.");
344         return 1;
345     }
346 
347     /* Glue together all the arguments */
348     log = sdsempty();
349     for (j = 1; j < argc; j++) {
350         size_t len;
351         char *s;
352 
353         s = (char*)lua_tolstring(lua,(-argc)+j,&len);
354         if (s) {
355             if (j != 1) log = sdscatlen(log," ",1);
356             log = sdscatlen(log,s,len);
357         }
358     }
359     redisLogRaw(level,log);
360     sdsfree(log);
361     return 0;
362 }
363 
364 void luaMaskCountHook(lua_State *lua, lua_Debug *ar) {
365     long long elapsed;
366     REDIS_NOTUSED(ar);
367     REDIS_NOTUSED(lua);
368 
369     elapsed = (ustime()/1000) - server.lua_time_start;
370     if (elapsed >= server.lua_time_limit && server.lua_timedout == 0) {
371         redisLog(REDIS_WARNING,"Lua slow script detected: still in execution after %lld milliseconds. You can try killing the script using the SCRIPT KILL command.",elapsed);
372         server.lua_timedout = 1;
373         /* Once the script timeouts we reenter the event loop to permit others
374          * to call SCRIPT KILL or SHUTDOWN NOSAVE if needed. For this reason
375          * we need to mask the client executing the script from the event loop.
376          * If we don't do that the client may disconnect and could no longer be
377          * here when the EVAL command will return. */
378          aeDeleteFileEvent(server.el, server.lua_caller->fd, AE_READABLE);
379     }
380     if (server.lua_timedout)
381         aeProcessEvents(server.el, AE_FILE_EVENTS|AE_DONT_WAIT);
382     if (server.lua_kill) {
383         redisLog(REDIS_WARNING,"Lua script killed by user with SCRIPT KILL.");
384         lua_pushstring(lua,"Script killed by user with SCRIPT KILL...");
385         lua_error(lua);
386     }
387 }
388 
389 void luaLoadLib(lua_State *lua, const char *libname, lua_CFunction luafunc) {
390   lua_pushcfunction(lua, luafunc);
391   lua_pushstring(lua, libname);
392   lua_call(lua, 1, 0);
393 }
394 
395 LUALIB_API int (luaopen_cjson) (lua_State *L);
396 LUALIB_API int (luaopen_struct) (lua_State *L);
397 LUALIB_API int (luaopen_cmsgpack) (lua_State *L);
398 
399 void luaLoadLibraries(lua_State *lua) {
400     luaLoadLib(lua, "", luaopen_base);
401     luaLoadLib(lua, LUA_TABLIBNAME, luaopen_table);
402     luaLoadLib(lua, LUA_STRLIBNAME, luaopen_string);
403     luaLoadLib(lua, LUA_MATHLIBNAME, luaopen_math);
404     luaLoadLib(lua, LUA_DBLIBNAME, luaopen_debug);
405     luaLoadLib(lua, "cjson", luaopen_cjson);
406     luaLoadLib(lua, "struct", luaopen_struct);
407     luaLoadLib(lua, "cmsgpack", luaopen_cmsgpack);
408 
409 #if 0 /* Stuff that we don't load currently, for sandboxing concerns. */
410     luaLoadLib(lua, LUA_LOADLIBNAME, luaopen_package);
411     luaLoadLib(lua, LUA_OSLIBNAME, luaopen_os);
412 #endif
413 }
414 
415 /* Initialize the scripting environment.
416  * It is possible to call this function to reset the scripting environment
417  * assuming that we call scriptingRelease() before.
418  * See scriptingReset() for more information. */
419 void scriptingInit(void) {
420     lua_State *lua = lua_open();
421     luaLoadLibraries(lua);
422 
423     /* Initialize a dictionary we use to map SHAs to scripts.
424      * This is useful for replication, as we need to replicate EVALSHA
425      * as EVAL, so we need to remember the associated script. */
426     server.lua_scripts = dictCreate(&dbDictType,NULL);
427 
428     /* Register the redis commands table and fields */
429     lua_newtable(lua);
430 
431     /* redis.call */
432     lua_pushstring(lua,"call");
433     lua_pushcfunction(lua,luaRedisCallCommand);
434     lua_settable(lua,-3);
435 
436     /* redis.pcall */
437     lua_pushstring(lua,"pcall");
438     lua_pushcfunction(lua,luaRedisPCallCommand);
439     lua_settable(lua,-3);
440 
441     /* redis.log and log levels. */
442     lua_pushstring(lua,"log");
443     lua_pushcfunction(lua,luaLogCommand);
444     lua_settable(lua,-3);
445 
446     lua_pushstring(lua,"LOG_DEBUG");
447     lua_pushnumber(lua,REDIS_DEBUG);
448     lua_settable(lua,-3);
449 
450     lua_pushstring(lua,"LOG_VERBOSE");
451     lua_pushnumber(lua,REDIS_VERBOSE);
452     lua_settable(lua,-3);
453 
454     lua_pushstring(lua,"LOG_NOTICE");
455     lua_pushnumber(lua,REDIS_NOTICE);
456     lua_settable(lua,-3);
457 
458     lua_pushstring(lua,"LOG_WARNING");
459     lua_pushnumber(lua,REDIS_WARNING);
460     lua_settable(lua,-3);
461 
462     /* redis.sha1hex */
463     lua_pushstring(lua, "sha1hex");
464     lua_pushcfunction(lua, luaRedisSha1hexCommand);
465     lua_settable(lua, -3);
466 
467     /* Finally set the table as 'redis' global var. */
468     lua_setglobal(lua,"redis");
469 
470     /* Replace math.random and math.randomseed with our implementations. */
471     lua_getglobal(lua,"math");
472 
473     lua_pushstring(lua,"random");
474     lua_pushcfunction(lua,redis_math_random);
475     lua_settable(lua,-3);
476 
477     lua_pushstring(lua,"randomseed");
478     lua_pushcfunction(lua,redis_math_randomseed);
479     lua_settable(lua,-3);
480 
481     lua_setglobal(lua,"math");
482 
483     /* Add a helper funciton that we use to sort the multi bulk output of non
484      * deterministic commands, when containing 'false' elements. */
485     {
486         char *compare_func =    "function __redis__compare_helper(a,b)\n"
487                                 "  if a == false then a = '' end\n"
488                                 "  if b == false then b = '' end\n"
489                                 "  return a<b\n"
490                                 "end\n";
491         luaL_loadbuffer(lua,compare_func,strlen(compare_func),"cmp_func_def");
492         lua_pcall(lua,0,0,0);
493     }
494 
495     /* Create the (non connected) client that we use to execute Redis commands
496      * inside the Lua interpreter.
497      * Note: there is no need to create it again when this function is called
498      * by scriptingReset(). */
499     if (server.lua_client == NULL) {
500         server.lua_client = createClient(-1);
501         server.lua_client->flags |= REDIS_LUA_CLIENT;
502     }
503 
504     server.lua = lua;
505 }
506 
507 /* Release resources related to Lua scripting.
508  * This function is used in order to reset the scripting environment. */
509 void scriptingRelease(void) {
510     dictRelease(server.lua_scripts);
511     lua_close(server.lua);
512 }
513 
514 void scriptingReset(void) {
515     scriptingRelease();
516     scriptingInit();
517 }
518 
519 /* Perform the SHA1 of the input string. We use this both for hasing script
520  * bodies in order to obtain the Lua function name, and in the implementation
521  * of redis.sha1().
522  *
523  * 'digest' should point to a 41 bytes buffer: 40 for SHA1 converted into an
524  * hexadecimal number, plus 1 byte for null term. */
525 void sha1hex(char *digest, char *script, size_t len) {
526     SHA1_CTX ctx;
527     unsigned char hash[20];
528     char *cset = "0123456789abcdef";
529     int j;
530 
531     SHA1Init(&ctx);
532     SHA1Update(&ctx,(unsigned char*)script,len);
533     SHA1Final(hash,&ctx);
534 
535     for (j = 0; j < 20; j++) {
536         digest[j*2] = cset[((hash[j]&0xF0)>>4)];
537         digest[j*2+1] = cset[(hash[j]&0xF)];
538     }
539     digest[40] = '\0';
540 }
541 
542 void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
543     int t = lua_type(lua,-1);
544 
545     switch(t) {
546     case LUA_TSTRING:
547         addReplyBulkCBuffer(c,(char*)lua_tostring(lua,-1),lua_strlen(lua,-1));
548         break;
549     case LUA_TBOOLEAN:
550         addReply(c,lua_toboolean(lua,-1) ? shared.cone : shared.nullbulk);
551         break;
552     case LUA_TNUMBER:
553         addReplyLongLong(c,(long long)lua_tonumber(lua,-1));
554         break;
555     case LUA_TTABLE:
556         /* We need to check if it is an array, an error, or a status reply.
557          * Error are returned as a single element table with 'err' field.
558          * Status replies are returned as single elment table with 'ok' field */
559         lua_pushstring(lua,"err");
560         lua_gettable(lua,-2);
561         t = lua_type(lua,-1);
562         if (t == LUA_TSTRING) {
563             sds err = sdsnew(lua_tostring(lua,-1));
564             sdsmapchars(err,"\r\n","  ",2);
565             addReplySds(c,sdscatprintf(sdsempty(),"-%s\r\n",err));
566             sdsfree(err);
567             lua_pop(lua,2);
568             return;
569         }
570 
571         lua_pop(lua,1);
572         lua_pushstring(lua,"ok");
573         lua_gettable(lua,-2);
574         t = lua_type(lua,-1);
575         if (t == LUA_TSTRING) {
576             sds ok = sdsnew(lua_tostring(lua,-1));
577             sdsmapchars(ok,"\r\n","  ",2);
578             addReplySds(c,sdscatprintf(sdsempty(),"+%s\r\n",ok));
579             sdsfree(ok);
580             lua_pop(lua,1);
581         } else {
582             void *replylen = addDeferredMultiBulkLength(c);
583             int j = 1, mbulklen = 0;
584 
585             lua_pop(lua,1); /* Discard the 'ok' field value we popped */
586             while(1) {
587                 lua_pushnumber(lua,j++);
588                 lua_gettable(lua,-2);
589                 t = lua_type(lua,-1);
590                 if (t == LUA_TNIL) {
591                     lua_pop(lua,1);
592                     break;
593                 }
594                 luaReplyToRedisReply(c, lua);
595                 mbulklen++;
596             }
597             setDeferredMultiBulkLength(c,replylen,mbulklen);
598         }
599         break;
600     default:
601         addReply(c,shared.nullbulk);
602     }
603     lua_pop(lua,1);
604 }
605 
606 /* Set an array of Redis String Objects as a Lua array (table) stored into a
607  * global variable. */
608 void luaSetGlobalArray(lua_State *lua, char *var, robj **elev, int elec) {
609     int j;
610 
611     lua_newtable(lua);
612     for (j = 0; j < elec; j++) {
613         lua_pushlstring(lua,(char*)elev[j]->ptr,sdslen(elev[j]->ptr));
614         lua_rawseti(lua,-2,j+1);
615     }
616     lua_setglobal(lua,var);
617 }
618 
619 /* Define a lua function with the specified function name and body.
620  * The function name musts be a 2 characters long string, since all the
621  * functions we defined in the Lua context are in the form:
622  *
623  *   f_<hex sha1 sum>
624  *
625  * On success REDIS_OK is returned, and nothing is left on the Lua stack.
626  * On error REDIS_ERR is returned and an appropriate error is set in the
627  * client context. */
628 int luaCreateFunction(redisClient *c, lua_State *lua, char *funcname, robj *body) {
629     sds funcdef = sdsempty();
630 
631     funcdef = sdscat(funcdef,"function ");
632     funcdef = sdscatlen(funcdef,funcname,42);
633     funcdef = sdscatlen(funcdef,"() ",3);
634     funcdef = sdscatlen(funcdef,body->ptr,sdslen(body->ptr));
635     funcdef = sdscatlen(funcdef," end",4);
636 
637     if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) {
638         addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
639             lua_tostring(lua,-1));
640         lua_pop(lua,1);
641         sdsfree(funcdef);
642         return REDIS_ERR;
643     }
644     sdsfree(funcdef);
645     if (lua_pcall(lua,0,0,0)) {
646         addReplyErrorFormat(c,"Error running script (new function): %s\n",
647             lua_tostring(lua,-1));
648         lua_pop(lua,1);
649         return REDIS_ERR;
650     }
651 
652     /* We also save a SHA1 -> Original script map in a dictionary
653      * so that we can replicate / write in the AOF all the
654      * EVALSHA commands as EVAL using the original script. */
655     {
656         int retval = dictAdd(server.lua_scripts,
657                              sdsnewlen(funcname+2,40),body);
658         redisAssertWithInfo(c,NULL,retval == DICT_OK);
659         incrRefCount(body);
660     }
661     return REDIS_OK;
662 }
663 
664 void evalGenericCommand(redisClient *c, int evalsha) {
665     lua_State *lua = server.lua;
666     char funcname[43];
667     long long numkeys;
668 
669     /* We want the same PRNG sequence at every call so that our PRNG is
670      * not affected by external state. */
671     redisSrand48(0);
672 
673     /* We set this flag to zero to remember that so far no random command
674      * was called. This way we can allow the user to call commands like
675      * SRANDMEMBER or RANDOMKEY from Lua scripts as far as no write command
676      * is called (otherwise the replication and AOF would end with non
677      * deterministic sequences).
678      *
679      * Thanks to this flag we'll raise an error every time a write command
680      * is called after a random command was used. */
681     server.lua_random_dirty = 0;
682     server.lua_write_dirty = 0;
683 
684     /* Get the number of arguments that are keys */
685     if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK)
686         return;
687     if (numkeys > (c->argc - 3)) {
688         addReplyError(c,"Number of keys can't be greater than number of args");
689         return;
690     }
691 
692     /* We obtain the script SHA1, then check if this function is already
693      * defined into the Lua state */
694     funcname[0] = 'f';
695     funcname[1] = '_';
696     if (!evalsha) {
697         /* Hash the code if this is an EVAL call */
698         sha1hex(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
699     } else {
700         /* We already have the SHA if it is a EVALSHA */
701         int j;
702         char *sha = c->argv[1]->ptr;
703 
704         for (j = 0; j < 40; j++)
705             funcname[j+2] = tolower(sha[j]);
706         funcname[42] = '\0';
707     }
708 
709     /* Try to lookup the Lua function */
710     lua_getglobal(lua, funcname);
711     if (lua_isnil(lua,1)) {
712         lua_pop(lua,1); /* remove the nil from the stack */
713         /* Function not defined... let's define it if we have the
714          * body of the funciton. If this is an EVALSHA call we can just
715          * return an error. */
716         if (evalsha) {
717             addReply(c, shared.noscripterr);
718             return;
719         }
720         if (luaCreateFunction(c,lua,funcname,c->argv[1]) == REDIS_ERR) return;
721         /* Now the following is guaranteed to return non nil */
722         lua_getglobal(lua, funcname);
723         redisAssert(!lua_isnil(lua,1));
724     }
725 
726     /* Populate the argv and keys table accordingly to the arguments that
727      * EVAL received. */
728     luaSetGlobalArray(lua,"KEYS",c->argv+3,numkeys);
729     luaSetGlobalArray(lua,"ARGV",c->argv+3+numkeys,c->argc-3-numkeys);
730 
731     /* Select the right DB in the context of the Lua client */
732     selectDb(server.lua_client,c->db->id);
733 
734     /* Set an hook in order to be able to stop the script execution if it
735      * is running for too much time.
736      * We set the hook only if the time limit is enabled as the hook will
737      * make the Lua script execution slower. */
738     if (server.lua_time_limit > 0 && server.masterhost == NULL) {
739         lua_sethook(lua,luaMaskCountHook,LUA_MASKCOUNT,100000);
740     } else {
741         lua_sethook(lua,luaMaskCountHook,0,0);
742     }
743 
744     /* At this point whatever this script was never seen before or if it was
745      * already defined, we can call it. We have zero arguments and expect
746      * a single return value. */
747     server.lua_caller = c;
748     server.lua_time_start = ustime()/1000;
749     server.lua_kill = 0;
750     if (lua_pcall(lua,0,1,0)) {
751         if (server.lua_timedout) {
752             server.lua_timedout = 0;
753             /* Restore the readable handler that was unregistered when the
754              * script timeout was detected. */
755             aeCreateFileEvent(server.el,c->fd,AE_READABLE,
756                               readQueryFromClient,c);
757         }
758         server.lua_caller = NULL;
759         selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
760         addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
761             funcname, lua_tostring(lua,-1));
762         lua_pop(lua,1);
763         lua_gc(lua,LUA_GCCOLLECT,0);
764         return;
765     }
766     server.lua_timedout = 0;
767     server.lua_caller = NULL;
768     selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
769     luaReplyToRedisReply(c,lua);
770     lua_gc(lua,LUA_GCSTEP,1);
771 
772     /* If we have slaves attached we want to replicate this command as
773      * EVAL instead of EVALSHA. We do this also in the AOF as currently there
774      * is no easy way to propagate a command in a different way in the AOF
775      * and in the replication link.
776      *
777      * IMPROVEMENT POSSIBLE:
778      * 1) Replicate this command as EVALSHA in the AOF.
779      * 2) Remember what slave already received a given script, and replicate
780      *    the EVALSHA against this slaves when possible.
781      */
782     if (evalsha) {
783         robj *script = dictFetchValue(server.lua_scripts,c->argv[1]->ptr);
784 
785         redisAssertWithInfo(c,NULL,script != NULL);
786         rewriteClientCommandArgument(c,0,
787             resetRefCount(createStringObject("EVAL",4)));
788         rewriteClientCommandArgument(c,1,script);
789     }
790 }
791 
792 void evalCommand(redisClient *c) {
793     evalGenericCommand(c,0);
794 }
795 
796 void evalShaCommand(redisClient *c) {
797     if (sdslen(c->argv[1]->ptr) != 40) {
798         /* We know that a match is not possible if the provided SHA is
799          * not the right length. So we return an error ASAP, this way
800          * evalGenericCommand() can be implemented without string length
801          * sanity check */
802         addReply(c, shared.noscripterr);
803         return;
804     }
805     evalGenericCommand(c,1);
806 }
807 
808 /* We replace math.random() with our implementation that is not affected
809  * by specific libc random() implementations and will output the same sequence
810  * (for the same seed) in every arch. */
811 
812 /* The following implementation is the one shipped with Lua itself but with
813  * rand() replaced by redisLrand48(). */
814 int redis_math_random (lua_State *L) {
815   /* the `%' avoids the (rare) case of r==1, and is needed also because on
816      some systems (SunOS!) `rand()' may return a value larger than RAND_MAX */
817   lua_Number r = (lua_Number)(redisLrand48()%REDIS_LRAND48_MAX) /
818                                 (lua_Number)REDIS_LRAND48_MAX;
819   switch (lua_gettop(L)) {  /* check number of arguments */
820     case 0: {  /* no arguments */
821       lua_pushnumber(L, r);  /* Number between 0 and 1 */
822       break;
823     }
824     case 1: {  /* only upper limit */
825       int u = luaL_checkint(L, 1);
826       luaL_argcheck(L, 1<=u, 1, "interval is empty");
827       lua_pushnumber(L, floor(r*u)+1);  /* int between 1 and `u' */
828       break;
829     }
830     case 2: {  /* lower and upper limits */
831       int l = luaL_checkint(L, 1);
832       int u = luaL_checkint(L, 2);
833       luaL_argcheck(L, l<=u, 2, "interval is empty");
834       lua_pushnumber(L, floor(r*(u-l+1))+l);  /* int between `l' and `u' */
835       break;
836     }
837     default: return luaL_error(L, "wrong number of arguments");
838   }
839   return 1;
840 }
841 
842 int redis_math_randomseed (lua_State *L) {
843   redisSrand48(luaL_checkint(L, 1));
844   return 0;
845 }
846 
847 /* ---------------------------------------------------------------------------
848  * SCRIPT command for script environment introspection and control
849  * ------------------------------------------------------------------------- */
850 
851 void scriptCommand(redisClient *c) {
852     if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"flush")) {
853         scriptingReset();
854         addReply(c,shared.ok);
855         server.dirty++; /* Replicating this command is a good idea. */
856     } else if (c->argc >= 2 && !strcasecmp(c->argv[1]->ptr,"exists")) {
857         int j;
858 
859         addReplyMultiBulkLen(c, c->argc-2);
860         for (j = 2; j < c->argc; j++) {
861             if (dictFind(server.lua_scripts,c->argv[j]->ptr))
862                 addReply(c,shared.cone);
863             else
864                 addReply(c,shared.czero);
865         }
866     } else if (c->argc == 3 && !strcasecmp(c->argv[1]->ptr,"load")) {
867         char funcname[43];
868         sds sha;
869 
870         funcname[0] = 'f';
871         funcname[1] = '_';
872         sha1hex(funcname+2,c->argv[2]->ptr,sdslen(c->argv[2]->ptr));
873         sha = sdsnewlen(funcname+2,40);
874         if (dictFind(server.lua_scripts,sha) == NULL) {
875             if (luaCreateFunction(c,server.lua,funcname,c->argv[2])
876                     == REDIS_ERR) {
877                 sdsfree(sha);
878                 return;
879             }
880         }
881         addReplyBulkCBuffer(c,funcname+2,40);
882         sdsfree(sha);
883     } else if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"kill")) {
884         if (server.lua_caller == NULL) {
885             addReplyError(c,"No scripts in execution right now.");
886         } else if (server.lua_write_dirty) {
887             addReplyError(c, "Sorry the script already executed write commands against the dataset. You can either wait the script termination or kill the server in an hard way using the SHUTDOWN NOSAVE command.");
888         } else {
889             server.lua_kill = 1;
890             addReply(c,shared.ok);
891         }
892     } else {
893         addReplyError(c, "Unknown SCRIPT subcommand or wrong # of args.");
894     }
895 }
896