122 lines
2.4 KiB
JavaScript
122 lines
2.4 KiB
JavaScript
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; |