Implementing TOTP-Based Two-Factor Authentication in Spring Boot

source: imperva.com

Two-factor authentication (2FA) is an essential security measure that adds an extra layer of protection to your applications. Time-based One-Time Password (TOTP) is one of the most commonly used methods for 2FA. In this article, we will discuss how to implement TOTP-based 2FA in a Spring Boot application. We will cover the necessary dependencies, configuration, and code examples to get you up and running.

What is TOTP?

TOTP stands for Time-based One-Time Password. It is a temporary, one-time password generated using the current time as a variable. TOTP is an extension of the HMAC-based One-Time Password (HOTP) algorithm, which generates passwords based on a counter value. TOTP uses the current time instead of a counter, making the passwords valid only for a short period (usually 30 seconds).

Why Use TOTP for 2FA?

TOTP is widely used for 2FA because it provides several benefits:

1. Enhanced Security: Adds an extra layer of security beyond just a username and password.
2. Short Lifespan: The OTP is valid for a short time, reducing the risk of it being used maliciously.
3. Compatibility: Supported by many 2FA apps such as Google Authenticator and Authy.

Prerequisites

Before we dive into the implementation, ensure you have the following:

• Java Development Kit (JDK) 8 or later
• Maven or Gradle for dependency management
• An Integrated Development Environment (IDE) like IntelliJ IDEA or Eclipse
• Basic knowledge of Spring Boot

Dependencies

To implement TOTP-based 2FA, we need the following dependencies:

• Spring Boot Starter Web: To create the RESTful APIs.
• Spring Boot Starter Security: To handle authentication and authorization.
• Google Authenticator Library: To generate and validate TOTP codes.

Add the following dependencies to your pom.xml file:

<dependencies>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-web</artifactId>
    </dependency>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-security</artifactId>
    </dependency>
    <dependency>
        <groupId>com.warrenstrange</groupId>
        <artifactId>googleauth</artifactId>
        <version>1.6.0</version>
    </dependency>
</dependencies>

Configuration

Spring Security Configuration

Create a SecurityConfig class to configure Spring Security:

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.annotation.web.configuration.EnableWebSecurity;
import org.springframework.security.config.annotation.web.configuration.WebSecurityConfigurerAdapter;
import org.springframework.security.crypto.bcrypt.BCryptPasswordEncoder;
import org.springframework.security.crypto.password.PasswordEncoder;

@Configuration
@EnableWebSecurity
public class SecurityConfig extends WebSecurityConfigurerAdapter {

    @Override
    protected void configure(HttpSecurity http) throws Exception {
        http
            .csrf().disable()
            .authorizeRequests()
            .antMatchers("/login", "/register").permitAll()
            .anyRequest().authenticated()
            .and()
            .formLogin()
            .loginPage("/login")
            .defaultSuccessUrl("/totp-setup", true)
            .and()
            .logout()
            .logoutUrl("/logout")
            .logoutSuccessUrl("/login");
    }

    @Bean
    public PasswordEncoder passwordEncoder() {
        return new BCryptPasswordEncoder();
    }
}

User Entity

Create a User entity to represent the users in your application:

import javax.persistence.*;

@Entity
public class User {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;

    private String username;
    private String password;
    private String secret;

    // Getters and Setters
}

User Repository

Create a UserRepository interface to interact with the database:

import org.springframework.data.jpa.repository.JpaRepository;

public interface UserRepository extends JpaRepository<User, Long> {
    User findByUsername(String username);
}

User Service

Create a UserService to handle user-related operations:

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.crypto.password.PasswordEncoder;
import org.springframework.stereotype.Service;

@Service
public class UserService {

    @Autowired
    private UserRepository userRepository;

    @Autowired
    private PasswordEncoder passwordEncoder;

    public User register(User user) {
        user.setPassword(passwordEncoder.encode(user.getPassword()));
        // Generate secret key for TOTP
        user.setSecret(generateSecretKey());
        return userRepository.save(user);
    }

    private String generateSecretKey() {
        return new com.warrenstrange.googleauth.GoogleAuthenticator().createCredentials().getKey();
    }
}

User Controller

Create a UserController to handle user registration and login:

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

@RestController
@RequestMapping("/api")
public class UserController {

    @Autowired
    private UserService userService;

    @PostMapping("/register")
    public User register(@RequestBody User user) {
        return userService.register(user);
    }

    @GetMapping("/totp-setup")
    public String setupTOTP() {
        User user = // Get the logged-in user from the security context
        String secret = user.getSecret();
        String qrUrl = getQRBarcodeURL(user.getUsername(), secret);
        return qrUrl;
    }

    private String getQRBarcodeURL(String user, String secret) {
        return "https://chart.googleapis.com/chart?chs=200x200&chld=M|0&cht=qr&chl=" + 
               "otpauth://totp/" + user + "?secret=" + secret;
    }
}

TOTP Validation

Create an endpoint to validate the TOTP code:

import org.springframework.web.bind.annotation.*;

@RestController
@RequestMapping("/api")
public class TOTPController {

    @Autowired
    private UserService userService;

    @PostMapping("/validate-totp")
    public boolean validateTOTP(@RequestBody TOTPRequest totpRequest) {
        User user = // Get the logged-in user from the security context
        return isValidTOTP(totpRequest.getCode(), user.getSecret());
    }

    private boolean isValidTOTP(int code, String secret) {
        com.warrenstrange.googleauth.GoogleAuthenticator gAuth = new com.warrenstrange.googleauth.GoogleAuthenticator();
        return gAuth.authorize(secret, code);
    }
}

class TOTPRequest {
    private int code;
    // Getters and Setters
}

Enhancing the Implementation

Now that we have a basic implementation, let’s explore some additional enhancements and considerations to make our TOTP-based 2FA more robust and user-friendly.

Enhancements

1. User Experience Improvements

While the basic implementation is functional, enhancing the user experience can make the process smoother and more intuitive.

Custom Login Page

Instead of using the default login page provided by Spring Security, you can create a custom login page. This allows you to provide better user experience and branding.

Create a login.html file under src/main/resources/templates:

<!DOCTYPE html>
<html xmlns:th="http://www.thymeleaf.org">
<head>
    <title>Login</title>
</head>
<body>
    <h1>Login</h1>
    <form th:action="@{/login}" method="post">
        <div>
            <label for="username">Username:</label>
            <input type="text" id="username" name="username">
        </div>
        <div>
            <label for="password">Password:</label>
            <input type="password" id="password" name="password">
        </div>
        <div>
            <button type="submit">Login</button>
        </div>
    </form>
</body>
</html>

Update the SecurityConfig to use the custom login page:

@Override
protected void configure(HttpSecurity http) throws Exception {
    http
        .csrf().disable()
        .authorizeRequests()
        .antMatchers("/login", "/register").permitAll()
        .anyRequest().authenticated()
        .and()
        .formLogin()
        .loginPage("/login")
        .defaultSuccessUrl("/totp-setup", true)
        .permitAll()
        .and()
        .logout()
        .logoutUrl("/logout")
        .logoutSuccessUrl("/login");
}

2. TOTP Backup Codes

Backup codes are a useful feature that allows users to authenticate if they lose access to their TOTP device. Generate a set of backup codes during registration and allow users to view and regenerate them.

Backup Codes Entity

Update the User entity to store backup codes:

import java.util.List;
import javax.persistence.ElementCollection;

@Entity
public class User {

    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;

    private String username;
    private String password;
    private String secret;

    @ElementCollection
    private List<String> backupCodes;

    // Getters and Setters
}

Generating Backup Codes

Generate backup codes during user registration:

import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.List;

@Service
public class UserService {

    @Autowired
    private UserRepository userRepository;

    @Autowired
    private PasswordEncoder passwordEncoder;

    public User register(User user) {
        user.setPassword(passwordEncoder.encode(user.getPassword()));
        user.setSecret(generateSecretKey());
        user.setBackupCodes(generateBackupCodes());
        return userRepository.save(user);
    }

    private String generateSecretKey() {
        return new com.warrenstrange.googleauth.GoogleAuthenticator().createCredentials().getKey();
    }

    private List<String> generateBackupCodes() {
        List<String> backupCodes = new ArrayList<>();
        SecureRandom random = new SecureRandom();
        for (int i = 0; i < 10; i++) {
            backupCodes.add(String.format("%06d", random.nextInt(1000000)));
        }
        return backupCodes;
    }
}

Displaying and Validating Backup Codes

Provide endpoints for users to view and regenerate their backup codes. Validate backup codes during the 2FA process.

@RestController
@RequestMapping("/api")
public class BackupCodeController {

    @Autowired
    private UserService userService;

    @GetMapping("/backup-codes")
    public List<String> getBackupCodes() {
        User user = // Get the logged-in user from the security context
        return user.getBackupCodes();
    }

    @PostMapping("/generate-backup-codes")
    public List<String> generateBackupCodes() {
        User user = // Get the logged-in user from the security context
        List<String> newBackupCodes = userService.generateBackupCodes();
        user.setBackupCodes(newBackupCodes);
        userService.saveUser(user);
        return newBackupCodes;
    }
}

@PostMapping("/validate-totp-or-backup")
public boolean validateTOTPOrBackup(@RequestBody TOTPRequest totpRequest) {
    User user = // Get the logged-in user from the security context
    boolean isValidTOTP = isValidTOTP(totpRequest.getCode(), user.getSecret());
    boolean isValidBackupCode = user.getBackupCodes().contains(String.valueOf(totpRequest.getCode()));
    
    if (isValidBackupCode) {
        // Remove the used backup code
        user.getBackupCodes().remove(String.valueOf(totpRequest.getCode()));
        userService.saveUser(user);
    }
    
    return isValidTOTP || isValidBackupCode;
}

3. User-Friendly Error Handling

Improve error handling to provide clear feedback to users. Create custom exception classes and handlers.

Custom Exceptions

Create a custom exception for failed 2FA validation:

public class InvalidTOTPException extends RuntimeException {
    public InvalidTOTPException(String message) {
        super(message);
    }
}

Exception Handler

Create an exception handler to manage custom exceptions:

import org.springframework.http.HttpStatus;
import org.springframework.web.bind.annotation.ExceptionHandler;
import org.springframework.web.bind.annotation.ResponseStatus;
import org.springframework.web.bind.annotation.RestControllerAdvice;

@RestControllerAdvice
public class GlobalExceptionHandler {

    @ExceptionHandler(InvalidTOTPException.class)
    @ResponseStatus(HttpStatus.UNAUTHORIZED)
    public String handleInvalidTOTPException(InvalidTOTPException ex) {
        return ex.getMessage();
    }
}

4. Logging and Monitoring

Add logging to track important events and potential security issues. Use Spring Boot’s built-in logging support.

Logging Configuration

Add logging configuration in application.properties:

logging.level.org.springframework.security=DEBUG
logging.level.com.yourpackage=INFO

Log Important Events

Log important events such as user registration, login, and TOTP validation:

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Service
public class UserService {

    private static final Logger logger = LoggerFactory.getLogger(UserService.class);

    @Autowired
    private UserRepository userRepository;

    @Autowired
    private PasswordEncoder passwordEncoder;

    public User register(User user) {
        user.setPassword(passwordEncoder.encode(user.getPassword()));
        user.setSecret(generateSecretKey());
        user.setBackupCodes(generateBackupCodes());
        User savedUser = userRepository.save(user);
        logger.info("User registered: {}", user.getUsername());
        return savedUser;
    }

    public boolean validateTOTP(int code, String secret) {
        boolean isValid = new GoogleAuthenticator().authorize(secret, code);
        logger.info("TOTP validation result: {}", isValid);
        return isValid;
    }
}

5. Implementing Rate Limiting

To prevent brute force attacks, implement rate limiting for TOTP validation attempts.

Rate Limiting Configuration

Use a library like Bucket4j to implement rate limiting. Add the dependency to your pom.xml:

<dependency>
    <groupId>com.github.vladimir-bukhtoyarov</groupId>
    <artifactId>bucket4j-core</artifactId>
    <version>7.3.0</version>
</dependency>

Rate Limiting Service

Create a rate limiting service:

import com.github.bucket4j.Bucket;
import com.github.bucket4j.Bandwidth;
import com.github.bucket4j.Refill;
import org.springframework.stereotype.Service;

import java.time.Duration;

@Service
public class RateLimitingService {

    private final Bucket bucket;

    public RateLimitingService() {
        Bandwidth limit = Bandwidth.classic(5, Refill.intervally(5, Duration.ofMinutes(1)));
        this.bucket = Bucket.builder().addLimit(limit).build();
    }

    public boolean tryConsume() {
        return bucket.tryConsume(1);
    }
}

Applying Rate Limiting

Apply rate limiting to the TOTP validation endpoint:

@RestController
@RequestMapping("/api")
public class TOTPController {

    @Autowired
    private UserService userService;

    @Autowired
    private RateLimitingService rateLimitingService;

    @PostMapping("/validate-totp")
    public boolean validateTOTP(@RequestBody TOTPRequest totpRequest) {
        if (!rateLimitingService.tryConsume()) {
            throw new RuntimeException("Too many attempts, please try again later.");
        }
        User user = // Get the logged-in user from the security context
        return isValidTOTP(totpRequest.getCode(), user.getSecret());
    }

    private boolean isValidTOTP(int code, String secret) {
        GoogleAuthenticator gAuth = new GoogleAuthenticator();
        return gAuth.authorize(secret, code);
    }
}

Conclusion

Implementing TOTP-based 2FA in a Spring Boot application enhances security by adding an extra layer of authentication. By following the steps outlined in this article, you can integrate TOTP 2FA into your Spring Boot application, providing users with a more secure authentication mechanism. We have covered the basics of setting up TOTP, enhancing user experience, handling errors, logging important events, and implementing rate limiting to prevent abuse.

Feel free to extend and customize this implementation according to your specific requirements. Ensuring that your application has a robust and user-friendly 2FA system is crucial for maintaining security in today’s threat landscape. Happy coding!