]>
Commit | Line | Data |
---|---|---|
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 | ||
19 | /* Take a Redis reply in the Redis protocol format and convert it into a | |
20 | * Lua type. Thanks to this function, and the introduction of not connected | |
21 | * clients, it is trvial to implement the redis() lua function. | |
22 | * | |
23 | * Basically we take the arguments, execute the Redis command in the context | |
24 | * of a non connected client, then take the generated reply and convert it | |
25 | * into a suitable Lua type. With this trick the scripting feature does not | |
26 | * need the introduction of a full Redis internals API. Basically the script | |
27 | * is like a normal client that bypasses all the slow I/O paths. | |
28 | * | |
29 | * Note: in this function we do not do any sanity check as the reply is | |
30 | * generated by Redis directly. This allows use to go faster. | |
31 | * The reply string can be altered during the parsing as it is discared | |
32 | * after the conversion is completed. | |
33 | * | |
34 | * Errors are returned as a table with a single 'err' field set to the | |
35 | * error string. | |
36 | */ | |
37 | ||
38 | char *redisProtocolToLuaType(lua_State *lua, char* reply) { | |
39 | char *p = reply; | |
40 | ||
41 | switch(*p) { | |
42 | case ':': | |
43 | p = redisProtocolToLuaType_Int(lua,reply); | |
44 | break; | |
45 | case '$': | |
46 | p = redisProtocolToLuaType_Bulk(lua,reply); | |
47 | break; | |
48 | case '+': | |
49 | p = redisProtocolToLuaType_Status(lua,reply); | |
50 | break; | |
51 | case '-': | |
52 | p = redisProtocolToLuaType_Error(lua,reply); | |
53 | break; | |
54 | case '*': | |
55 | p = redisProtocolToLuaType_MultiBulk(lua,reply); | |
56 | break; | |
57 | } | |
58 | return p; | |
59 | } | |
60 | ||
61 | char *redisProtocolToLuaType_Int(lua_State *lua, char *reply) { | |
62 | char *p = strchr(reply+1,'\r'); | |
63 | long long value; | |
64 | ||
65 | string2ll(reply+1,p-reply-1,&value); | |
66 | lua_pushnumber(lua,(lua_Number)value); | |
67 | return p+2; | |
68 | } | |
69 | ||
70 | char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply) { | |
71 | char *p = strchr(reply+1,'\r'); | |
72 | long long bulklen; | |
73 | ||
74 | string2ll(reply+1,p-reply-1,&bulklen); | |
75 | if (bulklen == -1) { | |
76 | lua_pushboolean(lua,0); | |
77 | return p+2; | |
78 | } else { | |
79 | lua_pushlstring(lua,p+2,bulklen); | |
80 | return p+2+bulklen+2; | |
81 | } | |
82 | } | |
83 | ||
84 | char *redisProtocolToLuaType_Status(lua_State *lua, char *reply) { | |
85 | char *p = strchr(reply+1,'\r'); | |
86 | ||
87 | lua_newtable(lua); | |
88 | lua_pushstring(lua,"ok"); | |
89 | lua_pushlstring(lua,reply+1,p-reply-1); | |
90 | lua_settable(lua,-3); | |
91 | return p+2; | |
92 | } | |
93 | ||
94 | char *redisProtocolToLuaType_Error(lua_State *lua, char *reply) { | |
95 | char *p = strchr(reply+1,'\r'); | |
96 | ||
97 | lua_newtable(lua); | |
98 | lua_pushstring(lua,"err"); | |
99 | lua_pushlstring(lua,reply+1,p-reply-1); | |
100 | lua_settable(lua,-3); | |
101 | return p+2; | |
102 | } | |
103 | ||
104 | char *redisProtocolToLuaType_MultiBulk(lua_State *lua, char *reply) { | |
105 | char *p = strchr(reply+1,'\r'); | |
106 | long long mbulklen; | |
107 | int j = 0; | |
108 | ||
109 | string2ll(reply+1,p-reply-1,&mbulklen); | |
110 | p += 2; | |
111 | if (mbulklen == -1) { | |
112 | lua_pushboolean(lua,0); | |
113 | return p; | |
114 | } | |
115 | lua_newtable(lua); | |
116 | for (j = 0; j < mbulklen; j++) { | |
117 | lua_pushnumber(lua,j+1); | |
118 | p = redisProtocolToLuaType(lua,p); | |
119 | lua_settable(lua,-3); | |
120 | } | |
121 | return p; | |
122 | } | |
123 | ||
124 | void luaPushError(lua_State *lua, char *error) { | |
125 | lua_newtable(lua); | |
126 | lua_pushstring(lua,"err"); | |
127 | lua_pushstring(lua, error); | |
128 | lua_settable(lua,-3); | |
129 | } | |
130 | ||
131 | int luaRedisGenericCommand(lua_State *lua, int raise_error) { | |
132 | int j, argc = lua_gettop(lua); | |
133 | struct redisCommand *cmd; | |
134 | robj **argv; | |
135 | redisClient *c = server.lua_client; | |
136 | sds reply; | |
137 | ||
138 | /* Build the arguments vector */ | |
139 | argv = zmalloc(sizeof(robj*)*argc); | |
140 | for (j = 0; j < argc; j++) { | |
141 | if (!lua_isstring(lua,j+1)) break; | |
142 | argv[j] = createStringObject((char*)lua_tostring(lua,j+1), | |
143 | lua_strlen(lua,j+1)); | |
144 | } | |
145 | ||
146 | /* Check if one of the arguments passed by the Lua script | |
147 | * is not a string or an integer (lua_isstring() return true for | |
148 | * integers as well). */ | |
149 | if (j != argc) { | |
150 | j--; | |
151 | while (j >= 0) { | |
152 | decrRefCount(argv[j]); | |
153 | j--; | |
154 | } | |
155 | zfree(argv); | |
156 | luaPushError(lua, | |
157 | "Lua redis() command arguments must be strings or integers"); | |
158 | return 1; | |
159 | } | |
160 | ||
161 | /* Setup our fake client for command execution */ | |
162 | c->argv = argv; | |
163 | c->argc = argc; | |
164 | ||
165 | /* Command lookup */ | |
166 | cmd = lookupCommand(argv[0]->ptr); | |
167 | if (!cmd || ((cmd->arity > 0 && cmd->arity != argc) || | |
168 | (argc < -cmd->arity))) | |
169 | { | |
170 | if (cmd) | |
171 | luaPushError(lua, | |
172 | "Wrong number of args calling Redis command From Lua script"); | |
173 | else | |
174 | luaPushError(lua,"Unknown Redis command called from Lua script"); | |
175 | goto cleanup; | |
176 | } | |
177 | ||
178 | if (cmd->flags & REDIS_CMD_NOSCRIPT) { | |
179 | luaPushError(lua, "This Redis command is not allowed from scripts"); | |
180 | goto cleanup; | |
181 | } | |
182 | ||
183 | if (cmd->flags & REDIS_CMD_WRITE && server.lua_random_dirty) { | |
184 | luaPushError(lua, | |
185 | "Write commands not allowed after non deterministic commands"); | |
186 | goto cleanup; | |
187 | } | |
188 | ||
189 | if (cmd->flags & REDIS_CMD_RANDOM) server.lua_random_dirty = 1; | |
190 | if (cmd->flags & REDIS_CMD_WRITE) server.lua_write_dirty = 1; | |
191 | ||
192 | /* Run the command */ | |
193 | cmd->proc(c); | |
194 | ||
195 | /* Convert the result of the Redis command into a suitable Lua type. | |
196 | * The first thing we need is to create a single string from the client | |
197 | * output buffers. */ | |
198 | reply = sdsempty(); | |
199 | if (c->bufpos) { | |
200 | reply = sdscatlen(reply,c->buf,c->bufpos); | |
201 | c->bufpos = 0; | |
202 | } | |
203 | while(listLength(c->reply)) { | |
204 | robj *o = listNodeValue(listFirst(c->reply)); | |
205 | ||
206 | reply = sdscatlen(reply,o->ptr,sdslen(o->ptr)); | |
207 | listDelNode(c->reply,listFirst(c->reply)); | |
208 | } | |
209 | if (raise_error && reply[0] != '-') raise_error = 0; | |
210 | redisProtocolToLuaType(lua,reply); | |
211 | sdsfree(reply); | |
212 | ||
213 | cleanup: | |
214 | /* Clean up. Command code may have changed argv/argc so we use the | |
215 | * argv/argc of the client instead of the local variables. */ | |
216 | for (j = 0; j < c->argc; j++) | |
217 | decrRefCount(c->argv[j]); | |
218 | zfree(c->argv); | |
219 | ||
220 | if (raise_error) { | |
221 | /* If we are here we should have an error in the stack, in the | |
222 | * form of a table with an "err" field. Extract the string to | |
223 | * return the plain error. */ | |
224 | lua_pushstring(lua,"err"); | |
225 | lua_gettable(lua,-2); | |
226 | return lua_error(lua); | |
227 | } | |
228 | return 1; | |
229 | } | |
230 | ||
231 | int luaRedisCallCommand(lua_State *lua) { | |
232 | return luaRedisGenericCommand(lua,1); | |
233 | } | |
234 | ||
235 | int luaRedisPCallCommand(lua_State *lua) { | |
236 | return luaRedisGenericCommand(lua,0); | |
237 | } | |
238 | ||
239 | int luaLogCommand(lua_State *lua) { | |
240 | int j, argc = lua_gettop(lua); | |
241 | int level; | |
242 | sds log; | |
243 | ||
244 | if (argc < 2) { | |
245 | luaPushError(lua, "redis.log() requires two arguments or more."); | |
246 | return 1; | |
247 | } else if (!lua_isnumber(lua,-argc)) { | |
248 | luaPushError(lua, "First argument must be a number (log level)."); | |
249 | return 1; | |
250 | } | |
251 | level = lua_tonumber(lua,-argc); | |
252 | if (level < REDIS_DEBUG || level > REDIS_WARNING) { | |
253 | luaPushError(lua, "Invalid debug level."); | |
254 | return 1; | |
255 | } | |
256 | ||
257 | /* Glue together all the arguments */ | |
258 | log = sdsempty(); | |
259 | for (j = 1; j < argc; j++) { | |
260 | size_t len; | |
261 | char *s; | |
262 | ||
263 | s = (char*)lua_tolstring(lua,(-argc)+j,&len); | |
264 | if (s) { | |
265 | if (j != 1) log = sdscatlen(log," ",1); | |
266 | log = sdscatlen(log,s,len); | |
267 | } | |
268 | } | |
269 | redisLogRaw(level,log); | |
270 | sdsfree(log); | |
271 | return 0; | |
272 | } | |
273 | ||
274 | void luaMaskCountHook(lua_State *lua, lua_Debug *ar) { | |
275 | long long elapsed; | |
276 | REDIS_NOTUSED(ar); | |
277 | REDIS_NOTUSED(lua); | |
278 | ||
279 | elapsed = (ustime()/1000) - server.lua_time_start; | |
280 | if (elapsed >= server.lua_time_limit && server.lua_timedout == 0) { | |
281 | 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); | |
282 | server.lua_timedout = 1; | |
283 | /* Once the script timeouts we reenter the event loop to permit others | |
284 | * to call SCRIPT KILL or SHUTDOWN NOSAVE if needed. For this reason | |
285 | * we need to mask the client executing the script from the event loop. | |
286 | * If we don't do that the client may disconnect and could no longer be | |
287 | * here when the EVAL command will return. */ | |
288 | aeDeleteFileEvent(server.el, server.lua_caller->fd, AE_READABLE); | |
289 | } | |
290 | if (server.lua_timedout) | |
291 | aeProcessEvents(server.el, AE_FILE_EVENTS|AE_DONT_WAIT); | |
292 | if (server.lua_kill) { | |
293 | redisLog(REDIS_WARNING,"Lua script killed by user with SCRIPT KILL."); | |
294 | lua_pushstring(lua,"Script killed by user with SCRIPT KILL..."); | |
295 | lua_error(lua); | |
296 | } | |
297 | } | |
298 | ||
299 | void luaLoadLib(lua_State *lua, const char *libname, lua_CFunction luafunc) { | |
300 | lua_pushcfunction(lua, luafunc); | |
301 | lua_pushstring(lua, libname); | |
302 | lua_call(lua, 1, 0); | |
303 | } | |
304 | ||
305 | LUALIB_API int (luaopen_cjson) (lua_State *L); | |
306 | ||
307 | void luaLoadLibraries(lua_State *lua) { | |
308 | luaLoadLib(lua, "", luaopen_base); | |
309 | luaLoadLib(lua, LUA_TABLIBNAME, luaopen_table); | |
310 | luaLoadLib(lua, LUA_STRLIBNAME, luaopen_string); | |
311 | luaLoadLib(lua, LUA_MATHLIBNAME, luaopen_math); | |
312 | luaLoadLib(lua, LUA_DBLIBNAME, luaopen_debug); | |
313 | luaLoadLib(lua, "cjson", luaopen_cjson); | |
314 | ||
315 | #if 0 /* Stuff that we don't load currently, for sandboxing concerns. */ | |
316 | luaLoadLib(lua, LUA_LOADLIBNAME, luaopen_package); | |
317 | luaLoadLib(lua, LUA_OSLIBNAME, luaopen_os); | |
318 | #endif | |
319 | } | |
320 | ||
321 | /* Initialize the scripting environment. | |
322 | * It is possible to call this function to reset the scripting environment | |
323 | * assuming that we call scriptingRelease() before. | |
324 | * See scriptingReset() for more information. */ | |
325 | void scriptingInit(void) { | |
326 | lua_State *lua = lua_open(); | |
327 | luaLoadLibraries(lua); | |
328 | ||
329 | /* Initialize a dictionary we use to map SHAs to scripts. | |
330 | * This is useful for replication, as we need to replicate EVALSHA | |
331 | * as EVAL, so we need to remember the associated script. */ | |
332 | server.lua_scripts = dictCreate(&dbDictType,NULL); | |
333 | ||
334 | /* Register the redis commands table and fields */ | |
335 | lua_newtable(lua); | |
336 | ||
337 | /* redis.call */ | |
338 | lua_pushstring(lua,"call"); | |
339 | lua_pushcfunction(lua,luaRedisCallCommand); | |
340 | lua_settable(lua,-3); | |
341 | ||
342 | /* redis.pcall */ | |
343 | lua_pushstring(lua,"pcall"); | |
344 | lua_pushcfunction(lua,luaRedisPCallCommand); | |
345 | lua_settable(lua,-3); | |
346 | ||
347 | /* redis.log and log levels. */ | |
348 | lua_pushstring(lua,"log"); | |
349 | lua_pushcfunction(lua,luaLogCommand); | |
350 | lua_settable(lua,-3); | |
351 | ||
352 | lua_pushstring(lua,"LOG_DEBUG"); | |
353 | lua_pushnumber(lua,REDIS_DEBUG); | |
354 | lua_settable(lua,-3); | |
355 | ||
356 | lua_pushstring(lua,"LOG_VERBOSE"); | |
357 | lua_pushnumber(lua,REDIS_VERBOSE); | |
358 | lua_settable(lua,-3); | |
359 | ||
360 | lua_pushstring(lua,"LOG_NOTICE"); | |
361 | lua_pushnumber(lua,REDIS_NOTICE); | |
362 | lua_settable(lua,-3); | |
363 | ||
364 | lua_pushstring(lua,"LOG_WARNING"); | |
365 | lua_pushnumber(lua,REDIS_WARNING); | |
366 | lua_settable(lua,-3); | |
367 | ||
368 | /* Finally set the table as 'redis' global var. */ | |
369 | lua_setglobal(lua,"redis"); | |
370 | ||
371 | /* Replace math.random and math.randomseed with our implementations. */ | |
372 | lua_getglobal(lua,"math"); | |
373 | ||
374 | lua_pushstring(lua,"random"); | |
375 | lua_pushcfunction(lua,redis_math_random); | |
376 | lua_settable(lua,-3); | |
377 | ||
378 | lua_pushstring(lua,"randomseed"); | |
379 | lua_pushcfunction(lua,redis_math_randomseed); | |
380 | lua_settable(lua,-3); | |
381 | ||
382 | lua_setglobal(lua,"math"); | |
383 | ||
384 | /* Create the (non connected) client that we use to execute Redis commands | |
385 | * inside the Lua interpreter. | |
386 | * Note: there is no need to create it again when this function is called | |
387 | * by scriptingReset(). */ | |
388 | if (server.lua_client == NULL) { | |
389 | server.lua_client = createClient(-1); | |
390 | server.lua_client->flags |= REDIS_LUA_CLIENT; | |
391 | } | |
392 | ||
393 | server.lua = lua; | |
394 | } | |
395 | ||
396 | /* Release resources related to Lua scripting. | |
397 | * This function is used in order to reset the scripting environment. */ | |
398 | void scriptingRelease(void) { | |
399 | dictRelease(server.lua_scripts); | |
400 | lua_close(server.lua); | |
401 | } | |
402 | ||
403 | void scriptingReset(void) { | |
404 | scriptingRelease(); | |
405 | scriptingInit(); | |
406 | } | |
407 | ||
408 | /* Hash the scripit into a SHA1 digest. We use this as Lua function name. | |
409 | * Digest should point to a 41 bytes buffer: 40 for SHA1 converted into an | |
410 | * hexadecimal number, plus 1 byte for null term. */ | |
411 | void hashScript(char *digest, char *script, size_t len) { | |
412 | SHA1_CTX ctx; | |
413 | unsigned char hash[20]; | |
414 | char *cset = "0123456789abcdef"; | |
415 | int j; | |
416 | ||
417 | SHA1Init(&ctx); | |
418 | SHA1Update(&ctx,(unsigned char*)script,len); | |
419 | SHA1Final(hash,&ctx); | |
420 | ||
421 | for (j = 0; j < 20; j++) { | |
422 | digest[j*2] = cset[((hash[j]&0xF0)>>4)]; | |
423 | digest[j*2+1] = cset[(hash[j]&0xF)]; | |
424 | } | |
425 | digest[40] = '\0'; | |
426 | } | |
427 | ||
428 | void luaReplyToRedisReply(redisClient *c, lua_State *lua) { | |
429 | int t = lua_type(lua,-1); | |
430 | ||
431 | switch(t) { | |
432 | case LUA_TSTRING: | |
433 | addReplyBulkCBuffer(c,(char*)lua_tostring(lua,-1),lua_strlen(lua,-1)); | |
434 | break; | |
435 | case LUA_TBOOLEAN: | |
436 | addReply(c,lua_toboolean(lua,-1) ? shared.cone : shared.nullbulk); | |
437 | break; | |
438 | case LUA_TNUMBER: | |
439 | addReplyLongLong(c,(long long)lua_tonumber(lua,-1)); | |
440 | break; | |
441 | case LUA_TTABLE: | |
442 | /* We need to check if it is an array, an error, or a status reply. | |
443 | * Error are returned as a single element table with 'err' field. | |
444 | * Status replies are returned as single elment table with 'ok' field */ | |
445 | lua_pushstring(lua,"err"); | |
446 | lua_gettable(lua,-2); | |
447 | t = lua_type(lua,-1); | |
448 | if (t == LUA_TSTRING) { | |
449 | sds err = sdsnew(lua_tostring(lua,-1)); | |
450 | sdsmapchars(err,"\r\n"," ",2); | |
451 | addReplySds(c,sdscatprintf(sdsempty(),"-%s\r\n",err)); | |
452 | sdsfree(err); | |
453 | lua_pop(lua,2); | |
454 | return; | |
455 | } | |
456 | ||
457 | lua_pop(lua,1); | |
458 | lua_pushstring(lua,"ok"); | |
459 | lua_gettable(lua,-2); | |
460 | t = lua_type(lua,-1); | |
461 | if (t == LUA_TSTRING) { | |
462 | sds ok = sdsnew(lua_tostring(lua,-1)); | |
463 | sdsmapchars(ok,"\r\n"," ",2); | |
464 | addReplySds(c,sdscatprintf(sdsempty(),"+%s\r\n",ok)); | |
465 | sdsfree(ok); | |
466 | lua_pop(lua,1); | |
467 | } else { | |
468 | void *replylen = addDeferredMultiBulkLength(c); | |
469 | int j = 1, mbulklen = 0; | |
470 | ||
471 | lua_pop(lua,1); /* Discard the 'ok' field value we popped */ | |
472 | while(1) { | |
473 | lua_pushnumber(lua,j++); | |
474 | lua_gettable(lua,-2); | |
475 | t = lua_type(lua,-1); | |
476 | if (t == LUA_TNIL) { | |
477 | lua_pop(lua,1); | |
478 | break; | |
479 | } | |
480 | luaReplyToRedisReply(c, lua); | |
481 | mbulklen++; | |
482 | } | |
483 | setDeferredMultiBulkLength(c,replylen,mbulklen); | |
484 | } | |
485 | break; | |
486 | default: | |
487 | addReply(c,shared.nullbulk); | |
488 | } | |
489 | lua_pop(lua,1); | |
490 | } | |
491 | ||
492 | /* Set an array of Redis String Objects as a Lua array (table) stored into a | |
493 | * global variable. */ | |
494 | void luaSetGlobalArray(lua_State *lua, char *var, robj **elev, int elec) { | |
495 | int j; | |
496 | ||
497 | lua_newtable(lua); | |
498 | for (j = 0; j < elec; j++) { | |
499 | lua_pushlstring(lua,(char*)elev[j]->ptr,sdslen(elev[j]->ptr)); | |
500 | lua_rawseti(lua,-2,j+1); | |
501 | } | |
502 | lua_setglobal(lua,var); | |
503 | } | |
504 | ||
505 | /* Define a lua function with the specified function name and body. | |
506 | * The function name musts be a 2 characters long string, since all the | |
507 | * functions we defined in the Lua context are in the form: | |
508 | * | |
509 | * f_<hex sha1 sum> | |
510 | * | |
511 | * On success REDIS_OK is returned, and nothing is left on the Lua stack. | |
512 | * On error REDIS_ERR is returned and an appropriate error is set in the | |
513 | * client context. */ | |
514 | int luaCreateFunction(redisClient *c, lua_State *lua, char *funcname, robj *body) { | |
515 | sds funcdef = sdsempty(); | |
516 | ||
517 | funcdef = sdscat(funcdef,"function "); | |
518 | funcdef = sdscatlen(funcdef,funcname,42); | |
519 | funcdef = sdscatlen(funcdef," ()\n",4); | |
520 | funcdef = sdscatlen(funcdef,body->ptr,sdslen(body->ptr)); | |
521 | funcdef = sdscatlen(funcdef,"\nend\n",5); | |
522 | ||
523 | if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) { | |
524 | addReplyErrorFormat(c,"Error compiling script (new function): %s\n", | |
525 | lua_tostring(lua,-1)); | |
526 | lua_pop(lua,1); | |
527 | sdsfree(funcdef); | |
528 | return REDIS_ERR; | |
529 | } | |
530 | sdsfree(funcdef); | |
531 | if (lua_pcall(lua,0,0,0)) { | |
532 | addReplyErrorFormat(c,"Error running script (new function): %s\n", | |
533 | lua_tostring(lua,-1)); | |
534 | lua_pop(lua,1); | |
535 | return REDIS_ERR; | |
536 | } | |
537 | ||
538 | /* We also save a SHA1 -> Original script map in a dictionary | |
539 | * so that we can replicate / write in the AOF all the | |
540 | * EVALSHA commands as EVAL using the original script. */ | |
541 | { | |
542 | int retval = dictAdd(server.lua_scripts, | |
543 | sdsnewlen(funcname+2,40),body); | |
544 | redisAssertWithInfo(c,NULL,retval == DICT_OK); | |
545 | incrRefCount(body); | |
546 | } | |
547 | return REDIS_OK; | |
548 | } | |
549 | ||
550 | void evalGenericCommand(redisClient *c, int evalsha) { | |
551 | lua_State *lua = server.lua; | |
552 | char funcname[43]; | |
553 | long long numkeys; | |
554 | ||
555 | /* We want the same PRNG sequence at every call so that our PRNG is | |
556 | * not affected by external state. */ | |
557 | redisSrand48(0); | |
558 | ||
559 | /* We set this flag to zero to remember that so far no random command | |
560 | * was called. This way we can allow the user to call commands like | |
561 | * SRANDMEMBER or RANDOMKEY from Lua scripts as far as no write command | |
562 | * is called (otherwise the replication and AOF would end with non | |
563 | * deterministic sequences). | |
564 | * | |
565 | * Thanks to this flag we'll raise an error every time a write command | |
566 | * is called after a random command was used. */ | |
567 | server.lua_random_dirty = 0; | |
568 | server.lua_write_dirty = 0; | |
569 | ||
570 | /* Get the number of arguments that are keys */ | |
571 | if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK) | |
572 | return; | |
573 | if (numkeys > (c->argc - 3)) { | |
574 | addReplyError(c,"Number of keys can't be greater than number of args"); | |
575 | return; | |
576 | } | |
577 | ||
578 | /* We obtain the script SHA1, then check if this function is already | |
579 | * defined into the Lua state */ | |
580 | funcname[0] = 'f'; | |
581 | funcname[1] = '_'; | |
582 | if (!evalsha) { | |
583 | /* Hash the code if this is an EVAL call */ | |
584 | hashScript(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr)); | |
585 | } else { | |
586 | /* We already have the SHA if it is a EVALSHA */ | |
587 | int j; | |
588 | char *sha = c->argv[1]->ptr; | |
589 | ||
590 | for (j = 0; j < 40; j++) | |
591 | funcname[j+2] = tolower(sha[j]); | |
592 | funcname[42] = '\0'; | |
593 | } | |
594 | ||
595 | /* Try to lookup the Lua function */ | |
596 | lua_getglobal(lua, funcname); | |
597 | if (lua_isnil(lua,1)) { | |
598 | lua_pop(lua,1); /* remove the nil from the stack */ | |
599 | /* Function not defined... let's define it if we have the | |
600 | * body of the funciton. If this is an EVALSHA call we can just | |
601 | * return an error. */ | |
602 | if (evalsha) { | |
603 | addReply(c, shared.noscripterr); | |
604 | return; | |
605 | } | |
606 | if (luaCreateFunction(c,lua,funcname,c->argv[1]) == REDIS_ERR) return; | |
607 | /* Now the following is guaranteed to return non nil */ | |
608 | lua_getglobal(lua, funcname); | |
609 | redisAssert(!lua_isnil(lua,1)); | |
610 | } | |
611 | ||
612 | /* Populate the argv and keys table accordingly to the arguments that | |
613 | * EVAL received. */ | |
614 | luaSetGlobalArray(lua,"KEYS",c->argv+3,numkeys); | |
615 | luaSetGlobalArray(lua,"ARGV",c->argv+3+numkeys,c->argc-3-numkeys); | |
616 | ||
617 | /* Select the right DB in the context of the Lua client */ | |
618 | selectDb(server.lua_client,c->db->id); | |
619 | ||
620 | /* Set an hook in order to be able to stop the script execution if it | |
621 | * is running for too much time. | |
622 | * We set the hook only if the time limit is enabled as the hook will | |
623 | * make the Lua script execution slower. */ | |
624 | if (server.lua_time_limit > 0 && server.masterhost == NULL) { | |
625 | lua_sethook(lua,luaMaskCountHook,LUA_MASKCOUNT,100000); | |
626 | } else { | |
627 | lua_sethook(lua,luaMaskCountHook,0,0); | |
628 | } | |
629 | ||
630 | /* At this point whatever this script was never seen before or if it was | |
631 | * already defined, we can call it. We have zero arguments and expect | |
632 | * a single return value. */ | |
633 | server.lua_caller = c; | |
634 | server.lua_time_start = ustime()/1000; | |
635 | server.lua_kill = 0; | |
636 | if (lua_pcall(lua,0,1,0)) { | |
637 | if (server.lua_timedout) { | |
638 | server.lua_timedout = 0; | |
639 | /* Restore the readable handler that was unregistered when the | |
640 | * script timeout was detected. */ | |
641 | aeCreateFileEvent(server.el,c->fd,AE_READABLE, | |
642 | readQueryFromClient,c); | |
643 | } | |
644 | server.lua_caller = NULL; | |
645 | selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */ | |
646 | addReplyErrorFormat(c,"Error running script (call to %s): %s\n", | |
647 | funcname, lua_tostring(lua,-1)); | |
648 | lua_pop(lua,1); | |
649 | lua_gc(lua,LUA_GCCOLLECT,0); | |
650 | return; | |
651 | } | |
652 | server.lua_timedout = 0; | |
653 | server.lua_caller = NULL; | |
654 | selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */ | |
655 | luaReplyToRedisReply(c,lua); | |
656 | lua_gc(lua,LUA_GCSTEP,1); | |
657 | ||
658 | /* If we have slaves attached we want to replicate this command as | |
659 | * EVAL instead of EVALSHA. We do this also in the AOF as currently there | |
660 | * is no easy way to propagate a command in a different way in the AOF | |
661 | * and in the replication link. | |
662 | * | |
663 | * IMPROVEMENT POSSIBLE: | |
664 | * 1) Replicate this command as EVALSHA in the AOF. | |
665 | * 2) Remember what slave already received a given script, and replicate | |
666 | * the EVALSHA against this slaves when possible. | |
667 | */ | |
668 | if (evalsha) { | |
669 | robj *script = dictFetchValue(server.lua_scripts,c->argv[1]->ptr); | |
670 | ||
671 | redisAssertWithInfo(c,NULL,script != NULL); | |
672 | rewriteClientCommandArgument(c,0, | |
673 | resetRefCount(createStringObject("EVAL",4))); | |
674 | rewriteClientCommandArgument(c,1,script); | |
675 | } | |
676 | } | |
677 | ||
678 | void evalCommand(redisClient *c) { | |
679 | evalGenericCommand(c,0); | |
680 | } | |
681 | ||
682 | void evalShaCommand(redisClient *c) { | |
683 | if (sdslen(c->argv[1]->ptr) != 40) { | |
684 | /* We know that a match is not possible if the provided SHA is | |
685 | * not the right length. So we return an error ASAP, this way | |
686 | * evalGenericCommand() can be implemented without string length | |
687 | * sanity check */ | |
688 | addReply(c, shared.noscripterr); | |
689 | return; | |
690 | } | |
691 | evalGenericCommand(c,1); | |
692 | } | |
693 | ||
694 | /* We replace math.random() with our implementation that is not affected | |
695 | * by specific libc random() implementations and will output the same sequence | |
696 | * (for the same seed) in every arch. */ | |
697 | ||
698 | /* The following implementation is the one shipped with Lua itself but with | |
699 | * rand() replaced by redisLrand48(). */ | |
700 | int redis_math_random (lua_State *L) { | |
701 | /* the `%' avoids the (rare) case of r==1, and is needed also because on | |
702 | some systems (SunOS!) `rand()' may return a value larger than RAND_MAX */ | |
703 | lua_Number r = (lua_Number)(redisLrand48()%REDIS_LRAND48_MAX) / | |
704 | (lua_Number)REDIS_LRAND48_MAX; | |
705 | switch (lua_gettop(L)) { /* check number of arguments */ | |
706 | case 0: { /* no arguments */ | |
707 | lua_pushnumber(L, r); /* Number between 0 and 1 */ | |
708 | break; | |
709 | } | |
710 | case 1: { /* only upper limit */ | |
711 | int u = luaL_checkint(L, 1); | |
712 | luaL_argcheck(L, 1<=u, 1, "interval is empty"); | |
713 | lua_pushnumber(L, floor(r*u)+1); /* int between 1 and `u' */ | |
714 | break; | |
715 | } | |
716 | case 2: { /* lower and upper limits */ | |
717 | int l = luaL_checkint(L, 1); | |
718 | int u = luaL_checkint(L, 2); | |
719 | luaL_argcheck(L, l<=u, 2, "interval is empty"); | |
720 | lua_pushnumber(L, floor(r*(u-l+1))+l); /* int between `l' and `u' */ | |
721 | break; | |
722 | } | |
723 | default: return luaL_error(L, "wrong number of arguments"); | |
724 | } | |
725 | return 1; | |
726 | } | |
727 | ||
728 | int redis_math_randomseed (lua_State *L) { | |
729 | redisSrand48(luaL_checkint(L, 1)); | |
730 | return 0; | |
731 | } | |
732 | ||
733 | /* --------------------------------------------------------------------------- | |
734 | * SCRIPT command for script environment introspection and control | |
735 | * ------------------------------------------------------------------------- */ | |
736 | ||
737 | void scriptCommand(redisClient *c) { | |
738 | if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"flush")) { | |
739 | scriptingReset(); | |
740 | addReply(c,shared.ok); | |
741 | server.dirty++; /* Replicating this command is a good idea. */ | |
742 | } else if (c->argc >= 2 && !strcasecmp(c->argv[1]->ptr,"exists")) { | |
743 | int j; | |
744 | ||
745 | addReplyMultiBulkLen(c, c->argc-2); | |
746 | for (j = 2; j < c->argc; j++) { | |
747 | if (dictFind(server.lua_scripts,c->argv[j]->ptr)) | |
748 | addReply(c,shared.cone); | |
749 | else | |
750 | addReply(c,shared.czero); | |
751 | } | |
752 | } else if (c->argc == 3 && !strcasecmp(c->argv[1]->ptr,"load")) { | |
753 | char funcname[43]; | |
754 | sds sha; | |
755 | ||
756 | funcname[0] = 'f'; | |
757 | funcname[1] = '_'; | |
758 | hashScript(funcname+2,c->argv[2]->ptr,sdslen(c->argv[2]->ptr)); | |
759 | sha = sdsnewlen(funcname+2,40); | |
760 | if (dictFind(server.lua_scripts,sha) == NULL) { | |
761 | if (luaCreateFunction(c,server.lua,funcname,c->argv[2]) | |
762 | == REDIS_ERR) { | |
763 | sdsfree(sha); | |
764 | return; | |
765 | } | |
766 | } | |
767 | addReplyBulkCBuffer(c,funcname+2,40); | |
768 | sdsfree(sha); | |
769 | } else if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"kill")) { | |
770 | if (server.lua_caller == NULL) { | |
771 | addReplyError(c,"No scripts in execution right now."); | |
772 | } else if (server.lua_write_dirty) { | |
773 | 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."); | |
774 | } else { | |
775 | server.lua_kill = 1; | |
776 | addReply(c,shared.ok); | |
777 | } | |
778 | } else { | |
779 | addReplyError(c, "Unknown SCRIPT subcommand or wrong # of args."); | |
780 | } | |
781 | } |