@johnhenry/wsh 0.29.0 → 0.31.0

package/src/server/index.mjs

@@ -15,7 +15,9 @@ import { SessionRegistry } from './sessions.mjs';
 import { RelayHub } from './relay.mjs';
 import { startWebTransport } from './webtransport.mjs';
 import { createConnectionFactory, STREAM_ANNOUNCE } from './connection.mjs';
+import { normalizeGateway } from './gateway.mjs';
 import { createRpcHost, mcpServerAdapter } from './rpc.mjs';
+import { bindSocket } from './socket.mjs';
 
 export { parseAuthorizedKeys, STREAM_ANNOUNCE, mcpServerAdapter };
 export { createReverseHost } from './reverse-host.mjs';
@@ -67,12 +69,16 @@ export { generateSelfSignedCertificate } from './self-signed.mjs';
  *   Handlers for application-defined messages: a control message whose `type` is a *string* (every protocol message
  *   has a numeric type, so these cannot collide) from an authenticated connection is passed to `extensions[type]`.
  *   A handler that throws is logged and ignored. Off unless given.
+ * @param {{ allow: string[], maxConnections?: number, connectTimeoutMs?: number, socks?: string | { host: string, port: number } }} [options.gateway]
+ *   Answer the gateway opcodes (OpenTcp, ResolveDns, GatewayData/Close): the client opens TCP connections and resolves
+ *   names from this host. Default deny: a destination must match `allow` (`"*"`, `"host"`, `"host:port"`). `socks` dials
+ *   through a SOCKS5 proxy (Tor) with the name unresolved. UDP and reverse tunnels are refused. Off unless given. See `gateway.mjs`.
  * @param {number} [options.bindTimeoutMs=3000]
  * @param {(line: string) => void} [options.onLog]
- * @returns {{ listen(): Promise<{address: string, port: number}>, close(): Promise<void>, address(): ({address: string, port: number} | null) }}
+ * @returns {{ listen(): Promise<{address: string, port: number}>, accept(socket: object, options?: object): Promise<{ close(): void }>, close(): Promise<void>, address(): ({address: string, port: number} | null) }}
  */
 export function createWshServer({
-  host = '127.0.0.1', port = 0, tls, auth, exec, pty, fs, hostKey, mcp, rpc, rpcMaxMessageBytes, rpcMaxInflight, sessions, sessionSecret, relay, extensions, webTransport, bindTimeoutMs = 3000, onLog = () => {},
+  host = '127.0.0.1', port = 0, tls, auth, exec, pty, fs, hostKey, mcp, rpc, rpcMaxMessageBytes, rpcMaxInflight, sessions, sessionSecret, relay, extensions, gateway, webTransport, bindTimeoutMs = 3000, onLog = () => {},
 } = {}) {
   let wss = null;
   let httpsServer = null;
@@ -94,17 +100,16 @@ export function createWshServer({
   let registry = null;
   let hub = null;
   let wt = null;
+  const gatewayCfg = gateway === undefined ? null : normalizeGateway(gateway);
   const sessionOptions = sessions ? (sessions === true ? {} : sessions) : null;
 
-  return {
-    async listen() {
-      if (wss) throw new Error('createWshServer: already listening');
-      let WebSocketServer;
-      try {
-        ({ WebSocketServer } = await import('ws'));
-      } catch (err) {
-        throw new Error('@johnhenry/wsh/server needs the "ws" package: npm install ws', { cause: err });
-      }
+  /** Sockets handed to `accept()` that are still open. */
+  const accepted = new Set();
+  /** The connection factory, built once on first use (`listen()` or `accept()`) and torn down by `close()`. */
+  let ready = null;
+
+  function start() {
+    ready ??= (async () => {
       const authorize = await buildAuthorizer(auth);
       host_ = hostKey ? await loadHostKey(hostKey) : null;
       if (sessionOptions) {
@@ -119,10 +124,25 @@ export function createWshServer({
       }
       if (relay) hub = new RelayHub({ ...relay, log: onLog });
       if (extensions !== undefined && (extensions === null || typeof extensions !== 'object')) throw new TypeError('createWshServer: extensions must be an object of handlers');
-      const attach = createConnectionFactory({
+      return createConnectionFactory({
         authorize, execRunner, execOptions, pty: ptyConfig, files, bindTimeoutMs, log: onLog,
-        methods, hostKey: host_, rateLimit, extensions, mcp: mcpHost, rpc: rpcHost, sessions: registry, relay: hub,
+        methods, hostKey: host_, rateLimit, extensions, mcp: mcpHost, rpc: rpcHost, gateway: gatewayCfg, sessions: registry, relay: hub,
       });
+    })();
+    ready.catch(() => { ready = null; });
+    return ready;
+  }
+
+  return {
+    async listen() {
+      if (wss) throw new Error('createWshServer: already listening');
+      let WebSocketServer;
+      try {
+        ({ WebSocketServer } = await import('ws'));
+      } catch (err) {
+        throw new Error('@johnhenry/wsh/server needs the "ws" package: npm install ws', { cause: err });
+      }
+      const attach = await start();
 
       if (tls) {
         const { createServer } = await import('node:https');
@@ -166,20 +186,45 @@ export function createWshServer({
       return bound;
     },
 
+    /**
+     * Serve a pre-upgraded transport: an already-accepted WebSocket-like duplex (see `socket.mjs`) that the embedder's
+     * own HTTP server, Worker or test upgraded. Same auth, features and session registry as `listen()`'s connections, and
+     * no listener of this server's own is needed -- `accept()` works without `listen()`, and with both.
+     * `remote: { address, headers }` is what the embedder knows of the peer (`address` keys the password-failure throttle).
+     * `close()` drops every accepted socket. Resolves once the connection is attached.
+     *
+     * @param {object} socket - `send`, `close`, `readyState`, and `on` or `addEventListener`
+     * @param {{ remote?: { address?: string, headers?: object } }} [options]
+     * @returns {Promise<{ close(): void }>}
+     */
+    async accept(socket, options = {}) {
+      const handle = await bindSocket(socket, start(), options);
+      const entry = { socket, handle };
+      accepted.add(entry);
+      const forget = () => accepted.delete(entry);
+      if (typeof socket.addEventListener === 'function') socket.addEventListener('close', forget);
+      else socket.on('close', forget);
+      return handle;
+    },
+
     async close() {
-      if (!wss) return;
+      for (const entry of [...accepted]) entry.handle.close();
+      accepted.clear();
       const server = wss;
       wss = null;
       bound = null;
-      for (const client of server.clients) client.terminate();
-      await wt?.close();
-      wt = null;
-      await new Promise((resolve) => server.close(resolve));
-      if (httpsServer) { const h = httpsServer; httpsServer = null; await new Promise((resolve) => h.close(resolve)); }
+      if (server) {
+        for (const client of server.clients) client.terminate();
+        await wt?.close();
+        wt = null;
+        await new Promise((resolve) => server.close(resolve));
+        if (httpsServer) { const h = httpsServer; httpsServer = null; await new Promise((resolve) => h.close(resolve)); }
+      }
       registry?.closeAll();
       registry = null;
       hub?.closeAll();
       hub = null;
+      ready = null;
     },
 
     address() {

package/src/server/pty.mjs

@@ -17,7 +17,7 @@ export function normalizePty(pty) {
   }
   return {
     spawn: pty.spawn,
-    shell: pty.shell ?? process.env.SHELL ?? '/bin/sh',
+    shell: pty.shell ?? globalThis.process?.env?.SHELL ?? '/bin/sh',
     cwd: pty.cwd,
     env: pty.env,
     term: pty.term ?? 'xterm-256color',

package/src/server/sessions.mjs

@@ -20,7 +20,8 @@ import { createHmac, randomBytes, randomUUID, timingSafeEqual } from 'node:crypt
 export const DEFAULT_RING_BYTES = 1024 * 1024;
 export const DEFAULT_DETACH_TTL_MS = 300_000;
 export const DEFAULT_MAX_DETACHED = 16;
-export const MAX_ATTACHMENTS = 16;
+import { MAX_ATTACHMENTS } from './constants.mjs';
+export { MAX_ATTACHMENTS };
 
 const NO_EXPIRY = Buffer.alloc(8, 0xff);
 

package/src/server/socket.mjs

@@ -0,0 +1,73 @@
+/**
+ * Serve a pre-upgraded transport: an already-accepted WebSocket-like duplex handed to a host by whoever did the
+ * upgrade (an HTTP server of the embedder's, a Worker, a message port wrapper, a test). Browser-safe -- it touches
+ * nothing but the object it is given -- and shared by `createWshServer().accept()` and `createMemoryHost()`.
+ *
+ * The duplex needs `send(bytes)`, `close(code?)`, `readyState` (1 = open) and either Node-style `on(event, fn)`
+ * (`ws`) or `addEventListener(event, fn)` (browsers, Workers, undici). Incoming frames must be binary; a text frame
+ * is ignored, as the QMux framing never carries one.
+ */
+
+const OPEN = 1;
+
+/** @param {unknown} data - a message payload as `ws` (Buffer, ArrayBuffer, Buffer[]) or a browser (ArrayBuffer, view, Blob) delivers it */
+async function toBytes(data) {
+  if (data instanceof Uint8Array) return new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
+  if (data instanceof ArrayBuffer) return new Uint8Array(data);
+  if (ArrayBuffer.isView(data)) return new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
+  if (Array.isArray(data)) {
+    const parts = await Promise.all(data.map(toBytes));
+    const out = new Uint8Array(parts.reduce((n, p) => n + p.byteLength, 0));
+    let off = 0;
+    for (const p of parts) { out.set(p, off); off += p.byteLength; }
+    return out;
+  }
+  if (data && typeof data.arrayBuffer === 'function') return new Uint8Array(await data.arrayBuffer());
+  return null;
+}
+
+/**
+ * @param {object} socket - the WebSocket-like duplex
+ * @param {Promise<(opts: object) => { receive: Function, close: Function }>} ready - resolves to the connection factory's `attach`
+ * @param {{ remote?: { address?: string, headers?: object } }} [options] - what the host knows about the peer
+ *   (`address` keys the password-failure throttle; the default key is shared by every socket without one)
+ * @returns {Promise<{ close(): void }>} resolves once the connection is attached; `close()` drops it
+ */
+export function bindSocket(socket, ready, { remote = {} } = {}) {
+  if (!socket || typeof socket.send !== 'function' || typeof socket.close !== 'function') {
+    throw new TypeError('accept: socket must be a WebSocket-like duplex with send() and close()');
+  }
+  const listen = typeof socket.addEventListener === 'function'
+    ? (event, fn) => socket.addEventListener(event, (ev) => fn(event === 'message' ? ev?.data : ev))
+    : typeof socket.on === 'function'
+      ? (event, fn) => socket.on(event, fn)
+      : null;
+  if (!listen) throw new TypeError('accept: socket must support addEventListener() or on()');
+  try { socket.binaryType = 'arraybuffer'; } catch { /* read-only on some implementations; toBytes copes with any */ }
+
+  // Listeners are attached now, before the host has finished starting, and the frames they see are
+  // chained behind `ready`: a peer that speaks first (a client's QMux handshake) is not dropped.
+  let conn = null;
+  let closed = false;
+  let chain = ready.then((attach) => {
+    if (closed) return;
+    conn = attach({
+      send: (bytes) => { if (socket.readyState === OPEN) socket.send(bytes); },
+      remote: { address: remote.address, headers: remote.headers ?? {} },
+      closeTransport: () => { try { socket.close(1000); } catch { /* already closed */ } },
+    });
+  });
+  const drop = () => { if (closed) return; closed = true; chain = chain.then(() => conn?.close()).catch(() => {}); };
+  listen('message', (data) => {
+    chain = chain.then(async () => {
+      if (closed || !conn) return;
+      const bytes = await toBytes(data);
+      if (bytes && !closed) conn.receive(bytes);
+    }).catch(() => {});
+  });
+  listen('close', drop);
+  listen('error', () => {});
+  return chain.then(() => ({
+    close() { drop(); try { socket.close(1000); } catch { /* already closed */ } },
+  }));
+}

package/src/transport-ws.mjs

@@ -97,6 +97,9 @@ class QMuxStreamAdapter {
   }
 }
 
+/** `WebSocket.OPEN` (spelled out so a substitute constructor need not define it). */
+const WS_OPEN = 1;
+
 // ── WebSocket transport ──────────────────────────────────────────────
 
 /**
@@ -133,6 +136,19 @@ export class WebSocketTransport extends WshTransport {
    */
   #inbox = new SerialQueue((raw) => this.#handleControlBytes(raw));
 
+  /** @type {typeof WebSocket} */
+  #WebSocket;
+
+  /**
+   * @param {object} [options]
+   * @param {typeof WebSocket} [options.WebSocket] - A WebSocket constructor to dial with instead of the global one: a
+   *   test double, or the in-memory host's (`createMemoryHost().WebSocket`, see `memory-host.mjs`).
+   */
+  constructor({ WebSocket: Impl } = {}) {
+    super();
+    this.#WebSocket = Impl;
+  }
+
   // ── Lifecycle ──────────────────────────────────────────────────────
 
   /** @override */
@@ -149,7 +165,7 @@ export class WebSocketTransport extends WshTransport {
         .replace(/^https:\/\//, 'wss://')
         .replace(/^http:\/\//, 'ws://');
 
-      const ws = new WebSocket(wsUrl);
+      const ws = new (this.#WebSocket ?? WebSocket)(wsUrl);
       ws.binaryType = 'arraybuffer';
       this.#ws = ws;
 
@@ -316,7 +332,7 @@ export class WebSocketTransport extends WshTransport {
    * @param {Uint8Array} data
    */
   #send(data) {
-    if (!this.#ws || this.#ws.readyState !== WebSocket.OPEN) {
+    if (!this.#ws || this.#ws.readyState !== WS_OPEN) {
       throw new Error('WebSocket is not open');
     }
     this.#ws.send(data);