【AWS API Gateway】WebSocket を使用した双方向通信

WebSocket とは

Webサーバとクライアントの間で双方向通信できるようにする技術です。通常はクライアントがリクエストをWebサーバに送り、それに応じてWebサーバがレスポンスを返しますが、WebSocketではサーバ側が任意のタイミングでクライアントにメッセージを送ることができます。

WebSocket API を作成する

バックエンドを Lambda 関数とする WebSocket API を作成します。先に Lambda 関数を作成する をやってもやらなくてもどちらでも構いません。

「WebSocket API」の「構築」をクリックします。

API 名、ルート選択式を入力します。ルート選択式とはルートキーを指定するための JSON パスです。ルートキーとは、ルートの名前です。ルートとは、API を呼び出したときの実行するアクションです。

カスタムルートを作成します。ルートキーには「sendmessage」を入力します。

Lambda 関数を指定します。指定する前に Lambda 関数を作成する必要があるので、作成していない場合は一旦適当な Lambda 関数を作成し、「前へ」→「次へ」を繰り返すと作成した Lambda 関数を選択することができます。

デプロイするステージ名を入力します。例えば開発環境であれば、development、テスト環境であれば test を入力します。

設定確認では「作成してデプロイ」をクリックします。

Lambda 関数を作成する

Lambda 関数は以下のような event を受け取ります。ここからクライアントメッセージを取り出し、加工してクライアントに返すような Lambda 関数を作成します。

{
  "requestContext": {
    "routeKey": "sendmessage",
    "messageId": "vvvvvvvvvvv=",
    "eventType": "MESSAGE",
    "extendedRequestId": "wwwwwwwwww=",
    "requestTime": "23/Jul/2022:07:09:00 +0000",
    "messageDirection": "IN",
    "stage": "production",
    "connectedAt": 1658559861848,
    "requestTimeEpoch": 1658560140781,
    "identity": {
      "sourceIp": "123.123.123.123"
    },
    "requestId": "xxxxxxxxx=",
    "domainName": "API_GATEWAY_ID.execute-api.ap-northeast-1.amazonaws.com",
    "connectionId": "yyyyyyyyyy=",
    "apiId": "API_GATEWAY_ID"
  },
  "body": "{'action': 'sendmessage', 'message': 'hello, everyone!'}",
  "isBase64Encoded": "False"
}

以下のコードを使用します。

import os
import json
import time
import boto3

def custom_function(msg):
    """
    メッセージを加工する関数
    Args:
        msg (str): 加工前メッセージ
    Returns:
        - (str): 加工後メッセージ
    """
    return 'Hello, ' + msg

def lambda_handler(event, context):
    # クライアントからのメッセージを取得する
    msg = json.loads(event['body'])['message']
    # デプロイステージ名を取得する
    stage = event['requestContext']['stage']
    # WebSocket URL のドメイン名を取得する
    domain = event['requestContext']['domainName']

    # ApiGatewayManagementApi クラスを初期化。endpoint_url を指定する。
    client = boto3.client(
        'apigatewaymanagementapi',
        endpoint_url=f'https://{domain}/{stage}'
    )

    # メッセージをカスタマイズする
    custom_msg = custom_function(msg)

    # クライアントにメッセージを送信する
    client.post_to_connection(
        Data=custom_msg,
        ConnectionId=event['requestContext']['connectionId']
    )

    return {
        'statusCode': 200,
        'body': msg
    }

Lambda 関数の IAM ロールに WebSocket API を呼び出す権限を付与します。IAM ロールのインラインポリシーに以下を追加します。

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": "execute-api:ManageConnections",
            "Resource": "arn:aws:execute-api:ap-northeast-1:ACCOUNT_ID:*/*/POST/@connections/*",
            "Effect": "Allow"
        }
    ]
}

Lambda 関数のタイムアウトは 1 分に変更しました。これはデフォルトの 3 秒では Lambda 関数でタイムアウトしたためです。

コマンドラインで動作確認

Windows の Powershell で動作確認をしました。wscat コマンドをインストールします。

npm install -g wscat

以下のように接続し、メッセージを送信します。

PS C:\> wscat -c wss://API_GATEWAY_ID.execute-api.ap-northeast-1.amazonaws.com/production
Connected (press CTRL+C to quit)
> {"action": "sendmessage", "message": "John"}
< Hello, John

次に API Gateway でタイムアウトしてもメッセージが返ってくることが確認します。API Gateway のコンソール画面でタイムアウトを 5 秒にし、Lambda 関数でそれ以上待機してタイムアウトするようにします。

def custom_function(msg):
    time.sleep(10)
    return 'Hello, ' + msg

コマンドラインでは API Gateway でタイムアウトしたことはわからないようです。次の「ブラウザで動作確認」の場合には API Gateway でタイムアウトしたことがわかります。

ブラウザで動作確認

API Gateway でタイムアウトした後でもメッセージが届いていることが確認できます。

ブラウザでの動作確認用ソースコードはこちらに置きました。

参考記事