FastAPIで任意のヘッダーをレスポンスに追加する方法


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

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

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


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

こちらはFastAPIの [公式ドキュメント](https://fastapi.tiangolo.com/advanced/response-headers/) にある方法です。

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

```python
from fastapi import Response

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

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

<img width="400" alt="response_before.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/29632/27c173e4-8156-5497-1f28-702d89f8d879.png">

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

<img width="400" alt="response_after.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/29632/cb1d43c1-5e2b-ded1-70c0-430b61d8f6e0.png">


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

[こちらドキュメントにある方法](https://fastapi.tiangolo.com/tutorial/middleware/#create-a-middleware) が利用できます。

```python
@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`](https://github.com/encode/starlette/blob/master/docs/middleware.md#basehttpmiddleware) というクラスを利用して記述することも可能です。

```python
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のインスタンスに対して設定してあげます。

```python
app = FastAPI()

app.add_middleware(MyHeadersMiddleware)
```

ちなみに、CORSに関しては専用の [`CORSMiddleware`](https://github.com/encode/starlette/blob/master/docs/middleware.md#corsmiddleware) が用意されていて、regexによるホワイトリストなどに対応しているので、上記のような実装をせずにこちらを利用するのがよいでしょう。
これに関しては [FastAPIのチュートリアル](https://fastapi.tiangolo.com/tutorial/cors/) にも詳しく書かれています。（`fastapi.middleware.cors.CORSMiddleware` クラスは実は単なる `starlette.middleware.cors.CORSMiddleware` のエイリアスです。）

# エラー処理に注意

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

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

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

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


# まとめ

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

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