Implementing Token Refresh Mechanism

Implementing Token Refresh Mechanism | CodeSignal Learn

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.

Frontend Implementation

Once we understand the role of access and refresh tokens, let's see how the token refresh mechanism is applied on the frontend. This will also give us a clear picture of what to expect from the backend implementation. On the frontend, we need to implement a function to refresh the access token and a wrapper to manage token expiration. This ensures a seamless user experience without frequent re-authentication.

Wrapper Function to Handle Token Expiration

If previously, fetching some url response required one-time query, now it becomes a two-step operation. Here we provide the wrapper function that will handle token expiration. This function will intercept requests and check if the access token has expired. If it has, the function will attempt to refresh the token: JavaScript // Wrapper function to handle token expiration const fetchWithToken = async (url, options) => {let response = await fetch(url, { ...options, credentials: 'include' }); if (response.status === 401) { // If unauthorized, try refreshing the token const success = await refreshAccessToken(); if (success) {response = await fetch(url, { ...options, credentials: 'include' });}} return response;}; This function makes an initial request and checks if the response status is 401 Unauthorized, indicating that the access token may have expired. If so, it calls the refreshAccessToken function to attempt a token refresh. If the refresh is successful, it retries the original request.

Join the 1M+ learners on CodeSignal

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

Function to Refresh the Access Token

Next, we implement the function responsible for refreshing the access token. This function will be called by the wrapper function when a token refresh is needed: JavaScript // Function to refresh the access token const refreshAccessToken = async () => {try {const response = await fetch('/api/auth/refresh', {method: 'POST', headers: { 'Content-Type': 'application/json' }, credentials: 'include' // Ensure cookies are sent with the request}); if (!response.ok) {throw new Error('Failed to refresh token');} // No need to handle accessToken directly if it's set via cookies return true;} catch (error) {console.error('Error refreshing token:', error); window.location.href = '/login'; // Redirect to login if refresh fails}}; // Function to refresh the access token const refreshAccessToken = async () => {try {const response = await fetch('/api/auth/refresh', {method: 'POST', headers: { 'Content-Type': 'application/json' }, credentials: 'include' // Ensure cookies are sent with the request}); if (!response.ok) {throw new Error('Failed to refresh token');} // No need to handle accessToken directly if it's set via cookies return true;} catch (error) {console.error('Error refreshing token:', error); window.location.href = '/login'; // Redirect to login if refresh fails}}; This function sends a request to the backend to obtain a new access token. If the refresh is successful, it returns true, allowing the wrapper function to retry the original request. If the refresh fails, it redirects to the login route. These snippets illustrate how to handle token expiration and refresh the access token on the frontend, setting the stage for the backend implementation that will support this functionality.

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://localhost:3000/api/auth/refresh \ -H "Content-Type: application/json" \ -H "Cookie: refresh_token=stolen_refresh_token"# Simulating an attack by using a stolen refresh token curl -X POST http://localhost:3000/api/auth/refresh \ -H "Content-Type: application/json" \ -H "Cookie: 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: Pythonimport jwt from datetime import datetime, timedelta import uuid from ..config import JWT_SECRET_KEY, JWT_REFRESH_SECRET # Set to store refresh tokens refresh_tokens = set() def generate_tokens(user_id: int): access_token = jwt.encode( { 'userId': user_id, 'type': 'access', 'exp': datetime.utcnow() + timedelta(minutes=15) }, JWT_SECRET_KEY, algorithm='HS256' ) refresh_token = jwt.encode( { 'userId': user_id, 'type': 'refresh', 'exp': datetime.utcnow() + timedelta(days=7) }, JWT_REFRESH_SECRET, algorithm='HS256' ) # Store refresh token securely refresh_tokens.add(refresh_token) return {"access_token": access_token, "refresh_token": refresh_token}import jwt from datetime import datetime, timedelta import uuid from ..config import JWT_SECRET_KEY, JWT_REFRESH_SECRET # Set to store refresh tokens refresh_tokens = set() def generate_tokens(user_id: int): access_token = jwt.encode( { 'userId': user_id, 'type': 'access', 'exp': datetime.utcnow() + timedelta(minutes=15) }, JWT_SECRET_KEY, algorithm='HS256' ) refresh_token = jwt.encode( { 'userId': user_id, 'type': 'refresh', 'exp': datetime.utcnow() + timedelta(days=7) }, JWT_REFRESH_SECRET, algorithm='HS256' ) # Store refresh token securely refresh_tokens.add(refresh_token) return {"access_token": access_token, "refresh_token": refresh_token} In this code, we generate new tokens and store the refresh token securely. This ensures that each token can be invalidated if necessary.

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: Python from fastapi import Request, HTTPException async def handle_refresh_token(request: Request): refresh_token_cookie = request.cookies.get('refresh_token') if not refresh_token_cookie: raise HTTPException(status_code=401, detail="Refresh token required") # ... rest of the code from fastapi import Request, HTTPException async def handle_refresh_token(request: Request): refresh_token_cookie = request.cookies.get('refresh_token') if not refresh_token_cookie: raise HTTPException(status_code=401, detail="Refresh token required") # ... rest of the code 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 and has not been tampered with: Python from jose import jwt, JWTError try: payload = jwt.decode(refresh_token_cookie, JWT_REFRESH_SECRET, algorithms=['HS256']) if payload.get('type') != 'refresh': raise ValueError("Invalid token type") # ... rest of the code except JWTError: raise HTTPException(status_code=403, detail="Invalid refresh token") from jose import jwt, JWTError try: payload = jwt.decode(refresh_token_cookie, JWT_REFRESH_SECRET, algorithms=['HS256']) if payload.get('type') != 'refresh': raise ValueError("Invalid token type") # ... rest of the code except JWTError: raise HTTPException(status_code=403, detail="Invalid refresh token") Here, we use the python-jose library to verify the refresh token against our secret key. We also check that the token type is refresh to ensure it's being used correctly. If the token is invalid, we respond with an error.

Integrating the Refresh Token Function

To integrate the refresh token function into your FastAPI application, you can set up a router: Python from fastapi import APIRouter from .auth import handle_refresh_token router = APIRouter() @router.post('/auth/refresh') async def refresh(request: Request): return await handle_refresh_token(request) from fastapi import APIRouter from .auth import handle_refresh_token router = APIRouter() @router.post('/auth/refresh') async def refresh(request: Request): return await handle_refresh_token(request) And then, in your main application file, include this router: Python from fastapi import FastAPI from .routes import auth_router app = FastAPI() app.include_router(auth_router, prefix="/api") # ... other routes and middleware from fastapi import FastAPI from .routes import auth_router app = FastAPI() app.include_router(auth_router, prefix="/api") # ... other routes and middleware In this setup, the refresh token functionality is associated with the /api/auth/refresh endpoint using a FastAPI router. This approach helps organize routes and middleware more modularly.

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. 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: Pythonfrom fastapi import Request, HTTPException refresh_tokens = set() # Set to store refresh tokens async def revoke_refresh_token(request: Request): data = await request.json() refresh_token = data.get('refreshToken') if not refresh_token: raise HTTPException(status_code=400, detail="Refresh token required") # Remove the token from the set if refresh_token in refresh_tokens: refresh_tokens.discard(refresh_token) return {"message": "Token revoked successfully"} else: raise HTTPException(status_code=404, detail="Token not found")from fastapi import Request, HTTPException refresh_tokens = set() # Set to store refresh tokens async def revoke_refresh_token(request: Request): data = await request.json() refresh_token = data.get('refreshToken') if not refresh_token: raise HTTPException(status_code=400, detail="Refresh token required") # Remove the token from the set if refresh_token in refresh_tokens: refresh_tokens.discard(refresh_token) return {"message": "Token revoked successfully"} else: raise HTTPException(status_code=404, detail="Token not found")

Integrating Token Revocation

To integrate token revocation into your application, you can set up an endpoint that allows users to revoke their tokens: Python from fastapi import APIRouter router = APIRouter() @router.post('/auth/logout') async def logout(request: Request): return await revoke_refresh_token(request) from fastapi import APIRouter router = APIRouter() @router.post('/auth/logout') async def logout(request: Request): return await revoke_refresh_token(request)

Introduction

Understanding Access and Refresh Tokens

Frontend Implementation

Wrapper Function to Handle Token Expiration

Join the 1M+ learners on CodeSignal

Function to Refresh the Access Token

Exploiting Token Vulnerabilities

Token Generation and Invalidation

Authentication Checks

Token Validation

Integrating the Refresh Token Function

Token Revocation

Integrating Token Revocation

Considerations

Conclusion and Next Steps