]> git.saurik.com Git - redis.git/blame - src/scripting.c
Make sure error and status replies emitted by Lua scripts can never have more than...
[redis.git] / src / scripting.c
CommitLineData
7585836e 1#include "redis.h"
2#include "sha1.h"
3
4#include <lua.h>
5#include <lauxlib.h>
6#include <lualib.h>
7229d60d 7#include <ctype.h>
7585836e 8
532e0f5d 9char *redisProtocolToLuaType_Int(lua_State *lua, char *reply);
10char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply);
11char *redisProtocolToLuaType_Status(lua_State *lua, char *reply);
3791000f 12char *redisProtocolToLuaType_Error(lua_State *lua, char *reply);
13char *redisProtocolToLuaType_MultiBulk(lua_State *lua, char *reply);
532e0f5d 14
15/* Take a Redis reply in the Redis protocol format and convert it into a
16 * Lua type. Thanks to this function, and the introduction of not connected
17 * clients, it is trvial to implement the redis() lua function.
18 *
19 * Basically we take the arguments, execute the Redis command in the context
20 * of a non connected client, then take the generated reply and convert it
21 * into a suitable Lua type. With this trick the scripting feature does not
22 * need the introduction of a full Redis internals API. Basically the script
23 * is like a normal client that bypasses all the slow I/O paths.
24 *
25 * Note: in this function we do not do any sanity check as the reply is
26 * generated by Redis directly. This allows use to go faster.
27 * The reply string can be altered during the parsing as it is discared
28 * after the conversion is completed.
29 *
30 * Errors are returned as a table with a single 'err' field set to the
31 * error string.
32 */
33
34char *redisProtocolToLuaType(lua_State *lua, char* reply) {
35 char *p = reply;
36
37 switch(*p) {
38 case ':':
39 p = redisProtocolToLuaType_Int(lua,reply);
40 break;
41 case '$':
42 p = redisProtocolToLuaType_Bulk(lua,reply);
43 break;
44 case '+':
45 p = redisProtocolToLuaType_Status(lua,reply);
46 break;
3791000f 47 case '-':
48 p = redisProtocolToLuaType_Error(lua,reply);
49 break;
50 case '*':
51 p = redisProtocolToLuaType_MultiBulk(lua,reply);
52 break;
532e0f5d 53 }
54 return p;
55}
56
57char *redisProtocolToLuaType_Int(lua_State *lua, char *reply) {
58 char *p = strchr(reply+1,'\r');
59 long long value;
60
61 string2ll(reply+1,p-reply-1,&value);
62 lua_pushnumber(lua,(lua_Number)value);
63 return p+2;
64}
65
66char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply) {
67 char *p = strchr(reply+1,'\r');
68 long long bulklen;
69
70 string2ll(reply+1,p-reply-1,&bulklen);
379789cc 71 if (bulklen == -1) {
82c6b825 72 lua_pushboolean(lua,0);
532e0f5d 73 return p+2;
74 } else {
75 lua_pushlstring(lua,p+2,bulklen);
76 return p+2+bulklen+2;
77 }
78}
79
80char *redisProtocolToLuaType_Status(lua_State *lua, char *reply) {
81 char *p = strchr(reply+1,'\r');
82
0d916763 83 lua_newtable(lua);
84 lua_pushstring(lua,"ok");
532e0f5d 85 lua_pushlstring(lua,reply+1,p-reply-1);
0d916763 86 lua_settable(lua,-3);
532e0f5d 87 return p+2;
88}
89
3791000f 90char *redisProtocolToLuaType_Error(lua_State *lua, char *reply) {
91 char *p = strchr(reply+1,'\r');
92
93 lua_newtable(lua);
94 lua_pushstring(lua,"err");
95 lua_pushlstring(lua,reply+1,p-reply-1);
96 lua_settable(lua,-3);
97 return p+2;
98}
99
100char *redisProtocolToLuaType_MultiBulk(lua_State *lua, char *reply) {
101 char *p = strchr(reply+1,'\r');
102 long long mbulklen;
103 int j = 0;
104
3791000f 105 string2ll(reply+1,p-reply-1,&mbulklen);
106 p += 2;
107 if (mbulklen == -1) {
82c6b825 108 lua_pushboolean(lua,0);
3791000f 109 return p;
110 }
3791000f 111 lua_newtable(lua);
112 for (j = 0; j < mbulklen; j++) {
10a6da7a 113 lua_pushnumber(lua,j+1);
3791000f 114 p = redisProtocolToLuaType(lua,p);
115 lua_settable(lua,-3);
116 }
117 return p;
118}
119
379789cc 120void luaPushError(lua_State *lua, char *error) {
121 lua_newtable(lua);
122 lua_pushstring(lua,"err");
123 lua_pushstring(lua, error);
124 lua_settable(lua,-3);
125}
126
0f1d64ca 127int luaRedisCommand(lua_State *lua) {
128 int j, argc = lua_gettop(lua);
129 struct redisCommand *cmd;
130 robj **argv;
131 redisClient *c = server.lua_client;
132 sds reply;
133
532e0f5d 134 /* Build the arguments vector */
0f1d64ca 135 argv = zmalloc(sizeof(robj*)*argc);
379789cc 136 for (j = 0; j < argc; j++) {
137 if (!lua_isstring(lua,j+1)) break;
532e0f5d 138 argv[j] = createStringObject((char*)lua_tostring(lua,j+1),
139 lua_strlen(lua,j+1));
379789cc 140 }
141
142 /* Check if one of the arguments passed by the Lua script
143 * is not a string or an integer (lua_isstring() return true for
144 * integers as well). */
145 if (j != argc) {
146 j--;
147 while (j >= 0) {
148 decrRefCount(argv[j]);
149 j--;
150 }
151 zfree(argv);
152 luaPushError(lua,
153 "Lua redis() command arguments must be strings or integers");
154 return 1;
155 }
0f1d64ca 156
157 /* Command lookup */
158 cmd = lookupCommand(argv[0]->ptr);
3791000f 159 if (!cmd || ((cmd->arity > 0 && cmd->arity != argc) ||
160 (argc < -cmd->arity)))
161 {
532e0f5d 162 for (j = 0; j < argc; j++) decrRefCount(argv[j]);
0f1d64ca 163 zfree(argv);
3791000f 164 if (cmd)
379789cc 165 luaPushError(lua,
3791000f 166 "Wrong number of args calling Redis command From Lua script");
167 else
379789cc 168 luaPushError(lua,"Unknown Redis command called from Lua script");
532e0f5d 169 return 1;
0f1d64ca 170 }
532e0f5d 171
0f1d64ca 172 /* Run the command in the context of a fake client */
173 c->argv = argv;
174 c->argc = argc;
175 cmd->proc(c);
176
177 /* Convert the result of the Redis command into a suitable Lua type.
178 * The first thing we need is to create a single string from the client
179 * output buffers. */
180 reply = sdsempty();
181 if (c->bufpos) {
7156f43c 182 reply = sdscatlen(reply,c->buf,c->bufpos);
0f1d64ca 183 c->bufpos = 0;
184 }
185 while(listLength(c->reply)) {
186 robj *o = listNodeValue(listFirst(c->reply));
187
09ab5591 188 reply = sdscatlen(reply,o->ptr,sdslen(o->ptr));
0f1d64ca 189 listDelNode(c->reply,listFirst(c->reply));
190 }
532e0f5d 191 redisProtocolToLuaType(lua,reply);
7156f43c 192 sdsfree(reply);
0f1d64ca 193
194 /* Clean up. Command code may have changed argv/argc so we use the
195 * argv/argc of the client instead of the local variables. */
196 for (j = 0; j < c->argc; j++)
197 decrRefCount(c->argv[j]);
198 zfree(c->argv);
199
200 return 1;
201}
202
288f811f 203int luaLogCommand(lua_State *lua) {
204 int j, argc = lua_gettop(lua);
205 int level;
206 sds log;
207
208 if (argc < 2) {
209 luaPushError(lua, "redis.log() requires two arguments or more.");
210 return 1;
211 } else if (!lua_isnumber(lua,-argc)) {
e927a246 212 luaPushError(lua, "First argument must be a number (log level).");
288f811f 213 return 1;
214 }
215 level = lua_tonumber(lua,-argc);
e927a246 216 if (level < REDIS_DEBUG || level > REDIS_WARNING) {
288f811f 217 luaPushError(lua, "Invalid debug level.");
218 return 1;
219 }
220
221 /* Glue together all the arguments */
222 log = sdsempty();
223 for (j = 1; j < argc; j++) {
224 size_t len;
225 char *s;
226
227 s = (char*)lua_tolstring(lua,(-argc)+j,&len);
228 if (s) {
229 if (j != 1) log = sdscatlen(log," ",1);
230 log = sdscatlen(log,s,len);
231 }
232 }
233 redisLogRaw(level,log);
234 sdsfree(log);
235 return 0;
236}
237
eeffcf38 238void luaMaskCountHook(lua_State *lua, lua_Debug *ar) {
239 long long elapsed;
240 REDIS_NOTUSED(ar);
241
eeffcf38 242 elapsed = (ustime()/1000) - server.lua_time_start;
243 if (elapsed >= server.lua_time_limit) {
87c21924 244 redisLog(REDIS_NOTICE,"Lua script aborted for max execution time after %lld milliseconds of running time.",elapsed);
245 lua_pushstring(lua,"Script aborted for max execution time.");
eeffcf38 246 lua_error(lua);
eeffcf38 247 }
248}
249
7585836e 250void scriptingInit(void) {
251 lua_State *lua = lua_open();
252 luaL_openlibs(lua);
0f1d64ca 253
288f811f 254 /* Register the redis commands table and fields */
255 lua_newtable(lua);
256
257 /* redis.call */
258 lua_pushstring(lua,"call");
0f1d64ca 259 lua_pushcfunction(lua,luaRedisCommand);
288f811f 260 lua_settable(lua,-3);
261
262 /* redis.log and log levels. */
263 lua_pushstring(lua,"log");
264 lua_pushcfunction(lua,luaLogCommand);
265 lua_settable(lua,-3);
266
267 lua_pushstring(lua,"LOG_DEBUG");
268 lua_pushnumber(lua,REDIS_DEBUG);
269 lua_settable(lua,-3);
270
271 lua_pushstring(lua,"LOG_VERBOSE");
272 lua_pushnumber(lua,REDIS_VERBOSE);
273 lua_settable(lua,-3);
274
275 lua_pushstring(lua,"LOG_NOTICE");
276 lua_pushnumber(lua,REDIS_NOTICE);
277 lua_settable(lua,-3);
278
279 lua_pushstring(lua,"LOG_WARNING");
280 lua_pushnumber(lua,REDIS_WARNING);
281 lua_settable(lua,-3);
282
283 /* Finally set the table as 'redis' global var. */
00b7541b 284 lua_setglobal(lua,"redis");
0f1d64ca 285
286 /* Create the (non connected) client that we use to execute Redis commands
287 * inside the Lua interpreter */
288 server.lua_client = createClient(-1);
7156f43c 289 server.lua_client->flags |= REDIS_LUA_CLIENT;
0f1d64ca 290
7585836e 291 server.lua = lua;
292}
293
294/* Hash the scripit into a SHA1 digest. We use this as Lua function name.
295 * Digest should point to a 41 bytes buffer: 40 for SHA1 converted into an
296 * hexadecimal number, plus 1 byte for null term. */
297void hashScript(char *digest, char *script, size_t len) {
298 SHA1_CTX ctx;
299 unsigned char hash[20];
300 char *cset = "0123456789abcdef";
301 int j;
302
303 SHA1Init(&ctx);
304 SHA1Update(&ctx,(unsigned char*)script,len);
305 SHA1Final(hash,&ctx);
306
307 for (j = 0; j < 20; j++) {
308 digest[j*2] = cset[((hash[j]&0xF0)>>4)];
309 digest[j*2+1] = cset[(hash[j]&0xF)];
310 }
311 digest[40] = '\0';
312}
313
314void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
82c6b825 315 int t = lua_type(lua,-1);
7585836e 316
317 switch(t) {
318 case LUA_TSTRING:
82c6b825 319 addReplyBulkCBuffer(c,(char*)lua_tostring(lua,-1),lua_strlen(lua,-1));
7585836e 320 break;
321 case LUA_TBOOLEAN:
82c6b825 322 addReply(c,lua_toboolean(lua,-1) ? shared.cone : shared.nullbulk);
7585836e 323 break;
324 case LUA_TNUMBER:
82c6b825 325 addReplyLongLong(c,(long long)lua_tonumber(lua,-1));
7585836e 326 break;
532e0f5d 327 case LUA_TTABLE:
0d916763 328 /* We need to check if it is an array, an error, or a status reply.
329 * Error are returned as a single element table with 'err' field.
330 * Status replies are returned as single elment table with 'ok' field */
532e0f5d 331 lua_pushstring(lua,"err");
332 lua_gettable(lua,-2);
333 t = lua_type(lua,-1);
334 if (t == LUA_TSTRING) {
3bb818df 335 sds err = sdsnew(lua_tostring(lua,-1));
336 sdsmapchars(err,"\r\n"," ",2);
337 addReplySds(c,sdscatprintf(sdsempty(),"-%s\r\n",err));
338 sdsfree(err);
0d916763 339 lua_pop(lua,2);
340 return;
341 }
342
343 lua_pop(lua,1);
344 lua_pushstring(lua,"ok");
345 lua_gettable(lua,-2);
346 t = lua_type(lua,-1);
347 if (t == LUA_TSTRING) {
3bb818df 348 sds ok = sdsnew(lua_tostring(lua,-1));
349 sdsmapchars(ok,"\r\n"," ",2);
350 addReplySds(c,sdscatprintf(sdsempty(),"+%s\r\n",ok));
351 sdsfree(ok);
532e0f5d 352 lua_pop(lua,1);
353 } else {
354 void *replylen = addDeferredMultiBulkLength(c);
355 int j = 1, mbulklen = 0;
356
0d916763 357 lua_pop(lua,1); /* Discard the 'ok' field value we popped */
532e0f5d 358 while(1) {
359 lua_pushnumber(lua,j++);
360 lua_gettable(lua,-2);
361 t = lua_type(lua,-1);
362 if (t == LUA_TNIL) {
363 lua_pop(lua,1);
364 break;
532e0f5d 365 }
82c6b825 366 luaReplyToRedisReply(c, lua);
367 mbulklen++;
532e0f5d 368 }
369 setDeferredMultiBulkLength(c,replylen,mbulklen);
370 }
371 break;
7585836e 372 default:
373 addReply(c,shared.nullbulk);
374 }
375 lua_pop(lua,1);
376}
377
4ae5b5e1 378/* Set an array of Redis String Objects as a Lua array (table) stored into a
379 * global variable. */
380void luaSetGlobalArray(lua_State *lua, char *var, robj **elev, int elec) {
381 int j;
382
383 lua_newtable(lua);
384 for (j = 0; j < elec; j++) {
385 lua_pushlstring(lua,(char*)elev[j]->ptr,sdslen(elev[j]->ptr));
386 lua_rawseti(lua,-2,j+1);
387 }
388 lua_setglobal(lua,var);
389}
390
7229d60d 391void evalGenericCommand(redisClient *c, int evalsha) {
7585836e 392 lua_State *lua = server.lua;
393 char funcname[43];
4ae5b5e1 394 long long numkeys;
395
396 /* Get the number of arguments that are keys */
397 if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK)
398 return;
399 if (numkeys > (c->argc - 3)) {
400 addReplyError(c,"Number of keys can't be greater than number of args");
401 return;
402 }
7585836e 403
404 /* We obtain the script SHA1, then check if this function is already
405 * defined into the Lua state */
406 funcname[0] = 'f';
407 funcname[1] = '_';
7229d60d 408 if (!evalsha) {
409 /* Hash the code if this is an EVAL call */
410 hashScript(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
411 } else {
412 /* We already have the SHA if it is a EVALSHA */
413 int j;
414 char *sha = c->argv[1]->ptr;
415
416 for (j = 0; j < 40; j++)
417 funcname[j+2] = tolower(sha[j]);
418 funcname[42] = '\0';
419 }
420
7585836e 421 lua_getglobal(lua, funcname);
422 if (lua_isnil(lua,1)) {
7229d60d 423 sds funcdef;
424
425 /* Function not defined... let's define it if we have the
426 * body of the funciton. If this is an EVALSHA call we can just
427 * return an error. */
428 if (evalsha) {
429 addReply(c, shared.noscripterr);
430 lua_pop(lua,1); /* remove the nil from the stack */
431 return;
432 }
433 funcdef = sdsempty();
7585836e 434
435 lua_pop(lua,1); /* remove the nil from the stack */
436 funcdef = sdscat(funcdef,"function ");
437 funcdef = sdscatlen(funcdef,funcname,42);
438 funcdef = sdscatlen(funcdef," ()\n",4);
439 funcdef = sdscatlen(funcdef,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
440 funcdef = sdscatlen(funcdef,"\nend\n",5);
379789cc 441 /* printf("Defining:\n%s\n",funcdef); */
7585836e 442
443 if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) {
444 addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
445 lua_tostring(lua,-1));
446 lua_pop(lua,1);
7156f43c 447 sdsfree(funcdef);
7585836e 448 return;
449 }
7156f43c 450 sdsfree(funcdef);
7585836e 451 if (lua_pcall(lua,0,0,0)) {
452 addReplyErrorFormat(c,"Error running script (new function): %s\n",
453 lua_tostring(lua,-1));
454 lua_pop(lua,1);
455 return;
456 }
457 lua_getglobal(lua, funcname);
458 }
4ae5b5e1 459
460 /* Populate the argv and keys table accordingly to the arguments that
461 * EVAL received. */
462 luaSetGlobalArray(lua,"KEYS",c->argv+3,numkeys);
463 luaSetGlobalArray(lua,"ARGV",c->argv+3+numkeys,c->argc-3-numkeys);
c2a7dd85 464
465 /* Select the right DB in the context of the Lua client */
466 selectDb(server.lua_client,c->db->id);
7585836e 467
da386cdf 468 /* Set an hook in order to be able to stop the script execution if it
469 * is running for too much time.
470 * We set the hook only if the time limit is enabled as the hook will
471 * make the Lua script execution slower. */
472 if (server.lua_time_limit > 0) {
473 lua_sethook(lua,luaMaskCountHook,LUA_MASKCOUNT,100000);
474 server.lua_time_start = ustime()/1000;
475 } else {
476 lua_sethook(lua,luaMaskCountHook,0,0);
477 }
478
7585836e 479 /* At this point whatever this script was never seen before or if it was
480 * already defined, we can call it. We have zero arguments and expect
481 * a single return value. */
482 if (lua_pcall(lua,0,1,0)) {
c2a7dd85 483 selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
7585836e 484 addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
485 funcname, lua_tostring(lua,-1));
486 lua_pop(lua,1);
40531be0 487 lua_gc(lua,LUA_GCCOLLECT,0);
7585836e 488 return;
489 }
c2a7dd85 490 selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
7585836e 491 luaReplyToRedisReply(c,lua);
40531be0 492 lua_gc(lua,LUA_GCSTEP,1);
7585836e 493}
7229d60d 494
495void evalCommand(redisClient *c) {
496 evalGenericCommand(c,0);
497}
498
499void evalShaCommand(redisClient *c) {
500 if (sdslen(c->argv[1]->ptr) != 40) {
501 /* We know that a match is not possible if the provided SHA is
502 * not the right length. So we return an error ASAP, this way
503 * evalGenericCommand() can be implemented without string length
504 * sanity check */
505 addReply(c, shared.noscripterr);
506 return;
507 }
508 evalGenericCommand(c,1);
509}