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