From d2c917db67faf1618de55714318672cf435fa9d9 Mon Sep 17 00:00:00 2001
From: jimingquan <mingquan.ji@vesoft.com>
Date: Fri, 17 Jun 2022 16:15:45 +0800
Subject: [PATCH] check god role exist when meta init (#4330)

* check god role exist when meta init

* return error if kv fail

Co-authored-by: Doodle <13706157+critical27@users.noreply.github.com>
---
 src/daemons/MetaDaemon.cpp       | 15 ++++++++++-----
 src/daemons/StandAloneDaemon.cpp | 15 ++++++++++-----
 src/meta/RootUserMan.h           | 20 ++++++++++++++++++++
 3 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/src/daemons/MetaDaemon.cpp b/src/daemons/MetaDaemon.cpp
index 5664bf2fac9..4c21c118a14 100644
--- a/src/daemons/MetaDaemon.cpp
+++ b/src/daemons/MetaDaemon.cpp
@@ -176,11 +176,16 @@ int main(int argc, char* argv[]) {
     }
     if (nebula::value(ret) == localhost) {
       LOG(INFO) << "Check and init root user";
-      if (!nebula::meta::RootUserMan::isUserExists(gKVStore.get())) {
-        if (!nebula::meta::RootUserMan::initRootUser(gKVStore.get())) {
-          LOG(ERROR) << "Init root user failed";
-          return EXIT_FAILURE;
-        }
+      auto checkRet = nebula::meta::RootUserMan::isGodExists(gKVStore.get());
+      if (!nebula::ok(checkRet)) {
+        auto retCode = nebula::error(checkRet);
+        LOG(ERROR) << "Parser God Role error:" << apache::thrift::util::enumNameSafe(retCode);
+        return EXIT_FAILURE;
+      }
+      auto existGod = nebula::value(checkRet);
+      if (!existGod && !nebula::meta::RootUserMan::initRootUser(gKVStore.get())) {
+        LOG(ERROR) << "Init root user failed";
+        return EXIT_FAILURE;
       }
     }
   }
diff --git a/src/daemons/StandAloneDaemon.cpp b/src/daemons/StandAloneDaemon.cpp
index 4b17749856b..5b519f726f5 100644
--- a/src/daemons/StandAloneDaemon.cpp
+++ b/src/daemons/StandAloneDaemon.cpp
@@ -234,11 +234,16 @@ int main(int argc, char *argv[]) {
       }
       if (nebula::value(ret) == metaLocalhost) {
         LOG(INFO) << "Check and init root user";
-        if (!nebula::meta::RootUserMan::isUserExists(gMetaKVStore.get())) {
-          if (!nebula::meta::RootUserMan::initRootUser(gMetaKVStore.get())) {
-            LOG(ERROR) << "Init root user failed";
-            return;
-          }
+        auto checkRet = nebula::meta::RootUserMan::isGodExists(gMetaKVStore.get());
+        if (!nebula::ok(checkRet)) {
+          auto retCode = nebula::error(checkRet);
+          LOG(ERROR) << "Parser God Role error:" << apache::thrift::util::enumNameSafe(retCode);
+          return;
+        }
+        auto existGod = nebula::value(checkRet);
+        if (!existGod && !nebula::meta::RootUserMan::initRootUser(gMetaKVStore.get())) {
+          LOG(ERROR) << "Init root user failed";
+          return;
         }
       }
     }
diff --git a/src/meta/RootUserMan.h b/src/meta/RootUserMan.h
index fbc45b1fd16..c8bde933203 100644
--- a/src/meta/RootUserMan.h
+++ b/src/meta/RootUserMan.h
@@ -20,6 +20,26 @@ namespace meta {
  * */
 class RootUserMan {
  public:
+  static ErrorOr<nebula::cpp2::ErrorCode, bool> isGodExists(kvstore::KVStore* kv) {
+    auto rolePrefix = MetaKeyUtils::roleSpacePrefix(kDefaultSpaceId);
+    std::unique_ptr<kvstore::KVIterator> iter;
+    auto code = kv->prefix(kDefaultSpaceId, kDefaultPartId, rolePrefix, &iter, false);
+    if (code != nebula::cpp2::ErrorCode::SUCCEEDED) {
+      return code;
+    }
+    while (iter->valid()) {
+      auto val = iter->val();
+      auto type = *reinterpret_cast<const meta::cpp2::RoleType*>(val.begin());
+      if (type == meta::cpp2::RoleType::GOD) {
+        LOG(INFO) << "God user exists";
+        return true;
+      }
+      iter->next();
+    }
+    LOG(INFO) << "God user is not exists";
+    return false;
+  }
+
   static bool isUserExists(kvstore::KVStore* kv) {
     auto userKey = MetaKeyUtils::userKey("root");
     std::string val;