Enhance login password check

src/clients/meta/MetaClient.cpp

@@ -175,8 +175,6 @@ bool MetaClient::loadUsersAndRoles() {
   }
   decltype(userRolesMap_) userRolesMap;
   decltype(userPasswordMap_) userPasswordMap;
-  decltype(userPasswordAttemptsRemain_) userPasswordAttemptsRemain;
-  decltype(userLoginLockTime_) userLoginLockTime;
   // List of username
   std::unordered_set<std::string> userNameList;
 
@@ -188,40 +186,37 @@ bool MetaClient::loadUsersAndRoles() {
     }
     userRolesMap[user.first] = rolesRet.value();
     userPasswordMap[user.first] = user.second;
-    userPasswordAttemptsRemain[user.first] = FLAGS_failed_login_attempts;
-    userLoginLockTime[user.first] = 0;
     userNameList.emplace(user.first);
   }
   {
     folly::RWSpinLock::WriteHolder holder(localCacheLock_);
     userRolesMap_ = std::move(userRolesMap);
     userPasswordMap_ = std::move(userPasswordMap);
 
-    // This method is called periodically by the heartbeat thread, but we don't want to reset the
-    // failed login attempts every time. Remove expired users from cache
-    for (auto& ele : userPasswordAttemptsRemain) {
-      if (userNameList.count(ele.first) == 0) {
-        userPasswordAttemptsRemain.erase(ele.first);
-      }
-    }
-    for (auto& ele : userLoginLockTime) {
-      if (userNameList.count(ele.first) == 0) {
-        userLoginLockTime.erase(ele.first);
+    // Remove expired users from cache
+    auto removeExpiredUser = [&](folly::ConcurrentHashMap<std::string, uint32>& userMap,
+                                 const std::unordered_set<std::string>& userList) {
+      for (auto& ele : userMap) {
+        if (userList.count(ele.first) == 0) {
+          userMap.erase(ele.first);
+        }
       }
-    }
+    };
+    removeExpiredUser(userPasswordAttemptsRemain_, userNameList);
+    removeExpiredUser(userLoginLockTime_, userNameList);
 
-    // If the user is not in the map, insert value with the default value
+    // This method is called periodically by the heartbeat thread, but we don't want to reset the
+    // failed login attempts every time.
     for (const auto& user : userNameList) {
-      if (userPasswordAttemptsRemain.count(user) == 0) {
-        userPasswordAttemptsRemain[user] = FLAGS_failed_login_attempts;
+      // If the user is not in the map, insert value with the default value
+      // Do nothing if the account is already in the map
+      if (userPasswordAttemptsRemain_.find(user) == userPasswordAttemptsRemain_.end()) {
+        userPasswordAttemptsRemain_.insert(user, FLAGS_failed_login_attempts);
       }
-      if (userLoginLockTime.count(user) == 0) {
-        userLoginLockTime[user] = 0;
+      if (userLoginLockTime_.find(user) == userLoginLockTime_.end()) {
+        userLoginLockTime_.insert(user, 0);
       }
     }
-
-    userPasswordAttemptsRemain_ = std::move(userPasswordAttemptsRemain);
-    userLoginLockTime_ = std::move(userLoginLockTime);
   }
   return true;
 }
@@ -2426,13 +2421,10 @@ Status MetaClient::authCheckFromCache(const std::string& account, const std::str
     return Status::Error("User not exist");
   }
 
-  folly::RWSpinLock::WriteHolder holder(localCacheLock_);
+  auto lockedSince = userLoginLockTime_[account];
+  auto passwordAttemtRemain = userPasswordAttemptsRemain_[account];
 
-  auto& lockedSince = userLoginLockTime_[account];
-  auto& passwordAttemtRemain = userPasswordAttemptsRemain_[account];
-  LOG(INFO) << "Thread id: " << std::this_thread::get_id()
-            << " ,passwordAttemtRemain: " << passwordAttemtRemain;
-  // lockedSince is non-zero means the account has been locked
+  // If lockedSince is non-zero, it means the account has been locked
   if (lockedSince != 0) {
     auto remainingLockTime =
         (lockedSince + FLAGS_password_lock_time_in_secs) - time::WallClock::fastNowInSec();
@@ -2446,8 +2438,9 @@ Status MetaClient::authCheckFromCache(const std::string& account, const std::str
           remainingLockTime);
     }
     // Clear lock state and reset attempts
-    lockedSince = 0;
-    passwordAttemtRemain = FLAGS_failed_login_attempts;
+    userLoginLockTime_.assign_if_equal(account, lockedSince, 0);
+    userPasswordAttemptsRemain_.assign_if_equal(
+        account, passwordAttemtRemain, FLAGS_failed_login_attempts);
   }
 
   if (iter->second != password) {
@@ -2458,27 +2451,29 @@ Status MetaClient::authCheckFromCache(const std::string& account, const std::str
 
     // If the password is not correct and passwordAttemtRemain > 0,
     // Allow another attemp
+    passwordAttemtRemain = userPasswordAttemptsRemain_[account];
     if (passwordAttemtRemain > 0) {
-      --passwordAttemtRemain;
-      if (passwordAttemtRemain == 0) {
+      auto newAttemtRemain = passwordAttemtRemain - 1;
+      userPasswordAttemptsRemain_.assign_if_equal(account, passwordAttemtRemain, newAttemtRemain);
+      if (newAttemtRemain == 0) {
         // If the remaining attemps is 0, failed to authenticate
         // Block user login
-        lockedSince = time::WallClock::fastNowInSec();
+        userLoginLockTime_.assign_if_equal(account, 0, time::WallClock::fastNowInSec());
         return Status::Error(
             "%d times consecutive incorrect passwords has been input, user name: %s has been "
             "locked, try again in %d seconds",
             FLAGS_failed_login_attempts,
             account.c_str(),
             FLAGS_password_lock_time_in_secs);
       }
-      LOG(ERROR) << "Invalid password, remaining attempts: " << passwordAttemtRemain;
-      return Status::Error("Invalid password, remaining attempts: %d", passwordAttemtRemain);
+      LOG(ERROR) << "Invalid password, remaining attempts: " << newAttemtRemain;
+      return Status::Error("Invalid password, remaining attempts: %d", newAttemtRemain);
     }
   }
 
-  // Reset password attempts
-  passwordAttemtRemain = FLAGS_failed_login_attempts;
-  lockedSince = 0;
+  // Authentication succeed, reset password attempts
+  userPasswordAttemptsRemain_.assign(account, FLAGS_failed_login_attempts);
+  userLoginLockTime_.assign(account, 0);
   return Status::OK();
 }
 

src/clients/meta/MetaClient.h

@@ -7,6 +7,7 @@
 #define CLIENTS_META_METACLIENT_H_
 
 #include <folly/RWSpinLock.h>
+#include <folly/concurrency/ConcurrentHashMap.h>
 #include <folly/container/F14Map.h>
 #include <folly/container/F14Set.h>
 #include <folly/executors/IOThreadPoolExecutor.h>
@@ -145,9 +146,9 @@ using UserRolesMap = std::unordered_map<std::string, std::vector<cpp2::RoleItem>
 // get user password by account
 using UserPasswordMap = std::unordered_map<std::string, std::string>;
 // Mapping of user name and remaining wrong password attempts
-using UserPasswordAttemptsRemain = std::unordered_map<std::string, uint32>;
+using UserPasswordAttemptsRemain = folly::ConcurrentHashMap<std::string, uint32>;
 // Mapping of user name and the timestamp when the account is locked
-using UserLoginLockTime = std::unordered_map<std::string, uint32>;
+using UserLoginLockTime = folly::ConcurrentHashMap<std::string, uint32>;
 
 // config cache, get config via module and name
 using MetaConfigMap =