ALBとCognitoとStreamlit構成で認証後のログアウト処理を実装する

2023.11.14

この記事を書いたメンバー：

大友佑介

こんにちは。大友（@yomon8)です。

結構前ですが、ALBとCognitoなどのIdPを連携させることで認証処理を実現できるようになりました。

以下の手順にあるような方法で簡易な認証処理はすぐに実装できます。Streamlitでもこれを取り入れようとしたのですが、ログアウト部分など少し工夫が必要だったので記事を書きます。

Application Load Balancer を使用してユーザーを認証する

認証処理概要

上記のドキュメントにもあるとおり、Cognitoとの連携はALBが行なってくれるので、Cognito側との検証済みの情報をALBのバックエンド（今回はStreamlit）で受け取れます。

主なものとして以下があります。

X-Amzn-Oidc-Identity	ユーザ特定に利用できるSubの値
X-Amzn-Oidc-Accesstoken	CognitoのアクセストークンUserInfoに投げて検証も可能
X-Amzn-Oidc-Data	ユーザークレーム

例えば以下のようなコードでStreamlitを動かしてCognitoと連携済みのEC2やECS上で動かしてみると、情報が取れます。この情報を使って認可処理に使うことも可能です。

import base64
import json
from typing import Any

import jwt
import requests
import streamlit as st
from streamlit.web.server.websocket_headers import _get_websocket_headers


def _decode_access_token(access_token: str) -> dict[str, Any]:
    decoded = jwt.decode(access_token, options={"verify_signature": False})
    return decoded

def _decode_oidc_data(jwt_data: str) -> dict[str, Any]:
    jwt_headers = jwt_data.split(".")[0]
    decoded_jwt_headers = base64.b64decode(jwt_headers)
    decoded_jwt_headers_str = decoded_jwt_headers.decode("utf-8")
    decoded_json = json.loads(decoded_jwt_headers_str)
    kid = decoded_json["kid"]
    url = "https://public-keys.auth.elb.ap-northeast-1.amazonaws.com/" + kid
    req = requests.get(url)
    pub_key = req.text
    payload = jwt.decode(jwt_data, pub_key, algorithms=["ES256"])
    return payload

def _get_user_info(access_token: str) -> dict:
    # xxxやregionは自身の環境の設定に修正してください
    cognito_domain_prefix = "xxxx"
    aws_region = "ap-northeast-1"
    user_pool_url = f"https://{cognito_domain_prefix}.auth.{aws_region}.amazoncognito.com/oauth2/userInfo"
    res = requests.get(
        user_pool_url,
        headers={"Authorization": f"Bearer {access_token}"},
    )
    return json.loads(res.text)

st.title("ALBから取得した認証情報")

# WebSocketのヘッダーを取得
headers = _get_websocket_headers()
if (
    not headers
    or "X-Amzn-Oidc-Accesstoken" not in headers
    or "X-Amzn-Oidc-Data" not in headers
):
    st.error("認証情報が取得できません")
    st.stop()
else:
    access_token = headers["X-Amzn-Oidc-Accesstoken"]
    oidc_data_jwt = headers["X-Amzn-Oidc-Data"]
    col1, col2, col3 = st.columns(3)
    with col1:
        oidc_data_claim = _decode_oidc_data(jwt_data=oidc_data_jwt)
        st.subheader("X-Amzn-Oidc-Data(Claim)")
        st.json(oidc_data_claim)
    with col2:
        access_token_claim = _decode_access_token(access_token=access_token)
        st.subheader("X-Amzn-Oidc-Accesstoken(Claim)")
        st.json(access_token_claim)
    with col3:
        access_token = headers["X-Amzn-Oidc-Accesstoken"]
        access_token_user_info = _get_user_info(access_token=access_token)
        st.subheader("X-Amzn-Oidc-Accesstoken(UserInfo)")
        st.json(access_token_user_info)

ログアウト処理

認証まではこの通りなのですが、Streamlitでのログアウト処理は少し工夫が必要でした。

AWSのドキュメントにもログアウト処理については書いてあるのですが、大きくは以下の2点の対応が必要です。

ALBの発行したセッションCookieを削除する
Cognitoのログアウトエンドポイントを叩く

両方とも少し工夫が必要ではあるのですが、特にCookieに関してはデフォルトでは以下のようになるのですが、見て分かる通り、HttpOnly属性が付いています。Streamlitからこれを扱うのが、かなり手間がかかります。

結果としてAWSのサービスの力を借りることにしました。

ここではその工夫について書いていきます。

Lambdaの実装

HttpOnlyのCookieを削除しつつ、ログアウトエンドポイントにリダイレクトLambdaを実装します。

import os
import urllib.parse

AWS_REGION = os.environ["AWS_REGION"]
APP_LOGOUT_URL = os.environ["APP_LOGOUT_URL"]
COGNITO_USER_POOL_DOMAIN_PREFIX = os.environ["COGNITO_USER_POOL_DOMAIN_PREFIX"]
COGNITO_CLIENT_ID = os.environ["COGNITO_CLIENT_ID"]

def lambda_handler(event, context):
    # ログアウト処理
    cognito_url = (
        f"https://{COGNITO_USER_POOL_DOMAIN_PREFIX}.auth.{AWS_REGION}.amazoncognito.com"
    )
    redirect_url = (
        f"{cognito_url}/logout?client_id="
        + COGNITO_CLIENT_ID
        + "&logout_uri="
        + urllib.parse.quote(APP_LOGOUT_URL, encoding="utf-8")
    )
    return {
        "statusCode": 302,
        "headers": {
            "Set-Cookie": 'AWSELBAuthSessionCookie-0=""; max-age=-1;HttpOnly;secure;path=/',
            "Set-CookiE": 'AWSELBAuthSessionCookie-1=""; max-age=-1;HttpOnly;secure;path=/',
            "Access-Control-Allow-Origin": APP_LOGOUT_URL,
            "Access-Control-Allow-Methods": "GET",
            "Location": redirect_url,
        },
    }

Lambdaの環境変数についても設定しておいてください。

COGNITO_USER_POOL_DOMAIN_PREFIX はCognitoで言うと以下の部分です。

COGNITO_CLIENT_IDとAPP_LOGOUT_URLに当たる部分はそれぞれ、アプリケーションクライアントの「クライアント ID」と「許可されているサインアウト URL」に設定したURLを指定してください。

ALBから特定のPathでLambdaを実行するように設定

以下の例の場合は /cognito/logoutのパスではLambdaを実行するようにしています。

Streamlitにログアウトボタンの追加

上記のStreamlitのコードに、以下のようにログアウトボタンを追加してみます。

logout_link = "/cognito/logout"
html_button_logout = f"""
    <a href='{logout_link}' class='button-logout' target='_self'>Log Out</a>
"""
st.sidebar.markdown(f"{html_button_logout}", unsafe_allow_html=True)