はじめに

サービスでチャット機能を提供するケースがあるかと思います。
その時に必要になるのが双方向通信であるWebsocketになります。
今回はAWSで公開されている API Gateway で Websocket API を構築して、トークルーム付きのチャット機能を実装してみます。
構成は API Gateway + Lambda + DynamoDB です。

また、Websocket接続を特定のユーザーだけが使用できるように制限したいケースもあるかと思います。
WebsocketではHTTPプロトコルで使われるCookie認証を使うことが難しいため、
チャット画面専用のトークンを用いて、Lamdbaオーソライザを間に挟んで認証をします。

対象者

チャット機能の導入を検討されている方
Amazon Web Servicesを使用される方

目次

構成
設定の流れ
Lamdbaの作成
DynamoDBの作成
API Gatewayの作成
動作確認
まとめ

構成

今回は以下の構成で設定していきます。

|サービス |用途
|---|---|---|
|Amazon API Gateway |ユーザーからのアクセスを元に後続のLamdbaに処理を割り振ります。また、Websocket管理を行います。
|AWS Lambda |Websocket接続から、メッセージ送信、切断処理を行います。
|Amazon DynamoDB |入退室管理を行います。Websocket接続時のconnection情報やroom情報を保持します。

設定の流れ

以下の流れで設定していきます。

Lamdbaの作成
1. OnConnect Lamdba
2. SendMessage Lamdba
3. OnDisconnect Lamdba
4. Lamdbaオーソライザ
DynamoDBの作成
API Gatewayの作成

Lamdbaの作成

OnConnect Lamdba

関数名: sample-onconnect
ランタイム: Node.js 14.x
アーキテクチャ: x86_64

index.js

var AWS = require("aws-sdk");
AWS.config.update({ region: process.env.AWS_REGION });
var db = new AWS.DynamoDB({ apiVersion: "2012-08-10" });

exports.handler = function (event, context, callback) {
  // connectionId,roomIdの取得
  const subprotocolHeader = event.headers['Sec-WebSocket-Protocol'];
  const subprotocols = subprotocolHeader.split(',');
  const secWebsocketProtocol = subprotocols[0];
  const connectionId = event.requestContext.connectionId
  const roomId = subprotocols[1]
  
  var putParams = {
    TableName: "sample-chat-dynamodb",
    Item: {
      connectionId: { S: connectionId },
      roomId: { S: roomId }
    }
  };

  db.putItem(putParams, function (err) {
    console.log(err);
    callback(null, {
      statusCode: err ? 500 : 200,
      body: err ? "Failed" : "Connected.",
      headers: {
        "Sec-WebSocket-Protocol" : secWebsocketProtocol
      }
    });
  });
};

SendMessage Lamdba

関数名: sample-sendmessage
ランタイム: Node.js 14.x
アーキテクチャ: x86_64

index.js

const AWS = require('aws-sdk');
AWS.config.update({ region: process.env.AWS_REGION });
var db = new AWS.DynamoDB({ apiVersion: "2012-08-10" });

exports.handler = async (event, context, callback) => {
  // メッセージ, connectionId, roomIdの取得 
  const postData = JSON.parse(event.body).data;
  const message = postData.message;
  const myConnectionId = event.requestContext.connectionId;
  const roomId = postData.roomId;
  console.log(message)
  
  const agma = new AWS.ApiGatewayManagementApi({
      apiVersion: "2018-11-29",
      endpoint: event.requestContext.domainName + "/" + event.requestContext.stage
  });
  
  // 同じroomのconnection検索
  const queryParams = {
    TableName: "sample-chat-dynamodb",
    KeyConditionExpression: "#roomId = :roomId",
    ExpressionAttributeNames: { "#roomId": "roomId" },
    ExpressionAttributeValues: { ":roomId": { "S": roomId} },
    IndexName: 'roomId-index'
  }
  const connectionData = await db.query(queryParams).promise();
  
  // メッセージのpush
  const postCalls = connectionData.Items.map(async ({ connectionId }) => {
    const target = connectionId.S.replace(/[\"]/g,"")
    console.log(target)
    
    try {
      if (myConnectionId !== target) {
        await agma.postToConnection({ ConnectionId: target, Data: message }).promise();
      }
      
    } catch (e) {
      const response = {
          statusCode: 500,
          body: JSON.stringify(e.message),
      };
    
      callback(null, response);
    }
  })
  
  const response = {
      statusCode: 200,
      body: JSON.stringify("success"),
  };

  callback(null, response);
};

OnDisconnect Lamdba

関数名: sample-disconnect
ランタイム: Node.js 14.x
アーキテクチャ: x86_64

index.js

var AWS = require("aws-sdk");
AWS.config.update({ region: process.env.AWS_REGION });
var db = new AWS.DynamoDB({ apiVersion: "2012-08-10" });

exports.handler = function (event, context, callback) {
  var deleteParams = {
    TableName: "sample-chat-dynamodb",
    Key: {
      connectionId: { S: event.requestContext.connectionId }
    }
  };

  // DynamoDBのテーブルから削除
  db.deleteItem(deleteParams, function (err) {
    callback(null, {
      statusCode: err ? 500 : 200,
      body: err ? "Failed" : "disconnected."
    });
  });
};

Lamdbaオーソライザ

関数名: sample-auth
ランタイム: Node.js 14.x
アーキテクチャ: x86_64

今回は簡略化するため、特定の文字列のトークンが送られた時に許可するようにします。

index.js

exports.handler = async function (event, context, callback) {
  const headers = event.headers;
  
  if (headers['Sec-WebSocket-Protocol'] != undefined) {
      const subprotocolHeader = headers['Sec-WebSocket-Protocol'];
      const subprotocols = subprotocolHeader.split(',');
      const token = subprotocols[0]
     
     // 要件に則ってここでtoken認証を追加
    if(token === "abcdefghijklmnopqrstuvwxyg") {
      return generatePolicy('user', 'Allow', event.methodArn)
    } else {
      callback("Unauthorized");
    }
  } else {
    callback("Unauthorized");
  }
}

var generatePolicy = function(principalId, effect, resource) {
  let authResponse = {};

  authResponse.principalId = principalId;
  if (effect && resource) {
    var policyDocument = {};
    policyDocument.Version = '2012-10-17';
    policyDocument.Statement = [];
    var statementOne = {};
    statementOne.Action = 'execute-api:Invoke';
    statementOne.Effect = effect;
    statementOne.Resource = resource;
    policyDocument.Statement[0] = statementOne;
    authResponse.policyDocument = policyDocument;
  }

  return authResponse;
}

ロールの設定

作成したLamdbaに同じロールを設定し、以下のポリシーをアタッチします。

AmazonDynamoDBFullAccess
ExecuteApi ※以下のようにビジュアルエディタから作成

DynamoDBの作成

テーブル名：sample-chat-dynamodb
パーティションキー : connection_id
グローバルインデックス：roomId-index　※roomIdから参加しているユーザーを検索できるようにするために必要です。
　パーティションキー：roomId (String)

API Gatewayの作成

APIの設定

WebSocket APIを選択して、設定を完了させる。
ルートとLamdba関数を紐づける。

|ルート |Lamdba関数
|---|---|---|
|$connect |sample-onconnect
|sendmessage |sample-sendMessage
|$disconnect |sample-disconnect
|$default |sample-onconnect

オーソライザの設定

以下のように設定します。

onConnectに設定したオーソライザを紐付けます。

API Gatewayのデプロイ

アクション > APIのデプロイから設定情報をデプロイする。

動作確認

ターミナルから以下のコマンドを実行することで接続ができることが確認できます。

wscat -c <デプロイ時のWebsocketURL> -s <token> -s <roomId>

トークンが異なる場合は、以下のように401が返却されます。

まとめ

認証付きのチャット機能を作成してみました。フロントからの接続については次の記事、もしくは本記事に後日追記したいと思います。