Node.js 中创建一个 WebSocket 服务器

安装依赖

npm install ws

创建websocket服务器

这里提供给一个简单的websocket服务器示例

// websocket-server.ts

import { Server as HTTPServer } from 'http'
import { Server as HTTPSServer } from 'https'
import { WebSocketServer, WebSocket } from 'ws'
import { parse } from 'url'

declare module 'ws' {
    interface WebSocket {
        isAlive?: boolean
        token?: string
    }
}

let wss: WebSocketServer | null = null

// 心跳间隔（单位：毫秒）
const HEARTBEAT_INTERVAL = 30000

/**
 * 初始化 WebSocket 服务
 */
export function initWebSocketServer(server: HTTPServer | HTTPSServer) {
    wss = new WebSocketServer({ server, path: '/ws' })

    wss.on('connection', async (ws: WebSocket, request: any) => {
        console.log('[wss] : 客户端已连接')

        const { query: { token } } = parse(request.url || '', true)

        if (!token) {
            console.log('[wss] : 无 token，断开连接')
            ws.close() // 无 token 时关闭连接
            return
        }

        // 将 token 存储到 WebSocket 连接中
        ws.token = token as string
        ws.isAlive = true
        console.log('[wss] : 客户端连接成功，token:', token)

        ws.on('pong', () => ws.isAlive = true)

        ws.on('message', (data) => {
            const message = data.toString()
            console.log('[wss] : 收到消息', message)
            ws.send(JSON.stringify({ value: `[wss] 回应消息：${message}` }))
        })

        ws.on('close', () => {
            console.log('[wss] : 客户端断开连接')
        })

        ws.send(JSON.stringify({ value: '你已经成功连接 WebSocket' }))
        broadcast({ value: `当前在线人数 ${wss!.clients.size}` })
    })

    // 启动心跳检测
    const interval = setInterval(() => {
        if (!wss) return
        wss.clients.forEach((ws) => {
            const client = ws as WebSocket & { isAlive?: boolean }

            if (client.isAlive === false) {
                console.log('[wss] : 发现无响应客户端，断开连接')
                return client.terminate()
            }

            client.isAlive = false
            client.ping()
        })
    }, HEARTBEAT_INTERVAL)

    // 在服务关闭时清理心跳定时器
    wss.on('close', () => {
        clearInterval(interval)
    })
}

/** 广播消息到所有客户端 */
export function broadcast(message: any) {
    if (!wss) {
        console.log('[wss] : 广播失败，wss 未初始化')
        return
    }

    const count = Array.from(wss.clients).filter(c => c.readyState === WebSocket.OPEN).length
    console.log('[wss] : 当前连接数（OPEN 状态）->', count)

    const stringifiedMsg = JSON.stringify(message)
    wss.clients.forEach((client) => {
        if (client.readyState === WebSocket.OPEN) {
            console.log('[wss] : 广播 ->', stringifiedMsg)
            client.send(stringifiedMsg)
        }
    })
}


/** 发送到特定客户端 */
export function deliver(token: string, message: any) {
    if (!wss) return

    const msg = JSON.stringify(message)

    let count = 0

    for (const client of wss.clients) {
        if (client.readyState === WebSocket.OPEN && client.token === token) {
            client.send(msg)
            count++
        }
    }

    console.log(`[wss] : 向 token(${token}) 推送 ${count} 个连接`)
}