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 { 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 Instance {
  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.)");
      };
    });
  }

  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, {});
      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) {
    if (!this.isServer || !this.wss) return;
    const packet = this.encodePacket(ptype, psub, payload);
    for (const ws of this.clients.keys()) {
      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);
  }

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