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