import { WebSocket, WebSocketServer } from 'ws'; import urlParse from 'url'; import { IncomingMessage } from 'node:http'; import { log } from '@/common/utils/log'; class WebsocketClientBase { private wsClient: WebSocket | undefined; constructor() { } send(msg: string) { if (this.wsClient && this.wsClient.readyState == WebSocket.OPEN) { this.wsClient.send(msg); } } onMessage(msg: string) { } } export class WebsocketServerBase { private ws: WebSocketServer | null = null; public token: string = ''; constructor() { } start(port: number) { try { this.ws = new WebSocketServer({ port }); } catch (e: any) { throw Error('ws服务启动失败, ' + e.toString()); } this.ws.on('connection', (wsClient, req) => { const url: string = req.url!.split('?').shift() || '/'; this.authorize(wsClient, req); this.onConnect(wsClient, url, req); wsClient.on('message', async (msg) => { this.onMessage(wsClient, url, msg.toString()); }); }); } stop() { this.ws && this.ws.close((err) => { log('ws server close failed!', err); }); this.ws = null; } restart(port: number) { this.stop(); this.start(port); } authorize(wsClient: WebSocket, req: IncomingMessage) { const url = req.url!.split('?').shift(); log('ws connect', url); let clientToken: string = ''; const authHeader = req.headers['authorization']; if (authHeader) { clientToken = authHeader.split('Bearer ').pop() || ''; log('receive ws header token', clientToken); } else { const parsedUrl = urlParse.parse(req.url || '/', true); const urlToken = parsedUrl.query.access_token; if (urlToken) { if (Array.isArray(urlToken)) { clientToken = urlToken[0]; } else { clientToken = urlToken; } log('receive ws url token', clientToken); } } if (this.token && clientToken != this.token) { this.authorizeFailed(wsClient); return wsClient.close(); } } authorizeFailed(wsClient: WebSocket) { } onConnect(wsClient: WebSocket, url: string, req: IncomingMessage) { } onMessage(wsClient: WebSocket, url: string, msg: string) { } sendHeart() { } }