Unified JWT Authentication

Welcome back! In the previous lesson, you learned how to use mock OAuth providers to simulate real-world authentication flows, and you explored how to link and unlink accounts securely. You also saw how important it is to have strong security practices, such as requiring confirmation before linking accounts and preventing duplicate links. These features help keep user accounts safe and flexible.

Now, let's take the next step in building a robust authentication system. In modern web applications, JWTs (JSON Web Tokens) are a popular way to securely manage user authentication and authorization. JWTs allow your server to issue a signed token to the client after a successful login, which the client can then use to access protected resources. This approach is especially important in OAuth systems, where users might sign in using different providers, such as Google or GitHub.

In this lesson, you will learn how to integrate JWT tokens into your OAuth system to manage authentication in a unified and secure way. By the end, you will understand how to generate JWTs for different authentication methods, verify them in middleware, and enforce permissions based on the user's authentication context. This will prepare you for hands-on practice and more advanced OAuth integrations.

Let's start by looking at how JWT tokens are created in a unified authentication system. The generate_unified_token function is responsible for generating a JWT that contains all the information needed to identify the user and their authentication method.

Here is the function:

from jose import jwt
from datetime import datetime, timedelta
import secrets

def generate_unified_token(user: User, auth_method: str) -> str:
    payload = {
        'userId': user.id,
        'authMethod': auth_method,
        'provider': getattr(user, 'provider', None)
    }
    
    # Add cryptographically secure unique token ID for enhanced security
    if auth_method != 'password':
        payload['jti'] = secrets.token_hex(16)
    
    expiry = (TOKEN_EXPIRY['password'] if auth_method == 'password' 
              else TOKEN_EXPIRY['oauth'])
    payload['exp'] = datetime.utcnow() + expiry
    payload['iat'] = datetime.utcnow()
    
    return jwt.encode(payload, JWT_SECRET_KEY, algorithm='HS256')

At the top of the file, we import the jose library's jwt module, which provides the tools we need to create and verify JWTs. The jwt.encode() method takes three arguments: the payload (the data you want to encode), a secret key (used to sign the token so it can be verified later), and the algorithm. This method returns a signed token string that can be sent to the client. The JWT_SECRET_KEY is a secret string stored securely on your server — it should never be exposed to clients. When you call jwt.encode(), it creates a cryptographically signed token that proves the payload hasn't been tampered with.

When you call this function, it creates a payload that includes the user's ID, the authentication method (such as password, google, or github), and the provider. If the authentication method is not password, it also adds a cryptographically secure unique token ID (jti) for extra security. The token's expiry time is set based on the authentication method — password tokens might last for one day, while OAuth tokens could last for a week.

For example, if a user logs in with a password, the payload might look like this:

{
  "userId": 42,
  "authMethod": "password",
  "provider": "local"
}

But if a user logs in with Google OAuth, the payload could look like this:

{
  "userId": 99,
  "authMethod": "google",
  "provider": "google",
  "jti": "a1b2c3d4e5f6789012345678901234ab"
}

This approach makes it easy to handle different authentication methods in a consistent way, while also allowing you to set different security rules for each method. The use of cryptographically secure random bytes for the jti field ensures that each token has a unique identifier that cannot be predicted or reproduced by malicious actors.

Once a JWT token is generated and sent to the client, you need a way to verify it on every request to protected endpoints. This is where the get_auth_context function comes in. Its job is to check the Authorization header for a valid JWT, decode it, and load the user's information.

Here's how the function works:

from fastapi import HTTPException, Request
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

async def get_auth_context(request: Request, db: AsyncSession) -> AuthContext:
    """Unified authentication middleware - extracts and validates JWT tokens"""
    auth_header = request.headers.get('authorization')
    
    if not auth_header or not auth_header.startswith('Bearer '):
        raise HTTPException(
            status_code=401,
            detail='Missing or invalid authorization header'
        )
    
    token = auth_header.split(' ')[1]
    
    try:
        decoded = jwt.decode(token, JWT_SECRET_KEY, algorithms=['HS256'])
        user_id = decoded.get('userId')
        auth_method = decoded.get('authMethod', 'password')
        
        result = await db.execute(select(User).where(User.id == user_id))
        user = result.scalar_one_or_none()
        
        if not user:
            raise HTTPException(status_code=404, detail='User not found')
        
        # Create authentication context
        auth_context = AuthContext(
            user=user,
            auth_method=auth_method,
            is_oauth=auth_method != 'password',
            permissions=determine_permissions(user, auth_method),
            token_info={
                'issued': datetime.fromtimestamp(decoded.get('iat', 0)),
                'expires': datetime.fromtimestamp(decoded.get('exp', 0)),
                'jti': decoded.get('jti')
            }
        )
        
        return auth_context
        
    except jwt.JWTError as e:
        if 'expired' in str(e).lower():
            raise HTTPException(status_code=401, detail='Token expired')
        raise HTTPException(status_code=401, detail='Invalid token')

When a request comes in, the function checks for a Bearer token in the header. If the token is missing or invalid, it raises an HTTPException. If the token is valid, it decodes the payload and fetches the user from the database. It then creates an AuthContext object, which includes the user, the authentication method, whether the user logged in with OAuth, their permissions, and token details like when it was issued and when it expires.

This AuthContext is returned from the function and can be used as a dependency in FastAPI route handlers. This makes it easy to enforce permissions and customize behavior based on how the user authenticated.

Not all users should have the same permissions. For example, an admin might be able to delete any resource, while a regular user can only delete their own. The determine_permissions function decides what permissions a user has based on their role and authentication method.

Here's how it works:

def determine_permissions(user: User, auth_method: str) -> List[str]:
    base_permissions = ['read', 'write']
    
    # OAuth users get additional permissions
    if auth_method != 'password':
        base_permissions.append('oauth_linked')
    
    # Admin users get all permissions
    if user.role == Role.ADMIN:
        base_permissions.extend(['admin', 'delete_any'])
    else:
        base_permissions.append('delete_own')
    
    return base_permissions

If a user logs in with OAuth, they get an extra permission called oauth_linked. If the user is an admin, they get admin and delete_any permissions. Regular users get delete_own.

To enforce these permissions, the require_permission function checks if the current user has the required permission before allowing access to a route. For example, if a route requires admin permission, only users with that permission can proceed.

Here's a simple example:

def require_permission(permission: str):
    """Decorator to require specific permission"""
    async def permission_checker(auth_context: AuthContext) -> AuthContext:
        if permission not in auth_context.permissions:
            print(f'Insufficient permissions: {auth_context.permissions}, required: {permission}')
            raise HTTPException(
                status_code=403,
                detail=f'Insufficient permissions. Required: {permission}'
            )
        return auth_context
    return permission_checker

If a user tries to access a protected route without the right permission, they will see a response like this:

{
  "detail": "Insufficient permissions. Required: admin"
}

In FastAPI, you can use this in a route like this:

@router.get("/protected")
async def protected_route(
    request: Request,
    db: AsyncSession = Depends(get_db)
):
    auth_context = await get_auth_context(request, db)
    await require_permission('admin')(auth_context)
    # Route logic here

This system makes it easy to control access based on both the user's role and how they authenticated.

Now that you have a way to generate and verify JWTs, you can use them to build useful endpoints for your application. The /status endpoint lets clients check the current authentication status, providing detailed information about the user, their permissions, and capabilities.

Here's how the /status endpoint works:

from fastapi import APIRouter, Depends, Request
from sqlalchemy.ext.asyncio import AsyncSession

router = APIRouter()

@router.get("/status")
async def get_status(
    request: Request,
    db: AsyncSession = Depends(get_db)
):
    """Get authentication status and capabilities"""
    auth_context = await get_auth_context(request, db)
    track_auth_usage(auth_context, request)
    
    user = auth_context.user
    
    return {
        "user": {
            "id": user.id,
            "username": user.username,
            "email": user.email,
            "provider": getattr(user, 'provider', None),
            "profileImage": user.profile_image
        },
        "authentication": {
            "method": auth_context.auth_method,
            "isOAuth": auth_context.is_oauth,
            "permissions": auth_context.permissions,
            "tokenInfo": {
                "issued": auth_context.token_info['issued'].isoformat(),
                "expires": auth_context.token_info['expires'].isoformat(),
                "jti": auth_context.token_info.get('jti')
            }
        },
        "capabilities": {
            "canLinkAccounts": not auth_context.is_oauth,
            "canCreateApiKeys": True,
            "canAccessAdmin": 'admin' in auth_context.permissions
        },
        "linkedAccounts": {
            "google": bool(user.google_id),
            "github": bool(user.github_id),
            "local": bool(user.password)
        }
    }

When a client calls this endpoint with a valid token, they get a response showing their user info, authentication method, permissions, and what actions they can perform. The function uses await get_auth_context(request, db) to extract and verify the JWT, which automatically handles authentication errors. For example:

{
  "user": {
    "id": 42,
    "username": "alice",
    "email": "alice@example.com",
    "provider": "local",
    "profileImage": null
  },
  "authentication": {
    "method": "password",
    "isOAuth": false,
    "permissions": ["read", "write", "delete_own"],
    "tokenInfo": {
      "issued": "2024-06-01T12:00:00.000000",
      "expires": "2024-06-02T12:00:00.000000"
    }
  },
  "capabilities": {
    "canLinkAccounts": true,
    "canCreateApiKeys": true,
    "canAccessAdmin": false
  },
  "linkedAccounts": {
    "google": false,
    "github": false,
    "local": true
  }
}

This endpoint is particularly useful for client applications that need to display user information, show or hide features based on permissions, or determine when a token is about to expire.

The /refresh endpoint allows clients to get a new token before the old one expires. This is useful for keeping users logged in without making them sign in again. Here's how it works:

@router.post("/refresh")
async def refresh_token(
    request: Request,
    db: AsyncSession = Depends(get_db)
):
    """Refresh JWT token"""
    auth_context = await get_auth_context(request, db)
    user = auth_context.user
    
    try:
        # Generate new token with same authentication method
        new_token = generate_unified_token(user, auth_context.auth_method)
        
        # Determine correct expiry based on authentication method
        expiry_key = 'password' if auth_context.auth_method == 'password' else 'oauth'
        
        return {
            "message": "Token refreshed successfully",
            "token": new_token,
            "authMethod": auth_context.auth_method,
            "expiresIn": str(TOKEN_EXPIRY[expiry_key])
        }
    except Exception as e:
        print(f'Token refresh error: {e}')
        raise HTTPException(status_code=500, detail='Failed to refresh token')

The endpoint uses await get_auth_context(request, db) to verify the current token and extract the authentication context. It then calls generate_unified_token to create a new token with the same authentication method. The expiry time is determined by checking if the user authenticated via OAuth or password. The try-except block handles any errors that might occur during token generation, ensuring the client receives a clear error message if something goes wrong.

For example, after calling /refresh, the response might be:

{
  "message": "Token refreshed successfully",
  "token": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...",
  "authMethod": "password",
  "expiresIn": "1 day, 0:00:00"
}

This endpoint is very useful in real-world applications where clients need to keep their sessions active. By refreshing tokens before they expire, you can provide a seamless user experience without forcing users to log in repeatedly.

In this lesson, you learned how to integrate JWT tokens into your OAuth authentication system. You saw how to generate unified tokens for different authentication methods, verify them using a dependency function, and attach an authentication context to each request. You also explored how to determine and enforce permissions, and how to use JWTs with endpoints that report authentication status and refresh tokens.

These concepts are essential for building secure and flexible authentication systems in modern web applications. You are now ready to put these ideas into practice in the CodeSignal IDE. In the next exercises, you will work directly with the code to generate tokens, protect routes, and manage permissions. Great job making it this far — these skills will serve you well as you continue to build advanced OAuth integrations!