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