@hedystia/ws 2.3.1

package/src/server.ts

@@ -0,0 +1,313 @@
+import { type WebSocket as NodeWebSocket, WebSocketServer as WSServer } from "ws";
+
+import type {
+  ServerWebSocket,
+  UpgradeOptions,
+  UpgradeRequest,
+  WebSocketHandlers,
+  WebSocketServerOptions,
+  WSData,
+  WSMessage,
+} from "./types";
+
+export type {
+  ServerWebSocket,
+  UpgradeOptions,
+  UpgradeRequest,
+  WebSocketHandlers,
+  WebSocketServerOptions,
+  WSData,
+  WSMessage,
+} from "./types";
+
+/**
+ * Runtime-agnostic WebSocket server.
+ *
+ * @remarks
+ * Internally backed by the [`ws`](https://github.com/websockets/ws) package
+ * which runs on Bun, Node.js and Deno (via `npm:` specifiers). The class
+ * does **not** create or own an HTTP server — callers feed it raw upgrade
+ * tuples coming from any HTTP runtime they prefer.
+ *
+ * It implements topic-based pub/sub on top of the per-connection
+ * `subscribe` / `unsubscribe` / `publish` API expected by Hedystia,
+ * matching the shape of `Bun.ServerWebSocket`.
+ *
+ * @typeParam Data - Shape of the user-attached `data` field
+ *
+ * @example
+ * ```ts
+ * import { createServer } from "node:http";
+ * import { WebSocketServer } from "@hedystia/ws/server";
+ *
+ * const wss = new WebSocketServer({
+ *   open: (ws) => ws.send("welcome"),
+ *   message: (ws, msg) => ws.publish("room", msg),
+ * });
+ *
+ * const http = createServer((_req, res) => res.end("ok"));
+ * http.on("upgrade", (req, socket, head) => {
+ *   wss.upgrade({ rawRequest: req, socket, head }, { data: { user: "anon" } });
+ * });
+ * http.listen(3000);
+ * ```
+ */
+export class WebSocketServer<Data extends WSData = WSData> {
+  private readonly handlers: WebSocketHandlers<Data>;
+  private readonly wss: WSServer;
+  private readonly topics = new Map<string, Set<NodeWebSocket>>();
+  private readonly socketTopics = new WeakMap<NodeWebSocket, Set<string>>();
+  private readonly allSockets = new Set<NodeWebSocket>();
+
+  /**
+   * Build a new WebSocket server.
+   *
+   * @param handlers - Lifecycle handlers ({@link WebSocketHandlers})
+   * @param options - Optional behavioural overrides ({@link WebSocketServerOptions})
+   *
+   * @example
+   * ```ts
+   * const wss = new WebSocketServer(
+   *   { message: (ws, msg) => ws.send(msg) },
+   *   { maxPayload: 1024 * 1024 },
+   * );
+   * ```
+   */
+  constructor(handlers: WebSocketHandlers<Data>, options: WebSocketServerOptions = {}) {
+    this.handlers = handlers;
+    this.wss = new WSServer({
+      noServer: true,
+      maxPayload: options.maxPayload,
+      perMessageDeflate: (options.perMessageDeflate ?? false) as any,
+    });
+  }
+
+  /**
+   * Upgrade a raw HTTP upgrade tuple to a WebSocket connection.
+   *
+   * @remarks
+   * The returned promise resolves to the connected {@link ServerWebSocket}
+   * once the handshake completes; rejection means the handshake failed.
+   *
+   * @param req - Upgrade tuple emitted by `node:http`'s `'upgrade'` event
+   * @param options - Optional initial `data` for the new connection
+   * @returns Promise that resolves with the established socket wrapper
+   *
+   * @throws {Error} When the underlying handshake throws synchronously
+   *
+   * @example
+   * ```ts
+   * http.on("upgrade", async (req, socket, head) => {
+   *   try {
+   *     await wss.upgrade({ rawRequest: req, socket, head });
+   *   } catch (err) {
+   *     console.error("Upgrade failed", err);
+   *     socket.destroy();
+   *   }
+   * });
+   * ```
+   */
+  upgrade(req: UpgradeRequest, options?: UpgradeOptions<Data>): Promise<ServerWebSocket<Data>> {
+    return new Promise((resolve, reject) => {
+      try {
+        this.wss.handleUpgrade(req.rawRequest, req.socket, req.head as Buffer, (socket) => {
+          const data = (options?.data ?? ({} as Data)) as Data;
+          const wrapped = this.wrap(socket, data, req.rawRequest);
+          this.bind(socket, wrapped);
+          resolve(wrapped);
+        });
+      } catch (err) {
+        reject(err);
+      }
+    });
+  }
+
+  /**
+   * Publish a message to all sockets currently subscribed to `topic`.
+   *
+   * @param topic - Topic name
+   * @param message - Payload to broadcast
+   * @param _compress - Reserved for future use; ignored under the `ws` adapter
+   * @returns Number of sockets that received the message.
+   *
+   * @example
+   * ```ts
+   * wss.publish("room", JSON.stringify({ kind: "ping" }));
+   * ```
+   */
+  publish(topic: string, message: WSMessage, _compress?: boolean): number {
+    const set = this.topics.get(topic);
+    if (!set || set.size === 0) {
+      return 0;
+    }
+    const payload = toSendable(message);
+    let count = 0;
+    for (const socket of set) {
+      if (socket.readyState === 1) {
+        socket.send(payload);
+        count++;
+      }
+    }
+    return count;
+  }
+
+  /**
+   * Close the server and optionally terminate all live sockets.
+   *
+   * @param closeActiveConnections - When `true`, calls `socket.terminate()`
+   *  on every live connection before shutting down.
+   */
+  close(closeActiveConnections = false): void {
+    if (closeActiveConnections) {
+      for (const socket of this.allSockets) {
+        try {
+          socket.terminate();
+        } catch {
+          /* ignore */
+        }
+      }
+      this.allSockets.clear();
+    }
+    this.wss.close();
+  }
+
+  /**
+   * Attach the per-socket lifecycle listeners (`message`, `close`, `error`).
+   *
+   * @internal
+   */
+  private bind(socket: NodeWebSocket, wrapped: ServerWebSocket<Data>): void {
+    this.allSockets.add(socket);
+
+    if (this.handlers.open) {
+      Promise.resolve(this.handlers.open(wrapped)).catch((err) =>
+        console.error("[ws] open handler error:", err),
+      );
+    }
+
+    socket.on("message", (raw, isBinary) => {
+      const message: WSMessage = isBinary
+        ? raw instanceof ArrayBuffer
+          ? new Uint8Array(raw)
+          : Array.isArray(raw)
+            ? Buffer.concat(raw)
+            : (raw as Buffer)
+        : raw.toString();
+      Promise.resolve(this.handlers.message(wrapped, message)).catch((err) =>
+        console.error("[ws] message handler error:", err),
+      );
+    });
+
+    socket.on("close", (code, reason) => {
+      const owned = this.socketTopics.get(socket);
+      if (owned) {
+        for (const topic of owned) {
+          this.topics.get(topic)?.delete(socket);
+        }
+        this.socketTopics.delete(socket);
+      }
+      this.allSockets.delete(socket);
+      if (this.handlers.close) {
+        Promise.resolve(this.handlers.close(wrapped, code, reason?.toString() ?? "")).catch((err) =>
+          console.error("[ws] close handler error:", err),
+        );
+      }
+    });
+
+    socket.on("error", (err) => {
+      if (this.handlers.error) {
+        Promise.resolve(this.handlers.error(wrapped, err)).catch((e) =>
+          console.error("[ws] error handler error:", e),
+        );
+      }
+    });
+  }
+
+  /**
+   * Build the {@link ServerWebSocket} wrapper exposed to user handlers.
+   *
+   * @internal
+   */
+  private wrap(socket: NodeWebSocket, data: Data, rawReq: any): ServerWebSocket<Data> {
+    const remoteAddress: string =
+      (rawReq?.socket?.remoteAddress as string) ||
+      (rawReq?.headers?.["x-forwarded-for"] as string) ||
+      "";
+
+    const subscribe = (topic: string) => {
+      let set = this.topics.get(topic);
+      if (!set) {
+        set = new Set();
+        this.topics.set(topic, set);
+      }
+      set.add(socket);
+      let owned = this.socketTopics.get(socket);
+      if (!owned) {
+        owned = new Set();
+        this.socketTopics.set(socket, owned);
+      }
+      owned.add(topic);
+    };
+
+    const unsubscribe = (topic: string) => {
+      this.topics.get(topic)?.delete(socket);
+      this.socketTopics.get(socket)?.delete(topic);
+    };
+
+    const publishToPeers = (topic: string, message: WSMessage) => {
+      const set = this.topics.get(topic);
+      if (!set) {
+        return;
+      }
+      const payload = toSendable(message);
+      for (const peer of set) {
+        if (peer !== socket && peer.readyState === 1) {
+          peer.send(payload);
+        }
+      }
+    };
+
+    const wrapper: ServerWebSocket<Data> = {
+      data,
+      get readyState() {
+        return socket.readyState;
+      },
+      remoteAddress,
+      send: (message, _compress) => {
+        const payload = toSendable(message);
+        socket.send(payload);
+        return typeof payload === "string"
+          ? Buffer.byteLength(payload)
+          : (payload as Buffer | Uint8Array).byteLength;
+      },
+      close: (code, reason) => {
+        socket.close(code, reason);
+      },
+      subscribe,
+      unsubscribe,
+      publish: publishToPeers,
+      isSubscribed: (topic) => !!this.socketTopics.get(socket)?.has(topic),
+      cork: (cb) => cb(wrapper),
+    };
+
+    return wrapper;
+  }
+}
+
+/**
+ * Coerce a {@link WSMessage} into something the `ws` package can transmit.
+ *
+ * @param message - User-supplied payload
+ * @returns A `string`, `Buffer` or `Uint8Array` ready to be sent
+ *
+ * @internal
+ */
+function toSendable(message: WSMessage): string | Buffer | Uint8Array {
+  if (typeof message === "string") {
+    return message;
+  }
+  if (message instanceof ArrayBuffer) {
+    return Buffer.from(message);
+  }
+  return message;
+}

package/src/types.ts

@@ -0,0 +1,226 @@
+/**
+ * Payload accepted by every `send`/`publish` method.
+ *
+ * @remarks
+ * Matches the WHATWG `WebSocket.send` signature plus the `Uint8Array`
+ * convenience accepted by Bun and the `ws` package.
+ */
+export type WSMessage = string | ArrayBuffer | Uint8Array;
+
+/**
+ * Bag of arbitrary, user-supplied state attached to a connection on
+ * upgrade and exposed to handlers as `ws.data`.
+ *
+ * @typeParam K - String key
+ * @typeParam V - Stored value
+ */
+export type WSData = Record<string, any>;
+
+/**
+ * The per-connection wrapper passed to every handler.
+ *
+ * @remarks
+ * The interface intentionally mirrors `Bun.ServerWebSocket` so that the
+ * same handler code works on Bun (native) and on Node.js (via
+ * {@link WebSocketServer}). Topic-based pub/sub is implemented in
+ * user-space when running outside Bun.
+ *
+ * @typeParam Data - Shape of the user-attached `data` field
+ *
+ * @example
+ * ```ts
+ * const handlers: WebSocketHandlers<{ user: string }> = {
+ *   open: (ws) => ws.subscribe(`user:${ws.data.user}`),
+ *   message: (ws, msg) => ws.publish(`user:${ws.data.user}`, msg),
+ * };
+ * ```
+ */
+export interface ServerWebSocket<Data extends WSData = WSData> {
+  /** User-supplied state attached to the socket on upgrade. */
+  readonly data: Data;
+  /** Standard WHATWG ready-state (`0` connecting, `1` open, `2` closing, `3` closed). */
+  readonly readyState: number;
+  /** Remote IP, taken from `X-Forwarded-For` when present. */
+  readonly remoteAddress: string;
+  /**
+   * Send a message to this socket only.
+   *
+   * @param message - Payload to send
+   * @param compress - Whether to compress (honoured on Bun, ignored on Node)
+   * @returns Number of bytes written (best-effort on Node)
+   */
+  send(message: WSMessage, compress?: boolean): number;
+  /**
+   * Close the connection.
+   *
+   * @param code - Close code (defaults to 1000)
+   * @param reason - Optional human-readable reason
+   */
+  close(code?: number, reason?: string): void;
+  /**
+   * Subscribe this socket to a topic so it receives subsequent
+   * {@link ServerWebSocket.publish | publish} or
+   * {@link WebSocketServer.publish | server.publish} broadcasts.
+   *
+   * @param topic - Topic name
+   */
+  subscribe(topic: string): void;
+  /**
+   * Unsubscribe this socket from a previously joined topic.
+   *
+   * @param topic - Topic name
+   */
+  unsubscribe(topic: string): void;
+  /**
+   * Broadcast a message to every other socket subscribed to `topic`.
+   *
+   * @remarks
+   * The sender is excluded by default — matching Bun's default behaviour.
+   *
+   * @param topic - Topic name
+   * @param message - Payload to broadcast
+   * @param compress - Whether to compress (honoured on Bun, ignored on Node)
+   */
+  publish(topic: string, message: WSMessage, compress?: boolean): void;
+  /**
+   * Check whether this socket is currently subscribed to `topic`.
+   *
+   * @param topic - Topic name
+   * @returns `true` when subscribed
+   */
+  isSubscribed(topic: string): boolean;
+  /**
+   * Batch multiple writes inside `cb`.
+   *
+   * @remarks
+   * On Bun this corresponds to `corked()`; on Node it is a no-op alias
+   * that simply invokes `cb(this)` synchronously.
+   *
+   * @param cb - Function invoked with the same socket
+   */
+  cork(cb: (ws: ServerWebSocket<Data>) => void): void;
+}
+
+/**
+ * Bun-style compression dictionary identifier.
+ *
+ * @remarks
+ * Used by Bun's `perMessageDeflate` configuration. The `@hedystia/ws`
+ * server forwards the value to the underlying implementation, which only
+ * Bun interprets natively; Node falls back to defaults when the value is
+ * not a recognised `ws` shape.
+ */
+export type Compressor =
+  | "disable"
+  | "shared"
+  | "dedicated"
+  | "3KB"
+  | "4KB"
+  | "8KB"
+  | "16KB"
+  | "32KB"
+  | "64KB"
+  | "128KB"
+  | "256KB";
+
+/**
+ * Per-message deflate configuration.
+ *
+ * @remarks
+ * Accepts either a boolean (`true` enables defaults, `false` disables it)
+ * or a free-form object whose shape is forwarded verbatim to the underlying
+ * implementation. Bun's {@link Compressor} strings (`"3KB"`, `"shared"`, …)
+ * and the [`ws`](https://github.com/websockets/ws) package's
+ * `PerMessageDeflateOptions` (`zlibDeflateOptions`, `threshold`, …) are
+ * both supported by simply matching whatever the runtime expects.
+ */
+export type PerMessageDeflate =
+  | boolean
+  | (Record<string, any> & {
+      compress?: boolean | Compressor;
+      decompress?: boolean | Compressor;
+    });
+
+/**
+ * Construction options for {@link WebSocketServer}.
+ */
+export interface WebSocketServerOptions {
+  /** Maximum allowed payload in bytes. Defaults to the underlying library's default. */
+  maxPayload?: number;
+  /** Per-message deflate configuration. */
+  perMessageDeflate?: PerMessageDeflate;
+}
+
+/**
+ * Lifecycle handlers passed to {@link WebSocketServer}.
+ *
+ * @typeParam Data - Shape of the user-attached `data` field
+ */
+export interface WebSocketHandlers<Data extends WSData = WSData> {
+  /** Called for every inbound message. */
+  message: (ws: ServerWebSocket<Data>, message: WSMessage) => void | Promise<void>;
+  /** Called once the handshake completes successfully. */
+  open?: (ws: ServerWebSocket<Data>) => void | Promise<void>;
+  /** Called after the connection closes (clean or otherwise). */
+  close?: (ws: ServerWebSocket<Data>, code: number, reason: string) => void | Promise<void>;
+  /** Called when the underlying transport raises an error. */
+  error?: (ws: ServerWebSocket<Data>, error: Error) => void | Promise<void>;
+  /**
+   * Called when back-pressure is relieved.
+   *
+   * @remarks
+   * Only fired by Bun; Node-backed servers never invoke it.
+   */
+  drain?: (ws: ServerWebSocket<Data>) => void | Promise<void>;
+}
+
+/**
+ * Options accepted by {@link createWebSocket}.
+ */
+export interface ClientWebSocketOptions {
+  /** Sub-protocols negotiated during the handshake. */
+  protocols?: string | string[];
+  /**
+   * Custom request headers.
+   *
+   * @remarks
+   * Honoured on Node via the `ws` package; runtimes that ship a WHATWG
+   * `WebSocket` (Bun, Deno, browsers, Node ≥ 22) ignore them — matching
+   * standard WebSocket semantics.
+   */
+  headers?: Record<string, string>;
+}
+
+/**
+ * Raw upgrade tuple consumed by {@link WebSocketServer.upgrade}.
+ *
+ * @remarks
+ * Mirrors what `node:http`'s `'upgrade'` event emits and what the `ws`
+ * package's `WebSocketServer.handleUpgrade` consumes.
+ */
+export interface UpgradeRequest {
+  /** Raw `IncomingMessage`-like object exposing `headers`, `method`, `url`. */
+  rawRequest: any;
+  /** Raw duplex socket (e.g. `node:net.Socket`). */
+  socket: any;
+  /** Initial buffer captured by the HTTP parser. */
+  head: Buffer | Uint8Array;
+}
+
+/**
+ * Options forwarded to {@link WebSocketServer.upgrade}.
+ *
+ * @typeParam Data - Shape of the user-attached `data` field
+ */
+export interface UpgradeOptions<Data extends WSData = WSData> {
+  /** Initial value of `ws.data` for the new connection. */
+  data?: Data;
+  /**
+   * Extra response headers.
+   *
+   * @remarks
+   * Forwarded to the underlying handshake when supported (Bun); ignored on
+   * Node-backed upgrades, which control headers via the `ws` package.
+   */
+  headers?: Record<string, string> | Headers;
+}

package/tsconfig.json

@@ -0,0 +1,3 @@
+{
+  "extends": "../../tsconfig.json"
+}

package/tsdown.config.ts

@@ -0,0 +1,12 @@
+import { defineConfig } from "tsdown";
+
+export default defineConfig({
+  entry: ["src/index.ts", "src/client.ts", "src/server.ts"],
+  format: ["cjs", "esm"],
+  dts: true,
+  sourcemap: true,
+  treeshake: true,
+  clean: true,
+  unbundle: true,
+  outputOptions: { exports: "named" },
+});