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