initial: LAN IM desktop app (Electron 32 + Vue 3 + TS + SQLite)

src/main/chat-server.ts

@@ -0,0 +1,116 @@
+// WebSocket 服务端: 监听来自其它设备的连接, 接收消息/ping/recall
+import { WebSocketServer, WebSocket } from 'ws'
+import { EventEmitter } from 'node:events'
+import type { WsFrame, DeviceInfo, MessageEnvelope } from './protocol'
+import { PROTOCOL_VERSION } from './protocol'
+
+export interface ChatServerEvents {
+  connect: (ws: WebSocket, peer: DeviceInfo) => void
+  disconnect: (deviceId: string) => void
+  message: (from: DeviceInfo, msg: MessageEnvelope) => void
+  recall: (from: DeviceInfo, messageId: string) => void
+  typing: (from: DeviceInfo) => void
+}
+
+export class ChatServer extends EventEmitter {
+  private wss: WebSocketServer
+  private port: number
+  // 一个 device 可能多个连接 (重连), 只信任最新 hello 的
+  private sockets = new Map<string, WebSocket>()
+
+  constructor(port: number) {
+    super()
+    this.port = port
+    this.wss = new WebSocketServer({ port, host: '0.0.0.0' })
+  }
+
+  getPort(): number {
+    return this.wss.address() as { port: number } | null as any
+  }
+
+  start(): Promise<number> {
+    return new Promise((resolve, reject) => {
+      this.wss.once('error', reject)
+      this.wss.once('listening', () => {
+        const addr = this.wss.address()
+        const realPort = typeof addr === 'object' && addr ? addr.port : this.port
+        this.port = realPort
+        this.wss.on('connection', (ws, req) => this.onConnection(ws, req))
+        resolve(realPort)
+      })
+    })
+  }
+
+  private onConnection(ws: WebSocket, _req: any) {
+    let peer: DeviceInfo | null = null
+
+    ws.on('message', (raw) => {
+      let frame: WsFrame
+      try {
+        frame = JSON.parse(raw.toString('utf8'))
+      } catch {
+        return
+      }
+      switch (frame.type) {
+        case 'hello':
+          if (frame.from && frame.from.deviceId) {
+            peer = frame.from
+            // 版本/协议不匹配 -> 关闭
+            if ((peer as any).version && (peer as any).version !== PROTOCOL_VERSION) {
+              ws.close(4001, 'protocol mismatch')
+              return
+            }
+            const old = this.sockets.get(peer.deviceId)
+            if (old && old !== ws) {
+              try { old.close(4000, 'replaced') } catch {}
+            }
+            this.sockets.set(peer.deviceId, ws)
+            this.emit('connect', ws, peer)
+          }
+          break
+        case 'ping':
+          this.sendTo(ws, { type: 'pong', ts: frame.ts })
+          break
+        case 'message':
+          if (peer) this.emit('message', peer, frame.payload)
+          break
+        case 'recall':
+          if (peer) this.emit('recall', peer, frame.messageId)
+          break
+        case 'typing':
+          if (peer) this.emit('typing', peer)
+          break
+      }
+    })
+
+    ws.on('close', () => {
+      if (peer) {
+        // 只有当这个 ws 仍是该 device 的代表 socket 才触发 disconnect
+        if (this.sockets.get(peer.deviceId) === ws) {
+          this.sockets.delete(peer.deviceId)
+          this.emit('disconnect', peer.deviceId)
+        }
+      }
+    })
+
+    ws.on('error', () => { /* swallow */ })
+  }
+
+  sendTo(ws: WebSocket, frame: WsFrame) {
+    if (ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify(frame))
+  }
+
+  sendToDevice(deviceId: string, frame: WsFrame): boolean {
+    const ws = this.sockets.get(deviceId)
+    if (!ws) return false
+    this.sendTo(ws, frame)
+    return true
+  }
+
+  async stop() {
+    for (const ws of this.sockets.values()) {
+      try { ws.close(1001, 'server shutdown') } catch {}
+    }
+    await new Promise<void>(r => this.wss.close(() => r()))
+  }
+}