feat(ws): support dynamic websocket resolver

unjs · Feb 25, 2024 · 7b006e2 · 7b006e2
1 parent f540c9d
commit 7b006e2
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 67 deletions.
diff --git a/README.md b/README.md
@@ -203,11 +203,11 @@ When enabled, listhen tries to listen to all network interfaces. You can also en
 
 - Default: `false`
 
-Enable experimental WebSocket support.
+Enable experimental WebSocket support using [unjs/crossws](https://crossws.unjs.io/) or node upgrade handler.
 
-Option can be a function for Node.js `upgrade` handler (`(req, head) => void`) or an Object to use [CrossWS Hooks](https://github.com/unjs/crossws).
+Option can be a function for Node.js `upgrade` handler (`(req, head) => void`) or an Object to use [CrossWS Hooks](https://crossws.unjs.io/guide/api).
 
-When using dev server CLI, you can easily use `--ws` and a named export called `webSocket` to define [CrossWS Hooks](https://github.com/unjs/crossws) with HMR support!
+When using dev server CLI, you can easily use `--ws` and a named export called `websocket` to define [CrossWS Hooks](https://github.com/unjs/crossws) with HMR support!
 
 ## License
 

diff --git a/playground/index.ts b/playground/index.ts
@@ -1,54 +1,23 @@
-import { createApp, eventHandler } from "h3";
+import { createApp, defineEventHandler } from "h3";
 import { defineWebSocketHooks } from "crossws";
 
 export const app = createApp();
 
 app.use(
   "/ws",
-  eventHandler(() => {
-    return getWebSocketTestPage();
-  }),
+  defineEventHandler(() =>
+    fetch(
+      "https://raw.githubusercontent.com/unjs/crossws/main/examples/h3/public/index.html",
+    ).then((r) => r.text()),
+  ),
 );
 
 app.use(
   "/",
-  eventHandler(() => ({ hello: "world!" })),
+  defineEventHandler(() => ({ hello: "world!" })),
 );
 
-function getWebSocketTestPage() {
-  return `
-  <!doctype html>
-  <head>
-    <title>WebSocket Test Page</title>
-  </head>
-  <body>
-    <div id="logs"></div>
-    <script type="module">
-      const url = \`ws://\${location.host}/_ws\`;
-      const logsEl = document.querySelector("#logs");
-      const log = (...args) => {
-        console.log("[ws]", ...args);
-        logsEl.innerHTML += \`<p>[\${new Date().toJSON()}] \${args.join(" ")}</p>\`;
-      };
-
-      log(\`Connecting to "\${url}""...\`);
-      const ws = new WebSocket(url);
-
-      ws.addEventListener("message", (event) => {
-        log("Message from server:", event.data);
-      });
-
-      log("Waiting for connection...");
-      await new Promise((resolve) => ws.addEventListener("open", resolve));
-
-      log("Sending ping...");
-      ws.send("ping");
-    </script>
-  </body>
-    `;
-}
-
-export const webSocket = defineWebSocketHooks({
+export const websocket = defineWebSocketHooks({
   open(peer) {
     console.log("[ws] open", peer);
     peer.send("Hello!");

diff --git a/src/listen.ts b/src/listen.ts
@@ -19,6 +19,7 @@ import type {
   HTTPSOptions,
   ListenURL,
   GetURLOptions,
+  WebSocketOptions,
 } from "./types";
 import {
   formatURL,
@@ -140,8 +141,14 @@ export async function listen(
       const nodeWSAdapter = await import("crossws/adapters/node").then(
         (r) => r.default || r,
       );
-      // @ts-expect-error TODO
-      const { handleUpgrade } = nodeWSAdapter(listhenOptions.ws);
+      const { $resolve, $options, ...hooks } =
+        listhenOptions.ws === true
+          ? ({} as WebSocketOptions)
+          : listhenOptions.ws;
+      const { handleUpgrade } = nodeWSAdapter(hooks, {
+        ...$options,
+        resolve: $resolve,
+      });
       server.on("upgrade", handleUpgrade);
     }
   }

diff --git a/src/server/dev.ts b/src/server/dev.ts
@@ -4,14 +4,14 @@ import { consola } from "consola";
 import { dirname, join, resolve } from "pathe";
 import type { ConsolaInstance } from "consola";
 import { resolve as _resolve } from "mlly";
-import type { WebSocketHooks } from "crossws";
+import type { WebSocketOptions, ListenOptions } from "../types";
 import { createResolver } from "./_resolver";
 
 export interface DevServerOptions {
   cwd?: string;
   staticDirs?: string[];
   logger?: ConsolaInstance;
-  ws?: boolean | Partial<WebSocketHooks>;
+  ws?: ListenOptions["ws"];
 }
 
 export async function createDevServer(
@@ -56,26 +56,18 @@ export async function createDevServer(
   // Create app instance
   const app = createApp();
 
-  // WebSocket
-  let _ws: Partial<WebSocketHooks> | undefined;
-  const webSocketHooks = Object.create(null);
-  if (options.ws) {
-    const createDynamicHook =
-      (name: string) =>
-      async (...args: any[]) => {
-        if (typeof options.ws === "object") {
-          await (options.ws as any)[name]?.(...args);
-        }
-        return (webSocketHooks as any)[name]?.(...args);
-      };
-    _ws = new Proxy(
-      {},
-      {
-        get(_, prop) {
-          return createDynamicHook(prop as string);
-        },
+  const webSocketHooks = Object.create(null); // Dynamically updated with HMR
+  let _ws: DevServerOptions["ws"] = options.ws;
+  if (_ws && typeof _ws !== "function") {
+    _ws = {
+      ...(options.ws as WebSocketOptions),
+      async $resolve(info) {
+        return {
+          ...webSocketHooks,
+          ...(await (options.ws as WebSocketOptions).$resolve?.(info)),
+        };
       },
-    );
+    };
   }
 
   // Register static asset handlers

diff --git a/src/types.ts b/src/types.ts
@@ -2,7 +2,12 @@ import type { IncomingMessage, Server } from "node:http";
 import type { Server as HTTPServer } from "node:https";
 import { AddressInfo } from "node:net";
 import type { GetPortInput } from "get-port-please";
-import type { WebSocketHooks } from "crossws";
+import type { UserHooks as WSHooks, CrossWSOptions } from "crossws";
+
+export interface WebSocketOptions extends WSHooks {
+  $resolve: CrossWSOptions["resolve"];
+  $options: CrossWSOptions;
+}
 
 export interface Certificate {
   key: string;
@@ -63,7 +68,7 @@ export interface ListenOptions {
    */
   ws?:
     | boolean
-    | Partial<WebSocketHooks>
+    | WebSocketOptions
     | ((req: IncomingMessage, head: Buffer) => void);
 }