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

# WebSocket とは

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

# WebSocket API を作成する

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

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

![1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/563db56f-4af7-71c8-cd44-739ca480c3fc.png)

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

![2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/3b4c1566-da57-fa28-6afd-cfbe902157c1.png)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/34c31446-1a93-06e8-3a83-cc509c983be4.png)

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

![3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/a2ebecba-71ea-d943-0636-97f7a1dfe621.png)

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

![4.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/8dc79d39-a337-cd97-f1fa-1d1523893fd5.png)

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

![5.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/c2a2d30c-6503-f868-d6aa-342b278f1fe8.png)

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

![6.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/14200795-a2fa-bb44-958c-32b9724161b9.png)

# Lambda 関数を作成する

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

```json
{
  "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"
}
```

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

```python
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 ロールのインラインポリシーに以下を追加します。

```json
{
    "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 関数でタイムアウトしたためです。

![9.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/9d2b649e-8489-5c92-edb5-7fc381474a4c.png)

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

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 関数でそれ以上待機してタイムアウトするようにします。

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

![10.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/8d0a64c0-c6b6-a7e3-ab89-47f143dfdc6f.png)

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

# ブラウザで動作確認

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

![Qiita-gif用 (6).gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/621414/31148c44-40d2-cfa4-195d-b9132bc3f4e0.gif)

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

https://github.com/quryu-1/websocket-demo/blob/main/App.js

# 参考記事

https://docs.aws.amazon.com/ja_jp/apigateway/latest/developerguide/apigateway-websocket-api-overview.html

https://docs.aws.amazon.com/ja_jp/apigateway/latest/developerguide/websocket-api-chat-app.html

https://docs.aws.amazon.com/ja_jp/apigateway/latest/developerguide/apigateway-how-to-call-websocket-api-wscat.html