Token Refresh Mechanism

Introduction

Welcome to the lesson on implementing a token refresh mechanism! In our previous lesson, we explored the basics of secure token-based authentication and the importance of using cookies and JWT expiration to enhance security. Now, we'll dive deeper into how we can extend user sessions securely using refresh tokens . This mechanism is crucial for maintaining seamless user experiences while ensuring robust security in web applications. Let's get started! 🚀

Understanding Access and Refresh Tokens

Now we know that access tokens, which are used to authenticate user requests, need to be short-lived. They have a limited lifespan to minimize the risk of misuse if compromised. However, this short lifespan can disrupt user sessions, requiring frequent re-authentication. This is where refresh tokens come into play. Refresh tokens are long-lived and can be used to obtain new access tokens without requiring the user to log in again. This combination allows for secure, uninterrupted user sessions.

Exploiting Token Vulnerabilities

Before we implement a secure token refresh mechanism, it's important to understand potential vulnerabilities. Attackers can exploit improperly managed tokens to gain unauthorized access. Let's see how an attacker might exploit a vulnerability in token handling: Shell# Simulating an attack by using a stolen refresh token curl -X POST http://api.pastebin.com/api/auth/refresh \ -H "Content-Type: application/json" \ -d '{"refresh_token": "stolen_refresh_token"}' In this example, an attacker uses a stolen refresh token to request a new access token. If the server doesn't properly validate the token, the attacker could gain unauthorized access. This highlights the importance of secure token management and demonstrates the critical need for secure token handling and validation.

Token Generation and Invalidation

The first step in implementing a token refresh mechanism is to generate new access and refresh tokens. It's also important to invalidate the old refresh token to prevent reuse. In Java, you can use the java-jwt library from Auth0 to generate and validate JWTs. Below is an example of how to generate access and refresh tokens and store refresh tokens in a thread-safe set: Javaimport com.auth0.jwt.JWT; import com.auth0.jwt.algorithms.Algorithm; import java.util.Date; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; public class TokenService { // ⚠️ Security Warning: Never hardcode secrets in production. Use environment variables or a secrets management system. The hardcoded values shown here are for demonstration purposes only. private static final Algorithm ACCESS_ALGORITHM = Algorithm.HMAC256("your-access-secret"); private static final Algorithm REFRESH_ALGORITHM = Algorithm.HMAC256("your-refresh-secret"); // Store refresh tokens in a thread-safe set // Note: In production environments, store this data in a distributed cache like Redis or a database to ensure consistency across multiple server instances. private static final Set<String> refreshTokens = ConcurrentHashMap.newKeySet(); public static String generateAccessToken(Integer userId) { return JWT.create() .withClaim("userId", userId) .withClaim("type", "access") .withExpiresAt(new Date(System.currentTimeMillis() + 15 * 60 * 1000)) // 15 minutes .sign(ACCESS_ALGORITHM); } public static String generateRefreshToken(Integer userId) { String refreshToken = JWT.create() .withClaim("userId", userId) .withClaim("type", "refresh") .withExpiresAt(new Date(System.currentTimeMillis() + 7 * 24 * 60 * 60 * 1000)) // 7 days .sign(REFRESH_ALGORITHM); refreshTokens.add(refreshToken); return refreshToken; } public static void invalidateRefreshToken(String refreshToken) { refreshTokens.remove(refreshToken); } public static boolean isRefreshTokenValid(String refreshToken) { return refreshTokens.contains(refreshToken); } } This code demonstrates how to generate access and refresh tokens, store refresh tokens, and invalidate them when necessary.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Authentication Checks

After generating the tokens, the next step is to ensure proper authentication. We need to verify that the user is authenticated before allowing them to refresh their tokens.

Below is an example using a Spring Boot controller to check for the presence of a refresh token in the request cookies:

Java

import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import jakarta.servlet.http.Cookie;
import jakarta.servlet.http.HttpServletRequest;

@RestController
@RequestMapping("/api/auth")
public class AuthController {

    @PostMapping("/refresh")
    public ResponseEntity<?> refreshToken(HttpServletRequest request) {
        String refreshToken = null;
        if (request.getCookies() != null) {
            for (Cookie cookie : request.getCookies()) {
                if ("refresh_token".equals(cookie.getName())) {
                    refreshToken = cookie.getValue();
                    break;
                }
            }
        }
        if (refreshToken == null) {
            return ResponseEntity.status(401).body("Refresh token required");
        }
        // ... rest of the code
        return ResponseEntity.ok().build();
    }
}

import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import jakarta.servlet.http.Cookie;
import jakarta.servlet.http.HttpServletRequest;

@RestController
@RequestMapping("/api/auth")
public class AuthController {

    @PostMapping("/refresh")
    public ResponseEntity<?> refreshToken(HttpServletRequest request) {
        String refreshToken = null;
        if (request.getCookies() != null) {
            for (Cookie cookie : request.getCookies()) {
                if ("refresh_token".equals(cookie.getName())) {
                    refreshToken = cookie.getValue();
                    break;
                }
            }
        }
        if (refreshToken == null) {
            return ResponseEntity.status(401).body("Refresh token required");
        }
        // ... rest of the code
        return ResponseEntity.ok().build();
    }
}

This code checks for the presence of a refresh token in the request cookies. If the token is missing, the server responds with an error, indicating that authentication is required.

Token Validation

Once we've ensured authentication, the next step is to validate the refresh token. This involves checking that the token is valid, has not been tampered with, and is of the correct type.

Here's how you can validate the refresh token using the java-jwt library:

Java

import com.auth0.jwt.JWT;
import com.auth0.jwt.algorithms.Algorithm;
import com.auth0.jwt.interfaces.DecodedJWT;
import jakarta.servlet.http.HttpServletResponse;

@PostMapping("/refresh")
public ResponseEntity<?> refreshToken(HttpServletRequest request, HttpServletResponse response) {
    String refreshToken = null;
    if (request.getCookies() != null) {
        for (Cookie cookie : request.getCookies()) {
            if ("refresh_token".equals(cookie.getName())) {
                refreshToken = cookie.getValue();
                break;
            }
        }
    }
    if (refreshToken == null) {
        return ResponseEntity.status(401).body("Refresh token required");
    }

    // Validate the refresh token
    try {
        DecodedJWT decoded = JWT.require(REFRESH_ALGORITHM)
                .build()
                .verify(refreshToken);

        String type = decoded.getClaim("type").asString();
        if (!"refresh".equals(type)) {
            throw new IllegalArgumentException("Invalid token type");
        }
        if (!TokenService.isRefreshTokenValid(refreshToken)) {
            return ResponseEntity.status(403).body("Refresh token is not valid");
        }

        // Generate new tokens
        Integer userId = decoded.getClaim("userId").asInt();
        String newAccessToken = TokenService.generateAccessToken(userId);
        String newRefreshToken = TokenService.generateRefreshToken(userId);

        // Invalidate the old refresh token
        TokenService.invalidateRefreshToken(refreshToken);

        // Set new tokens as cookies (example, adjust as needed)
        // ...

        return ResponseEntity.ok("Tokens refreshed");
    } catch (Exception e) {
        return ResponseEntity.status(403).body("Invalid refresh token");
    }
}

import com.auth0.jwt.JWT;
import com.auth0.jwt.algorithms.Algorithm;
import com.auth0.jwt.interfaces.DecodedJWT;
import jakarta.servlet.http.HttpServletResponse;

@PostMapping("/refresh")
public ResponseEntity<?> refreshToken(HttpServletRequest request, HttpServletResponse response) {
    String refreshToken = null;
    if (request.getCookies() != null) {
        for (Cookie cookie : request.getCookies()) {
            if ("refresh_token".equals(cookie.getName())) {
                refreshToken = cookie.getValue();
                break;
            }
        }
    }
    if (refreshToken == null) {
        return ResponseEntity.status(401).body("Refresh token required");
    }

    // Validate the refresh token
    try {
        DecodedJWT decoded = JWT.require(REFRESH_ALGORITHM)
                .build()
                .verify(refreshToken);

        String type = decoded.getClaim("type").asString();
        if (!"refresh".equals(type)) {
            throw new IllegalArgumentException("Invalid token type");
        }
        if (!TokenService.isRefreshTokenValid(refreshToken)) {
            return ResponseEntity.status(403).body("Refresh token is not valid");
        }

        // Generate new tokens
        Integer userId = decoded.getClaim("userId").asInt();
        String newAccessToken = TokenService.generateAccessToken(userId);
        String newRefreshToken = TokenService.generateRefreshToken(userId);

        // Invalidate the old refresh token
        TokenService.invalidateRefreshToken(refreshToken);

        // Set new tokens as cookies (example, adjust as needed)
        // ...

        return ResponseEntity.ok("Tokens refreshed");
    } catch (Exception e) {
        return ResponseEntity.status(403).body("Invalid refresh token");
    }
}

This code verifies the refresh token's signature, checks its type, and ensures it is still valid (not revoked). If the token is valid, it generates new tokens and invalidates the old refresh token.

Token Revocation

In a secure authentication system, it's important to have a mechanism for revoking refresh tokens. Token revocation allows you to invalidate tokens when necessary, such as when a user logs out, a token is suspected to be compromised, or a user wants to terminate a session.

To enable token revocation, refresh tokens should be stored in a persistent storage solution, such as a database or an in-memory data structure. This allows you to track and manage active tokens effectively.

When a token needs to be revoked, you can remove it from the storage or mark it as invalid. Here's an example of how you might implement token revocation in Java:

Java

import jakarta.servlet.http.Cookie;
import jakarta.servlet.http.HttpServletRequest;

@PostMapping("/logout")
public ResponseEntity<?> revokeRefreshToken(HttpServletRequest request) {
    String refreshToken = null;
    if (request.getCookies() != null) {
        for (Cookie cookie : request.getCookies()) {
            if ("refresh_token".equals(cookie.getName())) {
                refreshToken = cookie.getValue();
                break;
            }
        }
    }
    if (refreshToken == null) {
        return ResponseEntity.status(400).body("Refresh token required");
    }

    if (TokenService.isRefreshTokenValid(refreshToken)) {
        TokenService.invalidateRefreshToken(refreshToken);
        return ResponseEntity.ok("Token revoked successfully");
    } else {
        return ResponseEntity.status(404).body("Token not found");
    }
}

import jakarta.servlet.http.Cookie;
import jakarta.servlet.http.HttpServletRequest;

@PostMapping("/logout")
public ResponseEntity<?> revokeRefreshToken(HttpServletRequest request) {
    String refreshToken = null;
    if (request.getCookies() != null) {
        for (Cookie cookie : request.getCookies()) {
            if ("refresh_token".equals(cookie.getName())) {
                refreshToken = cookie.getValue();
                break;
            }
        }
    }
    if (refreshToken == null) {
        return ResponseEntity.status(400).body("Refresh token required");
    }

    if (TokenService.isRefreshTokenValid(refreshToken)) {
        TokenService.invalidateRefreshToken(refreshToken);
        return ResponseEntity.ok("Token revoked successfully");
    } else {
        return ResponseEntity.status(404).body("Token not found");
    }
}

This method checks for the refresh token in the request cookies, and if found, removes it from the set of valid tokens.

Introduction

Understanding Access and Refresh Tokens

Exploiting Token Vulnerabilities

Token Generation and Invalidation

Join the 1M+ learners on CodeSignal

Why Refresh Token Rotation Matters

Authentication Checks

Token Validation

Integrating the Refresh Token Function

Token Revocation

Integrating Token Revocation

Considerations

Conclusion and Next Steps