Zenn
🦃

AWS API GatewayのWebSocketのechoサーバをCDKで立てる

2021/05/08に公開
AWS
TypeScript
WebSocket
Amazon API Gateway
AWS CDK
tech

AWSのAPI GatewayでのWebSocketサーバを構成する勉強にechoサーバ、つまりClientがWebSocketメッセージを送信したらそのメッセージをそのまま返信するサーバを書いてみます。ググって引っ掛かる構成例がグループチャットサーバで、これはDynamoDBも必要だったりと勉強の1ステップ目としては結構難しいと感じたので、より簡単な例にしてみました。

全体は https://github.com/keshihoriuchi/samples/tree/master/aws/api-gateway-websocket-echo にあります。

CDK

以下がStackを一枚で記述するCDKのtsファイルです。

lib/api-gatewaty-websocket-echo-stack.ts
import { WebSocketApi, WebSocketStage } from "@aws-cdk/aws-apigatewayv2";
import { LambdaWebSocketIntegration } from "@aws-cdk/aws-apigatewayv2-integrations";
import { PolicyStatement } from "@aws-cdk/aws-iam";
import { NodejsFunction } from "@aws-cdk/aws-lambda-nodejs";
import { Stack, StackProps, Construct } from "@aws-cdk/core";

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

    const defaultHandler = new NodejsFunction(this, "WsEchoDefault", {
      entry: "lambda/default.ts",
    });

    const webSocketApi = new WebSocketApi(this, "WsEcho", {
      defaultRouteOptions: {
        integration: new LambdaWebSocketIntegration({
          handler: defaultHandler,
        }),
      },
    });
    const stage = new WebSocketStage(this, "DevStage", {
      webSocketApi,
      stageName: "dev",
      autoDeploy: true,
    });

    const wsManageConnPolicy = new PolicyStatement({
      actions: ["execute-api:ManageConnections"],
      resources: [
        this.formatArn({
          service: "execute-api",
          resourceName: `${stage.stageName}/POST/*`,
          resource: webSocketApi.apiId,
        }),
      ],
    });
    defaultHandler.addToRolePolicy(wsManageConnPolicy);
  }
}

Clientから何かメッセージが送られてきたら、WebSocketApiのdefaultRouteOptionsに指定したLambdaが呼び出されるようになります。

LambdaからClientにメッセージを送信するにはwsManageConnPolicyで定義しているPolicyStatementが必要です。ここで結構はまりました。CDK使うと生CfnよりIAM周りで悩まないで済む箇所も多いですがここは厳しいですね…。ただ記載時点でAWS::APIGatewayv2はEXPERIEMENTAL扱いなので改善されるかもしれません。

new NodejsFunction(this, "WsEchoDefault", { entry: "lambda/default.ts" }); は lambda/default.ts に置いたtsファイルをesbuildでjsにトランスパイルして、cdk bootstrapで作成されたS3バケットにアップロードしてくれます。このためには@aws-cdk/aws-lambda-nodejsパッケージに加えてesbuildパッケージもpackage.jsonに追加する必要があります。

Lambda

以下がLambdaでClientからメッセージを送信されるたびに呼び出されるjsの元となるtsファイルです。

lambda/default.ts
import { APIGatewayProxyEvent, APIGatewayProxyResultV2 } from "aws-lambda";
import { ApiGatewayManagementApi } from "aws-sdk";

exports.handler = async (
  event: APIGatewayProxyEvent
): Promise<APIGatewayProxyResultV2> => {
  try {
    const endpoint = `https://${event.requestContext.domainName}/${event.requestContext.stage}`;
    const apigw = new ApiGatewayManagementApi({ endpoint });
    const connId = event.requestContext.connectionId!;
    await apigw
      .postToConnection({
        ConnectionId: connId,
        Data: event.body!,
      })
      .promise();
  } catch (e) {
    console.error(e);
    return { statusCode: 500 };
  }
  return { statusCode: 200 };
};

aws-lambda、aws-sdkをpackage.jsonに追加しておくと型定義を参照してくれます。

event.requestContext.connectionIdから自分のconnectionIdを取得します。connectionIdはWebSocket接続が確立するたびにランダムに振られるIDのようです。今回は自分にメッセージを返送するだけなので自分のconnectionIdが分かればよいですが、グループチャットなどを作る場合は他の接続のconnectionIdをDynamoDBなどで記録、参照する必要があります。

await apigw.postToConnection({ConnectionId: connId, Data: event.body!}).promise() でメッセージを送信してきたClient自身に対して、そのClientが送信してきたメッセージをそのまま返信します。.promise()をawaitしないと先にreturnが呼ばれるので処理が打ち切られます。

なお、default.tsはトランスパイルされると以下のようなjsファイルになります。aws-sdkはLambdaランタイム上のモジュールを参照するのでバンドルされずrequireするようになります。

cdk.out/asset.*/index.js
var __create = Object.create;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __getProtoOf = Object.getPrototypeOf;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __markAsModule = (target) => __defProp(target, "__esModule", {value: true});
var __reExport = (target, module2, desc) => {
  if (module2 && typeof module2 === "object" || typeof module2 === "function") {
    for (let key of __getOwnPropNames(module2))
      if (!__hasOwnProp.call(target, key) && key !== "default")
        __defProp(target, key, {get: () => module2[key], enumerable: !(desc = __getOwnPropDesc(module2, key)) || desc.enumerable});
  }
  return target;
};
var __toModule = (module2) => {
  return __reExport(__markAsModule(__defProp(module2 != null ? __create(__getProtoOf(module2)) : {}, "default", module2 && module2.__esModule && "default" in module2 ? {get: () => module2.default, enumerable: true} : {value: module2, enumerable: true})), module2);
};

// default.ts
var import_aws_sdk = __toModule(require("aws-sdk"));
exports.handler = async (event) => {
  try {
    const endpoint = `https://${event.requestContext.domainName}/${event.requestContext.stage}`;
    const apigw = new import_aws_sdk.ApiGatewayManagementApi({endpoint});
    const connId = event.requestContext.connectionId;
    await apigw.postToConnection({
      ConnectionId: connId,
      Data: event.body
    }).promise();
  } catch (e) {
    console.error(e);
    return {statusCode: 500};
  }
  return {statusCode: 200};
};

動作確認

npx cdk deployでデプロイしたらwscatで動作確認してみます。 https://github.com/websockets/wscat

$ wscat -c wss://hako54h42l.execute-api.ap-northeast-1.amazonaws.com/dev
Connected (press CTRL+C to quit)
> foo
< foo
> bar
< bar

参考

Discussion