diff --git a/benchmark/BDN.benchmark/Lua/LuaScriptCacheOperations.cs b/benchmark/BDN.benchmark/Lua/LuaScriptCacheOperations.cs index 72558d01f1c..df9dbf50741 100644 --- a/benchmark/BDN.benchmark/Lua/LuaScriptCacheOperations.cs +++ b/benchmark/BDN.benchmark/Lua/LuaScriptCacheOperations.cs @@ -51,20 +51,28 @@ public void GlobalSetup() outerHitDigest = GC.AllocateUninitializedArray(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(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(SessionScriptCache.SHA1Len, pinned: true); sessionScriptCache.GetScriptDigest("foobar"u8, missDigest); + + static LuaScriptHandle CompileScript(ReadOnlySpan 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 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 _); } diff --git a/libs/server/Lua/LuaCommands.cs b/libs/server/Lua/LuaCommands.cs index 4a37d88e902..b60648b3367 100644 --- a/libs/server/Lua/LuaCommands.cs +++ b/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(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,38 +240,8 @@ 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(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(); } @@ -314,6 +249,32 @@ private bool NetworkScriptLoad() return true; } + private bool TryGetOrCreateScriptRunner(ReadOnlySpan source, ReadOnlySpan 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(SessionScriptCache.SHA1Len, pinned: true); + digest.CopyTo(digestCopy); + globalScriptKey = new(digestCopy); + } + + if (!storeWrapper.storeScriptCache.TryAdd(globalScriptKey, sessionScriptHandle)) + sessionScriptHandle.Dispose(); + + return true; + } + /// /// 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(); + } + /// /// Run a resolved script for the current session. /// diff --git a/libs/server/Lua/LuaRunner.Functions.cs b/libs/server/Lua/LuaRunner.Functions.cs index 23a6295df57..a82cd8c96ad 100644 --- a/libs/server/Lua/LuaRunner.Functions.cs +++ b/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(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(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(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(nint luaState, ref TResponse resp) return 0; } + private static unsafe void WriteOutOfMemoryError(ref TResponse resp) where TResponse : struct, IResponseAdapter + { + while (!RespWriteUtils.TryWriteError(CmdStrings.LUA_out_of_memory, ref resp.BufferCur, resp.BufferEnd)) + resp.SendAndReset(); + } + /// /// Entry point method for executing commands from a Lua Script /// diff --git a/libs/server/Lua/LuaRunner.Loader.cs b/libs/server/Lua/LuaRunner.Loader.cs index a9a95ef9e53..c7e23f061d5 100644 --- a/libs/server/Lua/LuaRunner.Loader.cs +++ b/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 PrepareLoaderBlockBytes(HashSet 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 PrepareLoaderBlockBytes(HashSet all /// /// These ops are faster to load into a runtime than parsing the whole source file again. /// - internal static byte[] CompileSource(ReadOnlySpan source) + internal static bool TryCompileSource(ReadOnlySpan source, out LuaScriptChunk compiledSource, out string error) { // This is equivalent to calling // @@ -496,23 +496,39 @@ internal static byte[] CompileSource(ReadOnlySpan 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"; + } } } } \ No newline at end of file diff --git a/libs/server/Lua/LuaRunner.cs b/libs/server/Lua/LuaRunner.cs index 5b68363b7af..6a60ec09fe4 100644 --- a/libs/server/Lua/LuaRunner.cs +++ b/libs/server/Lua/LuaRunner.cs @@ -160,7 +160,7 @@ public void SendAndReset() readonly LuaLoggingMode logMode; readonly HashSet allowedFunctions; - readonly ReadOnlyMemory source; + readonly LuaScriptChunk source; readonly ScratchBufferNetworkSender scratchBufferNetworkSender; readonly RespServerSession respServerSession; @@ -208,7 +208,7 @@ public unsafe LuaRunner( int? memLimitBytes, LuaLoggingMode logMode, HashSet allowedFunctions, - ReadOnlyMemory 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 name, delegat /// Creates a new runner with the source of the script /// 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) { } diff --git a/libs/server/Lua/LuaScriptHandle.cs b/libs/server/Lua/LuaScriptHandle.cs index caa9f066dfa..d93a688467c 100644 --- a/libs/server/Lua/LuaScriptHandle.cs +++ b/libs/server/Lua/LuaScriptHandle.cs @@ -5,6 +5,14 @@ namespace Garnet.server { + internal enum LuaScriptChunkKind : byte + { + Text, + GarnetGeneratedBinary + } + + internal readonly record struct LuaScriptChunk(ReadOnlyMemory Data, LuaScriptChunkKind Kind); + /// /// Used to track the lifetime a shared Lua script, which may end up backing multiple s. /// @@ -19,14 +27,13 @@ public sealed class LuaScriptHandle : IDisposable public bool IsDisposed { get; private set; } /// - /// Source (or compiled source) for the associated Lua script. + /// Internally generated bytecode for the associated Lua script. /// - public ReadOnlyMemory ScriptData { get; } + public ReadOnlyMemory ScriptData => Chunk.Data; - public LuaScriptHandle(ReadOnlyMemory scriptData) - { - ScriptData = scriptData; - } + internal LuaScriptChunk Chunk { get; } + + internal LuaScriptHandle(ReadOnlyMemory generatedBytecode) => Chunk = new(generatedBytecode, LuaScriptChunkKind.GarnetGeneratedBinary); /// public void Dispose() diff --git a/libs/server/Lua/LuaStateWrapper.cs b/libs/server/Lua/LuaStateWrapper.cs index 48f9fca6203..4ecfa34ed65 100644 --- a/libs/server/Lua/LuaStateWrapper.cs +++ b/libs/server/Lua/LuaStateWrapper.cs @@ -456,20 +456,14 @@ internal bool TrySetGlobal(ReadOnlySpan nullTerminatedGlobalName) } /// - /// This should be used for all LoadBuffers into Lua. - /// - /// Note that this is different from pushing a buffer, as the loaded buffer is compiled and executed. - /// - /// Maintains and to minimize p/invoke calls. + /// Load Garnet-generated bytecode into Lua. /// [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal LuaStatus LoadBuffer(ReadOnlySpan buffer) + internal LuaStatus LoadBinaryBuffer(ReadOnlySpan buffer) { AssertLuaStackNotFull(2); - // Note that https://www.lua.org/source/5.4/lauxlib.c.html#luaL_loadbufferx is implemented in terms of - // a PCall, so we don't have to worry about crashes. - var ret = NativeMethods.LoadBuffer(state, buffer); + var ret = NativeMethods.LoadBinaryBuffer(state, buffer); if (ret != LuaStatus.OK) { @@ -486,19 +480,16 @@ internal LuaStatus LoadBuffer(ReadOnlySpan buffer) } /// - /// This should be used for all LoadStrings into Lua. - /// - /// Note that this is different from pushing or loading buffer, as the loaded buffer is compiled but NOT executed. - /// - /// Maintains and to minimize p/invoke calls. + /// This should be used for compiling untrusted Lua source text. + /// + /// Binary chunks are rejected by the Lua runtime. /// - internal LuaStatus LoadString(ReadOnlySpan buffer) + internal LuaStatus LoadTextBuffer(ReadOnlySpan buffer) { AssertLuaStackNotFull(2); - // Note that https://www.lua.org/source/5.4/lauxlib.h.html#luaL_loadbuffer is implemented in terms of - // a PCall, so we don't have to worry about crashes. - var ret = NativeMethods.LoadString(state, buffer); + // Text-only mode rejects binary chunks at the untrusted input boundary. + var ret = NativeMethods.LoadTextBuffer(state, buffer); if (ret != LuaStatus.OK) { diff --git a/libs/server/Lua/NativeMethods.cs b/libs/server/Lua/NativeMethods.cs index 254b6c43e88..bd24f0b5a6b 100644 --- a/libs/server/Lua/NativeMethods.cs +++ b/libs/server/Lua/NativeMethods.cs @@ -45,13 +45,6 @@ internal static partial class NativeMethods [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])] private static partial LuaStatus luaL_loadbufferx(lua_State luaState, charptr_t buff, size_t sz, charptr_t name, charptr_t mode); - /// - /// see: https://www.lua.org/manual/5.4/manual.html#luaL_loadstring - /// - [LibraryImport(LuaLibraryName)] - [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])] - private static partial LuaStatus luaL_loadstring(lua_State lua_State, charptr_t buff); - /// /// see: https://www.lua.org/manual/5.4/manual.html#luaL_newstate /// @@ -388,28 +381,34 @@ internal static unsafe void PushBuffer(lua_State luaState, ReadOnlySpan st } /// - /// Push given span to stack, compiles it, and executes it. + /// Load Garnet-generated bytecode. /// /// Provided data is copied, and can be reused once this call returns. /// - internal static unsafe LuaStatus LoadBuffer(lua_State luaState, ReadOnlySpan str) + internal static unsafe LuaStatus LoadBinaryBuffer(lua_State luaState, ReadOnlySpan str) { + // Binary-only mode is reserved for bytecode generated internally by Garnet. + ReadOnlySpan mode = "b\0"u8; fixed (byte* ptr = str) + fixed (byte* modePtr = mode) { - return luaL_loadbufferx(luaState, (charptr_t)ptr, (size_t)str.Length, (charptr_t)UIntPtr.Zero, (charptr_t)UIntPtr.Zero); + return luaL_loadbufferx(luaState, (charptr_t)ptr, (size_t)str.Length, (charptr_t)UIntPtr.Zero, (charptr_t)modePtr); } } /// - /// Push given span to stack, and compiles it. - /// + /// Push given text span to stack and compile it. + /// /// Provided data is copied, and can be reused once this call returns. /// - internal static unsafe LuaStatus LoadString(lua_State luaState, ReadOnlySpan str) + internal static unsafe LuaStatus LoadTextBuffer(lua_State luaState, ReadOnlySpan str) { + // Text-only mode rejects externally supplied precompiled bytecode. + ReadOnlySpan mode = "t\0"u8; fixed (byte* ptr = str) + fixed (byte* modePtr = mode) { - return luaL_loadstring(luaState, (charptr_t)ptr); + return luaL_loadbufferx(luaState, (charptr_t)ptr, (size_t)str.Length, (charptr_t)UIntPtr.Zero, (charptr_t)modePtr); } } diff --git a/libs/server/Lua/SessionScriptCache.cs b/libs/server/Lua/SessionScriptCache.cs index 705a0b9deaf..cfe38ab2bf2 100644 --- a/libs/server/Lua/SessionScriptCache.cs +++ b/libs/server/Lua/SessionScriptCache.cs @@ -163,15 +163,15 @@ public bool TryGetFromDigest(ScriptHashKey digest, out LuaRunner scriptRunner, o } /// - /// Load script into the cache. + /// Compile Lua source text and load it into the session cache. /// /// If necessary, will be set so the allocation can be reused. /// - internal bool TryLoad( + internal bool TryGetOrCreateRunnerFromSource( RespServerSession session, ReadOnlySpan source, ScriptHashKey digest, - ref LuaScriptHandle luaScriptHandle, + out LuaScriptHandle luaScriptHandle, out LuaRunner runner, out ScriptHashKey? digestOnHeap ) @@ -183,11 +183,60 @@ out ScriptHashKey? digestOnHeap return true; } + luaScriptHandle = null; + return TryCompileSourceAndCreateRunner(session, source, digest, ref luaScriptHandle, out runner, out digestOnHeap); + } + + /// + /// Load a script previously stored in the global cache. + /// + internal bool TryGetOrCreateRunnerFromCachedScript(RespServerSession session, ScriptHashKey digest, LuaScriptHandle cachedScriptHandle, out LuaRunner runner) + { + if (TryGetFromDigest(digest, out runner, out _)) + return true; + + return TryGetOrCreateRunnerFromGeneratedBytecode(session, cachedScriptHandle.Chunk, digest, ref cachedScriptHandle, out runner, out _); + } + + private bool TryCompileSourceAndCreateRunner(RespServerSession session, ReadOnlySpan source, ScriptHashKey digest, ref LuaScriptHandle luaScriptHandle, out LuaRunner runner, out ScriptHashKey? digestOnHeap) + { + LuaScriptChunk generatedBytecode; + string error; try { - var compiledSource = LuaRunner.CompileSource(source); + if (LuaRunner.TryCompileSource(source, out generatedBytecode, out error)) + return TryGetOrCreateRunnerFromGeneratedBytecode(session, generatedBytecode, digest, ref luaScriptHandle, out runner, out digestOnHeap); + } + catch (Exception ex) + { + logger?.LogError(ex, "During Lua script compilation, an unexpected exception"); + runner = null; + digestOnHeap = null; + luaScriptHandle = null; + return false; + } + + session.WriteLuaCompilationError(error); + runner = null; + digestOnHeap = null; + return false; + } - runner = new LuaRunner(memoryManagementMode, memoryLimitBytes, logMode, allowedFunctions, compiledSource, storeWrapper.serverOptions.LuaTransactionMode, processor, scratchBufferNetworkSender, storeWrapper.redisProtocolVersion, logger); + /// + /// Load internally compiled script bytecode into the cache. + /// + private bool TryGetOrCreateRunnerFromGeneratedBytecode(RespServerSession session, LuaScriptChunk generatedBytecode, ScriptHashKey digest, ref LuaScriptHandle luaScriptHandle, out LuaRunner runner, out ScriptHashKey? digestOnHeap) + { + if (TryGetFromDigest(digest, out runner, out var existingLuaScriptHandle)) + { + luaScriptHandle = existingLuaScriptHandle; + digestOnHeap = null; + return true; + } + + try + { + runner = new LuaRunner(memoryManagementMode, memoryLimitBytes, logMode, allowedFunctions, generatedBytecode, storeWrapper.serverOptions.LuaTransactionMode, processor, scratchBufferNetworkSender, storeWrapper.redisProtocolVersion, logger); // If compilation fails, an error is written out if (runner.CompileForSession(session)) @@ -203,7 +252,7 @@ out ScriptHashKey? digestOnHeap ScriptHashKey storeKeyDigest = new(into); digestOnHeap = storeKeyDigest; - luaScriptHandle ??= new(compiledSource); + luaScriptHandle ??= new(generatedBytecode.Data); scriptCache.Add(storeKeyDigest, (runner, luaScriptHandle)); // On first script load, register for timeout notifications diff --git a/test/Garnet.fuzz/Targets/LuaScriptCompilation.cs b/test/Garnet.fuzz/Targets/LuaScriptCompilation.cs index 20c2e2052b4..36d8e402a17 100644 --- a/test/Garnet.fuzz/Targets/LuaScriptCompilation.cs +++ b/test/Garnet.fuzz/Targets/LuaScriptCompilation.cs @@ -34,7 +34,7 @@ public static void Fuzz(ReadOnlySpan input) { try { - using var runner = new LuaRunner(op.MemoryManagementMode, op.GetMemoryLimitBytes(), op.LogMode, op.AllowedFunctions, input.ToArray()); + using var runner = new LuaRunner(op.MemoryManagementMode, op.GetMemoryLimitBytes(), op.LogMode, op.AllowedFunctions, new LuaScriptChunk(input.ToArray(), LuaScriptChunkKind.Text)); runner.CompileForRunner(); _ = runner.RunForRunner([], []); diff --git a/test/standalone/Garnet.test.scripting/LuaScriptRunnerTests.cs b/test/standalone/Garnet.test.scripting/LuaScriptRunnerTests.cs index cf879a46ad9..d12f4ebab6a 100644 --- a/test/standalone/Garnet.test.scripting/LuaScriptRunnerTests.cs +++ b/test/standalone/Garnet.test.scripting/LuaScriptRunnerTests.cs @@ -153,6 +153,28 @@ public void CanLoadScript() } } + [Test] + public void TryCompileSourceRejectsBinaryAndInvalidInput() + { + // Valid source is compiled into bytecode explicitly tagged as generated by Garnet. + ClassicAssert.IsTrue(LuaRunner.TryCompileSource("return 1"u8, out var compiledSource, out var compileError)); + ClassicAssert.IsNull(compileError); + ClassicAssert.AreEqual(LuaScriptChunkKind.GarnetGeneratedBinary, compiledSource.Kind); + ClassicAssert.GreaterOrEqual(compiledSource.Data.Length, 4); + CollectionAssert.AreEqual(new byte[] { 0x1B, (byte)'L', (byte)'u', (byte)'a' }, compiledSource.Data.Span[..4].ToArray()); + + // Externally provided bytecode must not be accepted as source or returned as trusted output. + ClassicAssert.IsFalse(LuaRunner.TryCompileSource(compiledSource.Data.Span, out var rejectedBinary, out var binaryError)); + ClassicAssert.AreEqual(default(LuaScriptChunk), rejectedBinary); + StringAssert.Contains("binary chunk", binaryError); + + // Invalid source must report an error without preserving the original bytes as executable data. + var invalidSource = "return )"u8; + ClassicAssert.IsFalse(LuaRunner.TryCompileSource(invalidSource, out var rejectedSource, out var sourceError)); + ClassicAssert.AreEqual(default(LuaScriptChunk), rejectedSource); + StringAssert.Contains("unexpected symbol", sourceError); + } + [Test] public void CanRunScript() { diff --git a/test/standalone/Garnet.test.scripting/LuaScriptTests.cs b/test/standalone/Garnet.test.scripting/LuaScriptTests.cs index 30c18ff21a0..e3620de7100 100644 --- a/test/standalone/Garnet.test.scripting/LuaScriptTests.cs +++ b/test/standalone/Garnet.test.scripting/LuaScriptTests.cs @@ -1499,6 +1499,48 @@ public void Issue1079() ClassicAssert.AreEqual("hello lua", success); } + [Test] + public void ScriptInputsRejectPrecompiledLuaBytecode() + { + const string Key = "binary-chunk-key"; + + using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); + var db = redis.GetDatabase(); + + // Generate valid bytecode matching the server's exact Lua version. + ClassicAssert.IsTrue(LuaRunner.TryCompileSource("return redis.call('SET', KEYS[1], 'binary-chunk-executed')"u8, out var compiledScript, out var compileError)); + ClassicAssert.IsNull(compileError); + var binaryChunk = compiledScript.Data.ToArray(); + ClassicAssert.GreaterOrEqual(binaryChunk.Length, 4); + CollectionAssert.AreEqual(new byte[] { 0x1B, (byte)'L', (byte)'u', (byte)'a' }, binaryChunk.AsSpan(0, 4).ToArray()); + + // EVAL must not execute arbitrary bytecode supplied as the script body. + var exception = ClassicAssert.Throws(() => db.Execute("EVAL", [binaryChunk, 1, Key])); + StringAssert.Contains("binary chunk", exception.Message); + ClassicAssert.IsFalse(db.KeyExists(Key)); + ClassicAssert.AreNotEqual("binary-chunk-executed", (string)db.StringGet(Key)); + + // SCRIPT LOAD must reject the same bytes without adding them to the global cache. + var hash = Convert.ToHexString(SHA1.HashData(binaryChunk)).ToLowerInvariant(); + exception = ClassicAssert.Throws(() => db.Execute("SCRIPT", ["LOAD", binaryChunk])); + StringAssert.Contains("binary chunk", exception.Message); + + var exists = (RedisResult[])db.Execute("SCRIPT", ["EXISTS", hash]); + ClassicAssert.AreEqual(0, (int)exists[0]); + } + + [Test] + public void EvalUsesFullSourceLength() + { + using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig()); + var db = redis.GetDatabase(); + var source = Encoding.UTF8.GetBytes("return 1\0return 2"); + + // Exact-length loading must parse bytes after the NUL instead of truncating the script. + var exception = ClassicAssert.Throws(() => db.Execute("EVAL", [source, 0])); + StringAssert.StartsWith("Compilation error:", exception.Message); + } + [TestCase(2)] [TestCase(3)] public void LuaToResp2Conversions(int redisSetRespVersion)