Reject untrusted Lua bytecode while preserving script caching

benchmark/BDN.benchmark/Lua/LuaScriptCacheOperations.cs

@@ -51,20 +51,28 @@ public void GlobalSetup()
 
             outerHitDigest = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
             sessionScriptCache.GetScriptDigest("return 1"u8, outerHitDigest);
-            if (!storeWrapper.storeScriptCache.TryAdd(new(outerHitDigest), new("return 1"u8.ToArray())))
+            if (!storeWrapper.storeScriptCache.TryAdd(new(outerHitDigest), CompileScript("return 1"u8)))
             {
                 throw new InvalidOperationException("Should have been able to load into global cache");
             }
 
             innerHitDigest = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
             sessionScriptCache.GetScriptDigest("return 1 + 1"u8, innerHitDigest);
-            if (!storeWrapper.storeScriptCache.TryAdd(new(innerHitDigest), new("return 1 + 1"u8.ToArray())))
+            if (!storeWrapper.storeScriptCache.TryAdd(new(innerHitDigest), CompileScript("return 1 + 1"u8)))
             {
                 throw new InvalidOperationException("Should have been able to load into global cache");
             }
 
             missDigest = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
             sessionScriptCache.GetScriptDigest("foobar"u8, missDigest);
+
+            static LuaScriptHandle CompileScript(ReadOnlySpan<byte> source)
+            {
+                if (!LuaRunner.TryCompileSource(source, out var generatedBytecode, out var error))
+                    throw new InvalidOperationException(error);
+
+                return new(generatedBytecode.Data);
+            }
         }
 
         [GlobalCleanup]
@@ -81,8 +89,7 @@ public void IterationSetup()
             sessionScriptCache.Clear();
 
             // Make outer hit available for every iteration
-            LuaScriptHandle scriptHandle = null;
-            if (!sessionScriptCache.TryLoad(session, "return 1"u8, new(outerHitDigest), ref scriptHandle, out _, out _))
+            if (!sessionScriptCache.TryGetOrCreateRunnerFromSource(session, "return 1"u8, new(outerHitDigest), out _, out _, out _))
             {
                 throw new InvalidOperationException("Should have been able to load");
             }
@@ -147,10 +154,9 @@ private void LoadScript(Span<byte> digest)
             {
                 if (storeWrapper.storeScriptCache.TryGetValue(digestKey, out var scriptHandle))
                 {
-                    LuaScriptHandle newScriptHandle = null;
-                    if (!sessionScriptCache.TryLoad(session, scriptHandle.ScriptData.Span, digestKey, ref newScriptHandle, out runner, out _))
+                    if (!sessionScriptCache.TryGetOrCreateRunnerFromCachedScript(session, digestKey, scriptHandle, out runner))
                     {
-                        // TryLoad will have written an error out, it any
+                        // The loading error was already written, if any
 
                         _ = storeWrapper.storeScriptCache.TryRemove(digestKey, out _);
                     }

libs/server/Lua/LuaCommands.cs

@@ -48,9 +48,9 @@ private unsafe bool TryEVALSHA()
                 {
                     if (storeWrapper.storeScriptCache.TryGetValue(scriptKey, out var globalScriptHandle))
                     {
-                        if (!sessionScriptCache.TryLoad(this, globalScriptHandle.ScriptData.Span, scriptKey, ref globalScriptHandle, out runner, out _))
+                        if (!sessionScriptCache.TryGetOrCreateRunnerFromCachedScript(this, scriptKey, globalScriptHandle, out runner))
                         {
-                            // TryLoad will have written an error out, it any
+                            // The loading error was already written, if any
                             //
                             // Note we DON'T dispose the script handle because this is just the session cache
                             _ = storeWrapper.storeScriptCache.TryRemove(scriptKey, out _);
@@ -114,43 +114,8 @@ private unsafe bool TryEVAL()
             sessionScriptCache.GetScriptDigest(script.ReadOnlySpan, digest);
 
             var onStackScriptKey = new ScriptHashKey(digest);
-            _ = storeWrapper.storeScriptCache.TryGetValue(onStackScriptKey, out var globalScriptHandle);
-
-            var sessionScriptHandle = globalScriptHandle;
-
-            if (!sessionScriptCache.TryLoad(this, script.ReadOnlySpan, onStackScriptKey, ref sessionScriptHandle, out var runner, out var digestOnHeap))
-            {
-                // TryLoad will have written any errors out
-                return true;
-            }
-            else if (sessionScriptHandle != globalScriptHandle)
-            {
-                // Add script to the store dictionary IF we didn't already have it cached
-                //
-                // This may strike you as odd, but it is how Redis behaves
-                if (digestOnHeap == null)
-                {
-                    var newAlloc = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
-                    digest.CopyTo(newAlloc);
-                    if (!storeWrapper.storeScriptCache.TryAdd(new(newAlloc), sessionScriptHandle))
-                    {
-                        // Some other session loaded the script, toss our new handle
-                        //
-                        // Next time this script is run, it'll be pulled from the global cache
-                        sessionScriptHandle.Dispose();
-                    }
-                }
-                else
-                {
-                    if (!storeWrapper.storeScriptCache.TryAdd(digestOnHeap.Value, sessionScriptHandle))
-                    {
-                        // Some other session loaded the script, toss our new handle
-                        //
-                        // Next time this script is run, it'll be pulled from the global cache
-                        sessionScriptHandle.Dispose();
-                    }
-                }
-            }
+            if (!TryGetOrCreateScriptRunner(script.ReadOnlySpan, digest, onStackScriptKey, out var runner))
+                return true; // The loading error was written to the response.
 
             if (runner == null)
             {
@@ -275,45 +240,41 @@ private bool NetworkScriptLoad()
             sessionScriptCache.GetScriptDigest(source.Span, digest);
 
             var onStackScriptHashKey = new ScriptHashKey(digest);
-            _ = storeWrapper.storeScriptCache.TryGetValue(onStackScriptHashKey, out var globalScriptHandle);
-
-            var sessionScriptHandle = globalScriptHandle;
-            if (sessionScriptCache.TryLoad(this, source.ReadOnlySpan, onStackScriptHashKey, ref sessionScriptHandle, out _, out var digestOnHeap))
+            if (TryGetOrCreateScriptRunner(source.ReadOnlySpan, digest, onStackScriptHashKey, out _))
             {
-                // TryLoad will write any errors out
-
-                // Add script to the global store dictionary if not already in there
-                if (globalScriptHandle != sessionScriptHandle)
-                {
-                    if (digestOnHeap == null)
-                    {
-                        var newAlloc = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
-                        digest.CopyTo(newAlloc);
-                        if (!storeWrapper.storeScriptCache.TryAdd(new(newAlloc), sessionScriptHandle))
-                        {
-                            // Some other caller added the script already, our new handle is dead
-                            // but we'll load it from the shared cache on next invocation
-                            sessionScriptHandle.Dispose();
-                        }
-                    }
-                    else
-                    {
-                        if (!storeWrapper.storeScriptCache.TryAdd(digestOnHeap.Value, sessionScriptHandle))
-                        {
-                            // Some other caller added the script already, our new handle is dead
-                            // but we'll load it from the shared cache on next invocation
-                            sessionScriptHandle.Dispose();
-                        }
-                    }
-                }
-
                 while (!RespWriteUtils.TryWriteBulkString(digest, ref dcurr, dend))
                     SendAndReset();
             }
 
             return true;
         }
 
+        private bool TryGetOrCreateScriptRunner(ReadOnlySpan<byte> source, ReadOnlySpan<byte> digest, ScriptHashKey scriptKey, out LuaRunner runner)
+        {
+            if (storeWrapper.storeScriptCache.TryGetValue(scriptKey, out var globalScriptHandle))
+                return sessionScriptCache.TryGetOrCreateRunnerFromCachedScript(this, scriptKey, globalScriptHandle, out runner);
+
+            if (!sessionScriptCache.TryGetOrCreateRunnerFromSource(this, source, scriptKey, out var sessionScriptHandle, out runner, out var digestOnHeap))
+                return false;
+
+            ScriptHashKey globalScriptKey;
+            if (digestOnHeap != null)
+            {
+                globalScriptKey = digestOnHeap.Value;
+            }
+            else
+            {
+                var digestCopy = GC.AllocateUninitializedArray<byte>(SessionScriptCache.SHA1Len, pinned: true);
+                digest.CopyTo(digestCopy);
+                globalScriptKey = new(digestCopy);
+            }
+
+            if (!storeWrapper.storeScriptCache.TryAdd(globalScriptKey, sessionScriptHandle))
+                sessionScriptHandle.Dispose();
+
+            return true;
+        }
+
         /// <summary>
         /// Returns true if Lua is enabled.
         /// 
@@ -332,6 +293,12 @@ private bool CheckLuaEnabled()
             return true;
         }
 
+        internal void WriteLuaCompilationError(string error)
+        {
+            while (!RespWriteUtils.TryWriteError($"Compilation error: {error}", ref dcurr, dend))
+                SendAndReset();
+        }
+
         /// <summary>
         /// Run a resolved script for the current session.
         /// 

libs/server/Lua/LuaRunner.Functions.cs

@@ -852,7 +852,7 @@ internal int LoadString(nint luaStatePtr)
                 return LuaWrappedError(1, constStrs.InsufficientLuaStackSpace);
             }
 
-            var res = state.LoadString(buff);
+            var res = state.LoadTextBuffer(buff);
             if (res != LuaStatus.OK)
             {
                 state.ClearStack();
@@ -3071,8 +3071,8 @@ int luaArgCount
         private unsafe int CompileCommon<TResponse>(nint luaState, ref TResponse resp)
             where TResponse : struct, IResponseAdapter
         {
-            // 1 for function, 1 for code string
-            const int NeededStackSpace = 2;
+            // 1 for function, 1 for code string, 1 for load mode
+            const int NeededStackSpace = 3;
 
             Debug.Assert(functionRegistryIndex == -1, "Shouldn't compile multiple times");
 
@@ -3081,15 +3081,19 @@ private unsafe int CompileCommon<TResponse>(nint luaState, ref TResponse resp)
             Debug.Assert(state.TryEnsureMinimumStackCapacity(NeededStackSpace), "LUA_MIN_STACK should be high enough that this cannot happen");
 
             _ = state.RawGetInteger(LuaType.Function, (int)LuaRegistry.Index, loadSandboxedRegistryIndex);
-            if (!state.TryPushBuffer(source.Span))
+            if (!state.TryPushBuffer(source.Data.Span))
             {
-                while (!RespWriteUtils.TryWriteError(CmdStrings.LUA_out_of_memory, ref resp.BufferCur, resp.BufferEnd))
-                    resp.SendAndReset();
+                WriteOutOfMemoryError(ref resp);
+                return 0;
+            }
 
+            if (!state.TryPushBuffer(source.Kind == LuaScriptChunkKind.GarnetGeneratedBinary ? "b"u8 : "t"u8))
+            {
+                WriteOutOfMemoryError(ref resp);
                 return 0;
             }
 
-            var callRes = state.PCall(1, 2);
+            var callRes = state.PCall(2, 2);
 
             // On success the stack will have two things on it:
             //  1. The error (nil if not error)
@@ -3104,9 +3108,7 @@ private unsafe int CompileCommon<TResponse>(nint luaState, ref TResponse resp)
                 if (!state.TryRef(out functionRegistryIndex))
                 {
                     // Uh-oh, couldn't save the function under the registry
-                    while (!RespWriteUtils.TryWriteError(CmdStrings.LUA_out_of_memory, ref resp.BufferCur, resp.BufferEnd))
-                        resp.SendAndReset();
-
+                    WriteOutOfMemoryError(ref resp);
                     return 0;
                 }
             }
@@ -3131,6 +3133,12 @@ private unsafe int CompileCommon<TResponse>(nint luaState, ref TResponse resp)
             return 0;
         }
 
+        private static unsafe void WriteOutOfMemoryError<TResponse>(ref TResponse resp) where TResponse : struct, IResponseAdapter
+        {
+            while (!RespWriteUtils.TryWriteError(CmdStrings.LUA_out_of_memory, ref resp.BufferCur, resp.BufferEnd))
+                resp.SendAndReset();
+        }
+
         /// <summary>
         /// Entry point method for executing commands from a Lua Script
         /// </summary>

libs/server/Lua/LuaRunner.Loader.cs

@@ -314,8 +314,8 @@ function reset_keys_and_argv(fromKey, fromArgv)
 -- force new 'global' environment to be readonly
 recursively_readonly_table(sandbox_env)
 -- responsible for sandboxing user provided code
-function load_sandboxed(source)
-    local rawFunc, err = load(source, nil, nil, sandbox_env)
+function load_sandboxed(source, mode)
+    local rawFunc, err = load(source, nil, mode, sandbox_env)
 
     return err, rawFunc
 end
@@ -448,7 +448,7 @@ internal static ReadOnlyMemory<byte> PrepareLoaderBlockBytes(HashSet<string> all
 
                 compilingState.Remove(1);
 
-                if (compilingState.LoadString(Encoding.UTF8.GetBytes(finalLoaderBlock)) != LuaStatus.OK)
+                if (compilingState.LoadTextBuffer(Encoding.UTF8.GetBytes(finalLoaderBlock)) != LuaStatus.OK)
                 {
                     throw new InvalidOperationException("Compiling function should not fail");
                 }
@@ -479,7 +479,7 @@ internal static ReadOnlyMemory<byte> PrepareLoaderBlockBytes(HashSet<string> all
         /// 
         /// These ops are faster to load into a runtime than parsing the whole source file again.
         /// </summary>
-        internal static byte[] CompileSource(ReadOnlySpan<byte> source)
+        internal static bool TryCompileSource(ReadOnlySpan<byte> source, out LuaScriptChunk compiledSource, out string error)
         {
             // This is equivalent to calling 
             //
@@ -496,23 +496,39 @@ internal static byte[] CompileSource(ReadOnlySpan<byte> source)
 
             state.Remove(1);
 
-            if (state.LoadString(source) != LuaStatus.OK)
+            if (state.LoadTextBuffer(source) != LuaStatus.OK)
             {
-                // If we're going to fail, just keep the source as is - a future load attempt will fail it too
-                return source.ToArray();
+                compiledSource = default;
+                error = GetErrorFromStackTop(state);
+                return false;
             }
 
             state.PushBoolean(true);
 
             if (state.PCall(2, 1) != LuaStatus.OK)
             {
-                // If we're going to fail, just keep the source as is - a future load attempt will fail it too
-                return source.ToArray();
+                compiledSource = default;
+                error = GetErrorFromStackTop(state);
+                return false;
             }
 
             state.KnownStringToBuffer(1, out var ops);
 
-            return ops.ToArray();
+            compiledSource = new(ops.ToArray(), LuaScriptChunkKind.GarnetGeneratedBinary);
+            error = null;
+            return true;
+
+            static string GetErrorFromStackTop(LuaStateWrapper state)
+            {
+                var errorIndex = state.StackTop;
+                if (errorIndex >= 1 && state.Type(errorIndex) == LuaType.String)
+                {
+                    state.KnownStringToBuffer(errorIndex, out var errorBuffer);
+                    return Encoding.UTF8.GetString(errorBuffer);
+                }
+
+                return "cause unknown";
+            }
         }
     }
 }

libs/server/Lua/LuaRunner.cs

@@ -160,7 +160,7 @@ public void SendAndReset()
 
         readonly LuaLoggingMode logMode;
         readonly HashSet<string> allowedFunctions;
-        readonly ReadOnlyMemory<byte> source;
+        readonly LuaScriptChunk source;
         readonly ScratchBufferNetworkSender scratchBufferNetworkSender;
         readonly RespServerSession respServerSession;
 
@@ -208,7 +208,7 @@ public unsafe LuaRunner(
             int? memLimitBytes,
             LuaLoggingMode logMode,
             HashSet<string> allowedFunctions,
-            ReadOnlyMemory<byte> source,
+            LuaScriptChunk source,
             bool txnMode = false,
             RespServerSession respServerSession = null,
             ScratchBufferNetworkSender scratchBufferNetworkSender = null,
@@ -328,7 +328,7 @@ public unsafe LuaRunner(
                 throw new GarnetException("Insufficient space in Lua VM for redis version number global");
             }
 
-            var loadRes = state.LoadBuffer(PrepareLoaderBlockBytes(allowedFunctions, logger).Span);
+            var loadRes = state.LoadBinaryBuffer(PrepareLoaderBlockBytes(allowedFunctions, logger).Span);
             if (loadRes != LuaStatus.OK)
             {
                 if (state.StackTop == 1 && state.Type(1) == LuaType.String)
@@ -409,7 +409,7 @@ static void Register(ref LuaStateWrapper state, ReadOnlySpan<byte> name, delegat
         /// Creates a new runner with the source of the script
         /// </summary>
         public LuaRunner(LuaOptions options, string source, bool txnMode = false, RespServerSession respServerSession = null, ScratchBufferNetworkSender scratchBufferNetworkSender = null, string redisVersion = "0.0.0.0", ILogger logger = null)
-            : this(options.MemoryManagementMode, options.GetMemoryLimitBytes(), options.LogMode, options.AllowedFunctions, Encoding.UTF8.GetBytes(source), txnMode, respServerSession, scratchBufferNetworkSender, redisVersion, logger)
+            : this(options.MemoryManagementMode, options.GetMemoryLimitBytes(), options.LogMode, options.AllowedFunctions, new LuaScriptChunk(Encoding.UTF8.GetBytes(source), LuaScriptChunkKind.Text), txnMode, respServerSession, scratchBufferNetworkSender, redisVersion, logger)
         {
         }