Add auto-refresh / timeout of user auth token to prevent stale Token IDs

[dreaminglucid]

1. Extend JWT Payload

First, extend your JWT payload to include an expiration time (exp) and a refresh expiration time (ref_exp).

import datetime

def generate_jwt(userEmail):
    expiration_time = datetime.datetime.utcnow() + datetime.timedelta(days=1)
    refresh_exp_time = datetime.datetime.utcnow() + datetime.timedelta(days=7)
    payload = {
        "email": userEmail,
        "exp": expiration_time,
        "ref_exp": refresh_exp_time
    }
    return jwt.encode(payload, YOUR_SECRET, algorithm='RS256')

2. Middleware for Auto Refresh

Extend your handle_jwt_token function to also handle auto-refreshing of the token.

from flask import make_response

def handle_jwt_token(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            auth_header = request.headers.get("Authorization")
            if not auth_header:
                return jsonify({"error": "Token missing"}), 401

            token = auth_header.split(" ")[1]
            userEmail = extract_user_email_from_token(token)
            decoded_token = decode_and_verify_token(token)

            exp_time = decoded_token.get("exp")
            ref_exp_time = decoded_token.get("ref_exp")
            current_time = datetime.datetime.utcnow()

            # Check if refresh expiry time is also expired
            if current_time > ref_exp_time:
                return jsonify({"error": "Token and refresh expired"}), 401

            # Check if token expired but can be refreshed
            if current_time > exp_time:
                new_token = generate_jwt(userEmail)
                resp = make_response(func(*args, **kwargs, userEmail=userEmail))
                resp.headers['Authorization'] = f"Bearer {new_token.decode('UTF-8')}"
                return resp

            return func(*args, **kwargs, userEmail=userEmail)

        except jwt.ExpiredSignatureError:
            return jsonify({"error": "Token expired"}), 401
        except jwt.InvalidTokenError:
            return jsonify({"error": "Invalid token"}), 401
        except Exception as e:
            log(f"Unhandled exception: {traceback.format_exc()}", type="error")
            return jsonify({"error": "Internal server error"}), 500

    return wrapper

In this enhanced middleware:

• If the token is missing, it returns a 401 error.
• If the token is invalid, it returns a 401 error.
• If the token is expired but within the refresh window (ref_exp), it generates a new token and attaches it to the outgoing response header.
• If both the token and refresh token are expired, it returns a 401 error.

3. Client-Side Logic

On the client-side, when you receive a response, check for the new Authorization header. If present, replace the old token with the new one.

// Example in JavaScript fetch
const response = await fetch(`${API_URL}/api/dreams/export/pdf`, {
  method: 'GET',
  headers: {
    'Authorization': `Bearer ${oldToken}`
  }
});
const newToken = response.headers.get("Authorization");
if (newToken) {
  // Replace oldToken with newToken
  oldToken = newToken.split(" ")[1];
  }

This way, the server automatically refreshes the token when necessary, and the client swaps it transparently, ensuring a seamless user experience without compromising on security.