FastAPIで、脆弱性対策のためにレスポンスヘッダーを追加する必要がありました。

すべてのレスポンスに同じヘッダーを追加したかったのですが、 ~~FastAPIのドキュメントには記述がなく~~ (発見しました) 、当初path operation関数 (例: @app.post("/comment") デコレータ) ごとの設定だけで全体に設定する方法がわかりませんでした。
今回メインで紹介したいのはすべてのレスポンスに追加する方法ですが、ここではせっかくなので両方の方法について横断的に紹介します。

以下の例では、画像などを返却するAPIを想定して、 Cache-Control を設定してみます。
もちろん、 Content-Security-Policy だろうが X-Frame-Options だろうが、同様に設定可能です。

1. あるpath operation関数に対して追加する方法

こちらはFastAPIの公式ドキュメントにある方法です。

以下のように、path operation関数の引数に Response 型のクラスを追加します。

from fastapi import Response

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

これによって、もともとこうだったレスポンスに、

"Cache-Control" が追加されたのがわかります。

2. すべてのリソースに追加する方法

こちらドキュメントにある方法が利用できます。

@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

response.headers のキーを変更すれば、複数のヘッダーを設定することも可能です。

また、FastAPIが依存しているStarletteというライブラリの BaseHTTPMiddleware というクラスを利用して記述することも可能です。

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

この場合は、上記で作成したクラスを、main.pyなどで定義しているFastAPIのインスタンスに対して設定してあげます。

app = FastAPI()

app.add_middleware(MyHeadersMiddleware)

ちなみに、CORSに関しては専用の CORSMiddleware が用意されていて、regexによるホワイトリストなどに対応しているので、上記のような実装をせずにこちらを利用するのがよいでしょう。
これに関しては FastAPIのチュートリアルにも詳しく書かれています。（fastapi.middleware.cors.CORSMiddleware クラスは実は単なる starlette.middleware.cors.CORSMiddleware のエイリアスです。）

エラー処理に注意

Exceptionがraiseされたり、Pydanticのバリデーションによってエラーハンドリングされる場合、2の方法ではエラーケースでも問題なくヘッダーが追加されるのに対し、1の方法だと、設定したヘッダーが出力されません。

もちろん1の方法であっても、FastAPIの公式ドキュメントにあるように、以下のように直接エラーハンドリングをpath operation関数の中で処理してあげれば問題なく出力されます。

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

しかし、すべてのエラーケースに対してこれを処理するとなると処理漏れが発生するリスクがありますので、多くの場合 @app.exception_handler を定義して処理することになるかと思います。この中でpath operation関数ごとの処理を書いていくと複雑化すると思いますので、多くの場合は1の方法よりも2の方法を採用して、レスポンスの中身やエラー内容に応じてヘッダーを書き換えてあげるのが良いのではないかと考えます。

まとめ

以上、FastAPIのレスポンスに任意のHTTPヘッダーを追加する方法を２つ紹介しました。

FastAPIは業務でガッツリ使っていますので、今後もシリーズ化していくつか記事を書きたいと思います。