feat: support tablet auth and flag to disable auth

src/auth/refreshable_map.h

@@ -29,6 +29,8 @@ namespace openmldb::auth {
 template <typename Key, typename Value>
 class RefreshableMap {
  public:
+    RefreshableMap() : map_(std::make_shared<std::unordered_map<Key, Value>>()) {}
+
     std::optional<Value> Get(const Key& key) const {
         std::shared_lock<std::shared_mutex> lock(mutex_);
         if (auto it = map_->find(key); it != map_->end()) {

src/auth/user_access_manager.cc

@@ -23,9 +23,9 @@
 #include "nameserver/system_table.h"
 
 namespace openmldb::auth {
-UserAccessManager::UserAccessManager(IteratorFactory iterator_factory,
-                                     std::shared_ptr<nameserver::TableInfo> user_table_info)
-    : user_table_iterator_factory_(std::move(iterator_factory)), user_table_info_(user_table_info) {
+
+UserAccessManager::UserAccessManager(IteratorFactory iterator_factory)
+    : user_table_iterator_factory_(std::move(iterator_factory)) {
     StartSyncTask();
 }
 
@@ -49,26 +49,28 @@ void UserAccessManager::StopSyncTask() {
 }
 
 void UserAccessManager::SyncWithDB() {
-    auto new_user_map = std::make_unique<std::unordered_map<std::string, std::string>>();
-    auto it = user_table_iterator_factory_(::openmldb::nameserver::USER_INFO_NAME);
-    it->SeekToFirst();
-    while (it->Valid()) {
-        auto row = it->GetValue();
-        auto buf = it->GetValue().buf();
-        auto size = it->GetValue().size();
-        codec::RowView row_view(user_table_info_->column_desc(), buf, size);
-        std::string host, user, password;
-        row_view.GetStrValue(0, &host);
-        row_view.GetStrValue(1, &user);
-        row_view.GetStrValue(2, &password);
-        if (host == "%") {
-            new_user_map->emplace(user, password);
-        } else {
-            new_user_map->emplace(FormUserHost(user, host), password);
+    if (auto it_pair = user_table_iterator_factory_(::openmldb::nameserver::USER_INFO_NAME); it_pair) {
+        auto new_user_map = std::make_unique<std::unordered_map<std::string, std::string>>();
+        auto it = it_pair->first.get();
+        it->SeekToFirst();
+        while (it->Valid()) {
+            auto row = it->GetValue();
+            auto buf = it->GetValue().buf();
+            auto size = it->GetValue().size();
+            codec::RowView row_view(*it_pair->second.get(), buf, size);
+            std::string host, user, password;
+            row_view.GetStrValue(0, &host);
+            row_view.GetStrValue(1, &user);
+            row_view.GetStrValue(2, &password);
+            if (host == "%") {
+                new_user_map->emplace(user, password);
+            } else {
+                new_user_map->emplace(FormUserHost(user, host), password);
+            }
+            it->Next();
         }
-        it->Next();
+        user_map_.Refresh(std::move(new_user_map));
     }
-    user_map_.Refresh(std::move(new_user_map));
 }
 
 bool UserAccessManager::IsAuthenticated(const std::string& host, const std::string& user, const std::string& password) {

src/auth/user_access_manager.h

@@ -22,23 +22,25 @@
 #include <memory>
 #include <string>
 #include <thread>
+#include <utility>
 
 #include "catalog/distribute_iterator.h"
 #include "refreshable_map.h"
 
 namespace openmldb::auth {
 class UserAccessManager {
  public:
-    using IteratorFactory =
-        std::function<std::unique_ptr<::openmldb::catalog::FullTableIterator>(const std::string& table_name)>;
+    using IteratorFactory = std::function<std::optional<
+        std::pair<std::unique_ptr<::openmldb::catalog::FullTableIterator>, std::unique_ptr<openmldb::codec::Schema>>>(
+        const std::string& table_name)>;
+
+    explicit UserAccessManager(IteratorFactory iterator_factory);
 
-    UserAccessManager(IteratorFactory iterator_factory, std::shared_ptr<nameserver::TableInfo> user_table_info);
     ~UserAccessManager();
     bool IsAuthenticated(const std::string& host, const std::string& username, const std::string& password);
 
  private:
     IteratorFactory user_table_iterator_factory_;
-    std::shared_ptr<nameserver::TableInfo> user_table_info_;
     RefreshableMap<std::string, std::string> user_map_;
     std::atomic<bool> sync_task_running_{false};
     std::thread sync_task_thread_;

src/catalog/client_manager.cc

@@ -554,7 +554,7 @@ bool ClientManager::UpdateClient(const std::map<std::string, std::string>& endpo
     for (const auto& kv : endpoint_map) {
         auto it = real_endpoint_map_.find(kv.first);
         if (it == real_endpoint_map_.end()) {
-            auto wrapper = std::make_shared<TabletAccessor>(kv.first);
+            auto wrapper = std::make_shared<TabletAccessor>(kv.first, auth_token_);
             if (!wrapper->UpdateClient(kv.second)) {
                 LOG(WARNING) << "add client failed. name " << kv.first << ", endpoint " << kv.second;
                 continue;
@@ -583,7 +583,7 @@ bool ClientManager::UpdateClient(
     for (const auto& kv : tablet_clients) {
         auto it = real_endpoint_map_.find(kv.first);
         if (it == real_endpoint_map_.end()) {
-            auto wrapper = std::make_shared<TabletAccessor>(kv.first, kv.second);
+            auto wrapper = std::make_shared<TabletAccessor>(kv.first, kv.second, auth_token_);
             DLOG(INFO) << "add client. name " << kv.first << ", endpoint " << kv.second->GetRealEndpoint();
             clients_.emplace(kv.first, wrapper);
             real_endpoint_map_.emplace(kv.first, kv.second->GetRealEndpoint());

src/catalog/client_manager.h

@@ -132,17 +132,20 @@ class AsyncTablesHandler : public ::hybridse::vm::MemTableHandler {
 
 class TabletAccessor : public ::hybridse::vm::Tablet {
  public:
-    explicit TabletAccessor(const std::string& name) : name_(name), tablet_client_() {}
+    explicit TabletAccessor(const std::string& name,
+                            const openmldb::authn::AuthToken auth_token = openmldb::authn::ServiceToken{"default"})
+        : name_(name), tablet_client_(), auth_token_(auth_token) {}
 
-    TabletAccessor(const std::string& name, const std::shared_ptr<::openmldb::client::TabletClient>& client)
-        : name_(name), tablet_client_(client) {}
+    TabletAccessor(const std::string& name, const std::shared_ptr<::openmldb::client::TabletClient>& client,
+                   const openmldb::authn::AuthToken auth_token = openmldb::authn::ServiceToken{"default"})
+        : name_(name), tablet_client_(client), auth_token_(auth_token) {}
 
     std::shared_ptr<::openmldb::client::TabletClient> GetClient() {
         return std::atomic_load_explicit(&tablet_client_, std::memory_order_relaxed);
     }
 
     bool UpdateClient(const std::string& endpoint) {
-        auto client = std::make_shared<::openmldb::client::TabletClient>(name_, endpoint);
+        auto client = std::make_shared<::openmldb::client::TabletClient>(name_, endpoint, auth_token_);
         if (client->Init() != 0) {
             return false;
         }
@@ -170,6 +173,7 @@ class TabletAccessor : public ::hybridse::vm::Tablet {
  private:
     std::string name_;
     std::shared_ptr<::openmldb::client::TabletClient> tablet_client_;
+    const openmldb::authn::AuthToken auth_token_;
 };
 
 class TabletsAccessor : public ::hybridse::vm::Tablet {
@@ -243,7 +247,8 @@ class TableClientManager {
 
 class ClientManager {
  public:
-    ClientManager() : real_endpoint_map_(), clients_(), mu_(), rand_(0xdeadbeef) {}
+    explicit ClientManager(const openmldb::authn::AuthToken auth_token = openmldb::authn::ServiceToken{"default"})
+        : real_endpoint_map_(), clients_(), mu_(), rand_(0xdeadbeef), auth_token_(auth_token) {}
     std::shared_ptr<TabletAccessor> GetTablet(const std::string& name) const;
     std::shared_ptr<TabletAccessor> GetTablet() const;
     std::vector<std::shared_ptr<TabletAccessor>> GetAllTablet() const;
@@ -257,6 +262,7 @@ class ClientManager {
     std::unordered_map<std::string, std::shared_ptr<TabletAccessor>> clients_;
     mutable ::openmldb::base::SpinMutex mu_;
     mutable ::openmldb::base::Random rand_;
+    const openmldb::authn::AuthToken auth_token_;
 };
 
 }  // namespace catalog

src/client/tablet_client.cc

@@ -35,11 +35,14 @@ DECLARE_uint32(absolute_ttl_max);
 namespace openmldb {
 namespace client {
 
-TabletClient::TabletClient(const std::string& endpoint, const std::string& real_endpoint)
-    : Client(endpoint, real_endpoint), client_(real_endpoint.empty() ? endpoint : real_endpoint) {}
-
-TabletClient::TabletClient(const std::string& endpoint, const std::string& real_endpoint, bool use_sleep_policy)
-    : Client(endpoint, real_endpoint), client_(real_endpoint.empty() ? endpoint : real_endpoint, use_sleep_policy) {}
+TabletClient::TabletClient(const std::string& endpoint, const std::string& real_endpoint,
+                           const openmldb::authn::AuthToken auth_token)
+    : Client(endpoint, real_endpoint), client_(real_endpoint.empty() ? endpoint : real_endpoint, auth_token) {}
+
+TabletClient::TabletClient(const std::string& endpoint, const std::string& real_endpoint, bool use_sleep_policy,
+                           const openmldb::authn::AuthToken auth_token)
+    : Client(endpoint, real_endpoint),
+      client_(real_endpoint.empty() ? endpoint : real_endpoint, use_sleep_policy, auth_token) {}
 
 TabletClient::~TabletClient() {}
 

src/client/tablet_client.h

@@ -46,9 +46,11 @@ const uint32_t INVALID_REMOTE_TID = UINT32_MAX;
 
 class TabletClient : public Client {
  public:
-    TabletClient(const std::string& endpoint, const std::string& real_endpoint);
+    TabletClient(const std::string& endpoint, const std::string& real_endpoint,
+                 const openmldb::authn::AuthToken auth_token = openmldb::authn::ServiceToken{"default"});
 
-    TabletClient(const std::string& endpoint, const std::string& real_endpoint, bool use_sleep_policy);
+    TabletClient(const std::string& endpoint, const std::string& real_endpoint, bool use_sleep_policy,
+                 const openmldb::authn::AuthToken auth_token = openmldb::authn::ServiceToken{"default"});
 
     ~TabletClient();
 

src/cmd/openmldb.cc

@@ -145,19 +145,19 @@ void StartNameServer() {
         PDLOG(WARNING, "Fail to register name");
         exit(1);
     }
-    std::shared_ptr<::openmldb::nameserver::TableInfo> table_info;
-    if (!name_server->GetTableInfo(::openmldb::nameserver::USER_INFO_NAME, ::openmldb::nameserver::INTERNAL_DB,
-                                   &table_info)) {
-        PDLOG(WARNING, "Failed to get table info for user table");
-        exit(1);
-    }
-    openmldb::auth::UserAccessManager user_access_manager(name_server->GetSystemTableIterator(), table_info);
+
     brpc::ServerOptions options;
-    openmldb::authn::BRPCAuthenticator server_authenticator(
-        [&user_access_manager](const std::string& host, const std::string& username, const std::string& password) {
-            return user_access_manager.IsAuthenticated(host, username, password);
-        });
-    options.auth = &server_authenticator;
+    std::unique_ptr<openmldb::auth::UserAccessManager> user_access_manager;
+    std::unique_ptr<openmldb::authn::BRPCAuthenticator> server_authenticator;
+    if (!FLAGS_skip_grant_tables) {
+        user_access_manager =
+            std::make_unique<openmldb::auth::UserAccessManager>(name_server->GetSystemTableIterator());
+        server_authenticator = std::make_unique<openmldb::authn::BRPCAuthenticator>(
+            [&user_access_manager](const std::string& host, const std::string& username, const std::string& password) {
+                return user_access_manager->IsAuthenticated(host, username, password);
+            });
+        options.auth = server_authenticator.get();
+    }
 
     options.num_threads = FLAGS_thread_pool_size;
     brpc::Server server;
@@ -256,8 +256,17 @@ void StartTablet() {
         exit(1);
     }
     brpc::ServerOptions options;
-    openmldb::authn::BRPCAuthenticator server_authenticator;
-    options.auth = &server_authenticator;
+    std::unique_ptr<openmldb::auth::UserAccessManager> user_access_manager;
+    std::unique_ptr<openmldb::authn::BRPCAuthenticator> server_authenticator;
+
+    if (!FLAGS_skip_grant_tables) {
+        user_access_manager = std::make_unique<openmldb::auth::UserAccessManager>(tablet->GetSystemTableIterator());
+        server_authenticator = std::make_unique<openmldb::authn::BRPCAuthenticator>(
+            [&user_access_manager](const std::string& host, const std::string& username, const std::string& password) {
+                return user_access_manager->IsAuthenticated(host, username, password);
+            });
+        options.auth = server_authenticator.get();
+    }
     options.num_threads = FLAGS_thread_pool_size;
     brpc::Server server;
     if (server.AddService(tablet, brpc::SERVER_DOESNT_OWN_SERVICE) != 0) {

src/cmd/sql_cmd_test.cc

@@ -3545,7 +3545,7 @@ TEST_P(DBSDKTest, ShowComponents) {
 void ExpectShowTableStatusResult(const std::vector<std::vector<test::CellExpectInfo>>& expect,
                                  hybridse::sdk::ResultSet* rs, bool all_db = false, bool is_cluster = false) {
     static const std::vector<std::vector<test::CellExpectInfo>> SystemClusterTableStatus = {
-        {{}, "USER", "__INTERNAL_DB", "memory", {}, {}, {}, "1", "0", "1", "NULL", "NULL", "NULL", ""},
+        {{}, "USER", "__INTERNAL_DB", "memory", {}, {}, {}, "1", "0", "2", "NULL", "NULL", "NULL", ""},
         {{}, "PRE_AGG_META_INFO", "__INTERNAL_DB", "memory", {}, {}, {}, "1", "0", "1", "NULL", "NULL", "NULL", ""},
         {{}, "JOB_INFO", "__INTERNAL_DB", "memory", "0", {}, {}, "1", "0", "1", "NULL", "NULL", "NULL", ""},
         {{},

src/flags.cc

@@ -186,3 +186,4 @@ DEFINE_uint32(keep_log_file_num, 5, "Maximal info log files to be kept");
 DEFINE_int32(sync_job_timeout, 30 * 60 * 1000,
              "sync job timeout, unit is milliseconds, should <= server.channel_keep_alive_time in TaskManager");
 DEFINE_int32(deploy_job_max_wait_time_ms, 30 * 60 * 1000, "the max wait time of waiting deploy job");
+DEFINE_bool(skip_grant_tables, true, "skip the grant tables");

src/nameserver/name_server_impl.cc

@@ -1520,9 +1520,12 @@ bool NameServerImpl::Init(const std::string& zk_cluster, const std::string& zk_p
     task_vec_.resize(FLAGS_name_server_task_max_concurrency + FLAGS_name_server_task_concurrency_for_replica_cluster);
     task_thread_pool_.DelayTask(FLAGS_make_snapshot_check_interval,
                                 boost::bind(&NameServerImpl::SchedMakeSnapshot, this));
-    std::shared_ptr<::openmldb::nameserver::TableInfo> table_info;
-    while (!GetTableInfo(::openmldb::nameserver::USER_INFO_NAME, ::openmldb::nameserver::INTERNAL_DB, &table_info)) {
-        std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    if (!FLAGS_skip_grant_tables) {
+        std::shared_ptr<::openmldb::nameserver::TableInfo> table_info;
+        while (
+            !GetTableInfo(::openmldb::nameserver::USER_INFO_NAME, ::openmldb::nameserver::INTERNAL_DB, &table_info)) {
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        }
     }
     return true;
 }
@@ -5590,9 +5593,9 @@ void NameServerImpl::OnLocked() {
         PDLOG(WARNING, "recover failed");
     }
     CreateDatabaseOrExit(INTERNAL_DB);
-    if (db_table_info_[INTERNAL_DB].count(USER_INFO_NAME) == 0) {
+    if (!FLAGS_skip_grant_tables && db_table_info_[INTERNAL_DB].count(USER_INFO_NAME) == 0) {
         auto temp = FLAGS_system_table_replica_num;
-        FLAGS_system_table_replica_num = temp == 0 ? 1 : temp;
+        FLAGS_system_table_replica_num = tablets_.size();
         CreateSystemTableOrExit(SystemTableType::kUser);
         FLAGS_system_table_replica_num = temp;
         InsertUserRecord("%", "root", "1e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
@@ -9585,26 +9588,33 @@ void NameServerImpl::DropProcedure(RpcController* controller, const api::DropPro
     response->set_msg("ok");
 }
 
-std::function<std::unique_ptr<::openmldb::catalog::FullTableIterator>(const std::string& table_name)>
+std::function<std::optional<std::pair<std::unique_ptr<::openmldb::catalog::FullTableIterator>,
+                                      std::unique_ptr<openmldb::codec::Schema>>>(const std::string& table_name)>
 NameServerImpl::GetSystemTableIterator() {
-    return [this](const std::string& table_name) -> std::unique_ptr<::openmldb::catalog::FullTableIterator> {
+    return [this](const std::string& table_name)
+               -> std::optional<std::pair<std::unique_ptr<::openmldb::catalog::FullTableIterator>,
+                                          std::unique_ptr<openmldb::codec::Schema>>> {
         std::shared_ptr<TableInfo> table_info;
         if (!GetTableInfo(table_name, INTERNAL_DB, &table_info)) {
-            return nullptr;
+            return std::nullopt;
         }
         auto tid = table_info->tid();
         auto table_partition = table_info->table_partition(0);  // only one partition for system table
         for (int meta_idx = 0; meta_idx < table_partition.partition_meta_size(); meta_idx++) {
             if (table_partition.partition_meta(meta_idx).is_leader() &&
                 table_partition.partition_meta(meta_idx).is_alive()) {
                 auto endpoint = table_partition.partition_meta(meta_idx).endpoint();
-                auto table_ptr = GetTablet(endpoint);
+                auto tablet_ptr = GetTablet(endpoint);
+                if (tablet_ptr == nullptr) {
+                    return std::nullopt;
+                }
                 std::map<uint32_t, std::shared_ptr<::openmldb::client::TabletClient>> tablet_clients = {
-                    {0, table_ptr->client_}};
-                return std::make_unique<catalog::FullTableIterator>(tid, nullptr, tablet_clients);
+                    {0, tablet_ptr->client_}};
+                return {{std::make_unique<catalog::FullTableIterator>(tid, nullptr, tablet_clients),
+                         std::make_unique<::openmldb::codec::Schema>(table_info->column_desc())}};
             }
         }
-        return nullptr;
+        return std::nullopt;
     };
 }
 

src/nameserver/name_server_impl.h

@@ -358,7 +358,8 @@ class NameServerImpl : public NameServer {
     void DropProcedure(RpcController* controller, const api::DropProcedureRequest* request, GeneralResponse* response,
                        Closure* done);
 
-    std::function<std::unique_ptr<::openmldb::catalog::FullTableIterator>(const std::string& table_name)>
+    std::function<std::optional<std::pair<std::unique_ptr<::openmldb::catalog::FullTableIterator>,
+                                          std::unique_ptr<openmldb::codec::Schema>>>(const std::string& table_name)>
     GetSystemTableIterator();
 
     bool GetTableInfo(const std::string& table_name, const std::string& db_name,

src/nameserver/name_server_test.cc

@@ -856,7 +856,6 @@ TEST_F(NameServerImplTest, AddAndRemoveReplicaCluster) {
     vector<shared_ptr<NameServerImpl>*> ns_vector = {&m1_ns1, &m1_ns2};
     vector<shared_ptr<TabletImpl>*> tb_vector = {&m1_t1, &m1_t2};
     vector<string*> endpoints = {&m1_t1_ep, &m1_t2_ep};
-    ;
 
     int port = 9632;
 
@@ -1028,7 +1027,6 @@ TEST_F(NameServerImplTest, SyncTableReplicaCluster) {
     vector<shared_ptr<NameServerImpl>*> ns_vector = {&m1_ns1, &m1_ns2};
     vector<shared_ptr<TabletImpl>*> tb_vector = {&m1_t1, &m1_t2};
     vector<string*> endpoints = {&m1_t1_ep, &m1_t2_ep};
-    ;
 
     int port = 9632;
 

src/nameserver/system_table.h

@@ -26,6 +26,7 @@
 #include "proto/name_server.pb.h"
 
 DECLARE_uint32(system_table_replica_num);
+DECLARE_bool(skip_grant_tables);
 
 namespace openmldb {
 namespace nameserver {

src/nameserver/system_table_test.cc

@@ -48,8 +48,8 @@ class MockClosure : public ::google::protobuf::Closure {
 };
 class SystemTableTest : public ::testing::Test {
  public:
-    SystemTableTest() {}
-    ~SystemTableTest() {}
+    SystemTableTest() { FLAGS_skip_grant_tables = false; }
+    ~SystemTableTest() { FLAGS_skip_grant_tables = true; }
 };
 
 TEST_F(SystemTableTest, SystemTable) {