Improving the Chat Application

import {
  DynamoDBClient,
} from "@aws-sdk/client-dynamodb";
import {
  DynamoDBDocumentClient,
  PutCommand,
  GetCommand,
  DeleteCommand,
} from "@aws-sdk/lib-dynamodb";
import {
  ApiGatewayManagementApiClient,
  PostToConnectionCommand,
} from "@aws-sdk/client-apigatewaymanagementapi";

// Initialize DynamoDB Document Client
const ddbClient = new DynamoDBClient({});
const ddb = DynamoDBDocumentClient.from(ddbClient);

// DynamoDB Table Name
const TABLE_NAME = "WebSocketConnections";

/**
 * When connected to the WebSocket, add records in DynamoDB
 */
const onConnect = async (event) => {
  const connectionId = event.requestContext.connectionId;
  const userName = event.headers?.msgfrom || "unknown";
  const friendName = event.headers?.msgto || "unknown";

  // Store connection details
  await ddb.send(
    new PutCommand({
      TableName: TABLE_NAME,
      Item: {
        context: "Connection",
        id: connectionId,
        userName,
        friendName,
      },
    })
  );

  // Store user-to-connection mapping
  await ddb.send(
    new PutCommand({
      TableName: TABLE_NAME,
      Item: {
        context: "User",
        id: userName,
        friendName,
        connectionId,
      },
    })
  );
};

/**
 * When a user disconnects, remove connection records
 */
const onDisconnect = async (event) => {
  const connectionId = event.requestContext.connectionId;

  await ddb.send(
    new DeleteCommand({
      TableName: TABLE_NAME,
      Key: {
        context: "Connection",
        id: connectionId,
      },
    })
  );
};

/**
 * On receiving a message, forward it to the friend
 */
const onMessage = async (event) => {
  const connectionId = event.requestContext.connectionId;
  const apiId = event.requestContext.apiId;

  // Get the user's connection record
  const userResponse = await ddb.send(
    new GetCommand({
      TableName: TABLE_NAME,
      Key: {
        context: "Connection",
        id: connectionId,
      },
    })
  );

  const friendName = userResponse.Item?.friendName;
  if (!friendName) {
    console.log("Friend not found for connection:", connectionId);
    return;
  }

  // Get the friend's connection record
  const friendResponse = await ddb.send(
    new GetCommand({
      TableName: TABLE_NAME,
      Key: {
        context: "User",
        id: friendName,
      },
    })
  );

  const apigwManagementApi = new ApiGatewayManagementApiClient({
    endpoint: `https://${apiId}.execute-api.us-east-1.amazonaws.com/v1`,
  });

  const friendConnectionId = friendResponse.Item?.connectionId;

  try {
    // Send message to friend
    await apigwManagementApi.send(
      new PostToConnectionCommand({
        ConnectionId: friendConnectionId,
        Data: event.body || "",
      })
    );
  } catch (err) {
    console.log("Error posting to connection:", err);

    // Notify sender if friend is not reachable
    await apigwManagementApi.send(
      new PostToConnectionCommand({
        ConnectionId: connectionId,
        Data: "Your friend is not reachable. Connect from the other terminal and try again.",
      })
    );
  }
};

/**
 * Lambda handler
 */
export const handler = async (event) => {
  const routeKey = event.requestContext?.routeKey;

  if (routeKey === "$connect") {
    await onConnect(event);
  } else if (routeKey === "$disconnect") {
    await onDisconnect(event);
  } else if (routeKey === "$default") {
    await onMessage(event);
  }

  return { statusCode: 200, body: event.body || "{}" };
};