feat(server): Add support for aliasing commands

src/server/command_registry.cc

@@ -22,6 +22,9 @@ using namespace std;
 ABSL_FLAG(vector<string>, rename_command, {},
           "Change the name of commands, format is: <cmd1_name>=<cmd1_new_name>, "
           "<cmd2_name>=<cmd2_new_name>");
+ABSL_FLAG(vector<string>, command_alias, {},
+          "Add an alias for given commands, format is: <alias>=<original>, "
+          "<alias>=<original>");
 ABSL_FLAG(vector<string>, restricted_commands, {},
           "Commands restricted to connections on the admin port");
 
@@ -72,6 +75,43 @@ uint32_t ImplicitAclCategories(uint32_t mask) {
   return out;
 }
 
+absl::flat_hash_map<std::string, std::string> ParseCmdlineArgMap(
+    const absl::Flag<std::vector<std::string>>& flag, const bool allow_duplicates = false) {
+  const auto& mappings = absl::GetFlag(flag);
+  absl::flat_hash_map<std::string, std::string> parsed_mappings;
+  parsed_mappings.reserve(mappings.size());
+
+  for (const std::string& mapping : mappings) {
+    if (mapping.find('=') == std::string::npos) {
+      LOG(ERROR) << "Malformed command " << mapping << " for " << flag.Name()
+                 << ", expected key=value";
+      exit(1);
+    }
+
+    const std::vector<std::string> kv = absl::StrSplit(mapping, '=');
+    if (kv.size() > 2) {
+      LOG(ERROR) << "Malformed command " << mapping << " for " << flag.Name()
+                 << ", expected key=value";
+      exit(1);
+    }
+
+    const std::string key = absl::AsciiStrToUpper(std::move(kv[0]));
+    const std::string value = kv.size() == 1 ? "" : absl::AsciiStrToUpper(std::move(kv[1]));
+
+    if (key == value) {
+      LOG(ERROR) << "Invalid attempt to map " << key << " to itself in " << flag.Name();
+      exit(1);
+    }
+
+    const bool inserted = parsed_mappings.emplace(std::move(key), std::move(value)).second;
+    if (!allow_duplicates && !inserted) {
+      LOG(ERROR) << "Duplicate insert to " << flag.Name() << " not allowed";
+      exit(1);
+    }
+  }
+  return parsed_mappings;
+}
+
 }  // namespace
 
 CommandId::CommandId(const char* name, uint32_t mask, int8_t arity, int8_t first_key,
@@ -133,17 +173,8 @@ optional<facade::ErrorReply> CommandId::Validate(CmdArgList tail_args) const {
 }
 
 CommandRegistry::CommandRegistry() {
-  vector<string> rename_command = GetFlag(FLAGS_rename_command);
-
-  for (string command_data : rename_command) {
-    pair<string_view, string_view> kv = StrSplit(command_data, '=');
-    auto [_, inserted] =
-        cmd_rename_map_.emplace(AsciiStrToUpper(kv.first), AsciiStrToUpper(kv.second));
-    if (!inserted) {
-      LOG(ERROR) << "Invalid rename_command flag, trying to give 2 names to a command";
-      exit(1);
-    }
-  }
+  cmd_rename_map_ = ParseCmdlineArgMap(FLAGS_rename_command);
+  cmd_aliases_ = ParseCmdlineArgMap(FLAGS_command_alias, true);
 
   for (string name : GetFlag(FLAGS_restricted_commands)) {
     restricted_cmds_.emplace(AsciiStrToUpper(name));
@@ -165,8 +196,7 @@ CommandRegistry& CommandRegistry::operator<<(CommandId cmd) {
 
   absl::InlinedVector<std::string_view, 2> maybe_subcommand = StrSplit(cmd.name(), " ");
   const bool is_sub_command = maybe_subcommand.size() == 2;
-  auto it = cmd_rename_map_.find(maybe_subcommand.front());
-  if (it != cmd_rename_map_.end()) {
+  if (const auto it = cmd_rename_map_.find(maybe_subcommand.front()); it != cmd_rename_map_.end()) {
     if (it->second.empty()) {
       return *this;  // Incase of empty string we want to remove the command from registry.
     }

src/server/command_registry.h

@@ -169,7 +169,15 @@ class CommandRegistry {
   CommandRegistry& operator<<(CommandId cmd);
 
   const CommandId* Find(std::string_view cmd) const {
-    auto it = cmd_map_.find(cmd);
+    const auto it = cmd_map_.find(cmd);
+    if (it == cmd_map_.end()) {
+      if (const auto alias_it = cmd_aliases_.find(cmd); alias_it != cmd_aliases_.end()) {
+        if (alias_it->first == alias_it->second) {
+          return nullptr;
+        }
+        return Find(alias_it->second);
+      }
+    }
     return it == cmd_map_.end() ? nullptr : &it->second;
   }
 
@@ -215,6 +223,11 @@ class CommandRegistry {
  private:
   absl::flat_hash_map<std::string, CommandId> cmd_map_;
   absl::flat_hash_map<std::string, std::string> cmd_rename_map_;
+  // Stores a mapping from alias to original command. During the find operation, the first lookup is
+  // done in the cmd_map_, then in the alias map. This results in two lookups but only for commands
+  // which are not in original map, ie either typos or aliases. While it would be faster, we cannot
+  // store iterators into cmd_map_ here as they may be invalidated on rehashing.
+  absl::flat_hash_map<std::string, std::string> cmd_aliases_;
   absl::flat_hash_set<std::string> restricted_cmds_;
   absl::flat_hash_set<std::string> oomdeny_cmds_;
 

src/server/dragonfly_test.cc

@@ -25,6 +25,7 @@ ABSL_DECLARE_FLAG(float, mem_defrag_threshold);
 ABSL_DECLARE_FLAG(float, mem_defrag_waste_threshold);
 ABSL_DECLARE_FLAG(uint32_t, mem_defrag_check_sec_interval);
 ABSL_DECLARE_FLAG(std::vector<std::string>, rename_command);
+ABSL_DECLARE_FLAG(std::vector<std::string>, command_alias);
 ABSL_DECLARE_FLAG(bool, lua_resp2_legacy_float);
 ABSL_DECLARE_FLAG(double, eviction_memory_budget_threshold);
 
@@ -109,17 +110,13 @@ TEST_F(DflyEngineTest, Sds) {
 
 class DflyRenameCommandTest : public DflyEngineTest {
  protected:
-  DflyRenameCommandTest() : DflyEngineTest() {
+  DflyRenameCommandTest() {
     // rename flushall to myflushall, flushdb command will not be able to execute
     absl::SetFlag(
         &FLAGS_rename_command,
         std::vector<std::string>({"flushall=myflushall", "flushdb=", "ping=abcdefghijklmnop"}));
   }
-
-  void TearDown() {
-    absl::SetFlag(&FLAGS_rename_command, std::vector<std::string>({""}));
-    DflyEngineTest::TearDown();
-  }
+  absl::FlagSaver saver_;
 };
 
 TEST_F(DflyRenameCommandTest, RenameCommand) {
@@ -846,4 +843,28 @@ TEST_F(DflyEngineTest, CommandMetricLabels) {
   EXPECT_EQ(metrics.facade_stats.conn_stats.num_conns_other, 0);
 }
 
+class DflyCommandAliasTest : public DflyEngineTest {
+ protected:
+  DflyCommandAliasTest() {
+    // Test an interaction of rename and alias, where we rename and then add an alias on the rename
+    absl::SetFlag(&FLAGS_rename_command, {"ping=gnip"});
+    absl::SetFlag(&FLAGS_command_alias, {"___set=set", "___ping=gnip"});
+  }
+
+  absl::FlagSaver saver_;
+};
+
+TEST_F(DflyCommandAliasTest, Aliasing) {
+  EXPECT_EQ(Run({"SET", "foo", "bar"}), "OK");
+  EXPECT_EQ(Run({"___SET", "a", "b"}), "OK");
+  EXPECT_EQ(Run({"GET", "foo"}), "bar");
+  EXPECT_EQ(Run({"GET", "a"}), "b");
+  // test the alias
+  EXPECT_EQ(Run({"___ping"}), "PONG");
+  // test the rename
+  EXPECT_EQ(Run({"gnip"}), "PONG");
+  // the original command is not accessible
+  EXPECT_THAT(Run({"PING"}), ErrArg("unknown command `PING`"));
+}
+
 }  // namespace dfly