CDKでWebSocket APIを構築する — サーバーレスチャットのインフラ定義

REST API:                              WebSocket API:
┌─────────────┐                        ┌─────────────────────┐
│ API         │                        │ API                 │
│  ├─ /users  │ (パス)                 │  ├─ $connect        │ (接続時)
│  │  ├─ GET  │ (HTTPメソッド)         │  ├─ $disconnect     │ (切断時)
│  │  └─ POST │                        │  ├─ sendMessage     │ (カスタム)
│  └─ /items  │                        │  └─ $default        │ (その他)
│     └─ GET  │                        │                     │
└─────────────┘                        └─────────────────────┘

リクエスト: GET /users                 リクエスト: {"action":"sendMessage",...}
→ パスとメソッドでルーティング         → bodyのactionフィールドでルーティング

クライアント                     API Gateway                Lambda

── WebSocket接続要求 ──────▶ ── $connect ──────────▶ connect.ts
                                                    (JWT検証 + DB保存)

── {"action":"sendMessage",  ── sendMessage ───────▶ sendMessage.ts
    "chatRoomId":"room-xyz",                         (保存 + ブロードキャスト)
    "content":"こんにちは"} ▶

── (ブラウザ閉じる) ────────▶ ── $disconnect ────────▶ disconnect.ts
                                                    (DB削除)

── {"action":"unknown"} ───▶ ── $default ──────────▶ default.ts
                                                    (ログ出力)

import * as cdk from 'aws-cdk-lib';
import * as apigatewayv2 from 'aws-cdk-lib/aws-apigatewayv2';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as cognito from 'aws-cdk-lib/aws-cognito';
import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
import * as iam from 'aws-cdk-lib/aws-iam';
import { Construct } from 'constructs';
import * as path from 'path';

export interface WebSocketStackProps extends cdk.StackProps {
  envName: string;
  userPool: cognito.UserPool;
  userPoolClient: cognito.UserPoolClient;
  table: dynamodb.Table;
}

export class WebSocketStack extends cdk.Stack {
  public readonly webSocketApi: apigatewayv2.CfnApi;
  public readonly webSocketEndpoint: string;

  constructor(scope: Construct, id: string, props: WebSocketStackProps) {
    super(scope, id, props);

    // ... 以下で各リソースを定義
  }
}

this.webSocketApi = new apigatewayv2.CfnApi(this, 'WebSocketApi', {
  name: `Connect40-WebSocket-${props.envName}`,
  protocolType: 'WEBSOCKET',
  routeSelectionExpression: '$request.body.action',
});

// クライアントが送るJSON
{"action": "sendMessage", "chatRoomId": "room-xyz", "content": "こんにちは"}
//          ↑ この値が "sendMessage" ルートにマッチ

// Lambda共通ロール
const lambdaRole = new iam.Role(this, 'WebSocketLambdaRole', {
  assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
  managedPolicies: [
    iam.ManagedPolicy.fromAwsManagedPolicyName(
      'service-role/AWSLambdaBasicExecutionRole'
    ),
  ],
});

// DynamoDBの読み書き権限
props.table.grantReadWriteData(lambdaRole);

// ★ WebSocket接続への書き戻し権限
lambdaRole.addToPolicy(
  new iam.PolicyStatement({
    actions: ['execute-api:ManageConnections'],
    resources: [
      `arn:aws:execute-api:${this.region}:${this.account}:${this.webSocketApi.ref}/*`,
    ],
  })
);

通常のREST API:
  クライアント ──リクエスト──▶ Lambda ──レスポンス──▶ クライアント

WebSocket API:
  クライアントA ──メッセージ──▶ Lambda ──────────────────────────┐
                                 │                              │
                                 ├──▶ DynamoDBに保存            │
                                 │                              │
                                 ├──▶ クライアントAに送信 ◀──────┘
                                 ├──▶ クライアントBに送信  ← これにManageConnections権限が必要
                                 └──▶ クライアントCに送信

// Connect Handler
const connectFunction = new lambda.Function(this, 'ConnectFunction', {
  runtime: lambda.Runtime.NODEJS_20_X,
  handler: 'connect.handler',
  code: lambda.Code.fromAsset(
    path.join(__dirname, '../../../backend/functions/websocket')
  ),
  role: lambdaRole,
  environment: {
    TABLE_NAME: props.table.tableName,
    USER_POOL_ID: props.userPool.userPoolId,   // JWT検証用
    CLIENT_ID: props.userPoolClient.userPoolClientId,
  },
  timeout: cdk.Duration.seconds(30),
});

// Disconnect / SendMessage / Default も同様に定義
// （handler名とenvironmentが異なるだけ）

const connectIntegration = new apigatewayv2.CfnIntegration(
  this, 'ConnectIntegration', {
    apiId: this.webSocketApi.ref,
    integrationType: 'AWS_PROXY',
    integrationUri: `arn:aws:apigateway:${this.region}:lambda:path/2015-03-31/functions/${connectFunction.functionArn}/invocations`,
  }
);

// 4つのIntegrationを作成
connectIntegration     → connectFunction
disconnectIntegration  → disconnectFunction
sendMessageIntegration → sendMessageFunction
defaultIntegration     → defaultFunction

const connectRoute = new apigatewayv2.CfnRoute(this, 'ConnectRoute', {
  apiId: this.webSocketApi.ref,
  routeKey: '$connect',
  target: `integrations/${connectIntegration.ref}`,
});

const sendMessageRoute = new apigatewayv2.CfnRoute(this, 'SendMessageRoute', {
  apiId: this.webSocketApi.ref,
  routeKey: 'sendMessage',  // ← action の値と一致
  target: `integrations/${sendMessageIntegration.ref}`,
});

// $disconnect, $default も同様

connectFunction.grantInvoke(
  new iam.ServicePrincipal('apigateway.amazonaws.com')
);
disconnectFunction.grantInvoke(
  new iam.ServicePrincipal('apigateway.amazonaws.com')
);
sendMessageFunction.grantInvoke(
  new iam.ServicePrincipal('apigateway.amazonaws.com')
);
defaultFunction.grantInvoke(
  new iam.ServicePrincipal('apigateway.amazonaws.com')
);

// Deployment（ルート定義の「スナップショット」）
const deployment = new apigatewayv2.CfnDeployment(
  this, 'WebSocketDeployment', {
    apiId: this.webSocketApi.ref,
  }
);

// ルートが先に作成されていないとデプロイが空になる
deployment.addDependency(connectRoute);
deployment.addDependency(disconnectRoute);
deployment.addDependency(sendMessageRoute);
deployment.addDependency(defaultRoute);

// Stage（デプロイのエイリアス）
const stage = new apigatewayv2.CfnStage(this, 'WebSocketStage', {
  apiId: this.webSocketApi.ref,
  stageName: props.envName,  // 'dev' or 'prod'
  deploymentId: deployment.ref,
  defaultRouteSettings: {
    throttlingBurstLimit: 500,
    throttlingRateLimit: 1000,
  },
});

作成順序:
  API → Lambda → Integration → Route → Deployment → Stage
                                  ↑ addDependency で順序保証

this.webSocketEndpoint =
  `wss://${this.webSocketApi.ref}.execute-api.${this.region}.amazonaws.com/${props.envName}`;

new cdk.CfnOutput(this, 'WebSocketEndpoint', {
  value: this.webSocketEndpoint,
  exportName: `Connect40-WebSocketEndpoint-${props.envName}`,
});

WebSocketStack
│
├─ CfnApi (WebSocket API)
│   └─ routeSelectionExpression: '$request.body.action'
│
├─ IAM Role (Lambda共通)
│   ├─ AWSLambdaBasicExecutionRole
│   ├─ DynamoDB ReadWrite
│   └─ execute-api:ManageConnections
│
├─ Lambda Functions (×4)
│   ├─ connect.handler    ← USER_POOL_ID, CLIENT_ID
│   ├─ disconnect.handler
│   ├─ sendMessage.handler
│   └─ default.handler
│
├─ CfnIntegration (×4)
│   └─ AWS_PROXY → 各Lambda
│
├─ CfnRoute (×4)
│   ├─ $connect → connectIntegration
│   ├─ $disconnect → disconnectIntegration
│   ├─ sendMessage → sendMessageIntegration
│   └─ $default → defaultIntegration
│
├─ CfnDeployment
│   └─ depends on: 全ルート
│
└─ CfnStage
    └─ stageName: 'dev' / 'prod'