Building Custom Middleware in FastAPI

Picture middleware as your API’s secret agent, effortlessly intercepting incoming requests before they are processed and outgoing responses before returning them. It’s like having a personal assistant that adds custom functionality to your request-response cycle without disrupting the core framework. While FastAPI provides built-in middleware options, the true power lies in creating custom middleware tailored to your needs.

In this article you’ll see how to build custom middleware, enabling you to extend the functionality of your APIs in unique ways by building function-based and class-based middleware to modify request and response objects to suit your need and to handle rate limiting. You’ll also see how to write test cases and the best practices to follow.

Prerequisites

• Basic understanding of Python and FastAPI
• Web Development Basics: HTTP requests and responses, RESTful APIs, and middleware concepts
• Firefox Browser: Although you could use any browser of your choice, Firefox provides an excellent interface to visualize API response.
• Development Environment: Create a virtual environment with Python. Once the virtual environment is active, proceed to install FastAPI and Pytest.

python -m venv env && source env/bin/activate
pip install "fastapi[all]" pytest

Understanding Middleware in FastAPI

In FastAPI, middleware sits before your API endpoints, handling requests and responses. When a request arrives, it goes through the middleware layer before reaching the API endpoint. Likewise, when a response is ready, it goes through the middleware layer before being sent back to the client.

Middleware simplifies your API by offering a central and reusable way to add extra features. It helps separate concerns, keepingyour API endpoints focused on their core tasks while handling shared operations through middleware components.

Built-in Middleware in FastAPI and Their Functionalities

FastAPI comes with several built-in middleware that provides essential functionalities. These include but are not limited to:

1. CORSMiddleware: Includes necessary CORS headers in outbound responses to enable cross-origin requests from web browsers.
2. TrustedHostMiddleware: Validates the Host header of incoming requests to prevent potential HTTP Host Header attacks.
3. SessionMiddleware: Implements signed cookie-based HTTP sessions where session data is readable but not editable.
4. GZip Middleware: Compresses the response payloads to reduce bandwidth usage, resulting in faster transmission of data.

These built-in middleware components cater to common requirements and simplify the development process. However, custom middleware can extend these functionalities or introduce entirely new ones based on your specific needs.

Advantages of Creating Custom Middleware

Creating custom middleware offers several advantages when developing APIs:

• Flexibility: With custom middleware, you can tailor the behaviour of your API according to your specific use cases.
• Reusability: Supports reuse across multiple endpoints or even in different projects. This helps prevent unnecessary duplication in your code and keeps your code DRY.
• Code Separation: Custom middleware lets you keep your API endpoints clean and focused on their core responsibilities.
• Extensibility: With custom middleware, you can extend FastAPI’s capabilities beyond built-in options.

These benefits contribute to efficient and scalable API development, enabling you to build robust and customized APIs that align perfectly with your project objectives.

Getting Started With FastAPI: Setting Up a Simple API

FastAPI is designed with a gentle learning curve in mind and aims to be developer-friendly for API development, providing robust, production-ready code and automatic interactive documentation.

You’ll create a basic API that you can use to integrate middleware components.

In your project directory, create a file app.py with the following code:

# /app.py
from fastapi import FastAPI

app = FastAPI()

@app.get("/info")
async def hello():
    return {"message": "Hello, World!"}

@app.get("/apiv2/info")
async def hellov2():
    return {"message": "Hello, World from V2"}

The code above creates a FastAPI application with two endpoints. Each endpoint returns a JSONResponse when accessed throughthe corresponding URLs.

Run the server :

$ uvicorn app:app --reload
...
> INFO:     Started server process [8396]
> INFO:     Waiting for application startup.
> INFO:     Application startup complete.

The --reload flag tells the server to auto-reload on changes and is only for development use.

Visit http://127.0.0.1:8000/info on your browser to confirm the application is running.

Enhancing API Functionality with Custom Middleware

You’ll create a function-based middleware to modify the request/response object and a class-based middleware to implement basic rate-limiting for your API endpoints.

Function-based middleware: Request and Response Modification

When creating a function-based middleware, use the decorator @app.middleware("http") placed on top of the function to indicate that the function will act as middleware. This tells FastAPI to register the function as a middleware component.

The middleware function receives two parameters:

• request: Contains all the information and data associated with the incoming request.
• call_next: By invoking the call_next function, the middleware ensures that the request flows seamlessly to the appropriate path operation and awaits the resulting response.

Once the call_next function returns the response generated by the endpoint, the middleware can further modify the response.

Consider a case where you want to modify the incoming request URL before passing it on. Also, you’ll like to add a custom headerto the response if the return type is a StreamingResponse.

Add the following code to the app.py file;

# at the import level
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse

# after the `app` variable
@app.middleware("http")
async def modify_request_response_middleware(request: Request, call_next):
    # Intercept and modify the incoming request
    request.scope["path"] = str(request.url.path).replace("api", "apiv2")
    # Process the modified request
    response = await call_next(request)
    # Transform the outgoing response
    if isinstance(response, StreamingResponse):
        response.headers["X-Custom-Header"] = "Modified"
    return response

The code above implements a middleware function modify_request_response_middleware that modifies the incoming request by replacing “api” with “apiv2” in the URL path and adds a custom header to the response if StreamingResponse.

To see the middleware in action, visit http://127.0.0.1:8000/api/info

Although there is no endpoint with the path requested("/api/info"), the middleware handles this by modifying the URL to the latest version("/apiv2/info").

Also, the response headers contain the x-custom-header added through the middleware.

You’ve seen how to modify the request and response object. Next, you’ll see how to implement middleware using classes.

Class-based middleware: Implementing Rate-Limiting

Rate limiting is a crucial aspect of API performance and security. It controls and limits a client’s requests to an API within a specific period.

You’ll create a class-based middleware to handle rate limiting for your API.

To implement a class-based middleware;

• Ensure your class inherits from Starlett’s BaseHTTPMiddleware class.
• Override the async def dispatch(request, call_next) method to implement your middleware logic.
• To provide configuration options to the middleware class, make sure you override the __init__ method, with app as the initial argument and any other argument as optional keyword arguments.

Add the following code to the app.py file;

# at the import level
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from datetime import datetime, timedelta

# immediately after imports
class RateLimitingMiddleware(BaseHTTPMiddleware):
    # Rate limiting configurations
    RATE_LIMIT_DURATION = timedelta(minutes=1)
    RATE_LIMIT_REQUESTS = 3

    def __init__(self, app):
        super().__init__(app)
        # Dictionary to store request counts for each IP
        self.request_counts = {}

    async def dispatch(self, request, call_next):
        # Get the client's IP address
        client_ip = request.client.host

        # Check if IP is already present in request_counts
        request_count, last_request = self.request_counts.get(client_ip, (0, datetime.min))

        # Calculate the time elapsed since the last request
        elapsed_time = datetime.now() - last_request

        if elapsed_time > self.RATE_LIMIT_DURATION:
            # If the elapsed time is greater than the rate limit duration, reset the count
            request_count = 1
        else:
            if request_count >= self.RATE_LIMIT_REQUESTS:
                # If the request count exceeds the rate limit, return a JSON response with an error message
                return JSONResponse(
                    status_code=429,
                    content={"message": "Rate limit exceeded. Please try again later."}
                )
            request_count += 1

        # Update the request count and last request timestamp for the IP
        self.request_counts[client_ip] = (request_count, datetime.now())

        # Proceed with the request
        response = await call_next(request)
        return response

The code above defines a custom middleware class, RateLimitingMiddleware which implements rate-limiting functionality for the API. It tracks the number of requests made by each client’s IP address within a specified duration and limits the requests based on a predefined threshold. A JSON response with an error message is returned once the client exceeds the rate limit.

Let FastAPI know about the middleware; Set the RateLimitingMiddleware middleware after the modify_request_response_middleware middleware.

...

# after modify_request_response_middleware
app.add_middleware(RateLimitingMiddleware)

Using app.add_middleware, tells FastAPI to register RateLimitingMiddleware to your application middleware stack.

Now FastAPI is aware of RateLimitingMiddleware.

Notice how function-based middleware uses the @app.middleware("http") decorator while class-based middleware uses the add_middleware method provided by the FastAPI class.

Visit http://127.0.0.1:8000/api/info and refresh your browser multiple times.

After the third request, the endpoint becomes unreachable because you’ve exceeded the rate limit(3) and have to wait for sometime before being granted access again.

Ordering and Chaining Middleware

When working with multiple middleware components in FastAPI, you must define the order of execution and understand how there are chained together effectively.

Order of Execution for Multiple Middleware

The order in which middleware components are executed can significantly impact the behaviour of your API.

In FastAPI, middleware functions are executed in the order they are registered, and the registration order is reversed compared to how you define them. This means that the last middleware you define will be the first one to run, and the first middleware you define will be the last one to run in the request/response cycle.

Let’s consider three middleware components, middleware1, SomeClassBasedMiddleware and middleware2.

You define the order as follows:

@app.middleware("http")
async def middleware1(request, call_next):
    # Perform operations before middleware2
    response = await call_next(request)
    # Perform operations after middleware2
    return response

app.add_middleware(SomeClassBasedMiddleware)

@app.middleware("http")
async def middleware2(request, call_next):
    # Perform operations
    response = await call_next(request)
    # Perform operations
    return response

In this case, middleware2 will be executed before SomeClassBasedMiddleware and finally middleware1 will be executed.

This is worth knowing as improper ordering can lead to unexpected results.

Strategy for Chaining Middleware

Chaining middleware involves passing the request and response objects from one middleware component to the next in a sequence.

The common approach for chaining middleware is to use the call_next function. By calling await call_next(request) within a middleware component, you delegate the processing to the next middleware component in the chain. Once the subsequent middleware completes its operations, the control flows back to the previous middleware, allowing it to perform any post-processing tasks.

Testing Middleware

There are cases where an update to your code base causes the middleware to behave unexpectedly, or an update to the framework might result in compatibility issues.

Test cases check and ensure that the middleware gives the desired output in cases of update, ensuring the robustness and reliability of the middleware.

FastAPI’s integration with Starlette makes testing straightforward. It allows you to directly use Pytest for testing FastAPIapplications.

To write test cases for your API;

1. Import TestClient.
2. Initialize a TestClient with the app variable as an argument.
3. Define test functions starting with test_*, following Pytest conventions.
4. Utilize TestClient similar to how you would use the Python requests library.

Let’s write two simple test cases, one to assert the custom headers and JSON response and the other to test rate limiting.

Add the following code to the app.py file;

# import level
import time
from fastapi.testclient import TestClient

# at the end of the file
client = TestClient(app)

def test_modify_request_response_middleware():
    # Send a GET request to the hello endpoint
    response = client.get("/info")
    # Assert the response status code is 200
    assert response.status_code == 200
    # Assert the middleware has been applied
    assert response.headers.get("X-Custom-Header") == "Modified"
    # Assert the response content
    assert response.json() == {"message": "Hello, World!"}
    
def test_rate_limiting_middleware():
    time.sleep(1)
    response = client.get("/info")
    # Assert the response status code is 200
    assert response.status_code == 200

    time.sleep(1)
    response = client.get("/info")
    # Assert the response status code is 200
    assert response.status_code == 200

    time.sleep(1)
    response = client.get("/info")
    # Assert the response status code is 200
    assert response.status_code == 429

Run the test:

pytest app.py

The command tells Pytest to discover test cases in the app.py file.

$ pytest app.py
========================== test session starts ===========================
> platform linux -- Python 3.10.6, pytest-7.3.1, pluggy-1.0.0
> rootdir: /home/prince/semaphore/PrinceInyang/middleware
> plugins: anyio-3.7.0
> collected 2 item

app.py ..                                                            [100%]

========================== 2 passed in 1.78s =============================

The test cases passed because the middleware functions as expected.

Best Practices and Considerations

When working with custom middleware, it is necessary to adhere to best practices in other to achieve optimal reliability, performance, and security.

Performance Considerations for Custom Middleware

When developing custom middleware, it’s important to consider its impact on the overall performance of your FastAPI application.

• Keep Middleware Lightweight: Ensure that your middleware components are designed to be efficient and perform minimal processing. Avoid unnecessary computations or queries that could introduce delays in request processing.
• Optimize Middleware Order: Arrange your middleware components in an order that minimizes redundant operations. If one middleware component performs authentication, ensure that it is placed before any other middleware that also attempts authentication.
• Caching and Memorization: Make sure you leverage caching mechanisms and memorization techniques where appropriate to minimize redundant computations or database access.
• Benchmarking and Profiling: Use tools like cProfile or dedicated profiling libraries to analyze the execution time and memory usage of your middleware components.

Security Best Practices for Middleware Implementation

When creating custom middleware components for bigger applications, always put the following security best practices into consideration:

• Input Validation and Sanitization: Utilize FastAPI’s request validation features and input sanitization techniques to ensure data integrity.
• Authentication and Authorization: Use industry-standard protocols like OAuth 2.0 or JWT (JSON Web Tokens) to ensure secure user authentication.
• Role-Based Access Control (RBAC): Implement RBAC to manage access to different resources based on user roles. This ensures that only authorized users can perform specific actions and access restricted endpoints.

Avoiding Common Pitfalls

Below are some tips on how to avoid common pitfalls when working with custom middleware:

• Thorough Testing: Test your middleware components thoroughly in different scenarios, including edge cases and error conditions. This helps identify and address any unforeseen issues before deploying your API.
• Regular Code Review: Teams should conduct code reviews to ensure that the middleware adheres to best practices, follows coding standards, and meets security requirements.
• Documentation and Versioning: Maintain up-to-date documentation for your middleware, including its purpose, configuration options, and any changes or updates. Clearly communicate versioning information to ensure compatibility with API updates.

By following these best practices and considering potential pitfalls. implementing secure, performant, and reliable custom middleware for your FastAPI application is feasible.

Conclusion

In conclusion, custom middleware offers a powerful way to extend and enhance the functionality of APIs. You can add authentication, handle errors, transform data, and more by intercepting requests and responses.

Throughout this article, we’ve covered the fundamentals of middleware, implementation techniques, best practices, and real-world examples. With this knowledge, you can confidently implement custom middleware in FastAPI and customize your API development.