添加websocket连接

websocket/index.js

@@ -0,0 +1,122 @@
+const WebSocket = require("ws");
+const url = require("url");
+
+const { verifyToken } = require("../utils/jwt");
+
+const {
+  addConnection,
+  getConnection,
+  removeConnection,
+} = require("./connectionManager");
+
+// 创建WebSocket服务
+function initWebSocket(server) {
+  const wss = new WebSocket.Server({
+    server,
+  });
+
+  wss.on("connection", (ws, request) => {
+    console.log("有新的WebSocket连接");
+
+    // 获取URL中的token
+    const query = url.parse(request.url, true).query;
+    const token = query.token;
+
+    if (!token) {
+      console.log("没有提供token");
+
+      ws.close();
+      return;
+    }
+
+    // 验证JWT
+    let user;
+
+    try {
+      user = verifyToken(token);
+    } catch (error) {
+      console.log("JWT验证失败：", error.message);
+
+      ws.close();
+      return;
+    }
+
+    console.log("WebSocket用户：", user);
+
+    const username = user.username;
+
+    // 保存连接
+    addConnection(username, ws);
+
+    // 接收客户端消息
+    ws.on("message", (message) => {
+      try {
+        const data = JSON.parse(message.toString());
+
+        console.log("收到WebSocket消息：", data);
+
+        handleMessage(username, data);
+      } catch (error) {
+        console.error("WebSocket消息处理失败：", error);
+      }
+    });
+
+    // 连接关闭
+    ws.on("close", () => {
+      removeConnection(username);
+    });
+
+    // WebSocket错误
+    ws.on("error", (error) => {
+      console.error(`用户 ${username} WebSocket错误：`, error);
+    });
+  });
+
+  console.log("WebSocket服务已启动");
+
+  return wss;
+}
+
+// 处理WebSocket消息
+function handleMessage(username, data) {
+  switch (data.type) {
+    case "sendMessage":
+      handleSendMessage(username, data);
+      break;
+
+    default:
+      console.log("未知消息类型：", data.type);
+  }
+}
+
+// 处理发送消息
+function handleSendMessage(senderName, data) {
+  const receiverName = data.receiverName;
+  const content = data.content;
+
+  console.log(
+    `${senderName} → ${receiverName}：${content}`
+  );
+
+  // 找到接收者的WebSocket
+  const receiverWs = getConnection(receiverName);
+
+  if (!receiverWs) {
+    console.log(`${receiverName} 当前不在线`);
+    return;
+  }
+
+  // 推送消息
+  receiverWs.send(
+    JSON.stringify({
+      type: "newMessage",
+      message: {
+        senderName,
+        receiverName,
+        content,
+      },
+    })
+  );
+}
+
+module.exports = initWebSocket;