Add auth command

conf/pika.conf

@@ -16,6 +16,10 @@ write_buffer_size : 268435456
 timeout : 60
 # Requirepass
 requirepass :
+# Userpass
+userpass :
+# User Blacklist
+userblacklist :
 # Dump Prefix
 dump_prefix : 
 # daemonize  [yes | no]

include/pika_admin.h

@@ -38,4 +38,14 @@ class TrysyncCmd : public Cmd {
   virtual void DoInitial(PikaCmdArgsType &argvs, const CmdInfo* const ptr_info);
 };
 
+class AuthCmd : public Cmd {
+public:
+  AuthCmd() {
+  }
+  virtual void Do();
+private:
+  std::string pwd_;
+  virtual void DoInitial(PikaCmdArgsType &argvs, const CmdInfo* const ptr_info);
+};
+
 #endif

include/pika_client_conn.h

@@ -15,11 +15,26 @@ class PikaClientConn: public pink::RedisConn {
   PikaClientConn(int fd, std::string ip_port, pink::Thread *thread);
   virtual ~PikaClientConn();
   virtual int DealMessage();
-
 private:
   PikaWorkerThread* self_thread_;
   std::string DoCmd(const std::string& opt);
   std::string RestoreArgs();
+
+  // Auth related
+  class AuthStat {
+  public:
+    void Init();
+    bool IsAuthed(const std::string& opt);
+    bool ChecknUpdate(const std::string& arg);
+  private:
+    enum StatType {
+      kNoAuthed = 0,
+      kAdminAuthed,
+      kLimitAuthed,
+    };
+    StatType stat_;
+  };
+  AuthStat auth_stat_;
 };
 
 #endif

include/pika_command.h

@@ -13,6 +13,7 @@
 //Admin
 const std::string kCmdNameSlaveof = "slaveof";
 const std::string kCmdNameTrysync = "trysync";
+const std::string kCmdNameAuth = "auth";
 
 //Kv
 const std::string kCmdNameSet = "set";
@@ -163,19 +164,26 @@ class CmdRes {
     kOverFlow,
     kNotFound,
     kOutOfRange,
+    kInvalidPwd,
     kWrongNum,
     kErrOther,
   };
 
   CmdRes():ret_(kNone) {}
 
+  bool none() const {
+    return ret_ == kNone && message_.empty();
+  }
   bool ok() const {
     return ret_ == kOk || ret_ == kNone;
   }
   void clear() {
     message_.clear();
     ret_ = kNone;
   }
+  std::string raw_message() const {
+    return message_;
+  }
   std::string message() const {
     std::string result;
     switch (ret_) {
@@ -195,6 +203,8 @@ class CmdRes {
       return "-ERR no such key\r\n";
     case kOutOfRange:
       return "-ERR index out of range\r\n";
+    case kInvalidPwd:
+      return "-ERR invalid password\r\n";
     case kWrongNum:
       result = "-ERR wrong number of arguments for '";
       result.append(message_);

pikatests/tests/support/server.tcl

@@ -218,7 +218,8 @@ proc start_server {options {code undefined}} {
     if {$::valgrind} {
         set pid [exec valgrind --suppressions=src/valgrind.sup --show-reachable=no --show-possibly-lost=no --leak-check=full src/redis-server $config_file > $stdout 2> $stderr &]
     } else {
-        set pid [exec src/redis-server -c $config_file > $stdout 2> $stderr &]
+        #set pid [exec src/redis-server -c $config_file > $stdout 2> $stderr &]
+        set pid [exec src/redis-server $config_file > $stdout 2> $stderr &]
     }
 
     puts "Starting ---- "

src/pika_admin.cc

@@ -1,8 +1,10 @@
 #include "slash_string.h"
+#include "pika_conf.h"
 #include "pika_admin.h"
 #include "pika_server.h"
 
 extern PikaServer *g_pika_server;
+extern PikaConf *g_pika_conf;
 
 void SlaveofCmd::DoInitial(PikaCmdArgsType &argv, const CmdInfo* const ptr_info) {
   if (!ptr_info->CheckArg(argv.size())) {
@@ -131,3 +133,29 @@ void TrysyncCmd::Do() {
   }
 }
 
+void AuthCmd::DoInitial(PikaCmdArgsType &argv, const CmdInfo* const ptr_info) {
+  if (!ptr_info->CheckArg(argv.size())) {
+    res_.SetRes(CmdRes::kWrongNum, kCmdNameAuth);
+    return;
+  }
+  pwd_ = argv[1];
+}
+
+void AuthCmd::Do() {
+  std::string root_password(g_pika_conf->requirepass());
+  std::string user_password(g_pika_conf->userpass());
+  if (user_password.empty() && root_password.empty()) {
+    res_.SetRes(CmdRes::kErrOther, "Client sent AUTH, but no password is set");
+    return;
+  }
+
+  if (pwd_ == user_password) {
+    res_.SetRes(CmdRes::kOk, "USER");
+  }
+  if (pwd_ == root_password) {
+    res_.SetRes(CmdRes::kOk, "ROOT");
+  }
+  if (res_.none()) {
+    res_.SetRes(CmdRes::kInvalidPwd);
+  }
+}

src/pika_client_conn.cc

@@ -1,14 +1,19 @@
 #include <sstream>
+#include <vector>
+#include <algorithm>
 #include <glog/logging.h>
 #include "pika_server.h"
+#include "pika_conf.h"
 #include "pika_client_conn.h"
 
 extern PikaServer* g_pika_server;
+extern PikaConf* g_pika_conf;
 static const int RAW_ARGS_LEN = 1024 * 1024; 
 
 PikaClientConn::PikaClientConn(int fd, std::string ip_port, pink::Thread* thread) :
   RedisConn(fd, ip_port) {
   self_thread_ = dynamic_cast<PikaWorkerThread*>(thread);
+  auth_stat_.Init();
 }
 
 PikaClientConn::~PikaClientConn() {
@@ -33,14 +38,19 @@ std::string PikaClientConn::DoCmd(const std::string& opt) {
   if (!cinfo_ptr || !c_ptr) {
       return "-Err unknown or unsupported command \'" + opt + "\'\r\n";
   }
+
+  // Check authed
+  if (!auth_stat_.IsAuthed(opt)) {
+    LOG(INFO) << "(" << ip_port() << ")Authentication required, close connection";
+    return "-ERR NOAUTH Authentication required.\r\n";
+  }
 
   // Initial
   c_ptr->Initial(argv_, cinfo_ptr);
   if (!c_ptr->res().ok()) {
     return c_ptr->res().message();
   }
 
-  // TODO Check authed
   // Add read lock for no suspend command
   if (!cinfo_ptr->is_suspend()) {
     pthread_rwlock_rdlock(g_pika_server->rwlock());
@@ -64,66 +74,75 @@ std::string PikaClientConn::DoCmd(const std::string& opt) {
   if (!cinfo_ptr->is_suspend()) {
       pthread_rwlock_unlock(g_pika_server->rwlock());
   }
+
+  if (opt == kCmdNameAuth) {
+    if(!auth_stat_.ChecknUpdate(c_ptr->res().raw_message())) {
+      LOG(WARNING) << "(" << ip_port() << ")Wrong Password, close connection";
+    }
+  }
   return c_ptr->res().message();
 }
 
 int PikaClientConn::DealMessage() {
 
   self_thread_->PlusThreadQuerynum();
-  /* Sync cmd
-   * 1 logger_->Lock()
-   * 2 cmd->Do
-   * 3 AppendLog
-   * 4 logger_->Unlock
-   */
-
-  //logger_->Lock();
-  //TODO return value
+
   if (argv_.empty()) return -2;
   std::string opt = argv_[0];
   slash::StringToLower(opt);
-  //TODO add logger lock
   std::string res = DoCmd(opt);
-  // if (opt == "pikasync") {
-  //     AppendReply(ret);
-  //     mutex_.Unlock();
-  // } else if ((role_ != PIKA_MASTER  && !(role_ == PIKA_SLAVE && opt == "ping")) || (role_ == PIKA_MASTER && opt == "syncdb")) {
-  //     if (role_ == PIKA_SLAVE) {
-  //         MutexLock l(&mutex_);
-  //         AppendReply(ret);
-  //     } else {
-  //         AppendReply(ret);
-  //     }
-  // }
-  //  // TEST trysync
-  //  if (argv_[0] == "trysync") {
-  //    DLOG(INFO) << "recieve \"trysync\"";
-  //
-  //    // Test BinlogSender
-  //    SlaveItem slave;
-  //    slave.ip_port = "127.0.0.1:9922";
-  //    slave.port = 9922;
-  //
-  //    Status s = g_pika_server->AddBinlogSender(slave, 0, 0);
-  //    if (s.ok()) {
-  //      DLOG(INFO) << "AddBinlogSender ok";
-  //    } else {
-  //      DLOG(INFO) << "AddBinlogSender failed, " << s.ToString();
-  //    }
-  //
-  //  } else if (argv_[0] == "slaveof") {
-  //    DLOG(INFO) << "recieve \"slaveof\"";
-  //    DLOG(INFO) << "SetMaster " << g_pika_server->SetMaster("127.0.0.1", 9821);
-  //  } else {
-  //    // TEST Put
-  //    for (int i = 0; i < 10; i++) {
-  //      DLOG(INFO) << "Logger Put a msg:" << i << ";";
-  //      g_pika_server->logger_->Put(std::string("*3\r\n$3\r\nset\r\n$3\r\nkey\r\n$1\r\n1\r\n"));
-  //    }
-  //  }
-  //
+
   memcpy(wbuf_ + wbuf_len_, res.data(), res.size());
   wbuf_len_ += res.size();
   set_is_reply(true);
   return 0;
 }
+
+// Initial permission status
+void PikaClientConn::AuthStat::Init() {
+  // Check auth required
+  stat_ = std::string(g_pika_conf->userpass()) == "" ?
+    kLimitAuthed : kNoAuthed;
+  if (stat_ == kLimitAuthed 
+      && std::string(g_pika_conf->requirepass()) == "") {
+    stat_ = kAdminAuthed;
+  }
+}
+
+// Check permission for current command
+bool PikaClientConn::AuthStat::IsAuthed(const std::string& opt) {
+  //TODO consider slaveauth
+  if (opt == kCmdNameAuth) {
+    return true;
+  }
+  const std::vector<std::string>& blacklist = g_pika_conf->vuser_blacklist();
+  switch (stat_) {
+    case kNoAuthed:
+      return false;
+    case kAdminAuthed:
+      break;
+    case kLimitAuthed:
+      if (find(blacklist.begin(), blacklist.end(), opt) 
+          != blacklist.end()) {
+      return false;
+      }
+      break;
+    default:
+      LOG(WARNING) << "Invalid auth stat : " << static_cast<unsigned>(stat_);
+      return false;
+  }
+  return true;
+}
+
+// Update permission status
+bool PikaClientConn::AuthStat::ChecknUpdate(const std::string& message) {
+  // Situations to change auth status
+  if (message == "USER") {
+    stat_ = kLimitAuthed;
+  } else if (message == "ROOT"){
+    stat_ = kAdminAuthed;
+  } else {
+    return false;
+  }
+  return true;
+}

src/pika_command.cc

@@ -14,6 +14,8 @@ void InitCmdInfoTable() {
   ////Trysync
   CmdInfo* trysyncptr = new CmdInfo(kCmdNameTrysync, 5, kCmdFlagsRead | kCmdFlagsAdmin);
   cmd_infos.insert(std::pair<std::string, CmdInfo*>(kCmdNameTrysync, trysyncptr));
+  CmdInfo* authptr = new CmdInfo(kCmdNameAuth, 2, kCmdFlagsRead | kCmdFlagsAdmin);
+  cmd_infos.insert(std::pair<std::string, CmdInfo*>(kCmdNameAuth, authptr));
 
   //Kv
   ////SetCmd
@@ -209,6 +211,8 @@ void InitCmdTable(std::unordered_map<std::string, Cmd*> *cmd_table) {
   ////Trysync
   Cmd* trysyncptr = new TrysyncCmd();
   cmd_table->insert(std::pair<std::string, Cmd*>(kCmdNameTrysync, trysyncptr));
+  Cmd* authptr = new AuthCmd();
+  cmd_table->insert(std::pair<std::string, Cmd*>(kCmdNameAuth, authptr));
 
   //Kv
   ////SetCmd