From 0456edd3726cf8311beb604a4c1b36075c8f2ac7 Mon Sep 17 00:00:00 2001
From: tangruilin <tang.ruilin@foxmail.com>
Date: Thu, 20 Jun 2024 20:23:43 +0800
Subject: [PATCH] feat: support_echo_and_auth

Signed-off-by: tangruilin <tang.ruilin@foxmail.com>
---
 src/base_cmd.h           |  1 +
 src/cmd_admin.cc         | 36 ++++++++++++++++++++++++++++++++++++
 src/cmd_admin.h          | 22 ++++++++++++++++++++++
 src/cmd_table_manager.cc |  4 +++-
 src/config.h             |  2 ++
 5 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/base_cmd.h b/src/base_cmd.h
index 6bb77dfa2..32bc92ec6 100644
--- a/src/base_cmd.h
+++ b/src/base_cmd.h
@@ -22,6 +22,7 @@ namespace pikiwidb {
 
 // command definition
 // base cmd
+const std::string kCmdNameEcho = " echo ";
 const std::string kCmdNamePing = "ping";
 
 // key cmd
diff --git a/src/cmd_admin.cc b/src/cmd_admin.cc
index 2112943d7..2287b559a 100644
--- a/src/cmd_admin.cc
+++ b/src/cmd_admin.cc
@@ -6,6 +6,9 @@
  */
 
 #include "cmd_admin.h"
+#include <optional>
+#include "client.h"
+#include "config.h"
 #include "db.h"
 
 #include "braft/raft.h"
@@ -135,6 +138,39 @@ bool PingCmd::DoInitial(PClient* client) { return true; }
 
 void PingCmd::DoCmd(PClient* client) { client->SetRes(CmdRes::kPong, "PONG"); }
 
+EchoCmd::EchoCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsFast, kAclCategoryFast) {}
+
+bool EchoCmd::DoInitial(PClient* client) { return true; }
+
+void EchoCmd::DoCmd(PClient* client) {
+  if (client->argv_.size() != 2) {
+    return client->SetRes(CmdRes::kWrongNum, client->CmdName());
+  }
+
+  client->AppendString(client->argv_[1]);
+}
+
+AuthCmd::AuthCmd(const std::string& name, int16_t arity) : BaseCmd(name, arity, kCmdFlagsAdmin, kAclCategoryFast) {}
+
+bool AuthCmd::DoInitial(PClient* client) { return true; }
+
+void AuthCmd::DoCmd(PClient* client) {
+  if (client->argv_.size() != 2) {
+    return client->SetRes(CmdRes::kWrongNum, client->CmdName());
+  }
+
+  if (g_config.auth == std::nullopt || !g_config.auth.has_value()) {
+    return client->SetRes(CmdRes::kOK, "Authentication is empty");
+  }
+
+  auto& auth = g_config.auth;
+  if (client->argv_[1] == auth.value()) {
+    client->SetRes(CmdRes::kOK);
+  } else {
+    client->SetRes(CmdRes::kErrOther, "Authentication failed");
+  }
+}
+
 InfoCmd::InfoCmd(const std::string& name, int16_t arity)
     : BaseCmd(name, arity, kCmdFlagsAdmin | kCmdFlagsReadonly, kAclCategoryAdmin) {}
 
diff --git a/src/cmd_admin.h b/src/cmd_admin.h
index c78164093..6daf8ff65 100644
--- a/src/cmd_admin.h
+++ b/src/cmd_admin.h
@@ -112,6 +112,28 @@ class PingCmd : public BaseCmd {
   void DoCmd(PClient* client) override;
 };
 
+class EchoCmd : public BaseCmd {
+ public:
+  EchoCmd(const std::string& name, int16_t arity);
+
+ protected:
+  bool DoInitial(PClient* client) override;
+
+ private:
+  void DoCmd(PClient* client) override;
+};
+
+class AuthCmd : public BaseCmd {
+ public:
+  AuthCmd(const std::string& name, int16_t arity);
+
+ protected:
+  bool DoInitial(PClient* client) override;
+
+ private:
+  void DoCmd(PClient* client) override;
+};
+
 class InfoCmd : public BaseCmd {
  public:
   InfoCmd(const std::string& name, int16_t arity);
diff --git a/src/cmd_table_manager.cc b/src/cmd_table_manager.cc
index 5a335465b..80ddeff8f 100644
--- a/src/cmd_table_manager.cc
+++ b/src/cmd_table_manager.cc
@@ -52,11 +52,13 @@ void CmdTableManager::InitCmdTable() {
   ADD_COMMAND_GROUP(Config, -2);
   ADD_SUBCOMMAND(Config, Get, -3);
   ADD_SUBCOMMAND(Config, Set, -4);
-  ADD_COMMAND(Ping, 0);
+  ADD_COMMAND(Ping, 1);
   ADD_COMMAND_GROUP(Debug, -2);
   ADD_SUBCOMMAND(Debug, Help, 2);
   ADD_SUBCOMMAND(Debug, OOM, 2);
   ADD_SUBCOMMAND(Debug, Segfault, 2);
+  ADD_COMMAND(Echo, 2);
+  ADD_COMMAND(Auth, 2);
 
   // server
   ADD_COMMAND(Flushdb, 1);
diff --git a/src/config.h b/src/config.h
index 7b1196388..fcfe2227a 100644
--- a/src/config.h
+++ b/src/config.h
@@ -11,6 +11,7 @@
 #include <functional>
 #include <map>
 #include <memory>
+#include <optional>
 #include <shared_mutex>
 #include <string>
 #include <unordered_map>
@@ -173,6 +174,7 @@ class PConfig {
   std::atomic_int rocksdb_level0_stop_writes_trigger = 36;
   std::atomic_uint64_t rocksdb_ttl_second = 604800;       // default 86400 * 7
   std::atomic_uint64_t rocksdb_periodic_second = 259200;  // default 86400 * 3
+  std::optional<std::string> auth = std::nullopt;
 
   rocksdb::Options GetRocksDBOptions();