From aa424c81af4d2271967b2072f3059503d3426bcd Mon Sep 17 00:00:00 2001
From: Shahar Mike <chakaz@users.noreply.github.com>
Date: Thu, 8 Aug 2024 10:27:27 +0300
Subject: [PATCH] feat: Allow pre-declaring Lua SHAs to run with undeclared
 keys (#3465)

* feat: Allow pre-declaring Lua SHAs to run with undeclared keys

By using `--lua_undeclared_keys_shas=SHA,SHA,SHA` users can now specify
which scripts should run globally (undeclared keys) without explicit
support from the scripts themselves.

Fixes #2442
---
 src/server/multi_test.cc    | 32 ++++++++++++++++++++++++++++++++
 src/server/script_mgr.cc    | 14 ++++++++++++++
 src/server/server_family.cc |  1 +
 3 files changed, 47 insertions(+)

diff --git a/src/server/multi_test.cc b/src/server/multi_test.cc
index 9a9bea3b2..1d0ba7687 100644
--- a/src/server/multi_test.cc
+++ b/src/server/multi_test.cc
@@ -380,6 +380,7 @@ TEST_F(MultiTest, Eval) {
     GTEST_SKIP() << "Skipped Eval test because default_lua_flags is set";
     return;
   }
+  absl::FlagSaver saver;
   absl::SetFlag(&FLAGS_lua_allow_undeclared_auto_correct, true);
 
   RespExpr resp;
@@ -762,6 +763,37 @@ TEST_F(MultiTest, ScriptFlagsEmbedded) {
   EXPECT_THAT(Run({"eval", s2, "0"}), ErrArg("Invalid flag: this-is-an-error"));
 }
 
+TEST_F(MultiTest, UndeclaredKeyFlag) {
+  if (auto mode = absl::GetFlag(FLAGS_multi_exec_mode); mode != Transaction::LOCK_AHEAD) {
+    GTEST_SKIP() << "Skipped test because multi_exec_mode is not default";
+    return;
+  }
+
+  absl::FlagSaver fs;  // lua_undeclared_keys_shas changed via CONFIG cmd below
+
+  const char* script = "return redis.call('GET', 'random-key');";
+  Run({"set", "random-key", "works"});
+
+  // Get SHA for script in a persistent way
+  string sha = Run({"script", "load", script}).GetString();
+
+  // Make sure we can't run the script before setting the flag
+  EXPECT_THAT(Run({"evalsha", sha, "0"}), ErrArg("undeclared"));
+  EXPECT_THAT(Run({"eval", script, "0"}), ErrArg("undeclared"));
+
+  // Clear all Lua scripts so we can configure the cache
+  EXPECT_THAT(Run({"script", "flush"}), "OK");
+  EXPECT_THAT(Run({"script", "exists", sha}), IntArg(0));
+
+  EXPECT_THAT(
+      Run({"config", "set", "lua_undeclared_keys_shas", absl::StrCat(sha, ",NON-EXISTING-HASH")}),
+      "OK");
+
+  // Check eval finds script flags.
+  EXPECT_EQ(Run({"eval", script, "0"}), "works");
+  EXPECT_EQ(Run({"evalsha", sha, "0"}), "works");
+}
+
 // todo: ASAN fails heres on arm
 #ifndef SANITIZERS
 TEST_F(MultiTest, ScriptBadCommand) {
diff --git a/src/server/script_mgr.cc b/src/server/script_mgr.cc
index 1c950c149..1881685e9 100644
--- a/src/server/script_mgr.cc
+++ b/src/server/script_mgr.cc
@@ -38,6 +38,15 @@ ABSL_FLAG(bool, lua_allow_undeclared_auto_correct, false,
           "access undeclared keys, automaticaly set the script flag to be able to run with "
           "undeclared key.");
 
+ABSL_FLAG(
+    std::vector<std::string>, lua_undeclared_keys_shas,
+    std::vector<std::string>({
+        "351130589c64523cb98978dc32c64173a31244f3",  // Sidekiq, see #2442
+        "6ae15ef4678593dc61f991c9953722d67d822776",  // Sidekiq, see #2442
+    }),
+    "Comma-separated list of Lua script SHAs which are allowed to access undeclared keys. SHAs are "
+    "only looked at when loading the script, and new values do not affect already-loaded script.");
+
 namespace dfly {
 using namespace std;
 using namespace facade;
@@ -262,6 +271,11 @@ io::Result<string, GenericError> ScriptMgr::Insert(string_view body, Interpreter
     return params_opt.get_unexpected();
   auto params = params_opt->value_or(default_params_);
 
+  auto undeclared_shas = absl::GetFlag(FLAGS_lua_undeclared_keys_shas);
+  if (find(undeclared_shas.begin(), undeclared_shas.end(), sha) != undeclared_shas.end()) {
+    params.undeclared_keys = true;
+  }
+
   // If the script is atomic, check for possible squashing optimizations.
   // For non atomic modes, squashing increases the time locks are held, which
   // can decrease throughput with frequently accessed keys.
diff --git a/src/server/server_family.cc b/src/server/server_family.cc
index 71be48bda..ffe53a39f 100644
--- a/src/server/server_family.cc
+++ b/src/server/server_family.cc
@@ -824,6 +824,7 @@ void ServerFamily::Init(util::AcceptServer* acceptor, std::vector<facade::Listen
   config_registry.RegisterMutable("tls_ca_cert_file");
   config_registry.RegisterMutable("tls_ca_cert_dir");
   config_registry.RegisterMutable("replica_priority");
+  config_registry.RegisterMutable("lua_undeclared_keys_shas");
 
   pb_task_ = shard_set->pool()->GetNextProactor();
   if (pb_task_->GetKind() == ProactorBase::EPOLL) {