cnb/src/workspace/keep-live.ts

// WebSocket Keep-Alive Client Library
// 运行时检测：Bun 使用原生 WebSocket，Node.js 使用 ws 库
let WebSocketModule: any;

if (typeof Bun !== 'undefined') {
  // Bun 环境：使用原生 WebSocket
  WebSocketModule = { WebSocket: globalThis.WebSocket };
} else {
  // Node.js 环境：使用 ws 库
  WebSocketModule = await import('ws');
}

const WebSocket = WebSocketModule.WebSocket;

export interface KeepAliveConfig {
  wsUrl: string;
  cookie: string;
  reconnectInterval?: number;
  maxReconnectAttempts?: number;
  pingInterval?: number;
  onMessage?: (data: Buffer | string) => void;
  onConnect?: () => void;
  onDisconnect?: (code: number) => void;
  onError?: (error: Error) => void;
  onSign?: (data: { type: string; data: string; signedData: string }) => void;
  onExit?: (code?: number) => void;
  debug?: boolean;
}

export interface ParsedMessage {
  type: string;
  raw: Buffer;
  payload?: any;
}

type MessageHandler = (msg: ParsedMessage) => void;

export class WSKeepAlive {
  private ws: WebSocket | null = null;
  private config: Required<KeepAliveConfig>;
  private reconnectAttempts = 0;
  private pingTimer: NodeJS.Timeout | null = null;
  private messageHandlers: Set<MessageHandler> = new Set();
  private url: URL;
  private readonly isBun: boolean;

  constructor(config: KeepAliveConfig) {
    this.config = {
      wsUrl: config.wsUrl,
      cookie: config.cookie,
      reconnectInterval: config.reconnectInterval ?? 5000,
      maxReconnectAttempts: config.maxReconnectAttempts ?? 3,
      pingInterval: config.pingInterval ?? 30000,
      onMessage: config.onMessage ?? (() => { }),
      onConnect: config.onConnect ?? (() => { }),
      onDisconnect: config.onDisconnect ?? (() => { }),
      onError: config.onError ?? (() => { }),
      onSign: config.onSign ?? (() => { }),
      onExit: config.onExit ?? (() => { }),
      debug: config.debug ?? false,
    };
    this.url = new URL(this.config.wsUrl);
    this.isBun = typeof Bun !== 'undefined';
  }

  private log(message: string) {
    if (!this.config.debug) return;
    const timestamp = new Date().toISOString();
    const msg = `[${timestamp}] ${message}`;
    console.log(msg);
  }

  private parseMessage(data: Buffer): ParsedMessage | null {
    const result: ParsedMessage = { type: "unknown", raw: data };

    if (data.length < 14) {
      result.type = "raw";
      return result;
    }

    const prefix = data.slice(0, 13);
    const msgType = prefix[0];
    const jsonStart = data.indexOf(0x71); // 0x71 = 'q'

    if (jsonStart !== -1) {
      try {
        const jsonStr = data.slice(jsonStart + 1).toString();
        const payload = JSON.parse(jsonStr);
        result.type = `binary(0x${msgType.toString(16)})`;
        result.payload = payload;

        // 特殊处理 sign 类型
        if (payload.type === "sign" && this.config.onSign) {
          this.config.onSign(payload);
        }
        return result;
      } catch {
        result.type = "binary(json-parse-error)";
        return result;
      }
    }

    result.type = "raw";
    return result;
  }

  connect() {
    const { wsUrl, cookie, debug } = this.config;
    this.log(`Connecting to ${wsUrl}...`);

    this.ws = new WebSocket(wsUrl, {
      headers: {
        "Origin": this.url.origin,
        "Cookie": cookie,
        "Cache-Control": "no-cache",
        "Accept-Language": "zh-CN,zh;q=0.9",
        "Pragma": "no-cache",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        "Sec-WebSocket-Extensions": "permessage-deflate",
      }
    });

    if (this.isBun) {
      // Bun 环境：使用标准 Web API
      const ws = this.ws as any;
      ws.onopen = () => {
        debug && this.log("Connected!");
        this.reconnectAttempts = 0;
        this.config.onConnect();
        this.startPing();
      };

      ws.onmessage = async (event: MessageEvent) => {
        let data: Buffer | string;

        if (event.data instanceof Blob) {
          data = Buffer.from(await event.data.arrayBuffer());
        } else if (event.data instanceof ArrayBuffer) {
          data = Buffer.from(event.data);
        } else if (typeof event.data === 'string') {
          data = event.data;
        } else {
          data = Buffer.from(event.data);
        }

        this.handleMessage(data);
      };

      ws.onclose = (event: CloseEvent) => {
        debug && this.log(`Disconnected (code: ${event.code})`);
        this.stopPing();
        this.config.onDisconnect(event.code);
        this.handleReconnect();
      };

      ws.onerror = (event: Event) => {
        debug && this.log(`Error: ${event}`);
        this.config.onError(new Error("WebSocket error"));
      };
    } else {
      // Node.js (ws 库)：使用 EventEmitter 模式
      const ws = this.ws as any;
      ws.on("open", () => {
        debug && this.log("Connected!");
        this.reconnectAttempts = 0;
        this.config.onConnect();
        this.startPing();
      });

      ws.on("message", (data: any) => {
        this.handleMessage(data);
      });

      ws.on("close", (code: number) => {
        debug && this.log(`Disconnected (code: ${code})`);
        this.stopPing();
        this.config.onDisconnect(code);
        this.handleReconnect();
      });

      ws.on("error", (err: Error) => {
        debug && this.log(`Error: ${err.message}`);
        this.config.onError(err);
      });
    }
  }

  // 统一的消息处理方法
  private handleMessage(data: Buffer | string) {
    if (Buffer.isBuffer(data)) {
      const parsed = this.parseMessage(data);
      this.config.onMessage(parsed?.raw ?? data);

      this.messageHandlers.forEach(handler => {
        if (parsed) handler(parsed);
      });
    } else {
      this.config.onMessage(data);
    }
  }

  private startPing() {
    this.stopPing();
    this.pingTimer = setInterval(() => {
      if (this.ws && this.ws.readyState === WebSocket.OPEN) {
        // 使用 JSON 格式的 ping 消息，兼容 Bun 和 Node.js
        this.ws.send(JSON.stringify({ type: "ping", timestamp: Date.now() }));
        this.log("Sent ping");
      }
    }, this.config.pingInterval);
  }

  private stopPing() {
    if (this.pingTimer) {
      clearInterval(this.pingTimer);
      this.pingTimer = null;
    }
  }

  private handleReconnect() {
    if (this.reconnectAttempts >= this.config.maxReconnectAttempts) {
      this.log(`Max reconnect attempts (${this.config.maxReconnectAttempts}) reached. Giving up.`);
      this.config.onExit(1);
      return;
    }
    this.reconnectAttempts++;
    this.log(`Reconnecting in ${this.config.reconnectInterval}ms... (attempt ${this.reconnectAttempts}/${this.config.maxReconnectAttempts})`);
    setTimeout(() => this.connect(), this.config.reconnectInterval);
  }

  onMessage(handler: MessageHandler) {
    this.messageHandlers.add(handler);
    return () => this.messageHandlers.delete(handler);
  }

  disconnect() {
    this.stopPing();
    if (this.ws) {
      this.ws.close();
      this.config.onExit(0);
      this.ws = null;
    }
  }
}

// 便捷函数：快速创建并启动
export function createKeepAlive(config: KeepAliveConfig): WSKeepAlive {
  const client = new WSKeepAlive(config);
  client.connect();
  return client;
}