Query the archive to provide a complete KEYS list.

[redis.git] / src / scripting.c

diff --git a/src/scripting.c b/src/scripting.c
index be6247a907d984a2a5b79ccdc30807e0098430dd..d614f42a13f87d14d383604ff59c1224b8d70fa9 100644 (file)
--- a/src/scripting.c
+++ b/src/scripting.c
@@ -1,16 +1,50 @@
+/*
+ * Copyright (c) 2009-2012, Salvatore Sanfilippo <antirez at gmail dot com>
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ *   * Redistributions of source code must retain the above copyright notice,
+ *     this list of conditions and the following disclaimer.
+ *   * Redistributions in binary form must reproduce the above copyright
+ *     notice, this list of conditions and the following disclaimer in the
+ *     documentation and/or other materials provided with the distribution.
+ *   * Neither the name of Redis nor the names of its contributors may be used
+ *     to endorse or promote products derived from this software without
+ *     specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+ * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ * POSSIBILITY OF SUCH DAMAGE.
+ */
+

 #include "redis.h"
 #include "sha1.h"
 #include "redis.h"
 #include "sha1.h"

+#include "rand.h"

 
 #include <lua.h>
 #include <lauxlib.h>
 #include <lualib.h>
 #include <ctype.h>
 
 #include <lua.h>
 #include <lauxlib.h>
 #include <lualib.h>
 #include <ctype.h>

+#include <math.h>

 
 char *redisProtocolToLuaType_Int(lua_State *lua, char *reply);
 char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply);
 char *redisProtocolToLuaType_Status(lua_State *lua, char *reply);
 char *redisProtocolToLuaType_Error(lua_State *lua, char *reply);
 char *redisProtocolToLuaType_MultiBulk(lua_State *lua, char *reply);
 
 char *redisProtocolToLuaType_Int(lua_State *lua, char *reply);
 char *redisProtocolToLuaType_Bulk(lua_State *lua, char *reply);
 char *redisProtocolToLuaType_Status(lua_State *lua, char *reply);
 char *redisProtocolToLuaType_Error(lua_State *lua, char *reply);
 char *redisProtocolToLuaType_MultiBulk(lua_State *lua, char *reply);

+int redis_math_random (lua_State *L);
+int redis_math_randomseed (lua_State *L);
+void sha1hex(char *digest, char *script, size_t len);

 
 /* Take a Redis reply in the Redis protocol format and convert it into a
  * Lua type. Thanks to this function, and the introduction of not connected
 
 /* Take a Redis reply in the Redis protocol format and convert it into a
  * Lua type. Thanks to this function, and the introduction of not connected


@@ -23,7 +57,7 @@ char *redisProtocolToLuaType_MultiBulk(lua_State *lua, char *reply);
  * is like a normal client that bypasses all the slow I/O paths.
  *
  * Note: in this function we do not do any sanity check as the reply is
  * is like a normal client that bypasses all the slow I/O paths.
  *
  * Note: in this function we do not do any sanity check as the reply is

- * generated by Redis directly. This allows use to go faster.
+ * generated by Redis directly. This allows us to go faster.

  * The reply string can be altered during the parsing as it is discared
  * after the conversion is completed.
  *
  * The reply string can be altered during the parsing as it is discared
  * after the conversion is completed.
  *


@@ -124,13 +158,51 @@ void luaPushError(lua_State *lua, char *error) {
     lua_settable(lua,-3);
 }
 
     lua_settable(lua,-3);
 }
 

-int luaRedisCommand(lua_State *lua) {
+/* Sort the array currently in the stack. We do this to make the output
+ * of commands like KEYS or SMEMBERS something deterministic when called
+ * from Lua (to play well with AOf/replication).
+ *
+ * The array is sorted using table.sort itself, and assuming all the
+ * list elements are strings. */
+void luaSortArray(lua_State *lua) {
+    /* Initial Stack: array */
+    lua_getglobal(lua,"table");
+    lua_pushstring(lua,"sort");
+    lua_gettable(lua,-2);       /* Stack: array, table, table.sort */
+    lua_pushvalue(lua,-3);      /* Stack: array, table, table.sort, array */
+    if (lua_pcall(lua,1,0,0)) {
+        /* Stack: array, table, error */
+
+        /* We are not interested in the error, we assume that the problem is
+         * that there are 'false' elements inside the array, so we try
+         * again with a slower function but able to handle this case, that
+         * is: table.sort(table, __redis__compare_helper) */
+        lua_pop(lua,1);             /* Stack: array, table */
+        lua_pushstring(lua,"sort"); /* Stack: array, table, sort */
+        lua_gettable(lua,-2);       /* Stack: array, table, table.sort */
+        lua_pushvalue(lua,-3);      /* Stack: array, table, table.sort, array */
+        lua_getglobal(lua,"__redis__compare_helper");
+        /* Stack: array, table, table.sort, array, __redis__compare_helper */
+        lua_call(lua,2,0);
+    }
+    /* Stack: array (sorted), table */
+    lua_pop(lua,1);             /* Stack: array (sorted) */
+}
+
+int luaRedisGenericCommand(lua_State *lua, int raise_error) {

     int j, argc = lua_gettop(lua);
     struct redisCommand *cmd;
     robj **argv;
     redisClient *c = server.lua_client;
     sds reply;
 
     int j, argc = lua_gettop(lua);
     struct redisCommand *cmd;
     robj **argv;
     redisClient *c = server.lua_client;
     sds reply;
 

+    /* Require at least one argument */
+    if (argc == 0) {
+        luaPushError(lua,
+            "Please specify at least one argument for redis.call()");
+        return 1;
+    }
+

     /* Build the arguments vector */
     argv = zmalloc(sizeof(robj*)*argc);
     for (j = 0; j < argc; j++) {
     /* Build the arguments vector */
     argv = zmalloc(sizeof(robj*)*argc);
     for (j = 0; j < argc; j++) {


@@ -154,25 +226,70 @@ int luaRedisCommand(lua_State *lua) {
         return 1;
     }
 
         return 1;
     }
 

+    /* Setup our fake client for command execution */
+    c->argv = argv;
+    c->argc = argc;
+

     /* Command lookup */
     cmd = lookupCommand(argv[0]->ptr);
     if (!cmd || ((cmd->arity > 0 && cmd->arity != argc) ||
                    (argc < -cmd->arity)))
     {
     /* Command lookup */
     cmd = lookupCommand(argv[0]->ptr);
     if (!cmd || ((cmd->arity > 0 && cmd->arity != argc) ||
                    (argc < -cmd->arity)))
     {

-        for (j = 0; j < argc; j++) decrRefCount(argv[j]);
-        zfree(argv);

         if (cmd)
             luaPushError(lua,
                 "Wrong number of args calling Redis command From Lua script");
         else
             luaPushError(lua,"Unknown Redis command called from Lua script");
         if (cmd)
             luaPushError(lua,
                 "Wrong number of args calling Redis command From Lua script");
         else
             luaPushError(lua,"Unknown Redis command called from Lua script");

-        return 1;
+        goto cleanup;

     }
 
     }
 

-    /* Run the command in the context of a fake client */
-    c->argv = argv;
-    c->argc = argc;
-    cmd->proc(c);
+    /* There are commands that are not allowed inside scripts. */
+    if (cmd->flags & REDIS_CMD_NOSCRIPT) {
+        luaPushError(lua, "This Redis command is not allowed from scripts");
+        goto cleanup;
+    }
+
+    /* Write commands are forbidden against read-only slaves, or if a
+     * command marked as non-deterministic was already called in the context
+     * of this script. */
+    if (cmd->flags & REDIS_CMD_WRITE) {
+        if (server.lua_random_dirty) {
+            luaPushError(lua,
+                "Write commands not allowed after non deterministic commands");
+            goto cleanup;
+        } else if (server.masterhost && server.repl_slave_ro &&
+                   !(server.lua_caller->flags & REDIS_MASTER))
+        {
+            luaPushError(lua, shared.roslaveerr->ptr);
+            goto cleanup;
+        } else if (server.stop_writes_on_bgsave_err &&
+                   server.saveparamslen > 0 &&
+                   server.lastbgsave_status == REDIS_ERR)
+        {
+            luaPushError(lua, shared.bgsaveerr->ptr);
+            goto cleanup;
+        }
+    }
+
+    /* If we reached the memory limit configured via maxmemory, commands that
+     * could enlarge the memory usage are not allowed, but only if this is the
+     * first write in the context of this script, otherwise we can't stop
+     * in the middle. */
+    if (server.maxmemory && server.lua_write_dirty == 0 &&
+        (cmd->flags & REDIS_CMD_DENYOOM))
+    {
+        if (freeMemoryIfNeeded() == REDIS_ERR) {
+            luaPushError(lua, shared.oomerr->ptr);
+            goto cleanup;
+        }
+    }
+
+    if (cmd->flags & REDIS_CMD_RANDOM) server.lua_random_dirty = 1;
+    if (cmd->flags & REDIS_CMD_WRITE) server.lua_write_dirty = 1;
+
+    /* Run the command */
+    c->cmd = cmd;
+    call(c,REDIS_CALL_SLOWLOG | REDIS_CALL_STATS);

 
     /* Convert the result of the Redis command into a suitable Lua type.
      * The first thing we need is to create a single string from the client
 
     /* Convert the result of the Redis command into a suitable Lua type.
      * The first thing we need is to create a single string from the client


@@ -188,18 +305,90 @@ int luaRedisCommand(lua_State *lua) {
         reply = sdscatlen(reply,o->ptr,sdslen(o->ptr));
         listDelNode(c->reply,listFirst(c->reply));
     }
         reply = sdscatlen(reply,o->ptr,sdslen(o->ptr));
         listDelNode(c->reply,listFirst(c->reply));
     }

+    if (raise_error && reply[0] != '-') raise_error = 0;

     redisProtocolToLuaType(lua,reply);
     redisProtocolToLuaType(lua,reply);

+    /* Sort the output array if needed, assuming it is a non-null multi bulk
+     * reply as expected. */
+    if ((cmd->flags & REDIS_CMD_SORT_FOR_SCRIPT) &&
+        (reply[0] == '*' && reply[1] != '-')) {
+            luaSortArray(lua);
+    }

     sdsfree(reply);
     sdsfree(reply);

+    c->reply_bytes = 0;

 
 

+cleanup:

     /* Clean up. Command code may have changed argv/argc so we use the
      * argv/argc of the client instead of the local variables. */
     for (j = 0; j < c->argc; j++)
         decrRefCount(c->argv[j]);
     zfree(c->argv);
 
     /* Clean up. Command code may have changed argv/argc so we use the
      * argv/argc of the client instead of the local variables. */
     for (j = 0; j < c->argc; j++)
         decrRefCount(c->argv[j]);
     zfree(c->argv);
 

+    if (raise_error) {
+        /* If we are here we should have an error in the stack, in the
+         * form of a table with an "err" field. Extract the string to
+         * return the plain error. */
+        lua_pushstring(lua,"err");
+        lua_gettable(lua,-2);
+        return lua_error(lua);
+    }

     return 1;
 }
 
     return 1;
 }
 

+int luaRedisCallCommand(lua_State *lua) {
+    return luaRedisGenericCommand(lua,1);
+}
+
+int luaRedisPCallCommand(lua_State *lua) {
+    return luaRedisGenericCommand(lua,0);
+}
+
+/* This adds redis.sha1hex(string) to Lua scripts using the same hashing
+ * function used for sha1ing lua scripts. */
+int luaRedisSha1hexCommand(lua_State *lua) {
+    int argc = lua_gettop(lua);
+    char digest[41];
+    size_t len;
+    char *s;
+
+    if (argc != 1) {
+        luaPushError(lua, "wrong number of arguments");
+        return 1;
+    }
+
+    s = (char*)lua_tolstring(lua,1,&len);
+    sha1hex(digest,s,len);
+    lua_pushstring(lua,digest);
+    return 1;
+}
+
+/* Returns a table with a single field 'field' set to the string value
+ * passed as argument. This helper function is handy when returning
+ * a Redis Protocol error or status reply from Lua:
+ *
+ * return redis.error_reply("ERR Some Error")
+ * return redis.status_reply("ERR Some Error")
+ */
+int luaRedisReturnSingleFieldTable(lua_State *lua, char *field) {
+    if (lua_gettop(lua) != 1 || lua_type(lua,-1) != LUA_TSTRING) {
+        luaPushError(lua, "wrong number or type of arguments");
+        return 1;
+    }
+
+    lua_newtable(lua);
+    lua_pushstring(lua, field);
+    lua_pushvalue(lua, -3);
+    lua_settable(lua, -3);
+    return 1;
+}
+
+int luaRedisErrorReplyCommand(lua_State *lua) {
+    return luaRedisReturnSingleFieldTable(lua,"err");
+}
+
+int luaRedisStatusReplyCommand(lua_State *lua) {
+    return luaRedisReturnSingleFieldTable(lua,"ok");
+}
+

 int luaLogCommand(lua_State *lua) {
     int j, argc = lua_gettop(lua);
     int level;
 int luaLogCommand(lua_State *lua) {
     int j, argc = lua_gettop(lua);
     int level;


@@ -238,25 +427,124 @@ int luaLogCommand(lua_State *lua) {
 void luaMaskCountHook(lua_State *lua, lua_Debug *ar) {
     long long elapsed;
     REDIS_NOTUSED(ar);
 void luaMaskCountHook(lua_State *lua, lua_Debug *ar) {
     long long elapsed;
     REDIS_NOTUSED(ar);

+    REDIS_NOTUSED(lua);

 
     elapsed = (ustime()/1000) - server.lua_time_start;
 
     elapsed = (ustime()/1000) - server.lua_time_start;

-    if (elapsed >= server.lua_time_limit) {
-        redisLog(REDIS_NOTICE,"Lua script aborted for max execution time after %lld milliseconds of running time.",elapsed);
-        lua_pushstring(lua,"Script aborted for max execution time.");
+    if (elapsed >= server.lua_time_limit && server.lua_timedout == 0) {
+        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);
+        server.lua_timedout = 1;
+        /* Once the script timeouts we reenter the event loop to permit others
+         * to call SCRIPT KILL or SHUTDOWN NOSAVE if needed. For this reason
+         * we need to mask the client executing the script from the event loop.
+         * If we don't do that the client may disconnect and could no longer be
+         * here when the EVAL command will return. */
+         aeDeleteFileEvent(server.el, server.lua_caller->fd, AE_READABLE);
+    }
+    if (server.lua_timedout)
+        aeProcessEvents(server.el, AE_FILE_EVENTS|AE_DONT_WAIT);
+    if (server.lua_kill) {
+        redisLog(REDIS_WARNING,"Lua script killed by user with SCRIPT KILL.");
+        lua_pushstring(lua,"Script killed by user with SCRIPT KILL...");

         lua_error(lua);
     }
 }
 
         lua_error(lua);
     }
 }
 

+void luaLoadLib(lua_State *lua, const char *libname, lua_CFunction luafunc) {
+  lua_pushcfunction(lua, luafunc);
+  lua_pushstring(lua, libname);
+  lua_call(lua, 1, 0);
+}
+
+LUALIB_API int (luaopen_cjson) (lua_State *L);
+LUALIB_API int (luaopen_struct) (lua_State *L);
+LUALIB_API int (luaopen_cmsgpack) (lua_State *L);
+
+void luaLoadLibraries(lua_State *lua) {
+    luaLoadLib(lua, "", luaopen_base);
+    luaLoadLib(lua, LUA_TABLIBNAME, luaopen_table);
+    luaLoadLib(lua, LUA_STRLIBNAME, luaopen_string);
+    luaLoadLib(lua, LUA_MATHLIBNAME, luaopen_math);
+    luaLoadLib(lua, LUA_DBLIBNAME, luaopen_debug); 
+    luaLoadLib(lua, "cjson", luaopen_cjson);
+    luaLoadLib(lua, "struct", luaopen_struct);
+    luaLoadLib(lua, "cmsgpack", luaopen_cmsgpack);
+
+#if 0 /* Stuff that we don't load currently, for sandboxing concerns. */
+    luaLoadLib(lua, LUA_LOADLIBNAME, luaopen_package);
+    luaLoadLib(lua, LUA_OSLIBNAME, luaopen_os);
+#endif
+}
+
+/* Remove a functions that we don't want to expose to the Redis scripting
+ * environment. */
+void luaRemoveUnsupportedFunctions(lua_State *lua) {
+    lua_pushnil(lua);
+    lua_setglobal(lua,"loadfile");
+}
+
+/* This function installs metamethods in the global table _G that prevent
+ * the creation of globals accidentally.
+ *
+ * It should be the last to be called in the scripting engine initialization
+ * sequence, because it may interact with creation of globals. */
+void scriptingEnableGlobalsProtection(lua_State *lua) {
+    char *s[32];
+    sds code = sdsempty();
+    int j = 0;
+
+    /* strict.lua from: http://metalua.luaforge.net/src/lib/strict.lua.html.
+     * Modified to be adapted to Redis. */
+    s[j++]="local mt = {}\n";
+    s[j++]="setmetatable(_G, mt)\n";
+    s[j++]="mt.__newindex = function (t, n, v)\n";
+    s[j++]="  if debug.getinfo(2) then\n";
+    s[j++]="    local w = debug.getinfo(2, \"S\").what\n";
+    s[j++]="    if w ~= \"main\" and w ~= \"C\" then\n";
+    s[j++]="      error(\"Script attempted to create global variable '\"..tostring(n)..\"'\", 2)\n";
+    s[j++]="    end\n";
+    s[j++]="  end\n";
+    s[j++]="  rawset(t, n, v)\n";
+    s[j++]="end\n";
+    s[j++]="mt.__index = function (t, n)\n";
+    s[j++]="  if debug.getinfo(2) and debug.getinfo(2, \"S\").what ~= \"C\" then\n";
+    s[j++]="    error(\"Script attempted to access unexisting global variable '\"..tostring(n)..\"'\", 2)\n";
+    s[j++]="  end\n";
+    s[j++]="  return rawget(t, n)\n";
+    s[j++]="end\n";
+    s[j++]=NULL;
+
+    for (j = 0; s[j] != NULL; j++) code = sdscatlen(code,s[j],strlen(s[j]));
+    luaL_loadbuffer(lua,code,sdslen(code),"@enable_strict_lua");
+    lua_pcall(lua,0,0,0);
+    sdsfree(code);
+}
+
+/* Initialize the scripting environment.
+ * It is possible to call this function to reset the scripting environment
+ * assuming that we call scriptingRelease() before.
+ * See scriptingReset() for more information. */

 void scriptingInit(void) {
     lua_State *lua = lua_open();
 void scriptingInit(void) {
     lua_State *lua = lua_open();

-    luaL_openlibs(lua);
+
+    luaLoadLibraries(lua);
+    luaRemoveUnsupportedFunctions(lua);
+
+    /* Initialize a dictionary we use to map SHAs to scripts.
+     * This is useful for replication, as we need to replicate EVALSHA
+     * as EVAL, so we need to remember the associated script. */
+    server.lua_scripts = dictCreate(&shaScriptObjectDictType,NULL);

 
     /* Register the redis commands table and fields */
     lua_newtable(lua);
 
     /* redis.call */
     lua_pushstring(lua,"call");
 
     /* Register the redis commands table and fields */
     lua_newtable(lua);
 
     /* redis.call */
     lua_pushstring(lua,"call");

-    lua_pushcfunction(lua,luaRedisCommand);
+    lua_pushcfunction(lua,luaRedisCallCommand);
+    lua_settable(lua,-3);
+
+    /* redis.pcall */
+    lua_pushstring(lua,"pcall");
+    lua_pushcfunction(lua,luaRedisPCallCommand);

     lua_settable(lua,-3);
 
     /* redis.log and log levels. */
     lua_settable(lua,-3);
 
     /* redis.log and log levels. */


@@ -280,21 +568,83 @@ void scriptingInit(void) {
     lua_pushnumber(lua,REDIS_WARNING);
     lua_settable(lua,-3);
 
     lua_pushnumber(lua,REDIS_WARNING);
     lua_settable(lua,-3);
 

+    /* redis.sha1hex */
+    lua_pushstring(lua, "sha1hex");
+    lua_pushcfunction(lua, luaRedisSha1hexCommand);
+    lua_settable(lua, -3);
+
+    /* redis.error_reply and redis.status_reply */
+    lua_pushstring(lua, "error_reply");
+    lua_pushcfunction(lua, luaRedisErrorReplyCommand);
+    lua_settable(lua, -3);
+    lua_pushstring(lua, "status_reply");
+    lua_pushcfunction(lua, luaRedisStatusReplyCommand);
+    lua_settable(lua, -3);
+

     /* Finally set the table as 'redis' global var. */
     lua_setglobal(lua,"redis");
 
     /* Finally set the table as 'redis' global var. */
     lua_setglobal(lua,"redis");
 

+    /* Replace math.random and math.randomseed with our implementations. */
+    lua_getglobal(lua,"math");
+
+    lua_pushstring(lua,"random");
+    lua_pushcfunction(lua,redis_math_random);
+    lua_settable(lua,-3);
+
+    lua_pushstring(lua,"randomseed");
+    lua_pushcfunction(lua,redis_math_randomseed);
+    lua_settable(lua,-3);
+
+    lua_setglobal(lua,"math");
+
+    /* Add a helper funciton that we use to sort the multi bulk output of non
+     * deterministic commands, when containing 'false' elements. */
+    {
+        char *compare_func =    "function __redis__compare_helper(a,b)\n"
+                                "  if a == false then a = '' end\n"
+                                "  if b == false then b = '' end\n"
+                                "  return a<b\n"
+                                "end\n";
+        luaL_loadbuffer(lua,compare_func,strlen(compare_func),"@cmp_func_def");
+        lua_pcall(lua,0,0,0);
+    }
+

     /* Create the (non connected) client that we use to execute Redis commands
     /* Create the (non connected) client that we use to execute Redis commands

-     * inside the Lua interpreter */
-    server.lua_client = createClient(-1);
-    server.lua_client->flags |= REDIS_LUA_CLIENT;
+     * inside the Lua interpreter.
+     * Note: there is no need to create it again when this function is called
+     * by scriptingReset(). */
+    if (server.lua_client == NULL) {
+        server.lua_client = createClient(-1);
+        server.lua_client->flags |= REDIS_LUA_CLIENT;
+    }
+
+    /* Lua beginners ofter don't use "local", this is likely to introduce
+     * subtle bugs in their code. To prevent problems we protect accesses
+     * to global variables. */
+    scriptingEnableGlobalsProtection(lua);

 
     server.lua = lua;
 }
 
 
     server.lua = lua;
 }
 

-/* Hash the scripit into a SHA1 digest. We use this as Lua function name.
- * Digest should point to a 41 bytes buffer: 40 for SHA1 converted into an
+/* Release resources related to Lua scripting.
+ * This function is used in order to reset the scripting environment. */
+void scriptingRelease(void) {
+    dictRelease(server.lua_scripts);
+    lua_close(server.lua);
+}
+
+void scriptingReset(void) {
+    scriptingRelease();
+    scriptingInit();
+}
+
+/* Perform the SHA1 of the input string. We use this both for hasing script
+ * bodies in order to obtain the Lua function name, and in the implementation
+ * of redis.sha1().
+ *
+ * 'digest' should point to a 41 bytes buffer: 40 for SHA1 converted into an

  * hexadecimal number, plus 1 byte for null term. */
  * hexadecimal number, plus 1 byte for null term. */

-void hashScript(char *digest, char *script, size_t len) {
+void sha1hex(char *digest, char *script, size_t len) {

     SHA1_CTX ctx;
     unsigned char hash[20];
     char *cset = "0123456789abcdef";
     SHA1_CTX ctx;
     unsigned char hash[20];
     char *cset = "0123456789abcdef";


@@ -332,8 +682,10 @@ void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
         lua_gettable(lua,-2);
         t = lua_type(lua,-1);
         if (t == LUA_TSTRING) {
         lua_gettable(lua,-2);
         t = lua_type(lua,-1);
         if (t == LUA_TSTRING) {

-            addReplySds(c,sdscatprintf(sdsempty(),
-                    "-%s\r\n",(char*)lua_tostring(lua,-1)));
+            sds err = sdsnew(lua_tostring(lua,-1));
+            sdsmapchars(err,"\r\n","  ",2);
+            addReplySds(c,sdscatprintf(sdsempty(),"-%s\r\n",err));
+            sdsfree(err);

             lua_pop(lua,2);
             return;
         }
             lua_pop(lua,2);
             return;
         }


@@ -343,8 +695,10 @@ void luaReplyToRedisReply(redisClient *c, lua_State *lua) {
         lua_gettable(lua,-2);
         t = lua_type(lua,-1);
         if (t == LUA_TSTRING) {
         lua_gettable(lua,-2);
         t = lua_type(lua,-1);
         if (t == LUA_TSTRING) {

-            addReplySds(c,sdscatprintf(sdsempty(),
-                    "+%s\r\n",(char*)lua_tostring(lua,-1)));
+            sds ok = sdsnew(lua_tostring(lua,-1));
+            sdsmapchars(ok,"\r\n","  ",2);
+            addReplySds(c,sdscatprintf(sdsempty(),"+%s\r\n",ok));
+            sdsfree(ok);

             lua_pop(lua,1);
         } else {
             void *replylen = addDeferredMultiBulkLength(c);
             lua_pop(lua,1);
         } else {
             void *replylen = addDeferredMultiBulkLength(c);


@@ -384,10 +738,71 @@ void luaSetGlobalArray(lua_State *lua, char *var, robj **elev, int elec) {
     lua_setglobal(lua,var);
 }
 
     lua_setglobal(lua,var);
 }
 

+/* Define a lua function with the specified function name and body.
+ * The function name musts be a 2 characters long string, since all the
+ * functions we defined in the Lua context are in the form:
+ *
+ *   f_<hex sha1 sum>
+ *
+ * On success REDIS_OK is returned, and nothing is left on the Lua stack.
+ * On error REDIS_ERR is returned and an appropriate error is set in the
+ * client context. */
+int luaCreateFunction(redisClient *c, lua_State *lua, char *funcname, robj *body) {
+    sds funcdef = sdsempty();
+
+    funcdef = sdscat(funcdef,"function ");
+    funcdef = sdscatlen(funcdef,funcname,42);
+    funcdef = sdscatlen(funcdef,"() ",3);
+    funcdef = sdscatlen(funcdef,body->ptr,sdslen(body->ptr));
+    funcdef = sdscatlen(funcdef," end",4);
+
+    if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"@user_script")) {
+        addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
+            lua_tostring(lua,-1));
+        lua_pop(lua,1);
+        sdsfree(funcdef);
+        return REDIS_ERR;
+    }
+    sdsfree(funcdef);
+    if (lua_pcall(lua,0,0,0)) {
+        addReplyErrorFormat(c,"Error running script (new function): %s\n",
+            lua_tostring(lua,-1));
+        lua_pop(lua,1);
+        return REDIS_ERR;
+    }
+
+    /* We also save a SHA1 -> Original script map in a dictionary
+     * so that we can replicate / write in the AOF all the
+     * EVALSHA commands as EVAL using the original script. */
+    {
+        int retval = dictAdd(server.lua_scripts,
+                             sdsnewlen(funcname+2,40),body);
+        redisAssertWithInfo(c,NULL,retval == DICT_OK);
+        incrRefCount(body);
+    }
+    return REDIS_OK;
+}
+

 void evalGenericCommand(redisClient *c, int evalsha) {
     lua_State *lua = server.lua;
     char funcname[43];
     long long numkeys;
 void evalGenericCommand(redisClient *c, int evalsha) {
     lua_State *lua = server.lua;
     char funcname[43];
     long long numkeys;

+    int delhook = 0;
+
+    /* We want the same PRNG sequence at every call so that our PRNG is
+     * not affected by external state. */
+    redisSrand48(0);
+
+    /* We set this flag to zero to remember that so far no random command
+     * was called. This way we can allow the user to call commands like
+     * SRANDMEMBER or RANDOMKEY from Lua scripts as far as no write command
+     * is called (otherwise the replication and AOF would end with non
+     * deterministic sequences).
+     *
+     * Thanks to this flag we'll raise an error every time a write command
+     * is called after a random command was used. */
+    server.lua_random_dirty = 0;
+    server.lua_write_dirty = 0;

 
     /* Get the number of arguments that are keys */
     if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK)
 
     /* Get the number of arguments that are keys */
     if (getLongLongFromObjectOrReply(c,c->argv[2],&numkeys,NULL) != REDIS_OK)


@@ -403,7 +818,7 @@ void evalGenericCommand(redisClient *c, int evalsha) {
     funcname[1] = '_';
     if (!evalsha) {
         /* Hash the code if this is an EVAL call */
     funcname[1] = '_';
     if (!evalsha) {
         /* Hash the code if this is an EVAL call */

-        hashScript(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
+        sha1hex(funcname+2,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));

     } else {
         /* We already have the SHA if it is a EVALSHA */
         int j;
     } else {
         /* We already have the SHA if it is a EVALSHA */
         int j;


@@ -414,43 +829,21 @@ void evalGenericCommand(redisClient *c, int evalsha) {
         funcname[42] = '\0';
     }
 
         funcname[42] = '\0';
     }
 

+    /* Try to lookup the Lua function */

     lua_getglobal(lua, funcname);
     if (lua_isnil(lua,1)) {
     lua_getglobal(lua, funcname);
     if (lua_isnil(lua,1)) {

-        sds funcdef;
-      
+        lua_pop(lua,1); /* remove the nil from the stack */

         /* Function not defined... let's define it if we have the
          * body of the funciton. If this is an EVALSHA call we can just
          * return an error. */
         if (evalsha) {
             addReply(c, shared.noscripterr);
         /* Function not defined... let's define it if we have the
          * body of the funciton. If this is an EVALSHA call we can just
          * return an error. */
         if (evalsha) {
             addReply(c, shared.noscripterr);

-            lua_pop(lua,1); /* remove the nil from the stack */
-            return;
-        }
-        funcdef = sdsempty();
-
-        lua_pop(lua,1); /* remove the nil from the stack */
-        funcdef = sdscat(funcdef,"function ");
-        funcdef = sdscatlen(funcdef,funcname,42);
-        funcdef = sdscatlen(funcdef," ()\n",4);
-        funcdef = sdscatlen(funcdef,c->argv[1]->ptr,sdslen(c->argv[1]->ptr));
-        funcdef = sdscatlen(funcdef,"\nend\n",5);
-        /* printf("Defining:\n%s\n",funcdef); */
-
-        if (luaL_loadbuffer(lua,funcdef,sdslen(funcdef),"func definition")) {
-            addReplyErrorFormat(c,"Error compiling script (new function): %s\n",
-                lua_tostring(lua,-1));
-            lua_pop(lua,1);
-            sdsfree(funcdef);
-            return;
-        }
-        sdsfree(funcdef);
-        if (lua_pcall(lua,0,0,0)) {
-            addReplyErrorFormat(c,"Error running script (new function): %s\n",
-                lua_tostring(lua,-1));
-            lua_pop(lua,1);

             return;
         }
             return;
         }

+        if (luaCreateFunction(c,lua,funcname,c->argv[1]) == REDIS_ERR) return;
+        /* Now the following is guaranteed to return non nil */

         lua_getglobal(lua, funcname);
         lua_getglobal(lua, funcname);

+        redisAssert(!lua_isnil(lua,1));

     }
 
     /* Populate the argv and keys table accordingly to the arguments that
     }
 
     /* Populate the argv and keys table accordingly to the arguments that


@@ -465,17 +858,27 @@ void evalGenericCommand(redisClient *c, int evalsha) {
      * is running for too much time.
      * We set the hook only if the time limit is enabled as the hook will
      * make the Lua script execution slower. */
      * is running for too much time.
      * We set the hook only if the time limit is enabled as the hook will
      * make the Lua script execution slower. */

-    if (server.lua_time_limit > 0) {
+    server.lua_caller = c;
+    server.lua_time_start = ustime()/1000;
+    server.lua_kill = 0;
+    if (server.lua_time_limit > 0 && server.masterhost == NULL) {

         lua_sethook(lua,luaMaskCountHook,LUA_MASKCOUNT,100000);
         lua_sethook(lua,luaMaskCountHook,LUA_MASKCOUNT,100000);

-        server.lua_time_start = ustime()/1000;
-    } else {
-        lua_sethook(lua,luaMaskCountHook,0,0);
+        delhook = 1;

     }
 
     /* At this point whatever this script was never seen before or if it was
      * already defined, we can call it. We have zero arguments and expect
      * a single return value. */
     if (lua_pcall(lua,0,1,0)) {
     }
 
     /* At this point whatever this script was never seen before or if it was
      * already defined, we can call it. We have zero arguments and expect
      * a single return value. */
     if (lua_pcall(lua,0,1,0)) {

+        if (delhook) lua_sethook(lua,luaMaskCountHook,0,0); /* Disable hook */
+        if (server.lua_timedout) {
+            server.lua_timedout = 0;
+            /* Restore the readable handler that was unregistered when the
+             * script timeout was detected. */
+            aeCreateFileEvent(server.el,c->fd,AE_READABLE,
+                              readQueryFromClient,c);
+        }
+        server.lua_caller = NULL;

         selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
         addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
             funcname, lua_tostring(lua,-1));
         selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
         addReplyErrorFormat(c,"Error running script (call to %s): %s\n",
             funcname, lua_tostring(lua,-1));


@@ -483,9 +886,31 @@ void evalGenericCommand(redisClient *c, int evalsha) {
         lua_gc(lua,LUA_GCCOLLECT,0);
         return;
     }
         lua_gc(lua,LUA_GCCOLLECT,0);
         return;
     }

+    if (delhook) lua_sethook(lua,luaMaskCountHook,0,0); /* Disable hook */
+    server.lua_timedout = 0;
+    server.lua_caller = NULL;

     selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
     luaReplyToRedisReply(c,lua);
     lua_gc(lua,LUA_GCSTEP,1);
     selectDb(c,server.lua_client->db->id); /* set DB ID from Lua client */
     luaReplyToRedisReply(c,lua);
     lua_gc(lua,LUA_GCSTEP,1);

+
+    /* If we have slaves attached we want to replicate this command as
+     * EVAL instead of EVALSHA. We do this also in the AOF as currently there
+     * is no easy way to propagate a command in a different way in the AOF
+     * and in the replication link.
+     *
+     * IMPROVEMENT POSSIBLE:
+     * 1) Replicate this command as EVALSHA in the AOF.
+     * 2) Remember what slave already received a given script, and replicate
+     *    the EVALSHA against this slaves when possible.
+     */
+    if (evalsha) {
+        robj *script = dictFetchValue(server.lua_scripts,c->argv[1]->ptr);
+
+        redisAssertWithInfo(c,NULL,script != NULL);
+        rewriteClientCommandArgument(c,0,
+            resetRefCount(createStringObject("EVAL",4)));
+        rewriteClientCommandArgument(c,1,script);
+    }

 }
 
 void evalCommand(redisClient *c) {
 }
 
 void evalCommand(redisClient *c) {


@@ -503,3 +928,92 @@ void evalShaCommand(redisClient *c) {
     }
     evalGenericCommand(c,1);
 }
     }
     evalGenericCommand(c,1);
 }

+
+/* We replace math.random() with our implementation that is not affected
+ * by specific libc random() implementations and will output the same sequence
+ * (for the same seed) in every arch. */
+
+/* The following implementation is the one shipped with Lua itself but with
+ * rand() replaced by redisLrand48(). */
+int redis_math_random (lua_State *L) {
+  /* the `%' avoids the (rare) case of r==1, and is needed also because on
+     some systems (SunOS!) `rand()' may return a value larger than RAND_MAX */
+  lua_Number r = (lua_Number)(redisLrand48()%REDIS_LRAND48_MAX) /
+                                (lua_Number)REDIS_LRAND48_MAX;
+  switch (lua_gettop(L)) {  /* check number of arguments */
+    case 0: {  /* no arguments */
+      lua_pushnumber(L, r);  /* Number between 0 and 1 */
+      break;
+    }
+    case 1: {  /* only upper limit */
+      int u = luaL_checkint(L, 1);
+      luaL_argcheck(L, 1<=u, 1, "interval is empty");
+      lua_pushnumber(L, floor(r*u)+1);  /* int between 1 and `u' */
+      break;
+    }
+    case 2: {  /* lower and upper limits */
+      int l = luaL_checkint(L, 1);
+      int u = luaL_checkint(L, 2);
+      luaL_argcheck(L, l<=u, 2, "interval is empty");
+      lua_pushnumber(L, floor(r*(u-l+1))+l);  /* int between `l' and `u' */
+      break;
+    }
+    default: return luaL_error(L, "wrong number of arguments");
+  }
+  return 1;
+}
+
+int redis_math_randomseed (lua_State *L) {
+  redisSrand48(luaL_checkint(L, 1));
+  return 0;
+}
+
+/* ---------------------------------------------------------------------------
+ * SCRIPT command for script environment introspection and control
+ * ------------------------------------------------------------------------- */
+
+void scriptCommand(redisClient *c) {
+    if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"flush")) {
+        scriptingReset();
+        addReply(c,shared.ok);
+        server.dirty++; /* Replicating this command is a good idea. */
+    } else if (c->argc >= 2 && !strcasecmp(c->argv[1]->ptr,"exists")) {
+        int j;
+
+        addReplyMultiBulkLen(c, c->argc-2);
+        for (j = 2; j < c->argc; j++) {
+            if (dictFind(server.lua_scripts,c->argv[j]->ptr))
+                addReply(c,shared.cone);
+            else
+                addReply(c,shared.czero);
+        }
+    } else if (c->argc == 3 && !strcasecmp(c->argv[1]->ptr,"load")) {
+        char funcname[43];
+        sds sha;
+
+        funcname[0] = 'f';
+        funcname[1] = '_';
+        sha1hex(funcname+2,c->argv[2]->ptr,sdslen(c->argv[2]->ptr));
+        sha = sdsnewlen(funcname+2,40);
+        if (dictFind(server.lua_scripts,sha) == NULL) {
+            if (luaCreateFunction(c,server.lua,funcname,c->argv[2])
+                    == REDIS_ERR) {
+                sdsfree(sha);
+                return;
+            }
+        }
+        addReplyBulkCBuffer(c,funcname+2,40);
+        sdsfree(sha);
+    } else if (c->argc == 2 && !strcasecmp(c->argv[1]->ptr,"kill")) {
+        if (server.lua_caller == NULL) {
+            addReplySds(c,sdsnew("-NOTBUSY No scripts in execution right now.\r\n"));
+        } else if (server.lua_write_dirty) {
+            addReplySds(c,sdsnew("-UNKILLABLE 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.\r\n"));
+        } else {
+            server.lua_kill = 1;
+            addReply(c,shared.ok);
+        }
+    } else {
+        addReplyError(c, "Unknown SCRIPT subcommand or wrong # of args.");
+    }
+}