Tech blog Produced by FOURIER

AWS CDKでWebSocketを使ったサーバーレスチャットアプリを作る

Hirayama Hirayama 2022.09.08

はじめに

AWSにはAWS Cloud Development Kit(AWS CDK)というAWSのリソースをコードで定義することができる開発フレームワークがあります。

そして、AWS CDKを使ってAPI Gateway V2を定義することもでき、これを利用することでWebSocket APIを定義することができます。

ただ、API Gateway V2の機能はまだα版ということもあり、公式ドキュメントと実際の実装が異なっていたり、ネット上でもあまり情報が出てこないため、初めて使用する人にとっては難しい面があります。

そこで、本記事では、AWSの公式チュートリアルであるWebSocket API、Lambda、DynamoDB を使用したサーバーレスチャットアプリケーションの構築と同じ構成のサーバーレスアプリケーションをAWS CDKで構築することで、基本的なAPI Gateway V2の定義方法を解説したいと思います。

https://docs.aws.amazon.com/ja_jp/apigateway/latest/developerguide/websocket-api-chat-app.html 本記事で構築するアプリケーション(チュートリアルページより引用)

本記事で作成するコードは以下のリポジトリで公開しています。

前提条件

本記事作成時点でのAWS CDKのバージョンは2.40.0です。 AWS CDKは週に数回程度マイナーバージョンがアップデートされるくらい高頻度で更新されるので、変更内容等には注意してください。 また、API Gateway V2はα版なため、将来のバージョンアップにより、この記事で記載した内容では動作しなくなる場合があります。

最後に、本記事ではAWS CDKのGetting Startedの内容を読み終えた方向けに解説しています。まだお読みでない場合はこの記事をご覧になる前に一読することをお勧めいたします。

https://docs.aws.amazon.com/ja_jp/cdk/v2/guide/getting_started.html

WebSocketの仕組み

実際に始める前に、WebSocketの仕組みについて簡単に解説します。

WebSocketはクライアントとサーバー間で対話的に通信をするための仕組みで、WebSocketを使用すると、サーバーからクライアントに向けてデータの送信が可能になります。

API Gateway V2WebSocketはクライアントが送信するメッセージの内容でルートを変える事ができ、メッセージのルート分岐の変数名やルートごとのAWSリソースへの紐付けは自分で定義することができます。 また、API Gateway V2には最初から$connect$disconnect$defaltのルートが用意されており、それぞれ接続時、切断時、適切なルートがない場合のフォールバック時に実行するように設定することができます。

Step1:新規プロジェクト作成

まずは、以下のコマンドを入力して新しいプロジェクトを作成します。

mkdir websocket-api-chat-app-tutorial
cd websocket-api-chat-app-tutorial
cdk init app --language typescript

プロジェクトが作成されたら、API Gateway V2のα版パッケージをインストールします。

AWS CDKではaws-cdkパッケージにすべてのAWSリソースを定義するコードが含まれていますが、α版などの安定していないパッケージは個別にインストールする必要があります。

⚠️
インストールする際はaws-cdkのバージョンとAPI Gateway V2パッケージのバージョンが一致していることを必ず確認してください。
npm i @aws-cdk/aws-apigatewayv2-integrations-alpha@^2.40.0-alpha.0

最後に、Lambda関数を定義する際に使用する、aws-sdk@types/aws-sdk@types/aws-lambdaをインストールします。

npm i aws-sdk
npm i -D @types/aws-sdk @types/aws-lambda

Step2:Dynamo DBテーブルを定義

必要なセットアップが完了したら、まずはConnectionsTableというDynamoDBテーブルを定義します。

このテーブルは、WebSocketに接続しているクライアントの固有ID(connectionId)を保存するために使用されます。

lib/websocket-api-chat-app-tutorial.tsファイルを開き、コードを記述します。

import { aws_dynamodb, RemovalPolicy } from "aws-cdk-lib";
import * as cdk                        from 'aws-cdk-lib';
import { BillingMode }  from "aws-cdk-lib/aws-dynamodb";
import { Construct }    from 'constructs';

export class WebsocketApiChatAppTutorialStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    const table = new aws_dynamodb.Table(this, "ConnectionsTable", {
      billingMode: BillingMode.PROVISIONED,
      readCapacity: 5,
      writeCapacity: 5,
      removalPolicy: RemovalPolicy.DESTROY,
      partitionKey: { name: "connectionId", type: aws_dynamodb.AttributeType.STRING },
    });
  }
}

Step3:Lambda関数の処理を記述

次に、サーバー上で実行されるLambda関数の処理を記述します。

プロジェクト直下にlambdaディレクトリを作成し、以下のようにファイルを作成します。

ℹ️
Lambda関数の処理内容をTypeScriptで記述していますが、AWS上ではTypeScriptを直接実行できないため、デプロイ前にコンパイルする必要があります。コンパイル方法については後ほど解説します。

$connect route用Lambda関数

connect-handler.tsファイルを作成し、以下の内容を記述します。

このLambda関数はクライアント接続時に実行され、処理内容はクライアントのconnectionIdを取得し、Step2で定義したDynamoDBに追加する処理を行います。

この関数のポイントとして、9行目のTableNameにはStep2で定義したDynamoDBの名前を渡すのではなく、環境変数(process.env.table)を渡すようにします。(理由は後述)

import { APIGatewayProxyWebsocketHandlerV2 } from "aws-lambda";
import * as AWS                              from "aws-sdk";

const ddb = new AWS.DynamoDB.DocumentClient();

export const handler: APIGatewayProxyWebsocketHandlerV2 = async (event) => {
    return await ddb
        .put({
            TableName: process.env.CONNECTIONS_TABLE_NAME ?? "",
            Item: {
                connectionId: event.requestContext.connectionId,
            },
        })
        .promise()
        .then(() => ({
            statusCode: 200,
        }))
        .catch((e) => {
            console.error(e);
            return {
                statusCode: 500,
            };
        });
};

$disconnect route用Lambda関数

disconnect-handler.tsファイルを作成し、以下の内容を記述します。

このLambda関数はクライアント切断時に実行され、$connect routeでDynamoDBに保存したconnectionIdを削除します。

import { APIGatewayProxyWebsocketHandlerV2 } from "aws-lambda";
import * as AWS                              from "aws-sdk";

const ddb = new AWS.DynamoDB.DocumentClient();

export const handler: APIGatewayProxyWebsocketHandlerV2 = async (event) => {
    return await ddb
        .delete({
            TableName: process.env.CONNECTIONS_TABLE_NAME ?? "",
            Key: {
                connectionId: event.requestContext.connectionId,
            },
        })
        .promise()
        .then(() => ({
            statusCode: 200,
        }))
        .catch((e) => {
            console.log(e);
            return {
                statusCode: 500,
            };
        });
};

send-message route用Lambda関数

send-handler.tsファイルを作成し、以下の内容を記述します。

このLambda関数では、受信したメッセージを接続しているクライアントに送信します。

import { APIGatewayProxyWebsocketHandlerV2 } from "aws-lambda";
import * as AWS                              from "aws-sdk";

const ddb = new AWS.DynamoDB.DocumentClient();

export const handler: APIGatewayProxyWebsocketHandlerV2 = async (event) => {
    let connections;
    try {
        connections = (
            await ddb.scan({ TableName: process.env.table ?? "" })
                     .promise()
        ).Items as { connectionId: string }[];
    } catch (err) {
        return {
            statusCode: 500,
        };
    }

    const callbackAPI = new AWS.ApiGatewayManagementApi({
        apiVersion: "2018-11-29",
        endpoint: event.requestContext.domainName + "/" + event.requestContext.stage,
    });

    const message = JSON.parse(event.body ?? "{}").message;

    const sendMessages = connections.map(async ({ connectionId }) => {
        if (connectionId === event.requestContext.connectionId) return;

        await callbackAPI
            .postToConnection({ ConnectionId: connectionId, Data: message })
            .promise()
            .catch(e => console.error(e));
    });

    return await Promise
        .all(sendMessages)
        .then(() => ({
            statusCode: 200,
        }))
        .catch((e) => {
            console.error(e);
            return {
                statusCode: 500,
            };
        });
};

$default route用Lambda関数

最後に、default-handler.tsファイルを作成し、以下の内容を記述します。

このLambda関数はフォールバック用の関数で、予め定義されたルート(今回の場合send_routeのみ)のいずれにも当てはまらないルートが選択された場合に呼び出されます。今回は、誤ったルートを呼び出していることをクライアントに知らせる処理をします。

import { APIGatewayProxyWebsocketHandlerV2 } from "aws-lambda";
import * as AWS                              from "aws-sdk";

export const handler: APIGatewayProxyWebsocketHandlerV2 = async (event) => {
    let connectionId = event.requestContext.connectionId;

    const callbackAPI = new AWS.ApiGatewayManagementApi({
        apiVersion: "2018-11-29",
        endpoint: event.requestContext.domainName + "/" + event.requestContext.stage,
    });

    let connectionInfo: AWS.ApiGatewayManagementApi.GetConnectionResponse;
    try {
        connectionInfo = await callbackAPI
            .getConnection({
                ConnectionId: event.requestContext.connectionId,
            })
            .promise();
    } catch (e) {
        console.log(e);
    }

    const info = {
        ...connectionInfo!,
        connectionId,
    };
    await callbackAPI.postToConnection({
        ConnectionId: event.requestContext.connectionId,
        Data: "Use the send-message route to send a message. Your info:" + JSON.stringify(info),
    }).promise();

    return {
        statusCode: 200,
    };
};

Step4:Lambda関数を定義

定義例

Step3で記述したLambda関数の処理は、AWS上で実行される際の内容なため、現時点ではCDKのデプロするリソースとして定義されていません。

そのため、以下のようにLambda関数を定義します。

export class WebsocketApiChatAppTutorialStack extends cdk.Stack {
    // ...

    connectHandlerBuilder(table: Table) {
        const handler = new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketConnectHandler", {
            environment: {
                table: table.tableName,
            },
            runtime: Runtime.NODEJS_16_X,
            entry: "lambda/connect-handler.ts",
        });

        table.grantWriteData(handler);

        return handler;
    }
}

ポイント

このLambda関数の定義にはいくつかのポイントがあります。

ポイント1: NodejsFunctionでインスタンス化する

このクラスを使用することで、CDKデプロイ時にTypeScriptで書かれたLambda関数の処理を自動的にJavaScriptにコンパイルしてくれます。

ポイント2:environmentオブジェクトにテーブル名を渡す

environmentオブジェクトに渡した値は、Lambda関数内のprocess.envで参照することができるようになります。今回はtableキーにConnectionTableの名前を渡しています。

わざわざ環境変数としてテーブル名を渡す理由は、CDKのリソース名生成処理にあります。 CDKはデプロイする際にリソースの名前にランダムな英数字を挿入するため、コーディング中に正確な名前を知ることができません。そのため、定義時にテーブル名のリファレンスを渡すようにすることで、この問題を回避しています。

ポイント3: テーブルへのアクセスを許可する

基本的にLambda関数は初期状態では他のリソースへのアクセス許可を持っておらず、明示的に許可しないと実行時に権限エラーが発生します。

セキュリティの観点から、Lambda関数には必要最小限の権限を付与するようにします。

table.grantWriteData(handler);     // 書き込み許可
table.grantReadData(handler);      // 読み込み許可
table.grantReadWriteData(handler); // 読み書き許可

他のLambda関数も定義

定義例で挙げたconnect-handlerの定義と同様に他のLambda関数も定義していきます。

以下に、全てのLambda関数を定義した後のlib/websocket-api-chat-app-tutorial-stack.tsを示します。

import { aws_dynamodb, aws_lambda_nodejs, RemovalPolicy } from "aws-cdk-lib";
import * as cdk                                           from "aws-cdk-lib";
import { BillingMode, Table }                             from "aws-cdk-lib/aws-dynamodb";
import { Runtime }                                        from "aws-cdk-lib/aws-lambda";
import { Construct }                                      from "constructs";

export class WebsocketApiChatAppTutorialStack extends cdk.Stack {
    constructor(scope: Construct, id: string, props?: cdk.StackProps) {
        super(scope, id, props);

        const table = new aws_dynamodb.Table(this, "ConnectionsTable", {
            billingMode: BillingMode.PROVISIONED,
            readCapacity: 5,
            writeCapacity: 5,
            removalPolicy: RemovalPolicy.DESTROY,
            partitionKey: { name: "connectionId", type: aws_dynamodb.AttributeType.STRING },
        });
    }

    connectHandlerBuilder(table: Table) {
        const handler = new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketConnectHandler", {
            environment: {
                table: table.tableName,
            },
            runtime: Runtime.NODEJS_16_X,
            entry: "lambda/connect-handler.ts",
        });

        table.grantWriteData(handler);

        return handler;
    }

    disconnectHandlerBuilder(table: Table) {
        const handler = new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketDisconnectHandler", {
            environment: {
                table: table.tableName,
            },
            runtime: Runtime.NODEJS_16_X,
            entry: "lambda/disconnect-handler.ts",
        });

        table.grantWriteData(handler);

        return handler;
    }

    sendMessageHandlerBuilder(table: Table) {
        const handler = new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketSendHandler", {
            environment: {
                table: table.tableName,
            },
            runtime: Runtime.NODEJS_16_X,
            entry: "lambda/send-handler.ts",
        });

        table.grantReadWriteData(handler);

        return handler;
    }

    defaultHandlerBuilder() {
        return new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketDefaultHandler", {
            entry: "lambda/default-handler.ts",
            runtime: Runtime.NODEJS_16_X,
        });
    }
}

Step5:API Gateway V2を定義

ここまでで、API Gateway V2を定義するのに必要なリソースの定義は完了しました。

最後にAPI Gateway V2をスタックのコンストラクタに定義すれば、必要なリソースの定義は完了します。

以下に定義の内容を順番に解説し、最後に全体のコードを載せます。

1. API Gateway V2を定義

まずは以下のように、API Gateway V2を定義し、インスタンス化します。

import { WebSocketLambdaIntegration } from "@aws-cdk/aws-apigatewayv2-integrations-alpha";
import { WebSocketApi }               from "@aws-cdk/aws-apigatewayv2-alpha/lib/websocket";

//DynamoDBテーブル定義の後に以下を追記

// 関数定義
const connectHandler     = this.connectHandlerBuilder(table);
const disconnectHandler  = this.disconnectHandlerBuilder(table);
const sendMessageHandler = this.sendMessageHandlerBuilder(table);
const defaultHandler     = this.defaultHandlerBuilder();

// API Gateway V2を定義
const webSocketApi = new WebSocketApi(this, "MessageApi", {
    routeSelectionExpression: "$request.body.action",
    connectRouteOptions: {
        integration: new WebSocketLambdaIntegration("MessageApiConnectIntegration", connectHandler),
    },
    disconnectRouteOptions: {
        integration: new WebSocketLambdaIntegration("MessageApiDisconnectIntegration", disconnectHandler),
    },
    defaultRouteOptions: {
        integration: new WebSocketLambdaIntegration("MessageApiDefaultIntegration", defaultHandler),
    },
});

ここでのポイントは以下の2点です。

routeSelectionExpression(選択式)

選択式ではリクエストのどの部分を元にルートを選択する評価方法を決める事ができます。今回はリクエストボディのactionパラメータを元に選択するので、$request.body.actionとします。

connectionRouteOptions, disconnectRouteOptions, defaultRouteOptions

インスタンス化するときに、これらのデフォルトルートを定義します。

それぞれの値を確認すると、integrationキーにWebSocketLambdaIntegrationインスタンスがありますが、ここでWebSocket用のLambda統合を定義し、Step4で定義したLambda関数の定義と紐付けています。

Lambda統合は、API Gateway V2が受け取ったリクエストをLambda関数用にマッピングする機能です。

2. send-messageルートを定義

次に、API Gateway V2インスタンスに追加でsend-messageルートを追加します。

const sendMessageHandler = this.sendMessageHandlerBuilder(table);

webSocketApi.addRoute("send-message", {
    integration: new WebSocketLambdaIntegration("MessageApiSendIntegration", sendMessageHandler),
});

addRoute関数の第一引数がルートキーで、リクエストボディのactionパラメータの値がsend-messageだった場合、このルートが選択されます。

第二引数はオプションで、デフォルトルートのときの同じようにLamdba統合を定義し、Lambda関数と紐付けています。

3. Lambda関数にManage Connection権限を付与する

Manage Connectionとは、WebSocketクライントへのメッセージの送信や接続情報の取得、クライアントの切断をできるようにする権限で、これが与えられてないと、send-mesageやdefault routeで使用しているAWS.ApiGatewayManagementApiを使用することができません。

権限は以下のように付与できます。

webSocketApi.grantManageConnections(sendMessageHandler);
webSocketApi.grantManageConnections(defaultHandler);

4. ステージを定義

最後にステージを定義して、Apiを公開できるようにします。

import { WebSocketStage } from "@aws-cdk/aws-apigatewayv2-alpha/lib/websocket";

new WebSocketStage(this, "MessageApiProd", {
    webSocketApi,
    stageName: "prod",
    autoDeploy: true,
});

全体コード

import { WebSocketLambdaIntegration }                     from "@aws-cdk/aws-apigatewayv2-integrations-alpha";
import { aws_dynamodb, aws_lambda_nodejs, RemovalPolicy } from "aws-cdk-lib";
import * as cdk                                           from "aws-cdk-lib";
import { BillingMode, Table }                             from "aws-cdk-lib/aws-dynamodb";
import { Runtime }                                        from "aws-cdk-lib/aws-lambda";
import { Construct }                                      from "constructs";
import { WebSocketApi, WebSocketStage }                   from "@aws-cdk/aws-apigatewayv2-alpha/lib/websocket";

export class WebsocketApiChatAppTutorialStack extends cdk.Stack {
    constructor(scope: Construct, id: string, props?: cdk.StackProps) {
        super(scope, id, props);

        const table = new aws_dynamodb.Table(this, "ConnectionsTable", {
            billingMode: BillingMode.PROVISIONED,
            readCapacity: 5,
            writeCapacity: 5,
            removalPolicy: RemovalPolicy.DESTROY,
            partitionKey: { name: "connectionId", type: aws_dynamodb.AttributeType.STRING },
        });

        const connectHandler     = this.connectHandlerBuilder(table);
        const disconnectHandler  = this.disconnectHandlerBuilder(table);
        const sendMessageHandler = this.sendMessageHandlerBuilder(table);
        const defaultHandler     = this.defaultHandlerBuilder();

        const webSocketApi = new WebSocketApi(this, "MessageApi", {
            routeSelectionExpression: "$request.body.action",
            connectRouteOptions: {
                integration: new WebSocketLambdaIntegration("MessageApiConnectIntegration", connectHandler),
            },
            disconnectRouteOptions: {
                integration: new WebSocketLambdaIntegration("MessageApiDisconnectIntegration", disconnectHandler),
            },
            defaultRouteOptions: {
                integration: new WebSocketLambdaIntegration("MessageApiDefaultIntegration", defaultHandler),
            },
        });

        webSocketApi.addRoute("send-message", {
            integration: new WebSocketLambdaIntegration("MessageApiSendIntegration", sendMessageHandler),
        });

        webSocketApi.grantManageConnections(sendMessageHandler);
        webSocketApi.grantManageConnections(defaultHandler);

        new WebSocketStage(this, "MessageApiProd", {
            webSocketApi,
            stageName: "prod",
            autoDeploy: true,
        });
    }

    connectHandlerBuilder(table: Table) {
        const handler = new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketConnectHandler", {
            environment: {
                table: table.tableName,
            },
            runtime: Runtime.NODEJS_16_X,
            entry: "lambda/connect-handler.ts",
        });

        table.grantWriteData(handler);

        return handler;
    }

    disconnectHandlerBuilder(table: Table) {
        const handler = new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketDisconnectHandler", {
            environment: {
                table: table.tableName,
            },
            runtime: Runtime.NODEJS_16_X,
            entry: "lambda/disconnect-handler.ts",
        });

        table.grantWriteData(handler);

        return handler;
    }

    sendMessageHandlerBuilder(table: Table) {
        const handler = new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketSendHandler", {
            environment: {
                table: table.tableName,
            },
            runtime: Runtime.NODEJS_16_X,
            entry: "lambda/send-handler.ts",
        });

        table.grantReadWriteData(handler);

        return handler;
    }

    defaultHandlerBuilder() {
        return new aws_lambda_nodejs.NodejsFunction(this, "MessageWebSocketDefaultHandler", {
            entry: "lambda/default-handler.ts",
            runtime: Runtime.NODEJS_16_X,
        });
    }
}

Step6:デプロイ

ここまで完了したら、実際にAWS上にアプリケーションをデプロイします。

cdk deploy

コマンド実行後にビルドされ、成功すると、以下のように作成するリソースの一覧と確認が出ますが、yを入力して続けてください。

デプロイが完了すると以下のように表示されます。(今回は完了までに117秒かかりました)

AWSコンソールのCloudFormationにもデプロイされていることが確認できます。

エラーが起きた場合

もしビルド中にDocker関連でのエラーが発生していた場合は、npm i -g esbuildコマンドを実行して、esbuildをインストールしてみてください。

また、デプロイ先のCloudFormationにCDKToolkitスタックがないと、デプロイすることができません。cdk bootstrapコマンドでスタックを作成できるので、なかった場合は作成してください。

Step7:動作確認

実際にWebSocketが動作するか確認します。

まずは、クライアントツールとしてwscatをインストールします。

npm i -g wscat

次に、AWSコンソールのAPI Gatewayを開きMessageAPIを選択します。

選択後に左側のStagesタブをクリックし、prodステージを選択し、ステージエディターを開きます。

ステージエディターのページにWebSocket URLが表示されるので、これをメモします。

メモが終わったら、コンソールを2つ開き、それぞれWebSocket APIに接続します。

wscat -c wss://abcdef123.execute-api.ap-northeast-1.amazonaws.com/prod

片方のコンソールから以下のJSONデータを入力して送信すると、もう片方のコンソールに送信メッセージが表示されます。

{"action": "send-message", "message": "hello!"}

このように、クライアント同士でメッセージのやり取りができれば、問題なく動作しています。

まとめ

今回の記事ではWebSocketを利用したチャットアプリケーションをAWS CDKで構築しました。

API Gateway V2は様々なリソースと連携できるので、例えばDynamoDB Streamを使用して、テーブルに変更があった場合にWebSocketクライアントに通知するといったように、様々な活用が考えられます。

以上、この記事がどなたかの参考になれば幸いです。

Hirayama

Hirayama / Engineer

1997年生まれ、南伊豆出身。学生時代にC#で画像処理アプリケーションを作ったりしていました。業務では主にLaravelを使用してサーバーサイドのプログラミングをしています。趣味はドライブとシミュレーションゲーム。