]> git.saurik.com Git - redis.git/blob - src/scripting.c
Scripting: redis.NIL to return nil bulk replies.
[redis.git] / src / scripting.c
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 /* Require at least one argument */
171 if (argc == 0) {
172 luaPushError(lua,
173 "Please specify at least one argument for redis.call()");
174 return 1;
175 }
176
177 /* Build the arguments vector */
178 argv = zmalloc(sizeof(robj*)*argc);
179 for (j = 0; j < argc; j++) {
180 if (!lua_isstring(lua,j+1)) break;
181 argv[j] = createStringObject((char*)lua_tostring(lua,j+1),
182 lua_strlen(lua,j+1));
183 }
184
185 /* Check if one of the arguments passed by the Lua script
186 * is not a string or an integer (lua_isstring() return true for
187 * integers as well). */
188 if (j != argc) {
189 j--;
190 while (j >= 0) {
191 decrRefCount(argv[j]);
192 j--;
193 }
194 zfree(argv);
195 luaPushError(lua,
196 "Lua redis() command arguments must be strings or integers");
197 return 1;
198 }
199
200 /* Setup our fake client for command execution */
201 c->argv = argv;
202 c->argc = argc;
203
204 /* Command lookup */
205 cmd = lookupCommand(argv[0]->ptr);
206 if (!cmd || ((cmd->arity > 0 && cmd->arity != argc) ||
207 (argc < -cmd->arity)))
208 {
209 if (cmd)
210 luaPushError(lua,
211 "Wrong number of args calling Redis command From Lua script");
212 else
213 luaPushError(lua,"Unknown Redis command called from Lua script");
214 goto cleanup;
215 }
216
217 /* There are commands that are not allowed inside scripts. */
218 if (cmd->flags & REDIS_CMD_NOSCRIPT) {
219 luaPushError(lua, "This Redis command is not allowed from scripts");
220 goto cleanup;
221 }
222
223 /* Write commands are forbidden against read-only slaves, or if a
224 * command marked as non-deterministic was already called in the context
225 * of this script. */
226 if (cmd->flags & REDIS_CMD_WRITE) {
227 if (server.lua_random_dirty) {
228 luaPushError(lua,
229 "Write commands not allowed after non deterministic commands");
230 goto cleanup;
231 } else if (server.masterhost && server.repl_slave_ro &&
232 !(server.lua_caller->flags & REDIS_MASTER))
233 {
234 luaPushError(lua, shared.roslaveerr->ptr);
235 goto cleanup;
236 } else if (server.stop_writes_on_bgsave_err &&
237 server.saveparamslen > 0 &&
238 server.lastbgsave_status == REDIS_ERR)
239 {
240 luaPushError(lua, shared.bgsaveerr->ptr);
241 goto cleanup;
242 }
243 }
244
245 /* If we reached the memory limit configured via maxmemory, commands that
246 * could enlarge the memory usage are not allowed, but only if this is the
247 * first write in the context of this script, otherwise we can't stop
248 * in the middle. */
249 if (server.maxmemory && server.lua_write_dirty == 0 &&
250 (cmd->flags & REDIS_CMD_DENYOOM))
251 {
252 if (freeMemoryIfNeeded() == REDIS_ERR) {
253 luaPushError(lua, shared.oomerr->ptr);
254 goto cleanup;
255 }
256 }
257
258 if (cmd->flags & REDIS_CMD_RANDOM) server.lua_random_dirty = 1;
259 if (cmd->flags & REDIS_CMD_WRITE) server.lua_write_dirty = 1;
260
261 /* Run the command */
262 c->cmd = cmd;
263 call(c,REDIS_CALL_SLOWLOG | REDIS_CALL_STATS);
264
265 /* Convert the result of the Redis command into a suitable Lua type.
266 * The first thing we need is to create a single string from the client
267 * output buffers. */
268 reply = sdsempty();
269 if (c->bufpos) {
270 reply = sdscatlen(reply,c->buf,c->bufpos);
271 c->bufpos = 0;
272 }
273 while(listLength(c->reply)) {
274 robj *o = listNodeValue(listFirst(c->reply));
275
276 reply = sdscatlen(reply,o->ptr,sdslen(o->ptr));
277 listDelNode(c->reply,listFirst(c->reply));
278 }
279 if (raise_error && reply[0] != '-') raise_error = 0;
280 redisProtocolToLuaType(lua,reply);
281 /* Sort the output array if needed, assuming it is a non-null multi bulk
282 * reply as expected. */
283 if ((cmd->flags & REDIS_CMD_SORT_FOR_SCRIPT) &&
284 (reply[0] == '*' && reply[1] != '-')) {
285 luaSortArray(lua);
286 }
287 sdsfree(reply);
288 c->reply_bytes = 0;
289
290 cleanup:
291 /* Clean up. Command code may have changed argv/argc so we use the
292 * argv/argc of the client instead of the local variables. */
293 for (j = 0; j < c->argc; j++)
294 decrRefCount(c->argv[j]);
295 zfree(c->argv);
296
297 if (raise_error) {
298 /* If we are here we should have an error in the stack, in the
299 * form of a table with an "err" field. Extract the string to
300 * return the plain error. */
301 lua_pushstring(lua,"err");
302 lua_gettable(lua,-2);
303 return lua_error(lua);
304 }
305 return 1;
306 }
307
308 int luaRedisCallCommand(lua_State *lua) {
309 return luaRedisGenericCommand(lua,1);
310 }
311
312 int luaRedisPCallCommand(lua_State *lua) {
313 return luaRedisGenericCommand(lua,0);
314 }
315
316 /* This adds redis.sha1hex(string) to Lua scripts using the same hashing
317 * function used for sha1ing lua scripts. */
318 int luaRedisSha1hexCommand(lua_State *lua) {
319 int argc = lua_gettop(lua);
320 char digest[41];
321 size_t len;
322 char *s;
323
324 if (argc != 1) {
325 luaPushError(lua, "wrong number of arguments");
326 return 1;
327 }
328
329 s = (char*)lua_tolstring(lua,1,&len);
330 sha1hex(digest,s,len);
331 lua_pushstring(lua,digest);
332 return 1;
333 }
334
335 int luaLogCommand(lua_State *lua) {
336 int j, argc = lua_gettop(lua);
337 int level;
338 sds log;
339
340 if (argc < 2) {
341 luaPushError(lua, "redis.log() requires two arguments or more.");
342 return 1;
343 } else if (!lua_isnumber(lua,-argc)) {
344 luaPushError(lua, "First argument must be a number (log level).");
345 return 1;
346 }
347 level = lua_tonumber(lua,-argc);
348 if (level < REDIS_DEBUG || level > REDIS_WARNING) {
349 luaPushError(lua, "Invalid debug level.");
350 return 1;
351 }
352
353 /* Glue together all the arguments */
354 log = sdsempty();
355 for (j = 1; j < argc; j++) {
356 size_t len;
357 char *s;
358
359 s = (char*)lua_tolstring(lua,(-argc)+j,&len);
360 if (s) {
361 if (j != 1) log = sdscatlen(log," ",1);
362 log = sdscatlen(log,s,len);
363 }
364 }
365 redisLogRaw(level,log);
366 sdsfree(log);
367 return 0;
368 }
369
370 void luaMaskCountHook(lua_State *lua, lua_Debug *ar) {
371 long long elapsed;
372 REDIS_NOTUSED(ar);
373 REDIS_NOTUSED(lua);
374
375 elapsed = (ustime()/1000) - server.lua_time_start;
376 if (elapsed >= server.lua_time_limit && server.lua_timedout == 0) {
377 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);
378 server.lua_timedout = 1;
379 /* Once the script timeouts we reenter the event loop to permit others
380 * to call SCRIPT KILL or SHUTDOWN NOSAVE if needed. For this reason
381 * we need to mask the client executing the script from the event loop.
382 * If we don't do that the client may disconnect and could no longer be
383 * here when the EVAL command will return. */
384 aeDeleteFileEvent(server.el, server.lua_caller->fd, AE_READABLE);
385 }
386 if (server.lua_timedout)
387 aeProcessEvents(server.el, AE_FILE_EVENTS|AE_DONT_WAIT);
388 if (server.lua_kill) {
389 redisLog(REDIS_WARNING,"Lua script killed by user with SCRIPT KILL.");
390 lua_pushstring(lua,"Script killed by user with SCRIPT KILL...");
391 lua_error(lua);
392 }
393 }
394
395 void luaLoadLib(lua_State *lua, const char *libname, lua_CFunction luafunc) {
396 lua_pushcfunction(lua, luafunc);
397 lua_pushstring(lua, libname);
398 lua_call(lua, 1, 0);
399 }
400
401 LUALIB_API int (luaopen_cjson) (lua_State *L);
402 LUALIB_API int (luaopen_struct) (lua_State *L);
403 LUALIB_API int (luaopen_cmsgpack) (lua_State *L);
404
405 void luaLoadLibraries(lua_State *lua) {
406 luaLoadLib(lua, "", luaopen_base);
407 luaLoadLib(lua, LUA_TABLIBNAME, luaopen_table);
408 luaLoadLib(lua, LUA_STRLIBNAME, luaopen_string);
409 luaLoadLib(lua, LUA_MATHLIBNAME, luaopen_math);
410 luaLoadLib(lua, LUA_DBLIBNAME, luaopen_debug);
411 luaLoadLib(lua, "cjson", luaopen_cjson);
412 luaLoadLib(lua, "struct", luaopen_struct);
413 luaLoadLib(lua, "cmsgpack", luaopen_cmsgpack);
414
415 #if 0 /* Stuff that we don't load currently, for sandboxing concerns. */
416 luaLoadLib(lua, LUA_LOADLIBNAME, luaopen_package);
417 luaLoadLib(lua, LUA_OSLIBNAME, luaopen_os);
418 #endif
419 }
420
421 /* Remove a functions that we don't want to expose to the Redis scripting
422 * environment. */
423 void luaRemoveUnsupportedFunctions(lua_State *lua) {
424 lua_pushnil(lua);
425 lua_setglobal(lua,"loadfile");
426 }
427
428 /* This function installs metamethods in the global table _G that prevent
429 * the creation of globals accidentally.
430 *
431 * It should be the last to be called in the scripting engine initialization
432 * sequence, because it may interact with creation of globals. */
433 void scriptingEnableGlobalsProtection(lua_State *lua) {
434 char *s[32];
435 sds code = sdsempty();
436 int j = 0;
437
438 /* strict.lua from: http://metalua.luaforge.net/src/lib/strict.lua.html.
439 * Modified to be adapted to Redis. */
440 s[j++]="local mt = {}\n";
441 s[j++]="setmetatable(_G, mt)\n";
442 s[j++]="mt.__newindex = function (t, n, v)\n";
443 s[j++]=" if debug.getinfo(2) then\n";
444 s[j++]=" local w = debug.getinfo(2, \"S\").what\n";
445 s[j++]=" if w ~= \"main\" and w ~= \"C\" then\n";
446 s[j++]=" error(\"Script attempted to create global variable '\"..tostring(n)..\"'\", 2)\n";
447 s[j++]=" end\n";
448 s[j++]=" end\n";
449 s[j++]=" rawset(t, n, v)\n";
450 s[j++]="end\n";
451 s[j++]="mt.__index = function (t, n)\n";
452 s[j++]=" if debug.getinfo(2) and debug.getinfo(2, \"S\").what ~= \"C\" then\n";
453 s[j++]=" error(\"Script attempted to access unexisting global variable '\"..tostring(n)..\"'\", 2)\n";
454 s[j++]=" end\n";
455 s[j++]=" return rawget(t, n)\n";
456 s[j++]="end\n";
457 s[j++]=NULL;
458
459 for (j = 0; s[j] != NULL; j++) code = sdscatlen(code,s[j],strlen(s[j]));
460 luaL_loadbuffer(lua,code,sdslen(code),"@enable_strict_lua");
461 lua_pcall(lua,0,0,0);
462 sdsfree(code);
463 }
464
465 /* Initialize the scripting environment.
466 * It is possible to call this function to reset the scripting environment
467 * assuming that we call scriptingRelease() before.
468 * See scriptingReset() for more information. */
469 void scriptingInit(void) {
470 lua_State *lua = lua_open();
471
472 luaLoadLibraries(lua);
473 luaRemoveUnsupportedFunctions(lua);
474
475 /* Initialize a dictionary we use to map SHAs to scripts.
476 * This is useful for replication, as we need to replicate EVALSHA
477 * as EVAL, so we need to remember the associated script. */
478 server.lua_scripts = dictCreate(&dbDictType,NULL);
479
480 /* Register the redis commands table and fields */
481 lua_newtable(lua);
482
483 /* redis.call */
484 lua_pushstring(lua,"call");
485 lua_pushcfunction(lua,luaRedisCallCommand);
486 lua_settable(lua,-3);
487
488 /* redis.pcall */
489 lua_pushstring(lua,"pcall");
490 lua_pushcfunction(lua,luaRedisPCallCommand);
491 lua_settable(lua,-3);
492
493 /* redis.log and log levels. */
494 lua_pushstring(lua,"log");
495 lua_pushcfunction(lua,luaLogCommand);
496 lua_settable(lua,-3);
497
498 lua_pushstring(lua,"LOG_DEBUG");
499 lua_pushnumber(lua,REDIS_DEBUG);
500 lua_settable(lua,-3);
501
502 lua_pushstring(lua,"LOG_VERBOSE");
503 lua_pushnumber(lua,REDIS_VERBOSE);
504 lua_settable(lua,-3);
505
506 lua_pushstring(lua,"LOG_NOTICE");
507 lua_pushnumber(lua,REDIS_NOTICE);
508 lua_settable(lua,-3);
509
510 lua_pushstring(lua,"LOG_WARNING");
511 lua_pushnumber(lua,REDIS_WARNING);
512 lua_settable(lua,-3);
513
514 /* redis.sha1hex */
515 lua_pushstring(lua, "sha1hex");
516 lua_pushcfunction(lua, luaRedisSha1hexCommand);
517 lua_settable(lua, -3);
518
519 /* redis.NIL */
520 lua_pushstring(lua, "NIL");
521 lua_newtable(lua);
522 lua_pushstring(lua, "nilbulk");
523 lua_pushboolean(lua, 1);
524 lua_settable(lua, -3);
525 lua_settable(lua, -3);
526
527 /* Finally set the table as 'redis' global var. */
528 lua_setglobal(lua,"redis");
529
530 /* Replace math.random and math.randomseed with our implementations. */
531 lua_getglobal(lua,"math");
532
533 lua_pushstring(lua,"random");
534 lua_pushcfunction(lua,redis_math_random);
535 lua_settable(lua,-3);
536
537 lua_pushstring(lua,"randomseed");
538 lua_pushcfunction(lua,redis_math_randomseed);
539 lua_settable(lua,-3);
540
541 lua_setglobal(lua,"math");
542
543 /* Add a helper funciton that we use to sort the multi bulk output of non
544 * deterministic commands, when containing 'false' elements. */
545 {
546 char *compare_func = "function __redis__compare_helper(a,b)\n"
547 " if a == false then a = '' end\n"
548 " if b == false then b = '' end\n"
549 " return a<b\n"
550 "end\n";
551 luaL_loadbuffer(lua,compare_func,strlen(compare_func),"@cmp_func_def");
552 lua_pcall(lua,0,0,0);
553 }
554
555 /* Create the (non connected) client that we use to execute Redis commands
556 * inside the Lua interpreter.
557 * Note: there is no need to create it again when this function is called
558 * by scriptingReset(). */
559 if (server.lua_client == NULL) {
560 server.lua_client = createClient(-1);
561 server.lua_client->flags |= REDIS_LUA_CLIENT;
562 }
563
564 /* Lua beginners ofter don't use "local", this is likely to introduce
565 * subtle bugs in their code. To prevent problems we protect accesses
566 * to global variables. */
567 scriptingEnableGlobalsProtection(lua);
568
569 server.lua = lua;
570 }
571
572 /* Release resources related to Lua scripting.
573 * This function is used in order to reset the scripting environment. */
574 void scriptingRelease(void) {
575 dictRelease(server.lua_scripts);
576 lua_close(server.lua);
577 }
578
579 void scriptingReset(void) {
580 scriptingRelease();
581 scriptingInit();
582 }
583
584 /* Perform the SHA1 of the input string. We use this both for hasing script
585 * bodies in order to obtain the Lua function name, and in the implementation
586 * of redis.sha1().
587 *
588 * 'digest' should point to a 41 bytes buffer: 40 for SHA1 converted into an
589 * hexadecimal number, plus 1 byte for null term. */
590 void sha1hex(char *digest, char *script, size_t len) {
591 SHA1_CTX ctx;
592 unsigned char hash[20];
593 char *cset = "0123456789abcdef";
594 int j;
595
596 SHA1Init(&ctx);
597 SHA1Update(&ctx,(unsigned char*)script,len);
598 SHA1Final(hash,&ctx);
599
600 for (j = 0; j < 20; j++) {
601 digest[j*2] = cset[((hash[j]&0xF0)>>4)];
602 digest[j*2+1] = cset[(hash[j]&0xF)];
603 }
604 digest[40] = '\0';
605 }
606
607 void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
608 int t = lua_type(lua,-1);
609
610 switch(t) {
611 case LUA_TSTRING:
612 addReplyBulkCBuffer(c,(char*)lua_tostring(lua,-1),lua_strlen(lua,-1));
613 break;
614 case LUA_TBOOLEAN:
615 addReply(c,lua_toboolean(lua,-1) ? shared.cone : shared.nullbulk);
616 break;
617 case LUA_TNUMBER:
618 addReplyLongLong(c,(long long)lua_tonumber(lua,-1));
619 break;
620 case LUA_TTABLE:
621 /* The table can be an array or it may be in a special format that
622 * Lua uses to return special Redis protocol data types.
623 *
624 * 1) Errors are retuned as a single element table with 'err' field.
625 * 2) Status reply are returned as a single element table with 'ok'
626 * field.
627 * 3) A Redis nil bulk reply is returned as a single element table
628 * with 'nilbulk' field set to true.
629 *
630 * All the rest is considered just an array and is translated into
631 * a Redis multi bulk reply. */
632
633 /* Nil bulk reply */
634 lua_pushstring(lua,"nilbulk");
635 lua_gettable(lua,-2);
636 t = lua_type(lua,-1);
637 if (t == LUA_TBOOLEAN) {
638 addReply(c,shared.nullbulk);
639 lua_pop(lua,2);
640 return;
641 }
642 lua_pop(lua,1);
643
644 /* Error reply */
645 lua_pushstring(lua,"err");
646 lua_gettable(lua,-2);
647 t = lua_type(lua,-1);
648 if (t == LUA_TSTRING) {
649 sds err = sdsnew(lua_tostring(lua,-1));
650 sdsmapchars(err,"\r\n"," ",2);
651 addReplySds(c,sdscatprintf(sdsempty(),"-%s\r\n",err));
652 sdsfree(err);
653 lua_pop(lua,2);
654 return;
655 }
656 lua_pop(lua,1);
657
658 /* Status reply */
659 lua_pushstring(lua,"ok");
660 lua_gettable(lua,-2);
661 t = lua_type(lua,-1);
662 if (t == LUA_TSTRING) {
663 sds ok = sdsnew(lua_tostring(lua,-1));
664 sdsmapchars(ok,"\r\n"," ",2);
665 addReplySds(c,sdscatprintf(sdsempty(),"+%s\r\n",ok));
666 sdsfree(ok);
667 lua_pop(lua,1);
668 } else {
669 /* Multi bulk reply. */
670 void *replylen = addDeferredMultiBulkLength(c);
671 int j = 1, mbulklen = 0;
672
673 lua_pop(lua,1); /* Discard the 'ok' field value we popped */
674 while(1) {
675 lua_pushnumber(lua,j++);
676 lua_gettable(lua,-2);
677 t = lua_type(lua,-1);
678 if (t == LUA_TNIL) {
679 lua_pop(lua,1);
680 break;
681 }
682 luaReplyToRedisReply(c, lua);
683 mbulklen++;
684 }
685 setDeferredMultiBulkLength(c,replylen,mbulklen);
686 }
687 break;
688 default:
689 addReply(c,shared.nullbulk);
690 }
691 lua_pop(lua,1);
692 }
693
694 /* Set an array of Redis String Objects as a Lua array (table) stored into a
695 * global variable. */
696 void luaSetGlobalArray(lua_State *lua, char *var, robj **elev, int elec) {
697 int j;
698
699 lua_newtable(lua);
700 for (j = 0; j < elec; j++) {
701 lua_pushlstring(lua,(char*)elev[j]->ptr,sdslen(elev[j]->ptr));
702 lua_rawseti(lua,-2,j+1);
703 }
704 lua_setglobal(lua,var);
705 }
706
707 /* Define a lua function with the specified function name and body.
708 * The function name musts be a 2 characters long string, since all the
709 * functions we defined in the Lua context are in the form:
710 *
711 * f_<hex sha1 sum>
712 *
713 * On success REDIS_OK is returned, and nothing is left on the Lua stack.
714 * On error REDIS_ERR is returned and an appropriate error is set in the
715 * client context. */
716 int luaCreateFunction(redisClient *c, lua_State *lua, char *funcname, robj *body) {
717 sds funcdef = sdsempty();
718
719 funcdef = sdscat(funcdef,"function ");
720 funcdef = sdscatlen(funcdef,funcname,42);
721 funcdef = sdscatlen(funcdef,"() ",3);
722 funcdef = sdscatlen(funcdef,body->ptr,sdslen(body->ptr));
723 funcdef = sdscatlen(funcdef," end",4);
724
725 if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"@user_script")) {
726 addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
727 lua_tostring(lua,-1));
728 lua_pop(lua,1);
729 sdsfree(funcdef);
730 return REDIS_ERR;
731 }
732 sdsfree(funcdef);
733 if (lua_pcall(lua,0,0,0)) {
734 addReplyErrorFormat(c,"Error running script (new function): %s\n",
735 lua_tostring(lua,-1));
736 lua_pop(lua,1);
737 return REDIS_ERR;
738 }
739
740 /* We also save a SHA1 -> Original script map in a dictionary
741 * so that we can replicate / write in the AOF all the
742 * EVALSHA commands as EVAL using the original script. */
743 {
744 int retval = dictAdd(server.lua_scripts,
745 sdsnewlen(funcname+2,40),body);
746 redisAssertWithInfo(c,NULL,retval == DICT_OK);
747 incrRefCount(body);
748 }
749 return REDIS_OK;
750 }
751
752 void evalGenericCommand(redisClient *c, int evalsha) {
753 lua_State *lua = server.lua;
754 char funcname[43];
755 long long numkeys;
756 int delhook = 0;
757
758 /* We want the same PRNG sequence at every call so that our PRNG is
759 * not affected by external state. */
760 redisSrand48(0);
761
762 /* We set this flag to zero to remember that so far no random command
763 * was called. This way we can allow the user to call commands like
764 * SRANDMEMBER or RANDOMKEY from Lua scripts as far as no write command
765 * is called (otherwise the replication and AOF would end with non
766 * deterministic sequences).
767 *
768 * Thanks to this flag we'll raise an error every time a write command
769 * is called after a random command was used. */
770 server.lua_random_dirty = 0;
771 server.lua_write_dirty = 0;
772
773 /* Get the number of arguments that are keys */
774 if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK)
775 return;
776 if (numkeys > (c->argc - 3)) {
777 addReplyError(c,"Number of keys can't be greater than number of args");
778 return;
779 }
780
781 /* We obtain the script SHA1, then check if this function is already
782 * defined into the Lua state */
783 funcname[0] = 'f';
784 funcname[1] = '_';
785 if (!evalsha) {
786 /* Hash the code if this is an EVAL call */
787 sha1hex(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
788 } else {
789 /* We already have the SHA if it is a EVALSHA */
790 int j;
791 char *sha = c->argv[1]->ptr;
792
793 for (j = 0; j < 40; j++)
794 funcname[j+2] = tolower(sha[j]);
795 funcname[42] = '\0';
796 }
797
798 /* Try to lookup the Lua function */
799 lua_getglobal(lua, funcname);
800 if (lua_isnil(lua,1)) {
801 lua_pop(lua,1); /* remove the nil from the stack */
802 /* Function not defined... let's define it if we have the
803 * body of the funciton. If this is an EVALSHA call we can just
804 * return an error. */
805 if (evalsha) {
806 addReply(c, shared.noscripterr);
807 return;
808 }
809 if (luaCreateFunction(c,lua,funcname,c->argv[1]) == REDIS_ERR) return;
810 /* Now the following is guaranteed to return non nil */
811 lua_getglobal(lua, funcname);
812 redisAssert(!lua_isnil(lua,1));
813 }
814
815 /* Populate the argv and keys table accordingly to the arguments that
816 * EVAL received. */
817 luaSetGlobalArray(lua,"KEYS",c->argv+3,numkeys);
818 luaSetGlobalArray(lua,"ARGV",c->argv+3+numkeys,c->argc-3-numkeys);
819
820 /* Select the right DB in the context of the Lua client */
821 selectDb(server.lua_client,c->db->id);
822
823 /* Set an hook in order to be able to stop the script execution if it
824 * is running for too much time.
825 * We set the hook only if the time limit is enabled as the hook will
826 * make the Lua script execution slower. */
827 server.lua_caller = c;
828 server.lua_time_start = ustime()/1000;
829 server.lua_kill = 0;
830 if (server.lua_time_limit > 0 && server.masterhost == NULL) {
831 lua_sethook(lua,luaMaskCountHook,LUA_MASKCOUNT,100000);
832 delhook = 1;
833 }
834
835 /* At this point whatever this script was never seen before or if it was
836 * already defined, we can call it. We have zero arguments and expect
837 * a single return value. */
838 if (lua_pcall(lua,0,1,0)) {
839 if (delhook) lua_sethook(lua,luaMaskCountHook,0,0); /* Disable hook */
840 if (server.lua_timedout) {
841 server.lua_timedout = 0;
842 /* Restore the readable handler that was unregistered when the
843 * script timeout was detected. */
844 aeCreateFileEvent(server.el,c->fd,AE_READABLE,
845 readQueryFromClient,c);
846 }
847 server.lua_caller = NULL;
848 selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
849 addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
850 funcname, lua_tostring(lua,-1));
851 lua_pop(lua,1);
852 lua_gc(lua,LUA_GCCOLLECT,0);
853 return;
854 }
855 if (delhook) lua_sethook(lua,luaMaskCountHook,0,0); /* Disable hook */
856 server.lua_timedout = 0;
857 server.lua_caller = NULL;
858 selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
859 luaReplyToRedisReply(c,lua);
860 lua_gc(lua,LUA_GCSTEP,1);
861
862 /* If we have slaves attached we want to replicate this command as
863 * EVAL instead of EVALSHA. We do this also in the AOF as currently there
864 * is no easy way to propagate a command in a different way in the AOF
865 * and in the replication link.
866 *
867 * IMPROVEMENT POSSIBLE:
868 * 1) Replicate this command as EVALSHA in the AOF.
869 * 2) Remember what slave already received a given script, and replicate
870 * the EVALSHA against this slaves when possible.
871 */
872 if (evalsha) {
873 robj *script = dictFetchValue(server.lua_scripts,c->argv[1]->ptr);
874
875 redisAssertWithInfo(c,NULL,script != NULL);
876 rewriteClientCommandArgument(c,0,
877 resetRefCount(createStringObject("EVAL",4)));
878 rewriteClientCommandArgument(c,1,script);
879 }
880 }
881
882 void evalCommand(redisClient *c) {
883 evalGenericCommand(c,0);
884 }
885
886 void evalShaCommand(redisClient *c) {
887 if (sdslen(c->argv[1]->ptr) != 40) {
888 /* We know that a match is not possible if the provided SHA is
889 * not the right length. So we return an error ASAP, this way
890 * evalGenericCommand() can be implemented without string length
891 * sanity check */
892 addReply(c, shared.noscripterr);
893 return;
894 }
895 evalGenericCommand(c,1);
896 }
897
898 /* We replace math.random() with our implementation that is not affected
899 * by specific libc random() implementations and will output the same sequence
900 * (for the same seed) in every arch. */
901
902 /* The following implementation is the one shipped with Lua itself but with
903 * rand() replaced by redisLrand48(). */
904 int redis_math_random (lua_State *L) {
905 /* the `%' avoids the (rare) case of r==1, and is needed also because on
906 some systems (SunOS!) `rand()' may return a value larger than RAND_MAX */
907 lua_Number r = (lua_Number)(redisLrand48()%REDIS_LRAND48_MAX) /
908 (lua_Number)REDIS_LRAND48_MAX;
909 switch (lua_gettop(L)) { /* check number of arguments */
910 case 0: { /* no arguments */
911 lua_pushnumber(L, r); /* Number between 0 and 1 */
912 break;
913 }
914 case 1: { /* only upper limit */
915 int u = luaL_checkint(L, 1);
916 luaL_argcheck(L, 1<=u, 1, "interval is empty");
917 lua_pushnumber(L, floor(r*u)+1); /* int between 1 and `u' */
918 break;
919 }
920 case 2: { /* lower and upper limits */
921 int l = luaL_checkint(L, 1);
922 int u = luaL_checkint(L, 2);
923 luaL_argcheck(L, l<=u, 2, "interval is empty");
924 lua_pushnumber(L, floor(r*(u-l+1))+l); /* int between `l' and `u' */
925 break;
926 }
927 default: return luaL_error(L, "wrong number of arguments");
928 }
929 return 1;
930 }
931
932 int redis_math_randomseed (lua_State *L) {
933 redisSrand48(luaL_checkint(L, 1));
934 return 0;
935 }
936
937 /* ---------------------------------------------------------------------------
938 * SCRIPT command for script environment introspection and control
939 * ------------------------------------------------------------------------- */
940
941 void scriptCommand(redisClient *c) {
942 if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"flush")) {
943 scriptingReset();
944 addReply(c,shared.ok);
945 server.dirty++; /* Replicating this command is a good idea. */
946 } else if (c->argc >= 2 && !strcasecmp(c->argv[1]->ptr,"exists")) {
947 int j;
948
949 addReplyMultiBulkLen(c, c->argc-2);
950 for (j = 2; j < c->argc; j++) {
951 if (dictFind(server.lua_scripts,c->argv[j]->ptr))
952 addReply(c,shared.cone);
953 else
954 addReply(c,shared.czero);
955 }
956 } else if (c->argc == 3 && !strcasecmp(c->argv[1]->ptr,"load")) {
957 char funcname[43];
958 sds sha;
959
960 funcname[0] = 'f';
961 funcname[1] = '_';
962 sha1hex(funcname+2,c->argv[2]->ptr,sdslen(c->argv[2]->ptr));
963 sha = sdsnewlen(funcname+2,40);
964 if (dictFind(server.lua_scripts,sha) == NULL) {
965 if (luaCreateFunction(c,server.lua,funcname,c->argv[2])
966 == REDIS_ERR) {
967 sdsfree(sha);
968 return;
969 }
970 }
971 addReplyBulkCBuffer(c,funcname+2,40);
972 sdsfree(sha);
973 } else if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"kill")) {
974 if (server.lua_caller == NULL) {
975 addReplyError(c,"No scripts in execution right now.");
976 } else if (server.lua_write_dirty) {
977 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.");
978 } else {
979 server.lua_kill = 1;
980 addReply(c,shared.ok);
981 }
982 } else {
983 addReplyError(c, "Unknown SCRIPT subcommand or wrong # of args.");
984 }
985 }