LocalNetMsg/src/main/chat-server.ts

// 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
  read: (from: DeviceInfo, upToMessageId: string) => void
  // 任意 ws 帧 (含 terminal/forward/usb 等), 让上层路由分发
  frame: (from: DeviceInfo, frame: WsFrame) => 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' })
    // 连接处理必须在这里挂: ws 在构造时就 listen, 'listening' 事件可能早于 start() 挂监听而丢失,
    // 那种情况下 start() 走 "已在监听" 分支, 不会再走到下面的 on('connection')
    this.wss.on('connection', (ws, req) => this.onConnection(ws, req))
  }

  getPort(): number {
    const addr = this.wss.address()
    return typeof addr === 'object' && addr ? addr.port : this.port
  }

  start(): Promise<number> {
    return new Promise((resolve, reject) => {
      // ws 的 WebSocketServer 在构造时就 listen 了, 'listening' 事件可能早于本方法挂监听而丢失
      // (例如同一 tick 里先 await 了别的东西再调 start()) —— 这里先看是否已经在监听
      const addr = this.wss.address()
      if (addr && typeof addr === 'object' && addr.port) {
        this.port = addr.port
        resolve(addr.port)
        return
      }
      this.wss.once('error', reject)
      this.wss.once('listening', () => {
        const a = this.wss.address()
        const realPort = typeof a === 'object' && a ? a.port : this.port
        this.port = realPort
        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
      }
      // 通用: 任何帧都抛 'frame', 上层 (ipc.bindNetworkContext) 自行路由
      if (peer) this.emit('frame', peer, frame)
      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 'read':
          if (peer) this.emit('read', peer, frame.upTo)
          break
        case 'typing':
          if (peer) this.emit('typing', peer)
          break
        // terminal / forward / usb 由 'frame' 事件统一交给上层处理
      }
    })

    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
  }

  // 该设备是否有活着的入站连接 (对方主动连过来的 ws)
  // 与 chat-client 的 outgoing 相对: 任何一条通, 消息通道就是通的
  isConnected(deviceId: string): boolean {
    const ws = this.sockets.get(deviceId)
    return !!ws && ws.readyState === WebSocket.OPEN
  }

  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()))
  }
}