bo3-js/js/instances/NetworkService.js

let WebSocketImpl;
let WebSocketServer;
if (typeof window !== "undefined") {
    WebSocketImpl = WebSocket;
} else {
    // Dynamically require 'ws' only on server side
    const wsModule = await import("ws");
    WebSocketImpl = wsModule.default;
    WebSocketServer = wsModule.WebSocketServer;
}
import { Instance } from "./Instance.js";
import { BaseService } from "./BaseService.js";
import { BO3ScriptSignal } from "../core/BO3ScriptSignal.js";

/**
 * NetworkService
 *  - Handles both client & server networking
 *  - Binary framed packets: [uint32 length][ptype][psub][payload]
 *  - Fires PacketReceived signal for incoming packets
 */
export class NetworkService extends BaseService {
  constructor() {
    super("NetworkService");

    /** @type {boolean} */
    this.isServer = false;

    /** @type {WebSocketServer|null} */
    this.wss = null;

    /** @type {WebSocket|null} */
    this.client = null;

    /** @type {Map<WebSocket, any>} */
    this.clients = new Map();

    /** @type {Map<number, Function>} */
    this.handlers = new Map();

    /** @type {BO3ScriptSignal} */
    this.PacketReceived = new BO3ScriptSignal();

    // Total bytes received (for stats)
    this.totalBytes = 0;
    // Interval to clear totalBytes every second
    setInterval(() => {
      this.totalBytes = 0;
    }, 1000);
  }

  // ==========================
  // === Utility Functions ===
  // ==========================
  static pktKey(ptype, psub) {
    return (ptype << 8) | psub;
  }

  registerHandler(ptype, psub, fn) {
    const key = NetworkService.pktKey(ptype, psub);
    this.handlers.set(key, fn);
  }

  unregisterHandler(ptype, psub) {
    const key = NetworkService.pktKey(ptype, psub);
    this.handlers.delete(key);
  }

  invokeHandler(ptype, psub, payload, client = null) {
    const key = NetworkService.pktKey(ptype, psub);
    const handler = this.handlers.get(key);
    if (handler) handler(payload, client);
    this.PacketReceived.Fire({ ptype, psub, payload, client });
  }

  // ======================
  // === Packet Encode ===
  // ======================
  encodePacket(ptype, psub, payload = new Uint8Array()) {
    if (!(payload instanceof Uint8Array)) {
      if (typeof payload === "string") {
        payload = new TextEncoder().encode(payload);
      } else if (payload instanceof ArrayBuffer) {
        payload = new Uint8Array(payload);
      } else {
        throw new Error("Payload must be Uint8Array, ArrayBuffer, or string");
      }
    }
    const body = new Uint8Array(2 + payload.length);
    body[0] = ptype;
    body[1] = psub;
    body.set(payload, 2);

    const header = new ArrayBuffer(4);
    new DataView(header).setUint32(0, body.length);
    return new Uint8Array([...new Uint8Array(header), ...body]);
  }

  decodePacket(buffer) {
    const view = new DataView(buffer);
    const length = view.getUint32(0);
    const body = new Uint8Array(buffer, 4, length);
    const ptype = body[0];
    const psub = body[1];
    const payload = body.slice(2);
    return { ptype, psub, payload };
  }

  // ============================
  // === Client Functionality ===
  // ============================
  async connect(url) {
    this.isServer = false;
    return new Promise((resolve, reject) => {
      const ws = new WebSocket(url);
      this.client = ws;

      ws.binaryType = "arraybuffer";

      ws.onopen = () => {
        console.log(`[NetworkService] Connected to ${url}`);
        resolve();
      };

      ws.onmessage = (event) => {
        const buf = event.data instanceof ArrayBuffer ? event.data : event.data.buffer;
        this.totalBytes += Math.max(buf.byteLength,0); // safeguard
        const { ptype, psub, payload } = this.decodePacket(buf);
        this.invokeHandler(ptype, psub, payload);
      };

      ws.onerror = (err) => reject(err);
      ws.onclose = () => {
        console.warn("[NetworkService] Disconnected from server.")
        alert("Disconnected from server! If this was unintentional, please refresh the page. You might have been kicked by the server.\n\n(If you are the server host, check if the servers running.)");
      };
      // Register 0xFF kick handler
      this.registerHandler(0x00, 0xFF, (payload) => {
        let reason = "Kicked by server.";
        try {
          reason = new TextDecoder().decode(payload);
        } catch (e) {
          // ignore
        }
        console.warn("[NetworkService] Kicked by server:", reason);
        alert("You have been kicked by the server:\n\n" + reason);
      });
    });
  }

  send(ptype, psub, payload) {
    if (!this.client || this.client.readyState !== WebSocket.OPEN) return;
    const packet = this.encodePacket(ptype, psub, payload);
    this.client.send(packet);
  }

  // ============================
  // === Server Functionality ===
  // ============================
  listen(port = 3000, host = "0.0.0.0") {
    this.isServer = true;
    this.wss = new WebSocketServer({ port, host });
    console.log(`[NetworkService] Listening on ws://${host}:${port}`);

    this.wss.on("connection", (ws) => {
      this.clients.set(ws, {});
      ws.isAuthed = false; // Determine whether to send any data except auth packets
      console.log("[NetworkService] New client connected. Total clients:", this.clients.size);
      ws.binaryType = "arraybuffer";

      ws.on("message", (msg) => {
        const buf = msg instanceof ArrayBuffer ? msg : msg.buffer;
        const { ptype, psub, payload } = this.decodePacket(buf);
        this.invokeHandler(ptype, psub, payload, ws);
      });

      ws.on("close", () => {
        this.clients.delete(ws)
        console.log("[NetworkService] Client disconnected. Total clients:", this.clients.size);
      });
    });
  }

  broadcast(ptype, psub, payload,ignoreUnauthed,ignoreCheck) {
    if (!this.isServer || !this.wss) return;
    const packet = this.encodePacket(ptype, psub, payload);
    for (const ws of this.clients.keys()) {
      if(ignoreUnauthed && !ws.isAuthed) continue;
      if (ignoreCheck && ignoreCheck(ws)) continue;
      if (ws.readyState === WebSocket.OPEN) ws.send(packet);
    }
  }

  sendToClient(ws, ptype, psub, payload) {
    if (!this.isServer) return;
    const packet = this.encodePacket(ptype, psub, payload);
    if (ws.readyState === WebSocket.OPEN) ws.send(packet);
  }

  KickClientLater(ws,code=4000,reason="") {
    // Send a kick packet
    this.sendToClient(ws, 0x00, 0xFF, new TextEncoder().encode(reason));
    setTimeout(() => {
      if (ws.readyState === WebSocket.OPEN) {
        ws.close(code,reason);
      }
    }, 500); // Slight delay to ensure any pending packets are sent
  }


  Broadcast(ptype, psub, payload) {
    console.warn("[NetworkService] Broadcast is deprecated, use broadcast() instead.");
    this.broadcast(ptype, psub, payload);
  }

  // ============================
  // === Shutdown ===
  // ============================
  close() {
    if (this.isServer && this.wss) {
      for (const ws of this.clients.keys()) ws.close();
      this.wss.close();
    } else if (this.client) {
      this.client.close();
    }
  }
}