[PYTHON] How to add arbitrary headers to response with FastAPI

In FastAPI, it was necessary to add a response header as a countermeasure against vulnerabilities.

I wanted to add the same header to every response, but the ~~ FastAPI docs didn't mention it ~~ (found), and the initial path operation function (eg @ app.post ("/ comment") I didn't know how to set it as a whole with only the settings for each) decorator). The main thing I would like to introduce this time is how to add it to all responses, but since it is a big deal here, I will introduce both methods in a cross-cutting manner.

In the following example, let's set Cache-Control assuming an API that returns images and so on. Of course, whether it's Content-Security-Policy or X-Frame-Options, you can set it as well.

1. How to add to a path operation function

This is the method in the Fast API Official Documentation.

Add a class of type Response to the argument of the path operation function as shown below.

from fastapi import Response

@app.get("/image")
def image(response: Response)
    response.headers["Cache-Control"] = "no-cache, no-store"
    ...
    return image_instance

As a result, the response that was originally like this

You can see that "Cache-Control" has been added.

2. How to add to all resources

Method in this document is available.

@app.middleware("http")
async def add_my_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["Cache-Control"] = "no-cache, no-store"
    return response

It is possible to set multiple headers by changing the key of response.headers.

It can also be written using a class called BaseHTTPMiddleware in a library called Starlette on which FastAPI depends.

from starlette.middleware.base import BaseHTTPMiddleware

class MyHeadersMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        response = await call_next(request)
        response.headers["Cache-Control"] = "no-cache, no-store"
        return response

In this case, set the class created above for the FastAPI instance defined in main.py etc.

app = FastAPI()

app.add_middleware(MyHeadersMiddleware)

By the way, for CORS, a dedicated CORSMiddleware is prepared, and it supports whitelisting by regex, so it is better to use this without implementing as above. Let's do it. You can read more about this in the Fast API Tutorial (https://fastapi.tiangolo.com/tutorial/cors/). (The fastapi.middleware.cors.CORSMiddleware class is actually just an alias for starlette.middleware.cors.CORSMiddleware.)

Be careful of error handling

If an Exception is raised or error handling is performed by Pydantic validation, the header will be added without any problem even in the error case with method 2, but the set header will not be output with method 1.

Of course, even with method 1, if you handle the error handling directly in the path operation function as shown in the official document of FastAPI, it will be output without any problem.

    if error:
        return JSONResponse(
            {}, status_code=404, headers={"Cache-Control": "no-cache, no-store"}
        )

However, if you handle this for all error cases, there is a risk of processing omissions, so in most cases you will probably have to define @ app.exception_handler to handle it. I think that it will be complicated if you write the processing for each path operation function in this, so in many cases, adopt method 2 rather than method 1 and rewrite the header according to the content of the response and the content of the error. I think it would be better to give it to you.

Summary

So far, we have introduced two ways to add arbitrary HTTP headers to the Fast API response.

I use FastAPI for business, so I would like to write some articles in a series in the future.