Email 2FA provider, refactoring

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

@@ -28,9 +28,10 @@ import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.RestController;
 import org.thingsboard.common.util.JacksonUtil;
+import org.thingsboard.server.common.data.StringUtils;
 import org.thingsboard.server.common.data.exception.ThingsboardErrorCode;
 import org.thingsboard.server.common.data.exception.ThingsboardException;
-import org.thingsboard.server.service.security.auth.TokenOutdatingService;
+import org.thingsboard.server.common.msg.tools.TbRateLimits;
 import org.thingsboard.server.service.security.auth.mfa.TwoFactorAuthService;
 import org.thingsboard.server.service.security.auth.mfa.config.TwoFactorAuthSettings;
 import org.thingsboard.server.service.security.auth.mfa.config.account.TotpTwoFactorAuthAccountConfig;
@@ -43,24 +44,30 @@ import org.thingsboard.server.service.security.model.token.JwtTokenFactory;
 import javax.servlet.ServletOutputStream;
 import javax.servlet.http.HttpServletResponse;
 import javax.validation.Valid;
+import java.util.HashMap;
+import java.util.Map;
 
 /*
  *
  * TODO [viacheslav]:
- *  - 2FA should be mandatory when logging in and must be rolled out to all existing users when 2FA is activated.
+ *  - Configurable softlock after XX (3) attempts: XX (15) mins - on session level
- *  - Rate limits should be implemented to protect against brute force leaked accounts to prevent SMS cost explosion.
+ *  - Configurable hardlock (user blocking) after a total of XX (10) unsuccessful attempts - on user level
- *  - Configurable softlock after XX (3) attempts: XX (15) mins
- *  - Configurable hardlock (user blocking) after a total of XX (10) unsuccessful attempts.
- *  - The OTP token should only be valid for XX (5) minutes.
- *  - Disable 2FA only possible after successful 2FA auth - it is possible with simple password resest
- *  - 2FA entries should be secured against code injection by code validation.
- *  - Email 2FA provider
  *
  * FIXME [viacheslav]:
  *  - Tests for 2FA
  *  - Swagger documentation
  *
  * */
+// TODO [viacheslav]: maybe get rid of sessionId concept..
+
+/*
+ *
+ *
+ * TODO (later):
+ *  - 2FA entries should be secured against code injection by code validation
+ *  - ability to force users to use 2FA (maybe on log in, do not give them token pair but to give temporary
+ *      token to configure 2FA account config); also will need to make users configure 2FA during activation and password setup...
+ * */
 @RestController
 @RequestMapping("/api")
 @RequiredArgsConstructor
@@ -122,7 +129,7 @@ public class TwoFactorAuthController extends BaseController {
 
         boolean verificationSuccess = twoFactorAuthService.processByTwoFaProvider(user.getTenantId(), accountConfig.getProviderType(),
                 (provider, providerConfig) -> {
-                    return provider.checkVerificationCode(user, verificationCode, accountConfig);
+                    return provider.checkVerificationCode(user, verificationCode, providerConfig, accountConfig);
                 });
 
         if (verificationSuccess) {
@@ -146,27 +153,23 @@ public class TwoFactorAuthController extends BaseController {
     }
 
 
-    @PostMapping("/auth/2fa/verification/check")
+    private final Map<String, TbRateLimits> verificationCodeSendRateLimits = new HashMap<>();
+    private final Map<String, TbRateLimits> verificationCodeCheckRateLimits = new HashMap<>();
+
+    @PostMapping("/auth/2fa/verification/send")
     @PreAuthorize("hasAuthority('PRE_VERIFICATION_TOKEN')")
-    public JwtTokenPair checkTwoFaVerificationCode(@RequestParam String verificationCode) throws Exception {
+    public void sendTwoFaVerificationCode() throws Exception {
         SecurityUser user = getCurrentUser();
 
-        boolean verificationSuccess = twoFactorAuthService.processByTwoFaProvider(user.getTenantId(), user.getId(),
+        TwoFactorAuthSettings twoFaSettings = twoFactorAuthService.getTwoFaSettings(user.getTenantId()).get();
-                (provider, providerConfig, accountConfig) -> {
+        if (StringUtils.isNotEmpty(twoFaSettings.getVerificationCodeSendRateLimit())) {
-                    return provider.checkVerificationCode(user, verificationCode, accountConfig);
+            TbRateLimits rateLimits = verificationCodeSendRateLimits.computeIfAbsent(user.getSessionId(), sessionId -> {
-                });
+                return new TbRateLimits(twoFaSettings.getVerificationCodeSendRateLimit());
-
+            });
-        if (verificationSuccess) {
+            if (!rateLimits.tryConsume()) {
-            return tokenFactory.createTokenPair(user);
+                throw new ThingsboardException(ThingsboardErrorCode.TOO_MANY_REQUESTS);
-        } else {
+            }
-            throw new ThingsboardException("Verification code is incorrect", ThingsboardErrorCode.AUTHENTICATION);
         }
-    }
-
-    @PostMapping("/auth/2fa/verification/resend")
-    @PreAuthorize("hasAuthority('PRE_VERIFICATION_TOKEN')")
-    public void resendTwoFaVerificationCode() throws Exception {
-        SecurityUser user = getCurrentUser();
 
         twoFactorAuthService.processByTwoFaProvider(user.getTenantId(), user.getId(),
                 (provider, providerConfig, accountConfig) -> {
@@ -174,4 +177,34 @@ public class TwoFactorAuthController extends BaseController {
                 });
     }
 
+    @PostMapping("/auth/2fa/verification/check")
+    @PreAuthorize("hasAuthority('PRE_VERIFICATION_TOKEN')")
+    public JwtTokenPair checkTwoFaVerificationCode(@RequestParam String verificationCode) throws Exception {
+        SecurityUser user = getCurrentUser();
+
+
+
+        // FIXME [viacheslav]: rate limits for verification code check
+        boolean verificationSuccess = twoFactorAuthService.processByTwoFaProvider(user.getTenantId(), user.getId(),
+                (provider, providerConfig, accountConfig) -> {
+                    return provider.checkVerificationCode(user, verificationCode, providerConfig, accountConfig);
+                });
+
+
+        if (verificationSuccess) {
+            return tokenFactory.createTokenPair(user);
+        } else {
+            TwoFactorAuthSettings twoFaSettings = twoFactorAuthService.getTwoFaSettings(user.getTenantId()).get();
+            if (StringUtils.isNotEmpty(twoFaSettings.getVerificationCodeSendRateLimit())) {
+                TbRateLimits rateLimits = verificationCodeSendRateLimits.computeIfAbsent(user.getSessionId(), sessionId -> {
+                    return new TbRateLimits(twoFaSettings.getVerificationCodeSendRateLimit());
+                });
+                if (!rateLimits.tryConsume()) {
+                    throw new ThingsboardException(ThingsboardErrorCode.TOO_MANY_REQUESTS);
+                }
+            }
+            throw new ThingsboardException("Verification code is incorrect", ThingsboardErrorCode.AUTHENTICATION);
+        }
+    }
+
 }

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

@@ -0,0 +1,24 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.security.auth;
+
+import org.thingsboard.server.service.security.model.SecurityUser;
+
+public class MfaAuthenticationToken extends AbstractJwtAuthenticationToken {
+    public MfaAuthenticationToken(SecurityUser securityUser) {
+        super(securityUser);
+    }
+}

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

@@ -63,24 +63,6 @@ public class TwoFactorAuthService {
     protected static final String TWO_FACTOR_AUTH_SETTINGS_KEY = "twoFaSettings";
 
 
-    @Autowired
-    private void setProviders(Collection<TwoFactorAuthProvider<?, ?>> providers) {
-        providers.forEach(provider -> {
-            this.providers.put(provider.getType(), provider);
-        });
-    }
-
-    private <A extends TwoFactorAuthAccountConfig, C extends TwoFactorAuthProviderConfig> Optional<TwoFactorAuthProvider<C, A>> getTwoFaProvider(TwoFactorAuthProviderType providerType) {
-        return Optional.of((TwoFactorAuthProvider<C, A>) providers.get(providerType));
-    }
-
-    private <C extends TwoFactorAuthProviderConfig> Optional<C> getTwoFaProviderConfig(TenantId tenantId, TwoFactorAuthProviderType providerType) {
-        return getTwoFaSettings(tenantId)
-                .flatMap(twoFaSettings -> twoFaSettings.getProviderConfig(providerType))
-                .map(providerConfig -> (C) providerConfig);
-    }
-
-
     public <R> R processByTwoFaProvider(TenantId tenantId, TwoFactorAuthProviderType providerType, ThrowingBiFunction<TwoFactorAuthProvider<TwoFactorAuthProviderConfig, TwoFactorAuthAccountConfig>, TwoFactorAuthProviderConfig, R> function) throws Exception {
         TwoFactorAuthProviderConfig providerConfig = getTwoFaProviderConfig(tenantId, providerType)
                 .orElseThrow(() -> new ThingsboardException("2FA provider is not configured", ThingsboardErrorCode.BAD_REQUEST_PARAMS));
@@ -182,4 +164,22 @@ public class TwoFactorAuthService {
         }
     }
 
+
+    private <A extends TwoFactorAuthAccountConfig, C extends TwoFactorAuthProviderConfig> Optional<TwoFactorAuthProvider<C, A>> getTwoFaProvider(TwoFactorAuthProviderType providerType) {
+        return Optional.of((TwoFactorAuthProvider<C, A>) providers.get(providerType));
+    }
+
+    private <C extends TwoFactorAuthProviderConfig> Optional<C> getTwoFaProviderConfig(TenantId tenantId, TwoFactorAuthProviderType providerType) {
+        return getTwoFaSettings(tenantId)
+                .flatMap(twoFaSettings -> twoFaSettings.getProviderConfig(providerType))
+                .map(providerConfig -> (C) providerConfig);
+    }
+
+    @Autowired
+    private void setProviders(Collection<TwoFactorAuthProvider<?, ?>> providers) {
+        providers.forEach(provider -> {
+            this.providers.put(provider.getType(), provider);
+        });
+    }
+
 }

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

@@ -16,20 +16,33 @@
 package org.thingsboard.server.service.security.auth.mfa.config;
 
 import lombok.Data;
+import org.checkerframework.checker.index.qual.NonNegative;
 import org.thingsboard.server.service.security.auth.mfa.config.provider.TwoFactorAuthProviderConfig;
 import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProviderType;
 
 import javax.validation.Valid;
+import javax.validation.constraints.Min;
+import javax.validation.constraints.NotNull;
+import javax.validation.constraints.Pattern;
 import java.util.List;
 import java.util.Optional;
 
 @Data
 public class TwoFactorAuthSettings {
 
-    private boolean useSystemTwoFactorAuthSettings;
+    @NotNull
+    private Boolean useSystemTwoFactorAuthSettings;
     @Valid
     private List<TwoFactorAuthProviderConfig> providers;
 
+    @Pattern(regexp = "\\d+:\\d+")
+    private String verificationCodeSendRateLimit; // 1:60 - one time in a minute
+    @Pattern(regexp = "\\d+:\\d+")
+    private String verificationCodeCheckRateLimit; // soft lockout, on session level
+    @Min(0)
+    private Integer maxVerificationCodeSubmitAttemptsBeforeUserBlocking;
+
+
     public Optional<TwoFactorAuthProviderConfig> getProviderConfig(TwoFactorAuthProviderType providerType) {
         return Optional.ofNullable(providers)
                 .flatMap(providersConfigs -> providersConfigs.stream()

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

@@ -0,0 +1,34 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.security.auth.mfa.config.account;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProviderType;
+
+@EqualsAndHashCode(callSuper = true)
+@Data
+public class EmailTwoFactorAuthAccountConfig extends OtpBasedTwoFactorAuthAccountConfig {
+
+    private boolean useAccountEmail; // TODO [viacheslav]: validate
+    private String email;
+
+    @Override
+    public TwoFactorAuthProviderType getProviderType() {
+        return TwoFactorAuthProviderType.EMAIL;
+    }
+
+}

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

@@ -0,0 +1,19 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.security.auth.mfa.config.account;
+
+public abstract class OtpBasedTwoFactorAuthAccountConfig implements TwoFactorAuthAccountConfig {
+}

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

@@ -16,12 +16,14 @@
 package org.thingsboard.server.service.security.auth.mfa.config.account;
 
 import lombok.Data;
+import lombok.EqualsAndHashCode;
 import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProviderType;
 
 import javax.validation.constraints.NotBlank;
 
+@EqualsAndHashCode(callSuper = true)
 @Data
-public class SmsTwoFactorAuthAccountConfig implements TwoFactorAuthAccountConfig {
+public class SmsTwoFactorAuthAccountConfig extends OtpBasedTwoFactorAuthAccountConfig {
 
     @NotBlank
     private String phoneNumber;

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

@@ -0,0 +1,33 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.security.auth.mfa.config.provider;
+
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProviderType;
+
+@EqualsAndHashCode(callSuper = true)
+@Data
+public class EmailTwoFactorAuthProviderConfig extends OtpBasedTwoFactorAuthProviderConfig{
+
+    private String emailVerificationMessageTemplate; // FIXME [viacheslav]:
+
+    @Override
+    public TwoFactorAuthProviderType getProviderType() {
+        return TwoFactorAuthProviderType.EMAIL;
+    }
+
+}

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

@@ -0,0 +1,23 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.security.auth.mfa.config.provider;
+
+import lombok.Data;
+
+@Data
+public abstract class OtpBasedTwoFactorAuthProviderConfig implements TwoFactorAuthProviderConfig {
+    private Integer verificationCodeLifetime; // seconds
+}

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

@@ -16,13 +16,15 @@
 package org.thingsboard.server.service.security.auth.mfa.config.provider;
 
 import lombok.Data;
+import lombok.EqualsAndHashCode;
 import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProviderType;
 
 import javax.validation.constraints.NotBlank;
 import javax.validation.constraints.Pattern;
 
+@EqualsAndHashCode(callSuper = true)
 @Data
-public class SmsTwoFactorAuthProviderConfig implements TwoFactorAuthProviderConfig {
+public class SmsTwoFactorAuthProviderConfig extends OtpBasedTwoFactorAuthProviderConfig {
 
     @NotBlank
     @Pattern(regexp = ".*\\$\\{verificationCode}.*", message = "template must contain verification code")

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

@@ -27,7 +27,7 @@ public interface TwoFactorAuthProvider<C extends TwoFactorAuthProviderConfig, A
 
     default void prepareVerificationCode(SecurityUser user, C providerConfig, A accountConfig) throws ThingsboardException {}
 
-    boolean checkVerificationCode(SecurityUser user, String verificationCode, A accountConfig);
+    boolean checkVerificationCode(SecurityUser user, String verificationCode, C providerConfig, A accountConfig);
 
 
     TwoFactorAuthProviderType getType();

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

@@ -17,5 +17,6 @@ package org.thingsboard.server.service.security.auth.mfa.provider;
 
 public enum TwoFactorAuthProviderType {
     TOTP,
-    SMS
+    SMS,
+    EMAIL
 }

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

@@ -0,0 +1,67 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.security.auth.mfa.provider.impl;
+
+import org.springframework.cache.CacheManager;
+import org.springframework.stereotype.Service;
+import org.thingsboard.rule.engine.api.MailService;
+import org.thingsboard.server.common.data.User;
+import org.thingsboard.server.common.data.exception.ThingsboardException;
+import org.thingsboard.server.queue.util.TbCoreComponent;
+import org.thingsboard.server.service.security.auth.mfa.config.account.EmailTwoFactorAuthAccountConfig;
+import org.thingsboard.server.service.security.auth.mfa.config.provider.EmailTwoFactorAuthProviderConfig;
+import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProviderType;
+import org.thingsboard.server.service.security.model.SecurityUser;
+
+@Service
+@TbCoreComponent
+public class EmailTwoFactorAuthProvider extends OtpBasedTwoFactorAuthProvider<EmailTwoFactorAuthProviderConfig, EmailTwoFactorAuthAccountConfig> {
+
+    private final MailService mailService;
+
+    protected EmailTwoFactorAuthProvider(CacheManager cacheManager, MailService mailService) {
+        super(cacheManager);
+        this.mailService = mailService;
+    }
+
+
+    @Override
+    public EmailTwoFactorAuthAccountConfig generateNewAccountConfig(User user, EmailTwoFactorAuthProviderConfig providerConfig) {
+        EmailTwoFactorAuthAccountConfig accountConfig = new EmailTwoFactorAuthAccountConfig();
+        accountConfig.setUseAccountEmail(true);
+        return accountConfig;
+    }
+
+    @Override
+    protected void sendVerificationCode(SecurityUser user, String verificationCode, EmailTwoFactorAuthProviderConfig providerConfig, EmailTwoFactorAuthAccountConfig accountConfig) throws ThingsboardException {
+        String email;
+        if (accountConfig.isUseAccountEmail()) {
+            email = user.getEmail();
+        } else {
+            email = accountConfig.getEmail();
+        }
+
+        // FIXME [viacheslav]: mail template for 2FA verification
+        mailService.sendEmail(user.getTenantId(), email, "subject", "");
+    }
+
+
+    @Override
+    public TwoFactorAuthProviderType getType() {
+        return TwoFactorAuthProviderType.EMAIL;
+    }
+
+}

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

@@ -0,0 +1,70 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.security.auth.mfa.provider.impl;
+
+import lombok.Data;
+import org.apache.commons.lang3.RandomStringUtils;
+import org.springframework.cache.Cache;
+import org.springframework.cache.CacheManager;
+import org.thingsboard.server.common.data.CacheConstants;
+import org.thingsboard.server.common.data.exception.ThingsboardException;
+import org.thingsboard.server.service.security.auth.mfa.config.account.OtpBasedTwoFactorAuthAccountConfig;
+import org.thingsboard.server.service.security.auth.mfa.config.provider.OtpBasedTwoFactorAuthProviderConfig;
+import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProvider;
+import org.thingsboard.server.service.security.model.SecurityUser;
+
+import java.util.concurrent.TimeUnit;
+
+public abstract class OtpBasedTwoFactorAuthProvider<C extends OtpBasedTwoFactorAuthProviderConfig, A extends OtpBasedTwoFactorAuthAccountConfig> implements TwoFactorAuthProvider<C, A> {
+
+    private final Cache verificationCodesCache;
+
+    protected OtpBasedTwoFactorAuthProvider(CacheManager cacheManager) {
+        this.verificationCodesCache = cacheManager.getCache(CacheConstants.TWO_FA_VERIFICATION_CODES_CACHE);
+    }
+
+
+    @Override
+    public final void prepareVerificationCode(SecurityUser user, C providerConfig, A accountConfig) throws ThingsboardException {
+        String verificationCode = RandomStringUtils.randomNumeric(6);
+        verificationCodesCache.put(user.getSessionId(), new Otp(System.currentTimeMillis(), verificationCode));
+
+        sendVerificationCode(user, verificationCode, providerConfig, accountConfig);
+    }
+
+    protected abstract void sendVerificationCode(SecurityUser user, String verificationCode, C providerConfig, A accountConfig) throws ThingsboardException;
+
+
+    @Override
+    public final boolean checkVerificationCode(SecurityUser user, String verificationCode, C providerConfig, A accountConfig) {
+        Otp correctVerificationCode = verificationCodesCache.get(user.getSessionId(), Otp.class);
+        if (correctVerificationCode != null && verificationCode.equals(correctVerificationCode.getValue())) {
+            if (System.currentTimeMillis() - correctVerificationCode.getTimestamp() <= TimeUnit.SECONDS.toMillis(providerConfig.getVerificationCodeLifetime())) {
+                verificationCodesCache.evict(user.getSessionId());
+                return true;
+            }
+        }
+        return false;
+    }
+
+
+    @Data
+    private static class Otp {
+        private final long timestamp;
+        private final String value;
+    }
+
+}

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

@@ -15,21 +15,15 @@
  */
 package org.thingsboard.server.service.security.auth.mfa.provider.impl;
 
-import lombok.Data;
-import lombok.SneakyThrows;
-import org.apache.commons.lang3.RandomStringUtils;
-import org.springframework.cache.Cache;
 import org.springframework.cache.CacheManager;
 import org.springframework.stereotype.Service;
 import org.thingsboard.rule.engine.api.SmsService;
 import org.thingsboard.rule.engine.api.util.TbNodeUtils;
-import org.thingsboard.server.common.data.CacheConstants;
 import org.thingsboard.server.common.data.User;
 import org.thingsboard.server.common.data.exception.ThingsboardException;
 import org.thingsboard.server.queue.util.TbCoreComponent;
 import org.thingsboard.server.service.security.auth.mfa.config.account.SmsTwoFactorAuthAccountConfig;
 import org.thingsboard.server.service.security.auth.mfa.config.provider.SmsTwoFactorAuthProviderConfig;
-import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProvider;
 import org.thingsboard.server.service.security.auth.mfa.provider.TwoFactorAuthProviderType;
 import org.thingsboard.server.service.security.model.SecurityUser;
 
@@ -37,14 +31,13 @@ import java.util.Map;
 
 @Service
 @TbCoreComponent
-public class SmsTwoFactorAuthProvider implements TwoFactorAuthProvider<SmsTwoFactorAuthProviderConfig, SmsTwoFactorAuthAccountConfig> {
+public class SmsTwoFactorAuthProvider extends OtpBasedTwoFactorAuthProvider<SmsTwoFactorAuthProviderConfig, SmsTwoFactorAuthAccountConfig> {
 
     private final SmsService smsService;
-    private final Cache verificationCodesCache;
 
     public SmsTwoFactorAuthProvider(SmsService smsService, CacheManager cacheManager) {
+        super(cacheManager);
         this.smsService = smsService;
-        this.verificationCodesCache = cacheManager.getCache(CacheConstants.TWO_FA_VERIFICATION_CODES_CACHE);
     }
 
 
@@ -54,42 +47,21 @@ public class SmsTwoFactorAuthProvider implements TwoFactorAuthProvider<SmsTwoFac
     }
 
     @Override
-    public void prepareVerificationCode(SecurityUser user, SmsTwoFactorAuthProviderConfig providerConfig, SmsTwoFactorAuthAccountConfig accountConfig) throws ThingsboardException {
+    protected void sendVerificationCode(SecurityUser user, String verificationCode, SmsTwoFactorAuthProviderConfig providerConfig, SmsTwoFactorAuthAccountConfig accountConfig) throws ThingsboardException {
-        String verificationCode = RandomStringUtils.randomNumeric(6);
+        Map<String, String> messageData = Map.of(
-        verificationCodesCache.put(user.getSessionId(), new VerificationCode(System.currentTimeMillis(), verificationCode));
-
-        String phoneNumber = accountConfig.getPhoneNumber();
-
-        Map<String, String> data = Map.of(
                 "verificationCode", verificationCode,
                 "userEmail", user.getEmail()
         );
-        String message = TbNodeUtils.processTemplate(providerConfig.getSmsVerificationMessageTemplate(), data);
+        String message = TbNodeUtils.processTemplate(providerConfig.getSmsVerificationMessageTemplate(), messageData);
+        String phoneNumber = accountConfig.getPhoneNumber();
 
         smsService.sendSms(user.getTenantId(), user.getCustomerId(), new String[]{phoneNumber}, message);
     }
 
-    @Override
-    public boolean checkVerificationCode(SecurityUser user, String verificationCode, SmsTwoFactorAuthAccountConfig accountConfig) {
-        VerificationCode correctVerificationCode = verificationCodesCache.get(user.getSessionId(), VerificationCode.class);
-        if (correctVerificationCode != null && verificationCode.equals(correctVerificationCode.getValue())) {
-            verificationCodesCache.evict(user.getSessionId());
-            return true;
-        } else {
-            return false;
-        }
-    }
 
     @Override
     public TwoFactorAuthProviderType getType() {
         return TwoFactorAuthProviderType.SMS;
     }
 
-
-    @Data
-    private static class VerificationCode {
-        private final long timestamp;
-        private final String value;
-    }
-
 }

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

@@ -45,7 +45,7 @@ public class TotpTwoFactorAuthProvider implements TwoFactorAuthProvider<TotpTwoF
     }
 
     @Override
-    public final boolean checkVerificationCode(SecurityUser user, String verificationCode, TotpTwoFactorAuthAccountConfig accountConfig) {
+    public final boolean checkVerificationCode(SecurityUser user, String verificationCode, TotpTwoFactorAuthProviderConfig providerConfig, TotpTwoFactorAuthAccountConfig accountConfig) {
         String secretKey = UriComponentsBuilder.fromUriString(accountConfig.getAuthUrl()).build().getQueryParams().getFirst("secret");
         return new Totp(secretKey).verify(verificationCode);
     }
@@ -65,6 +65,7 @@ public class TotpTwoFactorAuthProvider implements TwoFactorAuthProvider<TotpTwoF
         return Base32.encode(RandomUtils.nextBytes(20));
     }
 
+
     @Override
     public TwoFactorAuthProviderType getType() {
         return TwoFactorAuthProviderType.TOTP;

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

@@ -39,6 +39,8 @@ import org.thingsboard.server.common.data.security.UserCredentials;
 import org.thingsboard.server.dao.audit.AuditLogService;
 import org.thingsboard.server.dao.customer.CustomerService;
 import org.thingsboard.server.dao.user.UserService;
+import org.thingsboard.server.service.security.auth.MfaAuthenticationToken;
+import org.thingsboard.server.service.security.auth.mfa.TwoFactorAuthService;
 import org.thingsboard.server.service.security.model.SecurityUser;
 import org.thingsboard.server.service.security.model.UserPrincipal;
 import org.thingsboard.server.service.security.system.SystemSecurityService;
@@ -55,16 +57,19 @@ public class RestAuthenticationProvider implements AuthenticationProvider {
     private final UserService userService;
     private final CustomerService customerService;
     private final AuditLogService auditLogService;
+    private final TwoFactorAuthService twoFactorAuthService;
 
     @Autowired
     public RestAuthenticationProvider(final UserService userService,
                                       final CustomerService customerService,
                                       final SystemSecurityService systemSecurityService,
-                                      final AuditLogService auditLogService) {
+                                      final AuditLogService auditLogService,
+                                      TwoFactorAuthService twoFactorAuthService) {
         this.userService = userService;
         this.customerService = customerService;
         this.systemSecurityService = systemSecurityService;
         this.auditLogService = auditLogService;
+        this.twoFactorAuthService = twoFactorAuthService;
     }
 
     @Override
@@ -77,17 +82,24 @@ public class RestAuthenticationProvider implements AuthenticationProvider {
         }
 
         UserPrincipal userPrincipal =  (UserPrincipal) principal;
+        SecurityUser securityUser;
         if (userPrincipal.getType() == UserPrincipal.Type.USER_NAME) {
             String username = userPrincipal.getValue();
             String password = (String) authentication.getCredentials();
-            return authenticateByUsernameAndPassword(authentication, userPrincipal, username, password);
+            securityUser = authenticateByUsernameAndPassword(authentication, userPrincipal, username, password);
+            if (twoFactorAuthService.getTwoFaAccountConfig(securityUser.getTenantId(), securityUser.getId()).isPresent()) {
+                return new MfaAuthenticationToken(securityUser);
+            }
+            logLoginAction((User) authentication.getPrincipal(), authentication, ActionType.LOGIN, null);
         } else {
             String publicId = userPrincipal.getValue();
-            return authenticateByPublicId(userPrincipal, publicId);
+            securityUser = authenticateByPublicId(userPrincipal, publicId);
         }
+
+        return new UsernamePasswordAuthenticationToken(securityUser, null, securityUser.getAuthorities());
     }
 
-    private Authentication authenticateByUsernameAndPassword(Authentication authentication, UserPrincipal userPrincipal, String username, String password) {
+    private SecurityUser authenticateByUsernameAndPassword(Authentication authentication, UserPrincipal userPrincipal, String username, String password) {
         User user = userService.findUserByEmail(TenantId.SYS_TENANT_ID, username);
         if (user == null) {
             throw new UsernameNotFoundException("User not found: " + username);
@@ -110,17 +122,14 @@ public class RestAuthenticationProvider implements AuthenticationProvider {
             if (user.getAuthority() == null)
                 throw new InsufficientAuthenticationException("User has no authority assigned");
 
-            SecurityUser securityUser = new SecurityUser(user, userCredentials.isEnabled(), userPrincipal);
+            return new SecurityUser(user, userCredentials.isEnabled(), userPrincipal);
-            // FIXME [viacheslav]: must not yet log login action if 2FA is used !
-            logLoginAction(user, authentication, ActionType.LOGIN, null);
-            return new UsernamePasswordAuthenticationToken(securityUser, null, securityUser.getAuthorities());
         } catch (Exception e) {
             logLoginAction(user, authentication, ActionType.LOGIN, e);
             throw e;
         }
     }
 
-    private Authentication authenticateByPublicId(UserPrincipal userPrincipal, String publicId) {
+    private SecurityUser authenticateByPublicId(UserPrincipal userPrincipal, String publicId) {
         CustomerId customerId;
         try {
             customerId = new CustomerId(UUID.fromString(publicId));
@@ -142,9 +151,7 @@ public class RestAuthenticationProvider implements AuthenticationProvider {
         user.setFirstName("Public");
         user.setLastName("Public");
 
-        SecurityUser securityUser = new SecurityUser(user, true, userPrincipal);
+        return new SecurityUser(user, true, userPrincipal);
-
-        return new UsernamePasswordAuthenticationToken(securityUser, null, securityUser.getAuthorities());
     }
 
     @Override

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

@@ -20,14 +20,14 @@ import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
-import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.web.WebAttributes;
 import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
 import org.springframework.stereotype.Component;
-import org.thingsboard.server.common.data.security.Authority;
+import org.thingsboard.server.service.security.auth.MfaAuthenticationToken;
+import org.thingsboard.server.service.security.auth.jwt.RefreshTokenRepository;
 import org.thingsboard.server.service.security.auth.mfa.TwoFactorAuthService;
-import org.thingsboard.server.service.security.auth.mfa.config.account.TwoFactorAuthAccountConfig;
+import org.thingsboard.server.service.security.auth.mfa.config.TwoFactorAuthSettings;
 import org.thingsboard.server.service.security.model.JwtTokenPair;
 import org.thingsboard.server.service.security.model.SecurityUser;
 import org.thingsboard.server.service.security.model.token.JwtTokenFactory;
@@ -37,7 +37,6 @@ import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import javax.servlet.http.HttpSession;
 import java.io.IOException;
-import java.util.Optional;
 
 @Component(value = "defaultAuthenticationSuccessHandler")
 @RequiredArgsConstructor
@@ -46,48 +45,31 @@ public class RestAwareAuthenticationSuccessHandler implements AuthenticationSucc
     private final ObjectMapper mapper;
     private final JwtTokenFactory tokenFactory;
     private final TwoFactorAuthService twoFactorAuthService;
+    private final RefreshTokenRepository refreshTokenRepository;
 
     @Override
     public void onAuthenticationSuccess(HttpServletRequest request, HttpServletResponse response,
                                         Authentication authentication) throws IOException, ServletException {
         SecurityUser securityUser = (SecurityUser) authentication.getPrincipal();
+        JwtTokenPair tokenPair = new JwtTokenPair();
 
-        JwtTokenPair tokenPair;
+        if (authentication instanceof MfaAuthenticationToken) {
-        if (authentication instanceof UsernamePasswordAuthenticationToken) {
+            TwoFactorAuthSettings twoFaSettings = twoFactorAuthService.getTwoFaSettings(securityUser.getTenantId()).get();
-            // TODO [viacheslav]: or maybe create another AuthenticationProvider and put it after rest authentication provider ?
+            // FIXME [viacheslav]: define the logic: pre-verification token lifetime,
-            tokenPair = processTwoFa(securityUser);
+            tokenPair.setToken(tokenFactory.createTwoFaPreVerificationToken(securityUser, ).getToken());
+            tokenPair.setRefreshToken(null);
         } else {
-            tokenPair = tokenFactory.createTokenPair(securityUser);
+            tokenPair.setToken(tokenFactory.createAccessJwtToken(securityUser).getToken());
+            tokenPair.setRefreshToken(refreshTokenRepository.requestRefreshToken(securityUser).getToken());
         }
 
         response.setStatus(HttpStatus.OK.value());
         response.setContentType(MediaType.APPLICATION_JSON_VALUE);
-
         mapper.writeValue(response.getWriter(), tokenPair);
 
         clearAuthenticationAttributes(request);
     }
 
-    private JwtTokenPair processTwoFa(SecurityUser securityUser) {
-        Optional<TwoFactorAuthAccountConfig> twoFaAccountConfig = twoFactorAuthService.getTwoFaAccountConfig(securityUser.getTenantId(), securityUser.getId());
-        if (twoFaAccountConfig.isPresent()) {
-            try {
-                twoFactorAuthService.processByTwoFaProvider(securityUser.getTenantId(), twoFaAccountConfig.get().getProviderType(),
-                        (provider, providerConfig) -> {
-                            provider.prepareVerificationCode(securityUser, providerConfig, twoFaAccountConfig.get());
-                        });
-                JwtTokenPair tokenPair = new JwtTokenPair();
-                tokenPair.setToken(tokenFactory.createPreVerificationToken(securityUser).getToken());
-                tokenPair.setScope(Authority.PRE_VERIFICATION_TOKEN);
-                return tokenPair;
-            } catch (Exception e) {
-                // TODO [viacheslav]: write audit log
-                log.error("Failed to process 2FA for user {}. Falling back to plain auth", securityUser.getId(), e);
-            }
-        }
-        return tokenFactory.createTokenPair(securityUser);
-    }
-
     /**
      * Removes temporary authentication-related data which may have been stored
      * in the session during the authentication process..

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

@@ -166,8 +166,8 @@ public class JwtTokenFactory {
         return securityUser;
     }
 
-    public JwtToken createPreVerificationToken(SecurityUser user) {
+    public JwtToken createTwoFaPreVerificationToken(SecurityUser user, Integer expirationTime) {
-        String token = setUpToken(user, Collections.singletonList(Authority.PRE_VERIFICATION_TOKEN.name()), settings.getPreVerificationTokenExpirationTime())
+        String token = setUpToken(user, Collections.singletonList(Authority.PRE_VERIFICATION_TOKEN.name()), expirationTime)
                 .claim(TENANT_ID, user.getTenantId().toString())
                 .compact();
         return new AccessJwtToken(token);

application/src/test/java/org/thingsboard/server/controller/TwoFactorAuthTest.java

@@ -25,7 +25,6 @@ import org.springframework.web.util.UriComponents;
 import org.springframework.web.util.UriComponentsBuilder;
 import org.thingsboard.rule.engine.api.SmsService;
 import org.thingsboard.server.service.security.auth.mfa.config.TwoFactorAuthSettings;
-import org.thingsboard.server.service.security.auth.mfa.config.account.SmsTwoFactorAuthAccountConfig;
 import org.thingsboard.server.service.security.auth.mfa.config.account.TotpTwoFactorAuthAccountConfig;
 import org.thingsboard.server.service.security.auth.mfa.config.account.TwoFactorAuthAccountConfig;
 import org.thingsboard.server.service.security.auth.mfa.config.provider.SmsTwoFactorAuthProviderConfig;
@@ -40,12 +39,12 @@ import java.util.stream.Collectors;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
 
+// TODO [viacheslav]: test sessionId
 public abstract class TwoFactorAuthTest extends AbstractControllerTest {
 
     @SpyBean