API Gateway + Web Socket で Auth0 の認証をするカスタムオーソライザーをつくってみた

西田@大阪です。

API Gateway + WebSocket で Auth0 を使った認証を行うカスタムオーソライザーを作成してみました。今回は Auth0 の Vue.js のサンプルを参考に実装をおこなっていきます。全体のソースコードは github に Pushしています

サーバー側を作成

Serverless Framework(以下 sls) のプロジェクトを作成し必要なライブラリをインストールします

$ sls create --template aws-python3 --name wsauth0 --path wsauth0
$ pipenv install --python 3.8
$ pipenv install requests
$ pipenv install pyjwt
$ pipenv install cryptography
$ sls plugin install -n serverless-python-requirements

cryptographyにCの拡張が入ってるため serverless.yml に以下を追加する必要があります

custom:
  pythonRequirements:
    dockerizePip: true

カスタムオーソライザーの作成

カスタムオーソライザはAuthorizationヘッダに検証するトークンを設定するのが一般的ですが、ブラウザ上で JavaSctipt のWebSocket オブジェクトの仕様の制限上、任意のHTTPヘッダを追加することができません。検証するトークンをカスタムオーソライザーに渡す方法としては、クエリパラメーターかWebSocketオブジェクトのprotocolにオプションにトークンを設定し Sec-WebSocket-Protocol HTTPヘッダで渡す方法があります。今回はクエリパラメーターで渡す方法でトークンを渡します

// Sec-WebSocket-Protocol で渡す
WebSocket("wss://...", "xxx")

// クエリパラメーターで渡す
WebSocket("wss://...?token=xxx")

クライアントからの接続時にルーティングされる$connectハンドラーにカスタムオーソライザーを設定します。その際にidentitySourceを設定し、前述の Sec-WebSocket-Protocolからトークンを取得するように設定します

functions:
  # カスタムオーソライザー
  auth:
    handler: handler.auth_handler

  connectHandler:
    handler: handler.connect_handler
    events:
      - websocket:
          route: $connect
          authorizer:
            name: auth
            identitySource:
              - 'route.request.querystring.token'

次にカスタムオーソライザーを作成していきます

トークンの検証の際に使用する公開鍵を取得します。Auth0 の Settings > Advanced Settings > Endpoints で確認できる JSON Web Key Set のURLからjwkを取得しpem形式の公開鍵を作成しています ※ Auth0 の設定は後ほど行いますので仮の値等を設定してすすめる必要があります

import json
import os

import requests
from cryptography.hazmat.primitives import serialization
from jwt.algorithms import RSAAlgorithm

AUTH0_JWKS_URL = os.getenv('AUTH0_JWKS_URL')

jwks_json_str = json.dumps(json.loads(requests.get(AUTH0_JWKS_URL).text)['keys'][0])
public_key = RSAAlgorithm.from_jwk(jwks_json_str)
pem = public_key.public_bytes(encoding=serialization.Encoding.PEM,
                              format=serialization.PublicFormat.SubjectPublicKeyInfo)

カスタムオーソライザーのハンドラーです。token クエリパラメーターからクライアントから渡されたトークンを公開鍵をつかって検証し、検証に成功したらポリシーを生成しています

def auth_handler(event, context):
    auth_token = event['queryStringParameters']['token']

    try:
        principal_id = jwt_verify(auth_token, pem)
        policy = generate_policy(principal_id, 'Allow', event['methodArn'])
        return policy
    except Exception as e:
        raise Exception('Unauthorized')

jwtを検証するための関数です。audienceには AUTH0に APIs に設定された Audience に値を設定します ※ Auth0 の設定は後ほど行いますので仮の値等を設定してすすめる必要があります

import jwt

AUTH0_AUDIENCE = os.getenv('AUTH0_AUDIENCE')

def jwt_verify(auth_token, pub_key):
    payload = jwt.decode(auth_token, pub_key, algorithms=['RS256'], audience=AUTH0_AUDIENCE)
    return payload['sub']

※ Auth0 の設定自体は後ほど行います

必要なポリシーを生成して返す関数です

def generate_policy(principal_id, effect, resource):
    return {
        'principalId': principal_id,
        'policyDocument': {
            'Version': '2012-10-17',
            'Statement': [
                {
                    "Action": "execute-api:Invoke",
                    "Effect": effect,
                    "Resource": resource
                }
            ]
        }
    }

WebSocketのイベントハンドラーの設定

設定の内容についてはこちらを参照ください

API Gateway + WebSocket の基本的な使い方を調べてみた

provider:
  # ...
  websocketsApiRouteSelectionExpression: $request.body.action
  # ...
functions:
  # カスタムオーソライザー
  auth:
    handler: handler.auth.auth_handler

  connectHandler:
    handler: handler.handler.connect_handler
    # ...

  disconnectHandler:
    handler: handler.handler.disconnect_handler
    events:
      - websocket: $disconnect

  sendMessageHandler:
    handler: handler.handler.send_message_handler
    events:
      - websocket: sendMessage

WebSocketのイベントハンドラーを作成

• クライアントの接続時に dynamodb に connection_id を登録します
• クライアントの切断時に dynamodb から connection_id を削除します
• クライアントからメッセージが送信されてきたら、それを dynamodb に登録済みのすべての connection_id に送信します

クライアント接続時のハンドラーです

def connect_handler(event, context):
    connection_id = event["requestContext"]["connectionId"]

    join_member(connection_id)

    return {
        "statusCode": 200
    }

def join_member(connection_id):
    connections_table.put_item(
        Item={
            'connection_id': connection_id,
        }
    )

クライアント切断時のハンドラーです

def disconnect_handler(event, context):
    connection_id = event["requestContext"]["connectionId"]

    leave_member(connection_id)

    return {
        "statusCode": 200
    }

def leave_member(connection_id):
    connections_table.delete_item(
        Key={
            'connection_id': connection_id,
        }
    )

メッセージ送信時のハンドラーです。

現在登録されているコネクションの一覧をDynamoDBより取得し、クライアントより送信されてきたデータをすべてのコネクションに送信しています。すでに接続が切れているクライアントに送信してしまった際にエラーになるためハンドルするコードが入っています

def send_message_handler(event, context):
    members = get_members()
    apigw = get_apigw_management_client(event)

    data = json.loads(event['body'])['data']

    for member in members:
        try:
            apigw.post_to_connection(
                ConnectionId=member['connection_id'],
                Data=json.dumps({
                    "message": data['message']
                })
            )
        except Exception as e:
            print(e)

    return {
        "statusCode": 200
    }

接続済みのすべてのコネクションIDを返す関数です

def get_members():
    return connections_table.scan()['Items']

クライアントにメッセージを送信するための、API Gateway Manager のクライアントを返す関数です。リクエストされてきた内容から endpoint を設定してます

def get_apigw_management_client(event):
    domain = event["requestContext"]["domainName"]
    stage = event["requestContext"]["stage"]

    return boto3.client('apigatewaymanagementapi', endpoint_url=f'https://{domain}/{stage}')

sls をデプロイします

$ sls deploy

デプロイ時にコンソールに出力されるエンドポイントを保存します

Auth0 のAPI を作成します

APIs Overview

こちらを参考に Auth0 のコンソール上でAPIを作成します

Identifier に先程保存したWebSocketのエンドポイントを入力します

Vue.js のサンプルをダウンロード

Auth0 のコンパネから 「CREATE APPLICATION」をクリックし、 「Single Page Web Application」を作成してください

作成したアプリケーションの「Quick Start」から Vue.js を選択肢サンプルをダウンロードしてください

ダウンロードする際に表示されるダイアログに従って、Callback URL、Allowed Web Origins、Allowed Logout URLs を設定し、自分の Vue.js を動作させる環境に合わせた設定を行ってください

Vue.jsのサンプルからWebSocketに接続する

Auth0 の設定

auth_config.jsonに Auth0 で作成したアプリケーションのauth0のドメインとクライアントIDを設定します

{
  "domain": "xxx.auth0.com",
  "clientId": "xxxxxx"
}

src/auth/authWrapper.js で、Auth0Client生成時のオプションにAPI作成時のIdentifierをaudience として追加します

async created() {
  this.auth0Client = await createAuth0Client({
    domain: options.domain,
    client_id: options.clientId,
    audience: '{YOUR API IDENTIFIER}',
    redirect_uri: redirectUri
  });
// ...

Chatコンポーネントを作成します

components/Chat.vueを作成します

テンプレート部分です。WebSocketとのコネクションが確立されていなければ、「接続」というキャプションのボタンを表示し、確立されていればテキスト入力欄とサーバーからのメッセージを表示するエリアを表示します

<template>
  <div>
    <button v-if="!isConnected" @click="connect">接続</button>
    <div v-if="isConnected">
      <ul style="list-style: none" v-for="message in messages">
        <li>{{message}}</li>
      </ul>
      <input placeholder="Enter your message" style="width: 100%" v-model="inputMessage"/>
      <button @click="sendMessage">送信</button>
    </div>
  </div>
</template>

スクリプト部分です。WebSocketオブジェクト生成時にwssのエンドポイントにtokenパラメーターでauth0の情報にアクセスできる$authプラグインから取得できるトークンを付与して接続しています。sendMessageメソッドでは、サーバーに送信するJSONのactionメンバーにルーティング用の文字列としてsendMessageを設定しています

export default {
  name: "Chat",
  data() {
    return {
      "inputMessage": "",
      "ws": null,
      "messages": [],
    }
  },
  computed: {
    isConnected() {
      return this.ws !== null
    }
  },
  methods: {
    async connect() {
      const token = await this.$auth.getTokenSilently()
      const ws = new WebSocket(`${YOURWSSENDPOINT}?token=${token}`)

      ws.onopen = () => {
        this.ws = ws
      }

      ws.onmessage = message => {
        const data = JSON.parse(message.data)
        this.messages.push(data.message)
      }
    },
    async sendMessage() {
      const data = {
        "action": "sendMessage",
        "data": {
          "message": this.inputMessage
        }
      }

      this.ws.send(JSON.stringify(data))

      this.inputMessage = ""
    }
  }
}

さいごに

いかがでしたでしょうか？途中説明を省力している部分もあるので、github を参考にしていただければともいます。この記事が誰かの参考になれば幸いです

参考

Python PyJWT で Google OAuth 2.0 API の ID Token を検証 — ykrods note

Lambda REQUEST オーソライザーの関数の作成 - Amazon API Gateway