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