Node.js WebSocket 实时通信全栈实践

const WebSocket = require('ws');
const wss = new WebSocket.Server({ port: 8080 });

wss.on('connection', (ws) => {
    // 每 30 秒发送一次 Ping
    const pingInterval = setInterval(() => {
        if (ws.readyState === WebSocket.OPEN) {
            ws.ping(Date.now().toString());
        }
    }, 30000);

    ws.on('pong', (data) => {
        const pingTime = parseInt(data.toString());
        console.log(`RTT: ${Date.now() - pingTime}ms`);
    });

    ws.on('close', () => clearInterval(pingInterval));
});

const WebSocket = require('ws');

// 服务端
const wss = new WebSocket.Server({ port: 8080 });
const clients = new Set();

wss.on('connection', (ws, req) => {
    clients.add(ws);
    console.log(`Connected: ${req.socket.remoteAddress}`);

    ws.on('message', (data) => {
        // 广播给所有客户端
        for (const client of clients) {
            if (client.readyState === WebSocket.OPEN) {
                client.send(data.toString());
            }
        }
    });

    ws.on('close', () => clients.delete(ws));
});

// 客户端
const client = new WebSocket('ws://localhost:8080');
client.on('open', () => client.send('Hello Server'));
client.on('message', (data) => console.log('Received:', data.toString()));

// 服务端
const { Server } = require('socket.io');
const io = new Server(3000, { cors: { origin: '*' } });

io.on('connection', (socket) => {
    console.log('User connected:', socket.id);

    socket.on('chat message', (msg) => {
        io.emit('chat message', { user: socket.id, text: msg });
    });

    socket.on('disconnect', () => {
        console.log('User disconnected:', socket.id);
    });
});

// 客户端
import { io } from 'socket.io-client';
const socket = io('http://localhost:3000');

socket.on('connect', () => console.log('Connected:', socket.id));
socket.on('chat message', (data) => console.log(data.text));
socket.emit('chat message', 'Hello World');

io.on('connection', (socket) => {
    // 加入房间
    socket.join('room-42');
    socket.join('room-premium');

    // 向房间广播（不包含发送者）
    socket.to('room-42').emit('message', 'Hello room!');

    // 向房间广播（包含发送者）
    io.to('room-42').emit('announcement', 'Server update');

    // 向多个房间广播
    socket.to('room-42').to('room-43').emit('message', 'Multi-room');

    // 离开房间
    socket.leave('room-42');

    // 获取房间列表
    const rooms = Array.from(socket.rooms);
    console.log('My rooms:', rooms);  // 包含默认的 socket.id 房间

    // 获取房间内所有 Socket ID
    const roomSockets = io.sockets.adapter.rooms.get('room-42');
    console.log('Room size:', roomSockets?.size || 0);
});

// /chat 命名空间
const chatNs = io.of('/chat');
chatNs.use(authMiddleware);  // Namespace 级别中间件

chatNs.on('connection', (socket) => {
    socket.on('send', (msg) => {
        chatNs.emit('broadcast', msg);
    });
});

// /admin 命名空间
const adminNs = io.of('/admin');
adminNs.use(adminAuthMiddleware);
adminNs.on('connection', (socket) => {
    socket.on('notify-all', (msg) => {
        io.emit('global-notify', msg);  // 跨 Namespace 广播需使用根 io
    });
});

// 客户端连接指定 Namespace
const chatSocket = io('http://localhost:3000/chat');
const adminSocket = io('http://localhost:3000/admin');

const { createServer } = require('http');
const { Server } = require('socket.io');
const { createAdapter } = require('@socket.io/redis-adapter');
const { createClient } = require('redis');
const express = require('express');

const app = express();
const server = createServer(app);
const io = new Server(server, {
    cors: { origin: process.env.CLIENT_ORIGIN || 'http://localhost:5173' },
});

// 创建 Redis Pub/Sub 客户端
async function setupAdapter() {
    const pubClient = createClient({ url: process.env.REDIS_URL || 'redis://localhost:6379' });
    const subClient = pubClient.duplicate();

    await Promise.all([pubClient.connect(), subClient.connect()]);
    io.adapter(createAdapter(pubClient, subClient));
    console.log('Redis adapter connected');
}

setupAdapter().catch(console.error);

// 房间管理中间件
io.use(async (socket, next) => {
    const token = socket.handshake.auth.token;
    try {
        const user = await validateToken(token);  // JWT 验证
        socket.user = user;
        next();
    } catch (err) {
        next(new Error('Authentication error'));
    }
});

io.on('connection', (socket) => {
    console.log(`[${process.pid}] User connected: ${socket.user.id}`);

    // 加入用户专属房间（用于一对一消息）
    socket.join(`user:${socket.user.id}`);

    // 加入公开聊天室
    socket.on('join-room', (roomId) => {
        socket.join(roomId);
        socket.to(roomId).emit('user-joined', {
            userId: socket.user.id,
            username: socket.user.name,
        });
    });

    // 发送消息
    socket.on('send-message', ({ roomId, content }) => {
        const message = {
            id: generateId(),
            roomId,
            userId: socket.user.id,
            username: socket.user.name,
            content,
            timestamp: new Date().toISOString(),
        };

        // 持久化到数据库（异步，不阻塞广播）
        saveMessage(message).catch(console.error);

        // 广播给房间内所有人（跨服务器）
        io.to(roomId).emit('new-message', message);
    });

    // 一对一私聊
    socket.on('private-message', ({ toUserId, content }) => {
        const message = {
            id: generateId(),
            from: socket.user.id,
            to: toUserId,
            content,
            timestamp: new Date().toISOString(),
        };

        // 目标用户可能在任意服务器上，Redis Adapter 确保送达
        io.to(`user:${toUserId}`).emit('private-message', message);
        savePrivateMessage(message).catch(console.error);
    });

    // 离开房间
    socket.on('leave-room', (roomId) => {
        socket.leave(roomId);
        socket.to(roomId).emit('user-left', { userId: socket.user.id });
    });

    socket.on('disconnect', () => {
        console.log(`[${process.pid}] User disconnected: ${socket.user.id}`);
    });
});

// 获取房间成员（跨服务器接口）
app.get('/rooms/:roomId/members', async (req, res) => {
    const sockets = await io.in(req.params.roomId).fetchSockets();
    const members = sockets.map((s) => ({
        id: s.user.id,
        name: s.user.name,
    }));
    res.json({ count: members.length, members });
});

const PORT = process.env.PORT || 3000;
server.listen(PORT, () => console.log(`Server ${process.pid} on port ${PORT}`));

upstream websocket_backend {
    ip_hash;  # 或 hash $cookie_io  consistent;
    server 127.0.0.1:3001;
    server 127.0.0.1:3002;
    server 127.0.0.1:3003;
}

server {
    listen 80;
    location /socket.io/ {
        proxy_pass http://websocket_backend;
        proxy_http_version 1.1;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection "upgrade";
        proxy_set_header Host $host;
        proxy_set_header X-Real-IP $remote_addr;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
        proxy_read_timeout 86400;
    }
}

// 服务端中间件
io.use(async (socket, next) => {
    const token = socket.handshake.auth.token || socket.handshake.query.token;

    if (!token) {
        return next(new Error('Missing token'));
    }

    try {
        const decoded = jwt.verify(token, process.env.JWT_SECRET);
        const user = await User.findById(decoded.userId);

        if (!user || !user.isActive) {
            return next(new Error('Invalid user'));
        }

        socket.user = user;
        next();
    } catch (err) {
        next(new Error('Authentication failed'));
    }
});

// 客户端
const socket = io('http://localhost:3000', {
    auth: { token: localStorage.getItem('accessToken') }
});

// 连接错误处理
socket.on('connect_error', async (err) => {
    if (err.message === 'Authentication failed') {
        // 尝试刷新 Token
        const newToken = await refreshAccessToken();
        localStorage.setItem('accessToken', newToken);

        // 使用新 Token 重连
        socket.auth.token = newToken;
        socket.connect();
    }
});

// 服务端主动断开过期连接后刷新
socket.on('disconnect', (reason) => {
    if (reason === 'io server disconnect') {
        // 服务端强制断开，尝试重连并刷新令牌
        refreshAndReconnect();
    }
});

const session = require('express-session');
const sessionMiddleware = session({
    store: new RedisStore({ client: redisClient }),
    secret: process.env.SESSION_SECRET,
    resave: false,
    saveUninitialized: false,
});

// Socket.IO 共享 Express Session
io.use((socket, next) => {
    sessionMiddleware(socket.request, socket.request.res || {}, next);
});

io.use((socket, next) => {
    if (socket.request.session?.userId) {
        socket.userId = socket.request.session.userId;
        next();
    } else {
        next(new Error('Unauthorized'));
    }
});

const io = new Server(server, {
    pingInterval: 10000,  // 10 秒一次心跳
    pingTimeout: 5000,    // 5 秒未收到 Pong 则断开
});

const socket = io('http://localhost:3000', {
    reconnection: true,
    reconnectionAttempts: 10,        // 最大重连次数
    reconnectionDelay: 1000,          // 初始延迟 1 秒
    reconnectionDelayMax: 5000,       // 最大延迟 5 秒
    randomizationFactor: 0.5,         // 抖动因子
});

// 监听重连事件
socket.on('reconnect', (attempt) => {
    console.log(`Reconnected after ${attempt} attempts`);
    // 重新加入房间、同步未读消息
    socket.emit('join-room', currentRoomId);
    socket.emit('sync-messages', { roomId: currentRoomId, since: lastMessageTime });
});

socket.on('reconnect_error', (err) => {
    console.error('Reconnection failed:', err.message);
});

socket.on('reconnect_failed', () => {
    // 所有重连尝试均失败
    showOfflineMode();
});

const rateLimit = require('ws-rate-limit')('40s', 20);  // 40 秒内最多 20 条

// 或使用自定义实现
class RateLimiter {
    constructor(maxMessages = 20, windowMs = 40000) {
        this.max = maxMessages;
        this.window = windowMs;
        this.clients = new Map();
    }

    isAllowed(clientId) {
        const now = Date.now();
        const history = this.clients.get(clientId) || [];

        // 清理过期记录
        const valid = history.filter((t) => now - t < this.window);
        this.clients.set(clientId, valid);

        if (valid.length >= this.max) {
            return false;
        }

        valid.push(now);
        return true;
    }

    reset(clientId) {
        this.clients.delete(clientId);
    }
}

const limiter = new RateLimiter(30, 60000);  // 每分钟 30 条

io.on('connection', (socket) => {
    socket.on('send-message', (data) => {
        if (!limiter.isAllowed(socket.user.id)) {
            socket.emit('error', { code: 'RATE_LIMITED', message: '消息发送过于频繁' });
            return;
        }
        // 处理消息...
    });
});

// 每 IP 最大连接数
const connectionsPerIP = new Map();

io.on('connection', (socket) => {
    const ip = socket.handshake.address || socket.request.socket.remoteAddress;
    const count = (connectionsPerIP.get(ip) || 0) + 1;

    if (count > 5) {
        socket.disconnect(true);
        return;
    }

    connectionsPerIP.set(ip, count);

    socket.on('disconnect', () => {
        const current = connectionsPerIP.get(ip) || 1;
        if (current <= 1) connectionsPerIP.delete(ip);
        else connectionsPerIP.set(ip, current - 1);
    });
});

const Ajv = require('ajv');
const ajv = new Ajv();

const messageSchema = {
    type: 'object',
    properties: {
        roomId: { type: 'string', maxLength: 50 },
        content: { type: 'string', minLength: 1, maxLength: 2000 },
    },
    required: ['roomId', 'content'],
    additionalProperties: false,
};

const validate = ajv.compile(messageSchema);

socket.on('send-message', (data) => {
    if (!validate(data)) {
        socket.emit('validation-error', validate.errors);
        return;
    }
    // 业务处理...
});

const redis = require('ioredis');
const redisClient = new Redis(process.env.REDIS_URL);

async function queueOfflineMessage(userId, message) {
    const key = `offline:${userId}`;
    await redisClient.rpush(key, JSON.stringify({
        ...message,
        queuedAt: new Date().toISOString(),
    }));
    // 设置过期时间（30 天）
    await redisClient.expire(key, 60 * 60 * 24 * 30);
}

async function getOfflineMessages(userId) {
    const key = `offline:${userId}`;
    const messages = await redisClient.lrange(key, 0, -1);
    await redisClient.del(key);  // 消费后删除
    return messages.map((m) => JSON.parse(m));
}

// 发送消息时检查用户在线状态
async function sendOrQueueMessage(toUserId, message) {
    const sockets = await io.in(`user:${toUserId}`).fetchSockets();

    if (sockets.length > 0) {
        // 用户在线，直接投递
        io.to(`user:${toUserId}`).emit('private-message', message);
    } else {
        // 用户离线，入队
        await queueOfflineMessage(toUserId, message);
    }
}

// 用户上线时拉取离线消息
io.on('connection', async (socket) => {
    const offlineMessages = await getOfflineMessages(socket.user.id);
    if (offlineMessages.length > 0) {
        socket.emit('offline-messages', offlineMessages);
    }
});

// 客户端收到消息后发送 ack
socket.on('private-message', async (msg) => {
    displayMessage(msg);
    socket.emit('ack', { messageId: msg.id });
});

// 服务端记录已送达
socket.on('ack', async ({ messageId }) => {
    await markAsDelivered(messageId);
});

// 投递时防止重复（幂等键）
async function deliverOnce(messageId, userId, payload) {
    const lockKey = `deliver:${messageId}:${userId}`;
    const locked = await redisClient.set(lockKey, '1', 'EX', 3600, 'NX');
    if (locked === 'OK') {
        io.to(`user:${userId}`).emit('private-message', payload);
    }
}

// 通知服务模块
class NotificationService {
    constructor(io) {
        this.io = io;
    }

    async notify(userIds, notification) {
        const message = {
            id: generateId(),
            type: notification.type,
            title: notification.title,
            body: notification.body,
            data: notification.data,
            createdAt: new Date().toISOString(),
            read: false,
        };

        // 持久化
        await Notification.insertMany(
            userIds.map((userId) => ({ ...message, userId }))
        );

        // 实时推送
        for (const userId of userIds) {
            const sockets = await this.io.in(`user:${userId}`).fetchSockets();
            if (sockets.length > 0) {
                this.io.to(`user:${userId}`).emit('notification', message);
            } else {
                // 离线：推送至手机（FCM/APNs）或邮件队列
                await pushToMobile(userId, message);
            }
        }
    }

    // 广播全局公告
    async broadcastAnnouncement(announcement) {
        await GlobalAnnouncement.create(announcement);
        this.io.emit('announcement', announcement);
    }
}

// 简化版协同编辑：操作序列广播
const ot = require('ot.js');

io.of('/editor').on('connection', (socket) => {
    socket.on('join-doc', (docId) => {
        socket.join(`doc:${docId}`);

        // 发送当前文档状态
        const snapshot = getDocumentSnapshot(docId);
        socket.emit('doc-state', snapshot);
    });

    socket.on('operation', ({ docId, operation, revision }) => {
        // OT 变换：将用户操作与服务端历史合并
        const transformed = ot.type.transform(operation, getHistorySince(docId, revision));

        // 应用操作
        applyOperation(docId, transformed);

        // 广播变换后的操作（其他客户端据此更新）
        socket.to(`doc:${docId}`).emit('operation', {
            operation: transformed,
            authorId: socket.user.id,
        });
    });

    // 光标位置同步（非关键操作，使用 volatile）
    socket.on('cursor-move', ({ docId, position }) => {
        socket.volatile.to(`doc:${docId}`).emit('cursor-update', {
            userId: socket.user.id,
            username: socket.user.name,
            position,
        });
    });
});

// Lambda 处理 WebSocket 事件
exports.handler = async (event) => {
    const { routeKey, connectionId, body } = event.requestContext;

    switch (routeKey) {
        case '$connect':
            // 连接建立，存储 connectionId 与 userId 映射
            await DynamoDB.put({
                TableName: 'Connections',
                Item: {
                    connectionId,
                    userId: event.queryStringParameters.userId,
                    connectedAt: new Date().toISOString(),
                },
            });
            return { statusCode: 200 };

        case '$disconnect':
            await DynamoDB.delete({
                TableName: 'Connections',
                Key: { connectionId },
            });
            return { statusCode: 200 };

        case 'sendMessage':
            const { roomId, content } = JSON.parse(body);
            // 查询房间所有 connectionId
            const connections = await getRoomConnections(roomId);

            // 向所有连接推送消息
            await Promise.all(connections.map((conn) =>
                postToConnection(conn.connectionId, { roomId, content })
            ));
            return { statusCode: 200 };

        default:
            return { statusCode: 400 };
    }
};

// 使用 AWS SDK 向指定连接发送消息
const { ApiGatewayManagementApiClient, PostToConnectionCommand } = require('@aws-sdk/client-apigatewaymanagementapi');

const client = new ApiGatewayManagementApiClient({
    endpoint: process.env.WEBSOCKET_ENDPOINT,
});

async function postToConnection(connectionId, data) {
    await client.send(new PostToConnectionCommand({
        ConnectionId: connectionId,
        Data: Buffer.from(JSON.stringify(data)),
    }));
}

// 服务端 SSE
app.get('/events', (req, res) => {
    res.setHeader('Content-Type', 'text/event-stream');
    res.setHeader('Cache-Control', 'no-cache');
    res.setHeader('Connection', 'keep-alive');

    const clientId = req.query.clientId;
    const send = (data) => res.write(`data: ${JSON.stringify(data)}\n\n`);

    eventEmitter.on(`update:${clientId}`, send);

    req.on('close', () => {
        eventEmitter.off(`update:${clientId}`, send);
    });
});

// 客户端
const es = new EventSource('/events?clientId=123');
es.onmessage = (e) => console.log(JSON.parse(e.data));