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