]> git.saurik.com Git - redis.git/blob - src/scripting.c
added redis.sha1hex(string) as lua scripting function.
[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 /* Initialize the scripting environment.
416 * It is possible to call this function to reset the scripting environment
417 * assuming that we call scriptingRelease() before.
418 * See scriptingReset() for more information. */
419 void scriptingInit(void) {
420 lua_State *lua = lua_open();
421 luaLoadLibraries(lua);
422
423 /* Initialize a dictionary we use to map SHAs to scripts.
424 * This is useful for replication, as we need to replicate EVALSHA
425 * as EVAL, so we need to remember the associated script. */
426 server.lua_scripts = dictCreate(&dbDictType,NULL);
427
428 /* Register the redis commands table and fields */
429 lua_newtable(lua);
430
431 /* redis.call */
432 lua_pushstring(lua,"call");
433 lua_pushcfunction(lua,luaRedisCallCommand);
434 lua_settable(lua,-3);
435
436 /* redis.pcall */
437 lua_pushstring(lua,"pcall");
438 lua_pushcfunction(lua,luaRedisPCallCommand);
439 lua_settable(lua,-3);
440
441 /* redis.log and log levels. */
442 lua_pushstring(lua,"log");
443 lua_pushcfunction(lua,luaLogCommand);
444 lua_settable(lua,-3);
445
446 lua_pushstring(lua,"LOG_DEBUG");
447 lua_pushnumber(lua,REDIS_DEBUG);
448 lua_settable(lua,-3);
449
450 lua_pushstring(lua,"LOG_VERBOSE");
451 lua_pushnumber(lua,REDIS_VERBOSE);
452 lua_settable(lua,-3);
453
454 lua_pushstring(lua,"LOG_NOTICE");
455 lua_pushnumber(lua,REDIS_NOTICE);
456 lua_settable(lua,-3);
457
458 lua_pushstring(lua,"LOG_WARNING");
459 lua_pushnumber(lua,REDIS_WARNING);
460 lua_settable(lua,-3);
461
462 /* redis.sha1hex */
463 lua_pushstring(lua, "sha1hex");
464 lua_pushcfunction(lua, luaRedisSha1hexCommand);
465 lua_settable(lua, -3);
466
467 /* Finally set the table as 'redis' global var. */
468 lua_setglobal(lua,"redis");
469
470 /* Replace math.random and math.randomseed with our implementations. */
471 lua_getglobal(lua,"math");
472
473 lua_pushstring(lua,"random");
474 lua_pushcfunction(lua,redis_math_random);
475 lua_settable(lua,-3);
476
477 lua_pushstring(lua,"randomseed");
478 lua_pushcfunction(lua,redis_math_randomseed);
479 lua_settable(lua,-3);
480
481 lua_setglobal(lua,"math");
482
483 /* Add a helper funciton that we use to sort the multi bulk output of non
484 * deterministic commands, when containing 'false' elements. */
485 {
486 char *compare_func = "function __redis__compare_helper(a,b)\n"
487 " if a == false then a = '' end\n"
488 " if b == false then b = '' end\n"
489 " return a<b\n"
490 "end\n";
491 luaL_loadbuffer(lua,compare_func,strlen(compare_func),"cmp_func_def");
492 lua_pcall(lua,0,0,0);
493 }
494
495 /* Create the (non connected) client that we use to execute Redis commands
496 * inside the Lua interpreter.
497 * Note: there is no need to create it again when this function is called
498 * by scriptingReset(). */
499 if (server.lua_client == NULL) {
500 server.lua_client = createClient(-1);
501 server.lua_client->flags |= REDIS_LUA_CLIENT;
502 }
503
504 server.lua = lua;
505 }
506
507 /* Release resources related to Lua scripting.
508 * This function is used in order to reset the scripting environment. */
509 void scriptingRelease(void) {
510 dictRelease(server.lua_scripts);
511 lua_close(server.lua);
512 }
513
514 void scriptingReset(void) {
515 scriptingRelease();
516 scriptingInit();
517 }
518
519 /* Perform the SHA1 of the input string. We use this both for hasing script
520 * bodies in order to obtain the Lua function name, and in the implementation
521 * of redis.sha1().
522 *
523 * 'digest' should point to a 41 bytes buffer: 40 for SHA1 converted into an
524 * hexadecimal number, plus 1 byte for null term. */
525 void sha1hex(char *digest, char *script, size_t len) {
526 SHA1_CTX ctx;
527 unsigned char hash[20];
528 char *cset = "0123456789abcdef";
529 int j;
530
531 SHA1Init(&ctx);
532 SHA1Update(&ctx,(unsigned char*)script,len);
533 SHA1Final(hash,&ctx);
534
535 for (j = 0; j < 20; j++) {
536 digest[j*2] = cset[((hash[j]&0xF0)>>4)];
537 digest[j*2+1] = cset[(hash[j]&0xF)];
538 }
539 digest[40] = '\0';
540 }
541
542 void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
543 int t = lua_type(lua,-1);
544
545 switch(t) {
546 case LUA_TSTRING:
547 addReplyBulkCBuffer(c,(char*)lua_tostring(lua,-1),lua_strlen(lua,-1));
548 break;
549 case LUA_TBOOLEAN:
550 addReply(c,lua_toboolean(lua,-1) ? shared.cone : shared.nullbulk);
551 break;
552 case LUA_TNUMBER:
553 addReplyLongLong(c,(long long)lua_tonumber(lua,-1));
554 break;
555 case LUA_TTABLE:
556 /* We need to check if it is an array, an error, or a status reply.
557 * Error are returned as a single element table with 'err' field.
558 * Status replies are returned as single elment table with 'ok' field */
559 lua_pushstring(lua,"err");
560 lua_gettable(lua,-2);
561 t = lua_type(lua,-1);
562 if (t == LUA_TSTRING) {
563 sds err = sdsnew(lua_tostring(lua,-1));
564 sdsmapchars(err,"\r\n"," ",2);
565 addReplySds(c,sdscatprintf(sdsempty(),"-%s\r\n",err));
566 sdsfree(err);
567 lua_pop(lua,2);
568 return;
569 }
570
571 lua_pop(lua,1);
572 lua_pushstring(lua,"ok");
573 lua_gettable(lua,-2);
574 t = lua_type(lua,-1);
575 if (t == LUA_TSTRING) {
576 sds ok = sdsnew(lua_tostring(lua,-1));
577 sdsmapchars(ok,"\r\n"," ",2);
578 addReplySds(c,sdscatprintf(sdsempty(),"+%s\r\n",ok));
579 sdsfree(ok);
580 lua_pop(lua,1);
581 } else {
582 void *replylen = addDeferredMultiBulkLength(c);
583 int j = 1, mbulklen = 0;
584
585 lua_pop(lua,1); /* Discard the 'ok' field value we popped */
586 while(1) {
587 lua_pushnumber(lua,j++);
588 lua_gettable(lua,-2);
589 t = lua_type(lua,-1);
590 if (t == LUA_TNIL) {
591 lua_pop(lua,1);
592 break;
593 }
594 luaReplyToRedisReply(c, lua);
595 mbulklen++;
596 }
597 setDeferredMultiBulkLength(c,replylen,mbulklen);
598 }
599 break;
600 default:
601 addReply(c,shared.nullbulk);
602 }
603 lua_pop(lua,1);
604 }
605
606 /* Set an array of Redis String Objects as a Lua array (table) stored into a
607 * global variable. */
608 void luaSetGlobalArray(lua_State *lua, char *var, robj **elev, int elec) {
609 int j;
610
611 lua_newtable(lua);
612 for (j = 0; j < elec; j++) {
613 lua_pushlstring(lua,(char*)elev[j]->ptr,sdslen(elev[j]->ptr));
614 lua_rawseti(lua,-2,j+1);
615 }
616 lua_setglobal(lua,var);
617 }
618
619 /* Define a lua function with the specified function name and body.
620 * The function name musts be a 2 characters long string, since all the
621 * functions we defined in the Lua context are in the form:
622 *
623 * f_<hex sha1 sum>
624 *
625 * On success REDIS_OK is returned, and nothing is left on the Lua stack.
626 * On error REDIS_ERR is returned and an appropriate error is set in the
627 * client context. */
628 int luaCreateFunction(redisClient *c, lua_State *lua, char *funcname, robj *body) {
629 sds funcdef = sdsempty();
630
631 funcdef = sdscat(funcdef,"function ");
632 funcdef = sdscatlen(funcdef,funcname,42);
633 funcdef = sdscatlen(funcdef,"() ",3);
634 funcdef = sdscatlen(funcdef,body->ptr,sdslen(body->ptr));
635 funcdef = sdscatlen(funcdef," end",4);
636
637 if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) {
638 addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
639 lua_tostring(lua,-1));
640 lua_pop(lua,1);
641 sdsfree(funcdef);
642 return REDIS_ERR;
643 }
644 sdsfree(funcdef);
645 if (lua_pcall(lua,0,0,0)) {
646 addReplyErrorFormat(c,"Error running script (new function): %s\n",
647 lua_tostring(lua,-1));
648 lua_pop(lua,1);
649 return REDIS_ERR;
650 }
651
652 /* We also save a SHA1 -> Original script map in a dictionary
653 * so that we can replicate / write in the AOF all the
654 * EVALSHA commands as EVAL using the original script. */
655 {
656 int retval = dictAdd(server.lua_scripts,
657 sdsnewlen(funcname+2,40),body);
658 redisAssertWithInfo(c,NULL,retval == DICT_OK);
659 incrRefCount(body);
660 }
661 return REDIS_OK;
662 }
663
664 void evalGenericCommand(redisClient *c, int evalsha) {
665 lua_State *lua = server.lua;
666 char funcname[43];
667 long long numkeys;
668
669 /* We want the same PRNG sequence at every call so that our PRNG is
670 * not affected by external state. */
671 redisSrand48(0);
672
673 /* We set this flag to zero to remember that so far no random command
674 * was called. This way we can allow the user to call commands like
675 * SRANDMEMBER or RANDOMKEY from Lua scripts as far as no write command
676 * is called (otherwise the replication and AOF would end with non
677 * deterministic sequences).
678 *
679 * Thanks to this flag we'll raise an error every time a write command
680 * is called after a random command was used. */
681 server.lua_random_dirty = 0;
682 server.lua_write_dirty = 0;
683
684 /* Get the number of arguments that are keys */
685 if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK)
686 return;
687 if (numkeys > (c->argc - 3)) {
688 addReplyError(c,"Number of keys can't be greater than number of args");
689 return;
690 }
691
692 /* We obtain the script SHA1, then check if this function is already
693 * defined into the Lua state */
694 funcname[0] = 'f';
695 funcname[1] = '_';
696 if (!evalsha) {
697 /* Hash the code if this is an EVAL call */
698 sha1hex(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
699 } else {
700 /* We already have the SHA if it is a EVALSHA */
701 int j;
702 char *sha = c->argv[1]->ptr;
703
704 for (j = 0; j < 40; j++)
705 funcname[j+2] = tolower(sha[j]);
706 funcname[42] = '\0';
707 }
708
709 /* Try to lookup the Lua function */
710 lua_getglobal(lua, funcname);
711 if (lua_isnil(lua,1)) {
712 lua_pop(lua,1); /* remove the nil from the stack */
713 /* Function not defined... let's define it if we have the
714 * body of the funciton. If this is an EVALSHA call we can just
715 * return an error. */
716 if (evalsha) {
717 addReply(c, shared.noscripterr);
718 return;
719 }
720 if (luaCreateFunction(c,lua,funcname,c->argv[1]) == REDIS_ERR) return;
721 /* Now the following is guaranteed to return non nil */
722 lua_getglobal(lua, funcname);
723 redisAssert(!lua_isnil(lua,1));
724 }
725
726 /* Populate the argv and keys table accordingly to the arguments that
727 * EVAL received. */
728 luaSetGlobalArray(lua,"KEYS",c->argv+3,numkeys);
729 luaSetGlobalArray(lua,"ARGV",c->argv+3+numkeys,c->argc-3-numkeys);
730
731 /* Select the right DB in the context of the Lua client */
732 selectDb(server.lua_client,c->db->id);
733
734 /* Set an hook in order to be able to stop the script execution if it
735 * is running for too much time.
736 * We set the hook only if the time limit is enabled as the hook will
737 * make the Lua script execution slower. */
738 if (server.lua_time_limit > 0 && server.masterhost == NULL) {
739 lua_sethook(lua,luaMaskCountHook,LUA_MASKCOUNT,100000);
740 } else {
741 lua_sethook(lua,luaMaskCountHook,0,0);
742 }
743
744 /* At this point whatever this script was never seen before or if it was
745 * already defined, we can call it. We have zero arguments and expect
746 * a single return value. */
747 server.lua_caller = c;
748 server.lua_time_start = ustime()/1000;
749 server.lua_kill = 0;
750 if (lua_pcall(lua,0,1,0)) {
751 if (server.lua_timedout) {
752 server.lua_timedout = 0;
753 /* Restore the readable handler that was unregistered when the
754 * script timeout was detected. */
755 aeCreateFileEvent(server.el,c->fd,AE_READABLE,
756 readQueryFromClient,c);
757 }
758 server.lua_caller = NULL;
759 selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
760 addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
761 funcname, lua_tostring(lua,-1));
762 lua_pop(lua,1);
763 lua_gc(lua,LUA_GCCOLLECT,0);
764 return;
765 }
766 server.lua_timedout = 0;
767 server.lua_caller = NULL;
768 selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
769 luaReplyToRedisReply(c,lua);
770 lua_gc(lua,LUA_GCSTEP,1);
771
772 /* If we have slaves attached we want to replicate this command as
773 * EVAL instead of EVALSHA. We do this also in the AOF as currently there
774 * is no easy way to propagate a command in a different way in the AOF
775 * and in the replication link.
776 *
777 * IMPROVEMENT POSSIBLE:
778 * 1) Replicate this command as EVALSHA in the AOF.
779 * 2) Remember what slave already received a given script, and replicate
780 * the EVALSHA against this slaves when possible.
781 */
782 if (evalsha) {
783 robj *script = dictFetchValue(server.lua_scripts,c->argv[1]->ptr);
784
785 redisAssertWithInfo(c,NULL,script != NULL);
786 rewriteClientCommandArgument(c,0,
787 resetRefCount(createStringObject("EVAL",4)));
788 rewriteClientCommandArgument(c,1,script);
789 }
790 }
791
792 void evalCommand(redisClient *c) {
793 evalGenericCommand(c,0);
794 }
795
796 void evalShaCommand(redisClient *c) {
797 if (sdslen(c->argv[1]->ptr) != 40) {
798 /* We know that a match is not possible if the provided SHA is
799 * not the right length. So we return an error ASAP, this way
800 * evalGenericCommand() can be implemented without string length
801 * sanity check */
802 addReply(c, shared.noscripterr);
803 return;
804 }
805 evalGenericCommand(c,1);
806 }
807
808 /* We replace math.random() with our implementation that is not affected
809 * by specific libc random() implementations and will output the same sequence
810 * (for the same seed) in every arch. */
811
812 /* The following implementation is the one shipped with Lua itself but with
813 * rand() replaced by redisLrand48(). */
814 int redis_math_random (lua_State *L) {
815 /* the `%' avoids the (rare) case of r==1, and is needed also because on
816 some systems (SunOS!) `rand()' may return a value larger than RAND_MAX */
817 lua_Number r = (lua_Number)(redisLrand48()%REDIS_LRAND48_MAX) /
818 (lua_Number)REDIS_LRAND48_MAX;
819 switch (lua_gettop(L)) { /* check number of arguments */
820 case 0: { /* no arguments */
821 lua_pushnumber(L, r); /* Number between 0 and 1 */
822 break;
823 }
824 case 1: { /* only upper limit */
825 int u = luaL_checkint(L, 1);
826 luaL_argcheck(L, 1<=u, 1, "interval is empty");
827 lua_pushnumber(L, floor(r*u)+1); /* int between 1 and `u' */
828 break;
829 }
830 case 2: { /* lower and upper limits */
831 int l = luaL_checkint(L, 1);
832 int u = luaL_checkint(L, 2);
833 luaL_argcheck(L, l<=u, 2, "interval is empty");
834 lua_pushnumber(L, floor(r*(u-l+1))+l); /* int between `l' and `u' */
835 break;
836 }
837 default: return luaL_error(L, "wrong number of arguments");
838 }
839 return 1;
840 }
841
842 int redis_math_randomseed (lua_State *L) {
843 redisSrand48(luaL_checkint(L, 1));
844 return 0;
845 }
846
847 /* ---------------------------------------------------------------------------
848 * SCRIPT command for script environment introspection and control
849 * ------------------------------------------------------------------------- */
850
851 void scriptCommand(redisClient *c) {
852 if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"flush")) {
853 scriptingReset();
854 addReply(c,shared.ok);
855 server.dirty++; /* Replicating this command is a good idea. */
856 } else if (c->argc >= 2 && !strcasecmp(c->argv[1]->ptr,"exists")) {
857 int j;
858
859 addReplyMultiBulkLen(c, c->argc-2);
860 for (j = 2; j < c->argc; j++) {
861 if (dictFind(server.lua_scripts,c->argv[j]->ptr))
862 addReply(c,shared.cone);
863 else
864 addReply(c,shared.czero);
865 }
866 } else if (c->argc == 3 && !strcasecmp(c->argv[1]->ptr,"load")) {
867 char funcname[43];
868 sds sha;
869
870 funcname[0] = 'f';
871 funcname[1] = '_';
872 sha1hex(funcname+2,c->argv[2]->ptr,sdslen(c->argv[2]->ptr));
873 sha = sdsnewlen(funcname+2,40);
874 if (dictFind(server.lua_scripts,sha) == NULL) {
875 if (luaCreateFunction(c,server.lua,funcname,c->argv[2])
876 == REDIS_ERR) {
877 sdsfree(sha);
878 return;
879 }
880 }
881 addReplyBulkCBuffer(c,funcname+2,40);
882 sdsfree(sha);
883 } else if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"kill")) {
884 if (server.lua_caller == NULL) {
885 addReplyError(c,"No scripts in execution right now.");
886 } else if (server.lua_write_dirty) {
887 addReplyError(c, "Sorry the script already executed write commands against the dataset. You can either wait the script termination or kill the server in an hard way using the SHUTDOWN NOSAVE command.");
888 } else {
889 server.lua_kill = 1;
890 addReply(c,shared.ok);
891 }
892 } else {
893 addReplyError(c, "Unknown SCRIPT subcommand or wrong # of args.");
894 }
895 }