From 7662e03d1f7bfe0f0b7c9e385256a41f3d8c1fc5 Mon Sep 17 00:00:00 2001
From: Vladislav <vlad@dragonflydb.io>
Date: Sun, 22 Jan 2023 18:56:19 +0300
Subject: [PATCH] fix(server): Fix Lua & Zset issue (#716)

* fix(server): Fix Lua & Zset issue

Signed-off-by: Vladislav Oleshko <vlad@dragonflydb.io>
---
 src/core/interpreter.cc        |  7 +++++++
 src/core/interpreter_test.cc   |  9 +++++++++
 src/server/zset_family.cc      | 33 +++++++++++++++++++++------------
 src/server/zset_family_test.cc |  8 ++++++++
 4 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/src/core/interpreter.cc b/src/core/interpreter.cc
index 01f08ca2b..0a06197e5 100644
--- a/src/core/interpreter.cc
+++ b/src/core/interpreter.cc
@@ -270,6 +270,13 @@ debug = nil
   lua_setglobal(lua, "loadfile");
   lua_pushnil(lua);
   lua_setglobal(lua, "dofile");
+
+  // unpack was a global function until Lua 5.2, but was moved into the table module.
+  // Register it globally to maintain compatibility.
+  lua_getglobal(lua, "table");
+  lua_getfield(lua, -1, "unpack");
+  lua_remove(lua, -2);
+  lua_setglobal(lua, "unpack");
 }
 
 // dest must have at least 41 chars.
diff --git a/src/core/interpreter_test.cc b/src/core/interpreter_test.cc
index e2864f15d..17aa2a712 100644
--- a/src/core/interpreter_test.cc
+++ b/src/core/interpreter_test.cc
@@ -321,4 +321,13 @@ TEST_F(InterpreterTest, Modules) {
   EXPECT_EQ("str(\x1\x2test)", ser_.res);
 }
 
+// Since Lua 5.2 global functions were moved to separate namespaces.
+// We need to register them globally to maintain 5.1 compatibility.
+TEST_F(InterpreterTest, OutdatedGlobals) {
+  // table.unpack is used in Laravel:
+  // https://github.com/laravel/framework/blob/6a5c2ec92200cc485983f26b284f7e78470b885f/src/Illuminate/Queue/LuaScripts.php#L118
+  EXPECT_TRUE(Execute("return unpack{1,2,3}"));
+  EXPECT_EQ("i(1)", ser_.res);
+}
+
 }  // namespace dfly
diff --git a/src/server/zset_family.cc b/src/server/zset_family.cc
index 1869982a3..45e718ca2 100644
--- a/src/server/zset_family.cc
+++ b/src/server/zset_family.cc
@@ -265,6 +265,12 @@ void IntervalVisitor::operator()(ZSetFamily::TopNScored sc) {
 }
 
 void IntervalVisitor::ActionRange(unsigned start, unsigned end) {
+  if (params_.limit == 0)
+    return;
+  // Calculate new start and end given offset and limit.
+  start += params_.offset;
+  end = static_cast<uint32_t>(min(1ULL * start + params_.limit - 1, 1ULL * end));
+
   container_utils::IterateSortedSet(
       zobj_,
       [this](container_utils::ContainerEntry ce, double score) {
@@ -1075,6 +1081,16 @@ void ZUnionFamilyInternal(CmdArgList args, bool store, ConnectionContext* cntx)
   }
 }
 
+bool ParseLimit(string_view offset_str, string_view limit_str, ZSetFamily::RangeParams* params) {
+  int64_t limit_arg;
+  if (!SimpleAtoi(offset_str, &params->offset) || !SimpleAtoi(limit_str, &limit_arg) ||
+      limit_arg > UINT32_MAX) {
+    return false;
+  }
+  params->limit = limit_arg < 0 ? UINT32_MAX : static_cast<uint32_t>(limit_arg);
+  return true;
+}
+
 }  // namespace
 
 void ZSetFamily::ZAdd(CmdArgList args, ConnectionContext* cntx) {
@@ -1376,9 +1392,7 @@ void ZSetFamily::ZRange(CmdArgList args, ConnectionContext* cntx) {
       if (i + 3 > args.size()) {
         return (*cntx)->SendError(kSyntaxErr);
       }
-      string_view os = ArgS(args, i + 1);
-      string_view cs = ArgS(args, i + 2);
-      if (!SimpleAtoi(os, &range_params.offset) || !SimpleAtoi(cs, &range_params.limit)) {
+      if (!ParseLimit(ArgS(args, i + 1), ArgS(args, i + 2), &range_params)) {
         return (*cntx)->SendError(kInvalidIntErr);
       }
       i += 2;
@@ -1441,11 +1455,9 @@ void ZSetFamily::ZRangeByLexInternal(CmdArgList args, bool reverse, ConnectionCo
     ToUpper(&args[4]);
     if (ArgS(args, 4) != "LIMIT")
       return (*cntx)->SendError(kSyntaxErr);
-    string_view os = ArgS(args, 5);
-    string_view cs = ArgS(args, 6);
-    if (!SimpleAtoi(os, &offset) || !SimpleAtoi(cs, &count)) {
+
+    if (!ParseLimit(ArgS(args, 5), ArgS(args, 6), &range_params))
       return (*cntx)->SendError(kInvalidIntErr);
-    }
   }
   range_params.offset = offset;
   range_params.limit = count;
@@ -1728,12 +1740,9 @@ bool ZSetFamily::ParseRangeByScoreParams(CmdArgList args, RangeParams* params) {
     } else if (cur_arg == "LIMIT") {
       if (i + 3 > args.size())
         return false;
-
-      string_view os = ArgS(args, i + 1);
-      string_view cs = ArgS(args, i + 2);
-
-      if (!SimpleAtoi(os, &params->offset) || !SimpleAtoi(cs, &params->limit))
+      if (!ParseLimit(ArgS(args, i + 1), ArgS(args, i + 2), params))
         return false;
+
       i += 2;
     } else {
       return false;
diff --git a/src/server/zset_family_test.cc b/src/server/zset_family_test.cc
index 6a76d337d..7bd42f6fc 100644
--- a/src/server/zset_family_test.cc
+++ b/src/server/zset_family_test.cc
@@ -86,6 +86,10 @@ TEST_F(ZSetFamilyTest, ZRangeRank) {
   ASSERT_THAT(resp, ArrLen(2));
   ASSERT_THAT(resp.GetVec(), ElementsAre("a", "1.1"));
 
+  resp = Run({"zrangebyscore", "x", "-inf", "+inf", "LIMIT", "0", "-1"});
+  ASSERT_THAT(resp, ArrLen(2));
+  ASSERT_THAT(resp.GetVec(), ElementsAre("a", "b"));
+
   resp = Run({"zrevrangebyscore", "x", "+inf", "-inf", "limit", "0", "5"});
   ASSERT_THAT(resp, ArgType(RespExpr::ARRAY));
   ASSERT_THAT(resp.GetVec(), ElementsAre("b", "a"));
@@ -190,6 +194,10 @@ TEST_F(ZSetFamilyTest, ZRange) {
   resp = Run({"zrange", "key", "-", "d", "BYLEX", "BYSCORE"});
   EXPECT_THAT(resp, ErrArg("BYSCORE and BYLEX options are not compatible"));
 
+  resp = Run({"zrange", "key", "0", "-1", "LIMIT", "3", "-1"});
+  ASSERT_THAT(resp, ArrLen(2));
+  ASSERT_THAT(resp.GetVec(), ElementsAre("c", "e"));
+
   Run({"zremrangebyscore", "key", "0", "4"});
 
   Run({