Refactor 2FA; add refilling setting to TbRateLimits

application/src/main/java/org/thingsboard/server/controller/TwoFactorAuthController.java

@@ -17,7 +17,6 @@ package org.thingsboard.server.controller;
 
 import lombok.RequiredArgsConstructor;
 import org.springframework.security.access.prepost.PreAuthorize;
-import org.springframework.security.core.Authentication;
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestParam;
@@ -25,6 +24,7 @@ import org.springframework.web.bind.annotation.RestController;
 import org.thingsboard.server.common.data.audit.ActionType;
 import org.thingsboard.server.common.data.exception.ThingsboardErrorCode;
 import org.thingsboard.server.common.data.exception.ThingsboardException;
+import org.thingsboard.server.dao.user.UserService;
 import org.thingsboard.server.queue.util.TbCoreComponent;
 import org.thingsboard.server.service.security.auth.mfa.TwoFactorAuthService;
 import org.thingsboard.server.service.security.auth.rest.RestAuthenticationDetails;
@@ -53,6 +53,7 @@ public class TwoFactorAuthController extends BaseController {
     private final TwoFactorAuthService twoFactorAuthService;
     private final JwtTokenFactory tokenFactory;
     private final SystemSecurityService systemSecurityService;
+    private final UserService userService;
 
 
     @PostMapping("/verification/send")
@@ -69,9 +70,10 @@ public class TwoFactorAuthController extends BaseController {
         boolean verificationSuccess = twoFactorAuthService.checkVerificationCode(user, verificationCode, true);
         if (verificationSuccess) {
             systemSecurityService.logLoginAction(user, new RestAuthenticationDetails(servletRequest), ActionType.LOGIN, null);
+            user = new SecurityUser(userService.findUserById(user.getTenantId(), user.getId()), true, user.getUserPrincipal());
             return tokenFactory.createTokenPair(user);
         } else {
-            ThingsboardException error = new ThingsboardException("Verification code is incorrect", ThingsboardErrorCode.AUTHENTICATION);
+            ThingsboardException error = new ThingsboardException("Verification code is incorrect", ThingsboardErrorCode.BAD_REQUEST_PARAMS);
             systemSecurityService.logLoginAction(user, new RestAuthenticationDetails(servletRequest), ActionType.LOGIN, error);
             throw error;
         }

application/src/main/java/org/thingsboard/server/service/security/auth/mfa/DefaultTwoFactorAuthService.java

@@ -16,9 +16,10 @@
 package org.thingsboard.server.service.security.auth.mfa;
 
 import lombok.RequiredArgsConstructor;
+import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.security.authentication.LockedException;
 import org.springframework.stereotype.Service;
-import org.thingsboard.server.common.data.StringUtils;
 import org.thingsboard.server.common.data.User;
 import org.thingsboard.server.common.data.exception.ThingsboardErrorCode;
 import org.thingsboard.server.common.data.exception.ThingsboardException;
@@ -76,7 +77,7 @@ public class DefaultTwoFactorAuthService implements TwoFactorAuthService {
         if (checkLimits) {
             if (StringUtils.isNotEmpty(twoFaSettings.getVerificationCodeSendRateLimit())) {
                 TbRateLimits rateLimits = verificationCodeSendingRateLimits.computeIfAbsent(securityUser.getId(), sessionId -> {
-                    return new TbRateLimits(twoFaSettings.getVerificationCodeSendRateLimit());
+                    return new TbRateLimits(twoFaSettings.getVerificationCodeSendRateLimit(), true);
                 });
                 if (!rateLimits.tryConsume()) {
                     throw new ThingsboardException("Too many verification code sending requests", ThingsboardErrorCode.TOO_MANY_REQUESTS);
@@ -107,19 +108,30 @@ public class DefaultTwoFactorAuthService implements TwoFactorAuthService {
         if (checkLimits) {
             if (StringUtils.isNotEmpty(twoFaSettings.getVerificationCodeCheckRateLimit())) {
                 TbRateLimits rateLimits = verificationCodeCheckingRateLimits.computeIfAbsent(securityUser.getId(), sessionId -> {
-                    return new TbRateLimits(twoFaSettings.getVerificationCodeCheckRateLimit());
+                    return new TbRateLimits(twoFaSettings.getVerificationCodeCheckRateLimit(), true);
                 });
                 if (!rateLimits.tryConsume()) {
                     throw new ThingsboardException("Too many verification code checking requests", ThingsboardErrorCode.TOO_MANY_REQUESTS);
                 }
             }
         }
-
         TwoFactorAuthProviderConfig providerConfig = twoFaSettings.getProviderConfig(accountConfig.getProviderType())
                 .orElseThrow(() -> PROVIDER_NOT_CONFIGURED_ERROR);
-        boolean verificationSuccess = getTwoFaProvider(accountConfig.getProviderType()).checkVerificationCode(securityUser, verificationCode, providerConfig, accountConfig);
+
+        boolean verificationSuccess;
+        if (StringUtils.isNumeric(verificationCode) && verificationCode.length() == 6) {
+            verificationSuccess = getTwoFaProvider(accountConfig.getProviderType()).checkVerificationCode(securityUser, verificationCode, providerConfig, accountConfig);
+        } else {
+            verificationSuccess = false;
+        }
         if (checkLimits) {
-            systemSecurityService.validateTwoFaVerification(securityUser, verificationSuccess, twoFaSettings);
+            try {
+                systemSecurityService.validateTwoFaVerification(securityUser, verificationSuccess, twoFaSettings);
+            } catch (LockedException e) {
+                verificationCodeCheckingRateLimits.remove(securityUser.getId());
+                verificationCodeSendingRateLimits.remove(securityUser.getId());
+                throw new ThingsboardException(e.getMessage(), ThingsboardErrorCode.AUTHENTICATION);
+            }
             if (verificationSuccess) {
                 verificationCodeCheckingRateLimits.remove(securityUser.getId());
                 verificationCodeSendingRateLimits.remove(securityUser.getId());

application/src/main/java/org/thingsboard/server/service/security/auth/mfa/config/TwoFactorAuthSettings.java

@@ -42,7 +42,7 @@ public class TwoFactorAuthSettings {
     @ApiModelProperty(example = "10")
     @Min(value = 0, message = "maximum number of verification failure before user lockout must be positive")
     private int maxVerificationFailuresBeforeUserLockout;
-    @ApiModelProperty(value = "in minutes", example = "60")
+    @ApiModelProperty(value = "in seconds", example = "3600 (60 minutes)")
     @Min(value = 1, message = "total amount of time allotted for verification must be greater than 0")
     private Integer totalAllowedTimeForVerification;
 

application/src/main/java/org/thingsboard/server/service/security/auth/mfa/provider/impl/OtpBasedTwoFactorAuthProvider.java

@@ -65,7 +65,6 @@ public abstract class OtpBasedTwoFactorAuthProvider<C extends OtpBasedTwoFactorA
         return false;
     }
 
-    // FIXME [viacheslav]: periodically clean up codes cache
 
     @Data
     public static class Otp {

application/src/main/java/org/thingsboard/server/service/security/auth/mfa/provider/impl/SmsTwoFactorAuthProvider.java

@@ -35,7 +35,7 @@ public class SmsTwoFactorAuthProvider extends OtpBasedTwoFactorAuthProvider<SmsT
 
     private final SmsService smsService;
 
-    public SmsTwoFactorAuthProvider(SmsService smsService, CacheManager cacheManager) {
+    public SmsTwoFactorAuthProvider(CacheManager cacheManager, SmsService smsService) {
         super(cacheManager);
         this.smsService = smsService;
     }

application/src/main/java/org/thingsboard/server/service/security/auth/rest/RestAwareAuthenticationSuccessHandler.java

@@ -54,8 +54,8 @@ public class RestAwareAuthenticationSuccessHandler implements AuthenticationSucc
         JwtTokenPair tokenPair = new JwtTokenPair();
 
         if (authentication instanceof MfaAuthenticationToken) {
-            int preVerificationTokenLifetime = (int) TimeUnit.MINUTES.toSeconds(twoFactorAuthConfigManager.getTwoFaSettings(securityUser.getTenantId(), true)
+            int preVerificationTokenLifetime = twoFactorAuthConfigManager.getTwoFaSettings(securityUser.getTenantId(), true)
-                    .flatMap(settings -> Optional.ofNullable(settings.getTotalAllowedTimeForVerification())).orElse(30));
+                    .flatMap(settings -> Optional.ofNullable(settings.getTotalAllowedTimeForVerification())).orElse((int) TimeUnit.MINUTES.toSeconds(30));
             tokenPair.setToken(tokenFactory.createPreVerificationToken(securityUser, preVerificationTokenLifetime).getToken());
             tokenPair.setRefreshToken(null);
             tokenPair.setScope(Authority.PRE_VERIFICATION_TOKEN);

application/src/main/java/org/thingsboard/server/service/security/model/token/JwtTokenFactory.java

@@ -115,21 +115,22 @@ public class JwtTokenFactory {
         } else if (securityUser.getAuthority() == Authority.SYS_ADMIN) {
             securityUser.setTenantId(TenantId.SYS_TENANT_ID);
         }
+        String customerId = claims.get(CUSTOMER_ID, String.class);
+        if (customerId != null) {
+            securityUser.setCustomerId(new CustomerId(UUID.fromString(customerId)));
+        }
 
+        UserPrincipal principal;
         if (securityUser.getAuthority() != Authority.PRE_VERIFICATION_TOKEN) {
             securityUser.setFirstName(claims.get(FIRST_NAME, String.class));
             securityUser.setLastName(claims.get(LAST_NAME, String.class));
             securityUser.setEnabled(claims.get(ENABLED, Boolean.class));
             boolean isPublic = claims.get(IS_PUBLIC, Boolean.class);
-            UserPrincipal principal = new UserPrincipal(isPublic ? UserPrincipal.Type.PUBLIC_ID : UserPrincipal.Type.USER_NAME, subject);
+            principal = new UserPrincipal(isPublic ? UserPrincipal.Type.PUBLIC_ID : UserPrincipal.Type.USER_NAME, subject);
-            securityUser.setUserPrincipal(principal);
-            String customerId = claims.get(CUSTOMER_ID, String.class);
-            if (customerId != null) {
-                securityUser.setCustomerId(new CustomerId(UUID.fromString(customerId)));
-            }
         } else {
-            securityUser.setUserPrincipal(new UserPrincipal(UserPrincipal.Type.USER_NAME, subject));
+            principal = new UserPrincipal(UserPrincipal.Type.USER_NAME, subject);
         }
+        securityUser.setUserPrincipal(principal);
 
         return securityUser;
     }
@@ -164,10 +165,12 @@ public class JwtTokenFactory {
     }
 
     public JwtToken createPreVerificationToken(SecurityUser user, Integer expirationTime) {
-        String token = setUpToken(user, Collections.singletonList(Authority.PRE_VERIFICATION_TOKEN.name()), expirationTime)
+        JwtBuilder jwtBuilder = setUpToken(user, Collections.singletonList(Authority.PRE_VERIFICATION_TOKEN.name()), expirationTime)
-                .claim(TENANT_ID, user.getTenantId().toString())
+                .claim(TENANT_ID, user.getTenantId().toString());
-                .compact();
+        if (user.getCustomerId() != null) {
-        return new AccessJwtToken(token);
+            jwtBuilder.claim(CUSTOMER_ID, user.getCustomerId().toString());
+        }
+        return new AccessJwtToken(jwtBuilder.compact());
     }
 
     private JwtBuilder setUpToken(SecurityUser securityUser, List<String> scopes, long expirationTime) {

common/message/src/main/java/org/thingsboard/server/common/msg/tools/TbRateLimits.java

@@ -17,6 +17,7 @@ package org.thingsboard.server.common.msg.tools;
 
 import io.github.bucket4j.Bandwidth;
 import io.github.bucket4j.Bucket4j;
+import io.github.bucket4j.Refill;
 import io.github.bucket4j.local.LocalBucket;
 import io.github.bucket4j.local.LocalBucketBuilder;
 
@@ -29,12 +30,17 @@ public class TbRateLimits {
     private final LocalBucket bucket;
 
     public TbRateLimits(String limitsConfiguration) {
+        this(limitsConfiguration, false);
+    }
+
+    public TbRateLimits(String limitsConfiguration, boolean refillIntervally) {
         LocalBucketBuilder builder = Bucket4j.builder();
         boolean initialized = false;
         for (String limitSrc : limitsConfiguration.split(",")) {
             long capacity = Long.parseLong(limitSrc.split(":")[0]);
             long duration = Long.parseLong(limitSrc.split(":")[1]);
-            builder.addLimit(Bandwidth.simple(capacity, Duration.ofSeconds(duration)));
+            Refill refill = refillIntervally ? Refill.intervally(capacity, Duration.ofSeconds(duration)) : Refill.greedy(capacity, Duration.ofSeconds(duration));
+            builder.addLimit(Bandwidth.classic(capacity, refill));
             initialized = true;
         }
         if (initialized) {
@@ -42,8 +48,6 @@ public class TbRateLimits {
         } else {
             throw new IllegalArgumentException("Failed to parse rate limits configuration: " + limitsConfiguration);
         }
-
-
     }
 
     public boolean tryConsume() {

dao/src/main/java/org/thingsboard/server/dao/user/UserServiceImpl.java

@@ -323,6 +323,7 @@ public class UserServiceImpl extends AbstractEntityService implements UserServic
         }
         ((ObjectNode) additionalInfo).put(LAST_LOGIN_TS, System.currentTimeMillis());
         user.setAdditionalInfo(additionalInfo);
+        saveUser(user);
     }
 
     @Override