[Enhancement]Lock account for timeout duration after failed login att…

…empts

Signed-off-by: Rohit Satardekar <[email protected]>

fe/fe-core/src/main/java/com/starrocks/authentication/AuthenticationMgr.java

@@ -50,6 +50,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 public class AuthenticationMgr {
@@ -62,6 +63,10 @@ public class AuthenticationMgr {
     // will be manually serialized one by one
     protected Map<UserIdentity, UserAuthenticationInfo> userToAuthenticationInfo;
 
+    private static final Map<String, Long> FAIL_TIME = new ConcurrentHashMap<>();
+    private static final Map<String, Integer> FAIL_COUNT = new ConcurrentHashMap<>();
+    private static final Map<String, Long> LOCKED_TIMES = new ConcurrentHashMap<>();
+
     private static class UserAuthInfoTreeMap extends TreeMap<UserIdentity, UserAuthenticationInfo> {
         public UserAuthInfoTreeMap() {
             super((o1, o2) -> {
@@ -248,7 +253,66 @@ private UserIdentity checkPasswordForNative(
     }
 
     public UserIdentity checkPassword(String remoteUser, String remoteHost, byte[] remotePasswd, byte[] randomString) {
-        return checkPasswordForNative(remoteUser, remoteHost, remotePasswd, randomString);
+        UserIdentity authenticatedUser = null;
+        String userAndHost = remoteUser.concat("@").concat(remoteHost);
+        for (Map.Entry<String, Long> attempt : FAIL_TIME.entrySet()) {
+            if (((System.currentTimeMillis() - attempt.getValue()) / 1000) > Config.password_lock_interval_seconds) {
+                if (!isUserLocked(attempt.getKey())) {
+                    clearFailedAttemptRecords(attempt.getKey());
+                } else {
+                    if (getRemainingLockedTime(attempt.getKey()) <= 0) {
+                        clearFailedAttemptRecords(attempt.getKey());
+                    }
+                }
+            }
+        }
+        if (!allowLoginAttempt(userAndHost)) {
+            return null;
+        }
+        authenticatedUser = checkPasswordForNative(remoteUser, remoteHost, remotePasswd, randomString);
+        if (Config.max_failed_login_attempts >= 0) {
+            if (authenticatedUser == null) {
+                recordFailedAttempt(userAndHost);
+            } else {
+                clearFailedAttemptRecords(userAndHost);
+            }
+        }
+        return authenticatedUser;
+    }
+
+    public Long getRemainingLockedTime(String userAndHost) {
+        return Config.password_lock_interval_seconds - ((System.currentTimeMillis() - LOCKED_TIMES.get(userAndHost)) / 1000);
+    }
+
+    private boolean allowLoginAttempt(String userAndHost) {
+        if (!FAIL_TIME.containsKey(userAndHost)) {
+            return true;
+        }
+        return !isUserLocked(userAndHost);
+    }
+
+    private void recordFailedAttempt(String userAndHost) {
+        if (!FAIL_TIME.containsKey(userAndHost)) {
+            FAIL_TIME.put(userAndHost, System.currentTimeMillis());
+        }
+        FAIL_COUNT.put(userAndHost,
+                FAIL_COUNT.getOrDefault(userAndHost, 0) + 1);
+        if (!LOCKED_TIMES.containsKey(userAndHost)
+                && FAIL_COUNT.get(userAndHost) >= Config.max_failed_login_attempts) {
+            LOCKED_TIMES.put(userAndHost, System.currentTimeMillis());
+        }
+    }
+
+    private void clearFailedAttemptRecords(String userAndHost) {
+        if (FAIL_TIME.containsKey(userAndHost)) {
+            FAIL_TIME.remove(userAndHost);
+            FAIL_COUNT.remove(userAndHost);
+            LOCKED_TIMES.remove(userAndHost);
+        }
+    }
+
+    public boolean isUserLocked(String userAndHost) {
+        return LOCKED_TIMES.containsKey(userAndHost);
     }
 
     public UserIdentity checkPlainPassword(String remoteUser, String remoteHost, String remotePasswd) {

fe/fe-core/src/main/java/com/starrocks/common/Config.java

@@ -2619,6 +2619,19 @@ public class Config extends ConfigBase {
      */
     @ConfField(mutable = true)
     public static boolean enable_password_reuse = true;
+
+    /**
+     * Number of attempts to login if failed. -1 means unlimited.
+     */
+    @ConfField(mutable = true)
+    public static int max_failed_login_attempts = -1;
+
+    /**
+     * Duration to lock the account after failed login attempts.
+     */
+    @ConfField(mutable = true)
+    public static long password_lock_interval_seconds = 600L;
+
     /**
      * If set to false, when the load is empty, success is returned.
      * Otherwise, `No partitions have data available for loading` is returned.

fe/fe-core/src/main/java/com/starrocks/common/ErrorCode.java

@@ -262,6 +262,8 @@ public enum ErrorCode {
                     ErrorCode.ERR_ACCESS_DENIED_HINT_MSG_FORMAT),
     ERR_ACCESS_DENIED_FOR_EXTERNAL_ACCESS_CONTROLLER(5204, new byte[] {'4', '2', '0', '0', '0'},
             "Access denied; you need (at least one of) the %s privilege(s) on %s%s for this operation."),
+    ERR_FAILED_ATTEMPT(5204, new byte[] {'H', 'Y', '0', '0', '0'},
+            "Access denied for user '%s'@'%s'. " + "Account is blocked for %d seconds due to %d consecutive failed logins."),
 
     /**
      * 5300 - 5399: Lock and Transaction

fe/fe-core/src/main/java/com/starrocks/http/BaseAction.java

@@ -317,8 +317,13 @@ public static UserIdentity checkPassword(ActionAuthorizationInfo authInfo)
                 globalStateMgr.getAuthenticationMgr().checkPlainPassword(
                         authInfo.fullUserName, authInfo.remoteIp, authInfo.password);
         if (currentUser == null) {
+            String userAndHost = authInfo.fullUserName.concat("@").concat(authInfo.remoteIp);
+            if (globalStateMgr.getAuthenticationMgr().isUserLocked(userAndHost)) {
+                throw new AccessDeniedException("Access denied for " + userAndHost + ". Locked for " +
+                    globalStateMgr.getAuthenticationMgr().getRemainingLockedTime(userAndHost) + " seconds.");
+            }
             throw new AccessDeniedException("Access denied for "
-                    + authInfo.fullUserName + "@" + authInfo.remoteIp);
+                    + userAndHost);
         }
         return currentUser;
     }

fe/fe-core/src/main/java/com/starrocks/mysql/MysqlProto.java

@@ -76,7 +76,14 @@ private static boolean authenticate(ConnectContext context, byte[] scramble, byt
         if (Config.enable_auth_check) {
             currentUser = authenticationManager.checkPassword(user, remoteIp, scramble, randomString);
             if (currentUser == null) {
-                ErrorReport.report(ErrorCode.ERR_AUTHENTICATION_FAIL, user, usePasswd);
+                String userAndHost = user.concat("@").concat(remoteIp);
+                if (authenticationManager.isUserLocked(userAndHost)) {
+                    ErrorReport.report(ErrorCode.ERR_FAILED_ATTEMPT, user, remoteIp,
+                            authenticationManager.getRemainingLockedTime(userAndHost),
+                            Config.max_failed_login_attempts);
+                } else {
+                    ErrorReport.report(ErrorCode.ERR_AUTHENTICATION_FAIL, user, usePasswd);
+                }
                 return false;
             }
         } else {

fe/fe-core/src/main/java/com/starrocks/service/FrontendServiceImpl.java

@@ -1132,6 +1132,11 @@ private void checkPasswordAndLoadPriv(String user, String passwd, String db, Str
         UserIdentity currentUser =
                 globalStateMgr.getAuthenticationMgr().checkPlainPassword(user, clientIp, passwd);
         if (currentUser == null) {
+            String userAndHost = user.concat("@").concat(clientIp);
+            if (globalStateMgr.getAuthenticationMgr().isUserLocked(userAndHost)) {
+                throw new AuthenticationException("Access denied for " + userAndHost + ". Locked for " +
+                    globalStateMgr.getAuthenticationMgr().getRemainingLockedTime(userAndHost) + "seconds.");
+            }
             throw new AuthenticationException("Access denied for " + user + "@" + clientIp);
         }
         // check INSERT action on table