WebSocket Lambdaの実装 — JWT認証・ブロードキャスト・切断処理

// ❌ これはできない（ブラウザのWebSocket APIの制限）
const ws = new WebSocket('wss://...', {
  headers: { Authorization: 'Bearer xxx' }  // サポートされていない
});

// ✅ クエリパラメータでトークンを渡す
const ws = new WebSocket('wss://...?token=eyJhbGciOi...');

import { APIGatewayProxyWebsocketHandlerV2 } from 'aws-lambda';
import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
import { DynamoDBDocumentClient, PutCommand } from '@aws-sdk/lib-dynamodb';
import { CognitoJwtVerifier } from 'aws-jwt-verify';

const client = new DynamoDBClient({});
const ddbDocClient = DynamoDBDocumentClient.from(client);

const TABLE_NAME = process.env.TABLE_NAME!;
const USER_POOL_ID = process.env.USER_POOL_ID!;
const CLIENT_ID = process.env.CLIENT_ID!;

// Cognito JWT検証器（コールドスタート時に1回だけ初期化）
const verifier = CognitoJwtVerifier.create({
  userPoolId: USER_POOL_ID,
  tokenUse: 'id',       // IDトークンを検証
  clientId: CLIENT_ID,
});

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

  // クエリパラメータからトークンを取得
  const queryStringParameters = (event as Record<string, unknown>)
    .queryStringParameters as Record<string, string> | undefined;
  const token = queryStringParameters?.token;

  if (!token) {
    return { statusCode: 401, body: 'Unauthorized: No token' };
  }

  let userId: string;
  try {
    const payload = await verifier.verify(token);
    userId = payload.sub;  // Cognito User PoolのユーザーID
  } catch (error) {
    return { statusCode: 401, body: 'Unauthorized: Invalid token' };
  }

const now = new Date().toISOString();

  // ① CONNECTION#レコード（connectionId → userId の正引き）
  await ddbDocClient.send(
    new PutCommand({
      TableName: TABLE_NAME,
      Item: {
        PK: `CONNECTION#${connectionId}`,
        SK: 'METADATA',
        Type: 'Connection',
        connectionId,
        userId,
        connectedAt: now,
        ttl: Math.floor(Date.now() / 1000) + 86400,
      },
    })
  );

  // ② USER#レコード（userId → connectionId の逆引き）
  await ddbDocClient.send(
    new PutCommand({
      TableName: TABLE_NAME,
      Item: {
        PK: `USER#${userId}`,
        SK: `CONNECTION#${connectionId}`,
        Type: 'UserConnection',
        connectionId,
        connectedAt: now,
        ttl: Math.floor(Date.now() / 1000) + 86400,
      },
    })
  );

  return { statusCode: 200, body: 'Connected' };
};

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

  // CONNECTION#レコードからuserIdを取得
  const queryResult = await ddbDocClient.send(
    new QueryCommand({
      TableName: TABLE_NAME,
      KeyConditionExpression: 'PK = :pk',
      ExpressionAttributeValues: {
        ':pk': `CONNECTION#${connectionId}`,
      },
    })
  );

  const connection = queryResult.Items?.[0];
  if (!connection) {
    return { statusCode: 200, body: 'OK' };
  }

  const userId = connection.userId;

  // 双方向レコードを両方削除
  await ddbDocClient.send(
    new DeleteCommand({
      TableName: TABLE_NAME,
      Key: { PK: `CONNECTION#${connectionId}`, SK: 'METADATA' },
    })
  );

  if (userId) {
    await ddbDocClient.send(
      new DeleteCommand({
        TableName: TABLE_NAME,
        Key: { PK: `USER#${userId}`, SK: `CONNECTION#${connectionId}` },
      })
    );
  }

  return { statusCode: 200, body: 'Disconnected' };
};

1. connectionIdからsenderIdを特定
2. 送信者の本人確認ステータスをチェック
3. チャットルームの存在と参加者検証
4. メッセージの長さバリデーション
5. DynamoDBにメッセージを保存
6. LASTMESSAGEキャッシュを更新
7. 全参加者のアクティブ接続を取得
8. 全接続にメッセージをブロードキャスト
9. 切断済み接続（GoneException）をクリーンアップ

import {
  ApiGatewayManagementApiClient,
  PostToConnectionCommand,
} from '@aws-sdk/client-apigatewaymanagementapi';

// endpointはイベントから動的に構築
const domain = event.requestContext.domainName;
const stage = event.requestContext.stage;

const apiGatewayClient = new ApiGatewayManagementApiClient({
  endpoint: `https://${domain}/${stage}`,
});

// 全参加者のアクティブ接続を収集
const connections: string[] = [];
for (const participantId of participantIds) {
  const userConnections = await ddbDocClient.send(
    new QueryCommand({
      TableName: TABLE_NAME,
      KeyConditionExpression: 'PK = :pk AND begins_with(SK, :sk)',
      ExpressionAttributeValues: {
        ':pk': `USER#${participantId}`,
        ':sk': 'CONNECTION#',
      },
    })
  );

  if (userConnections.Items) {
    connections.push(
      ...userConnections.Items.map((item) => item.connectionId as string)
    );
  }
}

// 全接続に並列送信
const messageData = JSON.stringify({
  type: 'message',
  data: {
    messageId,
    chatRoomId,
    senderId,
    content,
    messageType: 'user',
    createdAt: now,
    timestamp,
  },
});

await Promise.allSettled(
  connections.map(async (connId) => {
    try {
      await apiGatewayClient.send(
        new PostToConnectionCommand({
          ConnectionId: connId,
          Data: Buffer.from(messageData),
        })
      );
    } catch (error: unknown) {
      if (error instanceof Error && error.name === 'GoneException') {
        // 切断済み接続 → レコードをクリーンアップ
        await cleanupStaleConnection(connId);
      }
    }
  })
);

Promise.all:      connA(成功) → connB(失敗) → ❌ 中断（connCに届かない）
Promise.allSettled: connA(成功) → connB(失敗) → connC(成功) ← 全部試行

// GoneException発生時のクリーンアップ
async function cleanupStaleConnection(connId: string) {
  // CONNECTION#レコードからuserIdを取得
  const connRecord = await ddbDocClient.send(
    new GetCommand({
      TableName: TABLE_NAME,
      Key: { PK: `CONNECTION#${connId}`, SK: 'METADATA' },
    })
  );
  const connUserId = connRecord.Item?.userId as string | undefined;

  // 双方向レコードを両方削除
  const deletePromises: Promise<unknown>[] = [
    ddbDocClient.send(
      new DeleteCommand({
        TableName: TABLE_NAME,
        Key: { PK: `CONNECTION#${connId}`, SK: 'METADATA' },
      })
    ),
  ];

  if (connUserId) {
    deletePromises.push(
      ddbDocClient.send(
        new DeleteCommand({
          TableName: TABLE_NAME,
          Key: { PK: `USER#${connUserId}`, SK: `CONNECTION#${connId}` },
        })
      )
    );
  }

  await Promise.all(deletePromises);
}

export const handler: APIGatewayProxyWebsocketHandlerV2 = async (event) => {
  console.log('Default route handler called', event);

  return {
    statusCode: 200,
    body: JSON.stringify({
      message: 'Default route',
      route: event.requestContext.routeKey,
    }),
  };
};

const senderProfile = await ddbDocClient.send(
  new GetCommand({
    TableName: TABLE_NAME,
    Key: { PK: `USER#${senderId}`, SK: 'PROFILE' },
  })
);

if (!senderProfile.Item ||
    senderProfile.Item['verificationStatus'] !== 'approved') {
  // WebSocket経由でエラーメッセージを送信者に返す
  await apiGatewayClient.send(
    new PostToConnectionCommand({
      ConnectionId: connectionId,
      Data: Buffer.from(JSON.stringify({
        type: 'VERIFICATION_REQUIRED',
        message: '本人確認が必要です。',
      })),
    })
  );
  return { statusCode: 403, body: 'Verification required' };
}

const chatRoom = chatRoomQuery.Items?.[0];
const participantIds = chatRoom.participantIds as string[];

if (!participantIds.includes(senderId)) {
  return { statusCode: 403, body: 'Not a participant' };
}

if (content.length > 5000) {
  return { statusCode: 400, body: 'Message too long' };
}