2FA support for platform authentication

application/src/main/java/org/thingsboard/server/config/JwtSettings.java

@@ -15,12 +15,14 @@
  */
 package org.thingsboard.server.config;
 
+import lombok.Data;
 import org.springframework.boot.context.properties.ConfigurationProperties;
-import org.springframework.context.annotation.Configuration;
+import org.springframework.stereotype.Component;
 import org.thingsboard.server.common.data.security.model.JwtToken;
 
-@Configuration
+@Component
 @ConfigurationProperties(prefix = "security.jwt")
+@Data
 public class JwtSettings {
     /**
      * {@link JwtToken} will expire after this time.
@@ -42,34 +44,10 @@ public class JwtSettings {
      */
     private Integer refreshTokenExpTime;
 
-    public Integer getRefreshTokenExpTime() {
-        return refreshTokenExpTime;
-    }
-
-    public void setRefreshTokenExpTime(Integer refreshTokenExpTime) {
-        this.refreshTokenExpTime = refreshTokenExpTime;
-    }
-
-    public Integer getTokenExpirationTime() {
-        return tokenExpirationTime;
-    }
-
-    public void setTokenExpirationTime(Integer tokenExpirationTime) {
-        this.tokenExpirationTime = tokenExpirationTime;
-    }
-
-    public String getTokenIssuer() {
-        return tokenIssuer;
-    }
-    public void setTokenIssuer(String tokenIssuer) {
-        this.tokenIssuer = tokenIssuer;
-    }
-
-    public String getTokenSigningKey() {
-        return tokenSigningKey;
-    }
+    /**
+    * Issued when 2FA is being used.
+    * Valid only for 2FA verification code checking.
+    * */
+    private Integer preVerificationTokenExpirationTime;
 
-    public void setTokenSigningKey(String tokenSigningKey) {
-        this.tokenSigningKey = tokenSigningKey;
-    }
-}
+}

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

@@ -35,6 +35,7 @@ import org.thingsboard.server.service.security.auth.mfa.config.TwoFactorAuthSett
 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.provider.TwoFactorAuthProviderType;
+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;
 
@@ -125,4 +126,22 @@ public class TwoFactorAuthController extends BaseController {
         twoFactorAuthService.saveTwoFaSettings(getTenantId(), twoFactorAuthSettings);
     }
 
+
+    @PostMapping("/auth/2fa/verification/check")
+    @PreAuthorize("hasAuthority('PRE_VERIFICATION_TOKEN')")
+    public JwtTokenPair checkTwoFaVerificationCode(@RequestParam String verificationCode) throws ThingsboardException {
+        SecurityUser user = getCurrentUser();
+
+        boolean verificationSuccess = twoFactorAuthService.processByTwoFaProvider(user.getTenantId(), user.getId(),
+                (provider, providerConfig, accountConfig) -> {
+                    return provider.checkVerificationCode(user, verificationCode, accountConfig);
+                });
+
+        if (verificationSuccess) {
+            return tokenFactory.createTokenPair(user);
+        } else {
+            throw new ThingsboardException("Verification code is incorrect", ThingsboardErrorCode.AUTHENTICATION);
+        }
+    }
+
 }

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

@@ -66,6 +66,7 @@ public class RefreshTokenAuthenticationProvider implements AuthenticationProvide
         } else {
             securityUser = authenticateByPublicId(principal.getValue());
         }
+        securityUser.setSessionId(unsafeUser.getSessionId());
 
         if (tokenOutdatingService.isOutdated(rawAccessToken, securityUser.getId())) {
             throw new CredentialsExpiredException("Token is outdated");

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

@@ -16,15 +16,18 @@
 package org.thingsboard.server.service.security.auth.rest;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import org.springframework.beans.factory.annotation.Autowired;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 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.model.JwtToken;
 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.model.JwtTokenPair;
 import org.thingsboard.server.service.security.model.SecurityUser;
 import org.thingsboard.server.service.security.model.token.JwtTokenFactory;
 
@@ -33,37 +36,44 @@ import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import javax.servlet.http.HttpSession;
 import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
+import java.util.Optional;
 
 @Component(value = "defaultAuthenticationSuccessHandler")
+@RequiredArgsConstructor
+@Slf4j
 public class RestAwareAuthenticationSuccessHandler implements AuthenticationSuccessHandler {
     private final ObjectMapper mapper;
     private final JwtTokenFactory tokenFactory;
-    private final RefreshTokenRepository refreshTokenRepository;
-
-    @Autowired
-    public RestAwareAuthenticationSuccessHandler(final ObjectMapper mapper, final JwtTokenFactory tokenFactory, final RefreshTokenRepository refreshTokenRepository) {
-        this.mapper = mapper;
-        this.tokenFactory = tokenFactory;
-        this.refreshTokenRepository = refreshTokenRepository;
-    }
+    private final TwoFactorAuthService twoFactorAuthService;
 
     @Override
     public void onAuthenticationSuccess(HttpServletRequest request, HttpServletResponse response,
                                         Authentication authentication) throws IOException, ServletException {
         SecurityUser securityUser = (SecurityUser) authentication.getPrincipal();
 
-        JwtToken accessToken = tokenFactory.createAccessJwtToken(securityUser);
-        JwtToken refreshToken = refreshTokenRepository.requestRefreshToken(securityUser);
+        JwtTokenPair tokenPair;
 
-        Map<String, String> tokenMap = new HashMap<String, String>();
-        tokenMap.put("token", accessToken.getToken());
-        tokenMap.put("refreshToken", refreshToken.getToken());
+        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());
+                        });
+                tokenPair = new JwtTokenPair();
+                tokenPair.setToken(tokenFactory.createPreVerificationToken(securityUser).getToken());
+            } catch (Exception e) {
+                log.error("Failed to process 2FA for user {}. Falling back to plain auth", securityUser.getId(), e);
+                tokenPair = tokenFactory.createTokenPair(securityUser);
+            }
+        } else {
+            tokenPair = tokenFactory.createTokenPair(securityUser);
+        }
 
         response.setStatus(HttpStatus.OK.value());
         response.setContentType(MediaType.APPLICATION_JSON_VALUE);
-        mapper.writeValue(response.getWriter(), tokenMap);
+
+        mapper.writeValue(response.getWriter(), tokenPair);
 
         clearAuthenticationAttributes(request);
     }

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

@@ -19,10 +19,12 @@ import io.swagger.annotations.ApiModel;
 import io.swagger.annotations.ApiModelProperty;
 import lombok.AllArgsConstructor;
 import lombok.Data;
+import lombok.NoArgsConstructor;
 
 @ApiModel(value = "JWT Token Pair")
 @Data
 @AllArgsConstructor
+@NoArgsConstructor
 public class JwtTokenPair {
 
     @ApiModelProperty(position = 1, value = "The JWT Access Token. Used to perform API calls.", example = "AAB254FF67D..")

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

@@ -21,6 +21,7 @@ import org.thingsboard.server.common.data.User;
 import org.thingsboard.server.common.data.id.UserId;
 
 import java.util.Collection;
+import java.util.UUID;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -31,6 +32,7 @@ public class SecurityUser extends User {
     private Collection<GrantedAuthority> authorities;
     private boolean enabled;
     private UserPrincipal userPrincipal;
+    private String sessionId;
 
     public SecurityUser() {
         super();
@@ -44,6 +46,7 @@ public class SecurityUser extends User {
         super(user);
         this.enabled = enabled;
         this.userPrincipal = userPrincipal;
+        this.sessionId = UUID.randomUUID().toString();
     }
 
     public Collection<GrantedAuthority> getAuthorities() {
@@ -71,4 +74,12 @@ public class SecurityUser extends User {
         this.userPrincipal = userPrincipal;
     }
 
+    public String getSessionId() {
+        return sessionId;
+    }
+
+    public void setSessionId(String sessionId) {
+        this.sessionId = sessionId;
+    }
+
 }

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

@@ -15,25 +15,17 @@
  */
 package org.thingsboard.server.service.security.model.token;
 
-import com.fasterxml.jackson.annotation.JsonIgnore;
-import io.jsonwebtoken.Claims;
 import org.thingsboard.server.common.data.security.model.JwtToken;
 
 public final class AccessJwtToken implements JwtToken {
     private final String rawToken;
-    @JsonIgnore
-    private transient Claims claims;
 
-    protected AccessJwtToken(final String token, Claims claims) {
-        this.rawToken = token;
-        this.claims = claims;
+    public AccessJwtToken(String rawToken) {
+        this.rawToken = rawToken;
     }
 
     public String getToken() {
         return this.rawToken;
     }
 
-    public Claims getClaims() {
-        return claims;
-    }
 }

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

@@ -18,6 +18,7 @@ package org.thingsboard.server.service.security.model.token;
 import io.jsonwebtoken.Claims;
 import io.jsonwebtoken.ExpiredJwtException;
 import io.jsonwebtoken.Jws;
+import io.jsonwebtoken.JwtBuilder;
 import io.jsonwebtoken.Jwts;
 import io.jsonwebtoken.MalformedJwtException;
 import io.jsonwebtoken.SignatureAlgorithm;
@@ -36,6 +37,7 @@ import org.thingsboard.server.common.data.security.Authority;
 import org.thingsboard.server.common.data.security.model.JwtToken;
 import org.thingsboard.server.config.JwtSettings;
 import org.thingsboard.server.service.security.exception.JwtExpiredTokenException;
+import org.thingsboard.server.service.security.model.JwtTokenPair;
 import org.thingsboard.server.service.security.model.SecurityUser;
 import org.thingsboard.server.service.security.model.UserPrincipal;
 
@@ -58,6 +60,7 @@ public class JwtTokenFactory {
     private static final String IS_PUBLIC = "isPublic";
     private static final String TENANT_ID = "tenantId";
     private static final String CUSTOMER_ID = "customerId";
+    private static final String SESSION_ID = "sessionId";
 
     private final JwtSettings settings;
 
@@ -70,39 +73,28 @@ public class JwtTokenFactory {
      * Factory method for issuing new JWT Tokens.
      */
     public AccessJwtToken createAccessJwtToken(SecurityUser securityUser) {
-        if (StringUtils.isBlank(securityUser.getEmail()))
-            throw new IllegalArgumentException("Cannot create JWT Token without username/email");
-
-        if (securityUser.getAuthority() == null)
+        if (securityUser.getAuthority() == null) {
             throw new IllegalArgumentException("User doesn't have any privileges");
+        }
 
         UserPrincipal principal = securityUser.getUserPrincipal();
-        String subject = principal.getValue();
-        Claims claims = Jwts.claims().setSubject(subject);
-        claims.put(SCOPES, securityUser.getAuthorities().stream().map(GrantedAuthority::getAuthority).collect(Collectors.toList()));
-        claims.put(USER_ID, securityUser.getId().getId().toString());
-        claims.put(FIRST_NAME, securityUser.getFirstName());
-        claims.put(LAST_NAME, securityUser.getLastName());
-        claims.put(ENABLED, securityUser.isEnabled());
-        claims.put(IS_PUBLIC, principal.getType() == UserPrincipal.Type.PUBLIC_ID);
+
+        JwtBuilder jwtBuilder = setUpToken(securityUser, securityUser.getAuthorities().stream()
+                .map(GrantedAuthority::getAuthority).collect(Collectors.toList()), settings.getTokenExpirationTime());
+        jwtBuilder.claim(FIRST_NAME, securityUser.getFirstName())
+                .claim(LAST_NAME, securityUser.getLastName())
+                .claim(ENABLED, securityUser.isEnabled())
+                .claim(IS_PUBLIC, principal.getType() == UserPrincipal.Type.PUBLIC_ID);
         if (securityUser.getTenantId() != null) {
-            claims.put(TENANT_ID, securityUser.getTenantId().getId().toString());
+            jwtBuilder.claim(TENANT_ID, securityUser.getTenantId().getId().toString());
         }
         if (securityUser.getCustomerId() != null) {
-            claims.put(CUSTOMER_ID, securityUser.getCustomerId().getId().toString());
+            jwtBuilder.claim(CUSTOMER_ID, securityUser.getCustomerId().getId().toString());
         }
 
-        ZonedDateTime currentTime = ZonedDateTime.now();
+        String token = jwtBuilder.compact();
 
-        String token = Jwts.builder()
-                .setClaims(claims)
-                .setIssuer(settings.getTokenIssuer())
-                .setIssuedAt(Date.from(currentTime.toInstant()))
-                .setExpiration(Date.from(currentTime.plusSeconds(settings.getTokenExpirationTime()).toInstant()))
-                .signWith(SignatureAlgorithm.HS512, settings.getTokenSigningKey())
-                .compact();
-
-        return new AccessJwtToken(token, claims);
+        return new AccessJwtToken(token);
     }
 
     public SecurityUser parseAccessJwtToken(RawAccessJwtToken rawAccessToken) {
@@ -118,47 +110,40 @@ public class JwtTokenFactory {
         SecurityUser securityUser = new SecurityUser(new UserId(UUID.fromString(claims.get(USER_ID, String.class))));
         securityUser.setEmail(subject);
         securityUser.setAuthority(Authority.parse(scopes.get(0)));
-        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);
-        securityUser.setUserPrincipal(principal);
         String tenantId = claims.get(TENANT_ID, String.class);
         if (tenantId != null) {
             securityUser.setTenantId(TenantId.fromUUID(UUID.fromString(tenantId)));
+        } 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)));
+        securityUser.setSessionId(claims.get(SESSION_ID, String.class));
+
+        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);
+            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));
         }
 
         return securityUser;
     }
 
     public JwtToken createRefreshToken(SecurityUser securityUser) {
-        if (StringUtils.isBlank(securityUser.getEmail())) {
-            throw new IllegalArgumentException("Cannot create JWT Token without username/email");
-        }
-
-        ZonedDateTime currentTime = ZonedDateTime.now();
-
         UserPrincipal principal = securityUser.getUserPrincipal();
-        Claims claims = Jwts.claims().setSubject(principal.getValue());
-        claims.put(SCOPES, Collections.singletonList(Authority.REFRESH_TOKEN.name()));
-        claims.put(USER_ID, securityUser.getId().getId().toString());
-        claims.put(IS_PUBLIC, principal.getType() == UserPrincipal.Type.PUBLIC_ID);
 
-        String token = Jwts.builder()
-                .setClaims(claims)
-                .setIssuer(settings.getTokenIssuer())
-                .setId(UUID.randomUUID().toString())
-                .setIssuedAt(Date.from(currentTime.toInstant()))
-                .setExpiration(Date.from(currentTime.plusSeconds(settings.getRefreshTokenExpTime()).toInstant()))
-                .signWith(SignatureAlgorithm.HS512, settings.getTokenSigningKey())
-                .compact();
+        String token = setUpToken(securityUser, Collections.singletonList(Authority.REFRESH_TOKEN.name()), settings.getRefreshTokenExpTime())
+                .claim(IS_PUBLIC, principal.getType() == UserPrincipal.Type.PUBLIC_ID)
+                .setId(UUID.randomUUID().toString()).compact();
 
-        return new AccessJwtToken(token, claims);
+        return new AccessJwtToken(token);
     }
 
     public SecurityUser parseRefreshToken(RawAccessJwtToken rawAccessToken) {
@@ -177,9 +162,41 @@ public class JwtTokenFactory {
         UserPrincipal principal = new UserPrincipal(isPublic ? UserPrincipal.Type.PUBLIC_ID : UserPrincipal.Type.USER_NAME, subject);
         SecurityUser securityUser = new SecurityUser(new UserId(UUID.fromString(claims.get(USER_ID, String.class))));
         securityUser.setUserPrincipal(principal);
+        securityUser.setSessionId(claims.get(SESSION_ID, String.class));
         return securityUser;
     }
 
+    public JwtToken createPreVerificationToken(SecurityUser user) {
+        String token = setUpToken(user, Collections.singletonList(Authority.PRE_VERIFICATION_TOKEN.name()), settings.getPreVerificationTokenExpirationTime())
+                .claim(TENANT_ID, user.getTenantId().toString())
+                .compact();
+        return new AccessJwtToken(token);
+    }
+
+    private JwtBuilder setUpToken(SecurityUser securityUser, List<String> scopes, long expirationTime) {
+        if (StringUtils.isBlank(securityUser.getEmail())) {
+            throw new IllegalArgumentException("Cannot create JWT Token without username/email");
+        }
+
+        UserPrincipal principal = securityUser.getUserPrincipal();
+
+        Claims claims = Jwts.claims().setSubject(principal.getValue());
+        claims.put(USER_ID, securityUser.getId().getId().toString());
+        claims.put(SCOPES, scopes);
+        if (securityUser.getSessionId() != null) {
+            claims.put(SESSION_ID, securityUser.getSessionId());
+        }
+
+        ZonedDateTime currentTime = ZonedDateTime.now();
+
+        return Jwts.builder()
+                .setClaims(claims)
+                .setIssuer(settings.getTokenIssuer())
+                .setIssuedAt(Date.from(currentTime.toInstant()))
+                .setExpiration(Date.from(currentTime.plusSeconds(expirationTime).toInstant()))
+                .signWith(SignatureAlgorithm.HS512, settings.getTokenSigningKey());
+    }
+
     public Jws<Claims> parseTokenClaims(JwtToken token) {
         try {
             return Jwts.parser()
@@ -193,4 +210,11 @@ public class JwtTokenFactory {
             throw new JwtExpiredTokenException(token, "JWT Token expired", expiredEx);
         }
     }
+
+    public JwtTokenPair createTokenPair(SecurityUser securityUser) {
+        JwtToken accessToken = createAccessJwtToken(securityUser);
+        JwtToken refreshToken = createRefreshToken(securityUser);
+        return new JwtTokenPair(accessToken.getToken(), refreshToken.getToken());
+    }
+
 }

application/src/main/resources/thingsboard.yml

@@ -127,6 +127,8 @@ security:
   jwt:
     tokenExpirationTime: "${JWT_TOKEN_EXPIRATION_TIME:9000}" # Number of seconds (2.5 hours)
     refreshTokenExpTime: "${JWT_REFRESH_TOKEN_EXPIRATION_TIME:604800}" # Number of seconds (1 week)
+    # Number of seconds. Issued when 2FA is being used; valid only for checking 2FA verification code after which usual token pair is issued
+    preVerificationTokenExpirationTime: "${JWT_PRE_VERIFICATION_TOKEN_EXPIRATION_TIME:30}"
     tokenIssuer: "${JWT_TOKEN_ISSUER:thingsboard.io}"
     tokenSigningKey: "${JWT_TOKEN_SIGNING_KEY:thingsboardDefaultSigningKey}"
   # Enable/disable access to Tenant Administrators JWT token by System Administrator or Customer Users JWT token by Tenant Administrator
@@ -425,6 +427,9 @@ caffeine:
     edges:
       timeToLiveInMinutes: "${CACHE_SPECS_EDGES_TTL:1440}"
       maxSize: "${CACHE_SPECS_EDGES_MAX_SIZE:10000}"
+    twoFaVerificationCodes:
+      timeToLiveInMinutes: "1"
+      maxSize: "100000"
 
 redis:
   # standalone or cluster

common/data/src/main/java/org/thingsboard/server/common/data/security/Authority.java

@@ -16,11 +16,12 @@
 package org.thingsboard.server.common.data.security;
 
 public enum Authority {
-    
+
     SYS_ADMIN(0),
     TENANT_ADMIN(1),
     CUSTOMER_USER(2),
-    REFRESH_TOKEN(10);
+    REFRESH_TOKEN(10),
+    PRE_VERIFICATION_TOKEN(11);
 
     private int code;