How to Protect Your API with OAuth2, JWT, and AWS Cognito

Introduction

Securing your API is essential to ensure data integrity, prevent unauthorized access, and maintain user trust. With a growing number of attacks targeting APIs, implementing robust security mechanisms is no longer optional. In this blog, we’ll explore three powerful tools for API security: OAuth2, JWT (JSON Web Tokens), and AWS Cognito. Each approach offers unique benefits and trade-offs depending on your requirements.

We will build the api using fastapi, and store user credential in AWS Cognito. We will use OAuth2 to authenticate the user and generate a JWT token. The JWT token will be used to authorize the user to access the API.

Java Web Tokens (JWT)

JWT is a format for securely transmitting information (claims) between two parties. It is commonly used as a token for stateless authentication. It consists of three parts:

1. Header: Contains metadata about the token. Typically includes:

• alg: The signing algorithm (e.g., HS256, RS256).
• typ: The token type, which is "JWT".
• Example:

{
  "alg": "HS256",
  "typ": "JWT"
}

2. Payload: Contains claims, which are statements about the token (e.g., user ID, roles, or permissions).
   • Common claims:
     • sub: Subject of the token (user ID or email).
     • exp: Expiration time (in Unix timestamp).
     • iat: Issued at (when the token was created).
   • Example:

{
  "sub": "1234567890",
  "name": "John Doe",
  "admin": true,
  "exp": 1712345678
}

3. Signature: Ensures the integrity of the token by verifying that the header and payload haven’t been tampered with. It is generated using a secret key or a private key to encode the header and the payload

• Example:

HMACSHA256(
  base64UrlEncode(header) + "." + base64UrlEncode(payload),
  secret
)

A JWT is a single string formed by concatenating the three parts, separated by dots (.):

eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWUsImV4cCI6MTcxMjM0NTY3OH0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c

The header and the payload are encoded into a Base64Url format, which is a URL-safe version of Base64 encoding. The signature part is encoded using the algorithm specified in the header.

To verify the token, the server recalculates the signature by encoding the header and payload using the secret key or the public key, and compares it with the signature in the token. If they match, the token is considered valid.

OAuth2

OAuth2 is an authorization protocol that simplifies secure access to resources across different applications. OAuth2 allows you to share your data or use features in one app from another app without giving away your password. It's like a special key that lets other apps do specific things on your behalf.

OAuth2 often uses a token called JWT (JSON Web Token) to prove you're authorized. This token is like a temporary pass that says what you're allowed to do. It's usually short-lived and can be revoked at any time.

FastAPI

FastAPI is a modern web framework for building APIs with Python based on standard Python type hints. It is designed to be fast, easy to use, and highly performant. FastAPI leverages Python's type system to provide automatic data validation, serialization, and documentation. It also supports asynchronous programming, making it ideal for high-performance applications.

AWS Cognito

AWS Congito is a managed service that stores user credentials such as Email and password. It provides a secure and scalable solution for user authentication and authorization. Cognito supports multiple authentication methods, including username and password, social logins, and multi-factor authentication. It also integrates with other AWS services, such as API Gateway and Lambda, to provide a seamless authentication experience for your users.

In this blog, we will use FastAPI to build a simple API that requires authentication using OAuth2 and JWT. We will store user credentials in AWS Cognito and use the AWS SDK to authenticate users and generate JWT tokens. The JWT tokens will be used to authorize users to access the API endpoints.

When using AWS Cognito for authentication, you don't need to manually implement JWT (JSON Web Token) handling, as Cognito handles this for you. Cognito issues JWTs: When a user successfully authenticates, Cognito issues ID, access, and refresh tokens in JWT format.

However we will manually implement JWT handling in this blog to understand how it works. We will use the PyJWT library to generate and verify JWT tokens.

Implementation

Step 1: Install the required libraries

from fastapi import APIRouter, FastAPI, HTTPException, Depends
from pydantic import BaseModel
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm

import jwt
from datetime import datetime, timedelta, timezone
from typing import Optional
import boto3
import os

Step 2: Configure the AWS Cognito client and JWT settings

The following code snippet configures the AWS Cognito client (using boto3) and sets the JWT secret key, algorithm, and expiration time. Replace the SECRET_KEY, COGNITO_CLIENT_ID, and AWS_REGION with your own values. Typically, these values are stored securely in environment variables and loaded into your application using libraries like dotenv to enhance security and maintainability.

# Secret key to encode the JWT
SECRET_KEY = "MY_SECRET_KEY"
ALGORITHM = "HS256"
COGNITO_CLIENT_ID = "set your cognito client id"
AWS_REGION = 'set your aws region'

# AWS Cognito client
cognito_client = boto3.client('cognito-idp', AWS_REGION)

Step 3: Define the OAuth2 scheme and Pydantic models

The OAuth2 scheme is used to authenticate users and generate JWT tokens. The OAuth2PasswordBearer is part of the FastAPI framework's support for OAuth2 authentication. It is a specific type of OAuth2 flow where the client sends a token in the Authorization header of HTTP requests to access protected resources.

By using Depends(oauth2_scheme) in your route functions, FastAPI automatically extracts the token from the Authorization header of incoming requests. It then passes this token to the function, allowing you to perform further validation or processing.

# OAuth2 scheme
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

The Pydantic models define the structure of the user data, token data, and token response. Pydantic is a data validation library that allows you to define data models using Python type hints. FastAPI uses Pydantic models to automatically validate incoming requests and serialize outgoing responses.

# Pydantic models
class User(BaseModel):
    username: str
    email: str
    password: str
    mode: str

class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    username: Optional[str] = None

Step 4: Implement helper functions for JWT token handling

When the user is authenticated, a JWT token is generated using the create_access_token function. This function takes a dictionary of data (e.g., user ID) and an optional expiration time (in minutes) as input. It encodes the data into a JWT token using the secret key and algorithm specified earlier.

The get_current_user function is a dependency that extracts the token from the Authorization header and decodes it to retrieve the user data. If the token is invalid or expired, an HTTP 401 error is raised, indicating that the user is not authenticated.

# Router
router = FastAPI()

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.now(timezone.utc) + expires_delta
    else:
        expire = datetime.now(timezone.utc) + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=401,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = TokenData(username=username)
    except jwt.PyJWTError:
        raise credentials_exception
    return token_data

Step 5: Implement the API routes for user signup, confirmation, signin, signout, and a demo page

The following code snippet defines the API routes for user signup, confirmation, signin, signout, and a demo page. These routes handle user authentication and authorization using AWS Cognito and JWT tokens. The signup route registers a new user in AWS Cognito, while the confirm route verifies the user's confirmation code. The signin route authenticates the user and generates a JWT token, which is used to authorize access to the demopage route. The signout route revokes the user's token, logging them out of the system.

The depends() function is used to inject the get_current_user dependency into the demopage route, ensuring that only authenticated users can access the page.

@router.post("/signup", response_model=Token)
async def signup(user: User):
    try:
        response = cognito_client.sign_up(
            ClientId=COGNITO_CLIENT_ID,
            Username=user.email,
            Password=user.password,
            UserAttributes=[
                {
                    'Name': 'custom:username',
                    'Value': user.username
                },
                {
                    'Name': 'custom:mode',
                    'Value': user.mode
                },
            ],
        )
        access_token = create_access_token(data={"sub": user.email})
        return {"access_token": access_token, "token_type": "bearer"}
    except cognito_client.exceptions.UsernameExistsException:
        raise HTTPException(status_code=400, detail="Username already exists")

@router.post("/confirm")
async def confirm(username: str, confirmation_code: str):
    try:
        cognito_client.confirm_sign_up(
            ClientId=COGNITO_CLIENT_ID,
            Username=username,
            ConfirmationCode=confirmation_code,
        )
        return {"msg": "User confirmed successfully"}
    except cognito_client.exceptions.CodeMismatchException:
        raise HTTPException(status_code=400, detail="Invalid confirmation code")

@router.post("/signin", response_model=Token)
async def signin(form_data: OAuth2PasswordRequestForm = Depends()):
    try:
        response = cognito_client.initiate_auth(
            ClientId=COGNITO_CLIENT_ID,
            AuthFlow='USER_PASSWORD_AUTH',
            AuthParameters={
                'USERNAME': form_data.username,
                'PASSWORD': form_data.password,
            },
        )
        access_token = create_access_token(data={"sub": form_data.username})
        return {"access_token": access_token, "token_type": "bearer"}
    except cognito_client.exceptions.NotAuthorizedException:
        raise HTTPException(status_code=400, detail="Incorrect username or password")

@router.post("/signout")
async def signout(token: str = Depends(oauth2_scheme)):
    try:
        cognito_client.global_sign_out(
            AccessToken=token
        )
        return {"msg": "User signed out successfully"}
    except cognito_client.exceptions.NotAuthorizedException:
        raise HTTPException(status_code=400, detail="Invalid token")

@router.get("/demopage")
async def demopage(current_user: TokenData = Depends(get_current_user)):
    return {"msg": f"Hello, {current_user.username}! Welcome to the demo page."}

Step 6: Run the FastAPI application

Finally, you can run the FastAPI application by creating an instance of the FastAPI class and including the router defined earlier. The application will start a local server that listens for incoming HTTP requests on the specified port.

uvicorn main:router --reload

Open the browser and navigate to http://127.0.0.1:8000/docs#/ to access the Swagger UI, where you can interact with the API endpoints and test the authentication flow.

Before using the API, you need to setup AWS Cognito and configure the client ID. You can follow the steps in the AWS Cognito documentation to create a user pool, configure the app client, and set up the necessary permissions.

Afterwards start with signup and confirm the user, then signin to get the JWT token and use it to access the demo page.

Conclusion

In this blog, we explored how to protect your API using OAuth2, JWT, and AWS Cognito. We discussed the fundamentals of JWT tokens, OAuth2 authorization, and AWS Cognito user authentication. We implemented a simple API using FastAPI, AWS Cognito, and PyJWT to demonstrate the authentication and authorization flow. By combining these tools, you can build secure and scalable APIs that protect user data and prevent unauthorized access.

Part 2: How to Secure Your API with OAuth2, JWT, and AWS Cognito