@johnhenry/wsh 0.27.0 → 0.28.0

package/src/index.d.ts

@@ -616,6 +616,14 @@ export function reverseReject(opts?: {
 export function relayForward(opts?: {
   fromFingerprint?: string;
   inner?: Uint8Array;
+  /** Peer -> relay: the operator this is for, when the peer serves several (wsh #89). */
+  toFingerprint?: string;
+}): WshMessage;
+
+/** Relay -> peer: one operator of a multi-operator peer left (wsh #89). */
+export function reverseClose(opts?: {
+  targetFingerprint?: string;
+  reason?: string;
 }): WshMessage;
 
 /** Message types a relay may forward inside a RelayForward wrapper. */
@@ -851,6 +859,8 @@ export function keyExchange(opts?: {
   sessionId?: string;
   kemPublicKey?: Uint8Array;
   kemCiphertext?: Uint8Array;
+  /** Round 1 only: the sender's Ed25519 signature over the exchange (see the spec's KeyExchange). */
+  signature?: Uint8Array;
 }): WshMessage;
 
 export function encryptedFrame(opts?: {
@@ -1909,6 +1919,27 @@ export interface WshPeerInfo {
   verified: boolean;
 }
 
+/** Host feature: it answers `KeyExchange` and seals/opens `EncryptedFrame`s (a reverse host with `e2e` on). */
+export const E2E_FEATURE: 'e2e';
+/** Host feature: it signs its `KeyExchange` reply, so `initiateE2E(..., { verifyPeer })` can authenticate it. */
+export const E2E_SIGN_FEATURE: 'e2e-sign';
+
+/** The bytes a responder signs in `KeyExchange.signature`. */
+export function keyExchangeTranscript(t: {
+  sessionId: string;
+  /** The initiator's `algorithm` string. */
+  algorithm: string;
+  initiatorKey: Uint8Array;
+  responderKey: Uint8Array;
+}): Uint8Array;
+
+/** Verify a responder's `KeyExchange.signature` against its long-term Ed25519 public key. */
+export function verifyKeyExchangeSignature(
+  peerIdentity: Uint8Array | CryptoKey,
+  signature: Uint8Array | undefined,
+  transcript: Parameters<typeof keyExchangeTranscript>[0],
+): Promise<boolean>;
+
 /** Result from WshClient.initiateE2E(). */
 export interface WshE2EResult {
   /** Non-extractable AES-256-GCM key derived from the exchange. */
@@ -1921,6 +1952,8 @@ export interface WshE2EResult {
    * requested but the peer didn't support it (automatic fallback).
    */
   hybrid: boolean;
+  /** True when `verifyPeer` was given and the peer's signature over the exchange verified. */
+  peerAuthenticated: boolean;
 }
 
 /**
@@ -2005,7 +2038,13 @@ export class WshClient {
    * In reverse mode, the relay bridge forwards messages from the CLI peer
    * to this browser client.
    */
-  onRelayMessage: ((msg: WshMessage) => void) | null;
+  onRelayMessage: ((msg: WshMessage, fromFingerprint?: string) => void) | null;
+
+  /**
+   * Called on a peer that serves several operators when the relay reports that one left (`ReverseClose`,
+   * wsh #89): end that operator's bridge and clean up what belongs to it.
+   */
+  onReverseClose: ((msg: WshMessage) => void) | null;
 
   /**
    * Called when a gateway-subsystem control message arrives (opcodes 0x70-0x7f).
@@ -2169,7 +2208,7 @@ export class WshClient {
   /**
    * Send a control message over the authenticated relay connection.
    */
-  sendRelayControl(msg: WshMessage): Promise<void>;
+  sendRelayControl(msg: WshMessage, opts?: { to?: string }): Promise<void>;
 
   /**
    * Send a raw control message on the authenticated connection.
@@ -2202,7 +2241,15 @@ export class WshClient {
    * fallback to classical if the peer doesn't support it). Experimental
    * -- the derived key is not yet wired to any actual encryption.
    */
-  initiateE2E(sessionId: string, algorithm?: 'X25519' | 'X25519+ML-KEM-768', timeout?: number): Promise<WshE2EResult>;
+  initiateE2E(
+    sessionId: string,
+    algorithm?: 'X25519' | 'X25519+ML-KEM-768',
+    timeout?: number,
+    opts?: {
+      /** The peer's long-term Ed25519 public key; the exchange is refused (`E2E_PEER_UNAUTHENTICATED`) unless it signed its half. */
+      verifyPeer?: Uint8Array | CryptoKey;
+    },
+  ): Promise<WshE2EResult>;
 
   /**
    * Mark a peer fingerprint as an accepted reverse-connect bridge partner.

package/src/index.mjs

@@ -17,7 +17,7 @@ export {
   presence, controlChanged, metrics,
   mcpDiscover, mcpTools, mcpCall, mcpResult,
   reverseRegister, reverseList, reversePeers, reverseConnect, reverseAccept, reverseReject,
-  relayForward, RELAY_FORWARDABLE, isRelayForwardable,
+  relayForward, reverseClose, RELAY_FORWARDABLE, isRelayForwardable,
   openTcp, openUdp, resolveDns, gatewayOk, gatewayFail, gatewayClose,
   inboundOpen, inboundAccept, inboundReject, dnsResult,
   listenRequest, listenOk, listenFail, listenClose, gatewayData,
@@ -76,6 +76,9 @@ export { WshSession } from './session.mjs';
 export { WshVirtualSessionBackend, normalizeSessionData } from './virtual-session.mjs';
 export { WshClient, MCP_CALL_ID_FEATURE, STREAM_ANNOUNCE_FEATURE } from './client.mjs';
 
+// End-to-end key exchange (responder side, feature names, signed transcript) -- wsh #90
+export { E2E_FEATURE, E2E_SIGN_FEATURE, keyExchangeTranscript, verifyKeyExchangeSignature } from './e2e-exchange.mjs';
+
 // Key storage
 export { WshKeyStore } from './keystore.mjs';
 

package/src/messages.gen.mjs

@@ -63,6 +63,7 @@ export const MSG = Object.freeze({
   REVERSE_ACCEPT:    0x54,
   REVERSE_REJECT:    0x55,
   RELAY_FORWARD:     0x56,
+  REVERSE_CLOSE:     0x57,
 
   // Session
   SESSION_LIST:      0x5f,
@@ -510,12 +511,16 @@ export function reverseReject({ targetFingerprint, username, reason } = {}) {
   };
 }
 
-export function relayForward({ fromFingerprint, inner } = {}) {
-  return {
-    type: MSG.RELAY_FORWARD,
-    from_fingerprint: fromFingerprint,
-    inner,
-  };
+export function relayForward({ fromFingerprint, inner, toFingerprint } = {}) {
+  const msg = { type: MSG.RELAY_FORWARD, from_fingerprint: fromFingerprint, inner };
+  if (toFingerprint !== undefined) msg.to_fingerprint = toFingerprint;
+  return msg;
+}
+
+export function reverseClose({ targetFingerprint, reason } = {}) {
+  const msg = { type: MSG.REVERSE_CLOSE, target_fingerprint: targetFingerprint };
+  if (reason !== undefined) msg.reason = reason;
+  return msg;
 }
 
 export function sessionList({ sessions } = {}) {
@@ -762,11 +767,12 @@ export function copilotDetach({ sessionId, reason } = {}) {
   return msg;
 }
 
-export function keyExchange({ algorithm, publicKey, sessionId, kemPublicKey, kemCiphertext } = {}) {
+export function keyExchange({ algorithm, publicKey, sessionId, kemPublicKey, kemCiphertext, signature } = {}) {
   const msg = { type: MSG.KEY_EXCHANGE, algorithm, session_id: sessionId };
   if (publicKey !== undefined) msg.public_key = publicKey;
   if (kemPublicKey !== undefined) msg.kem_public_key = kemPublicKey;
   if (kemCiphertext !== undefined) msg.kem_ciphertext = kemCiphertext;
+  if (signature !== undefined) msg.signature = signature;
   return msg;
 }
 

package/src/server/connection.mjs

@@ -18,7 +18,7 @@ import { QMuxConnection } from '../qmux-connection.mjs';
 import { FrameDecoder, frameEncode } from '../cbor.mjs';
 import {
   MSG, serverHello, challenge, authOk, authFail, openOk, openFail, sessionData,
-  exit as exitMsg, close as closeMsg, pong, fileResult, fileChunk, mcpTools, mcpResult,
+  exit as exitMsg, close as closeMsg, encryptedFrame as encryptedFrameMsg, pong, fileResult, fileChunk, mcpTools, mcpResult,
   presence as presenceMsg, error as errorMsg, sessionList, detachOk, detachFail, reversePeers, reverseReject,
 } from '../messages.gen.mjs';
 import {
@@ -28,6 +28,9 @@ import { FILE_CHUNK_BYTES } from './fs.mjs';
 import { MAX_ATTACHMENTS } from './sessions.mjs';
 import { MCP_CALL_ID_FEATURE } from '../client.mjs';
 import { RpcChannel, RPC_FEATURE, RPC_MAX_MESSAGE_PREFIX, rpcProtocolFeature } from '../rpc.mjs';
+import { sealFrame, openFrame, ROLE_TAGS } from '../e2e-frame.mjs';
+import { E2EResponder, E2E_FEATURE, E2E_SIGN_FEATURE } from '../e2e-exchange.mjs';
+import { MULTI_OPERATOR_FEATURE } from './relay.mjs';
 import { HOST_KEY_PREFIX, HOST_KEY_SIG_PREFIX, findClientNonce, hostKeyProofMessage, toHex } from '../host-key.mjs';
 
 export const STREAM_ANNOUNCE = 'stream-announce';
@@ -42,6 +45,15 @@ export function hostFeatures(cfg, { dataStreams = true } = {}) {
   const features = dataStreams ? [STREAM_ANNOUNCE] : [];
   if (cfg.mcp) features.push(MCP_CALL_ID_FEATURE);
   if (cfg.files) features.push('file-transfer', 'file-write', 'file-rename');
+  // End-to-end encryption (#90) is the virtual-mode EncryptedFrame layer, so it is offered where there are no
+  // client-opened streams: over a relay bridge.
+  // A relay carries a signed KeyExchange untouched, and (when configured) several operators per peer: peers read
+  // these from its ServerHello before relying on either, since an older relay would drop the connection on them.
+  if (cfg.relay) {
+    features.push(E2E_SIGN_FEATURE);
+    if (cfg.relay.maxOperatorsPerPeer > 1) features.push(MULTI_OPERATOR_FEATURE);
+  }
+  if (cfg.e2e && !dataStreams) features.push(E2E_FEATURE, ...(cfg.e2e.signKey ? [E2E_SIGN_FEATURE] : []));
   if (cfg.rpc && dataStreams && cfg.rpc.protocols.size) {
     features.push(RPC_FEATURE, ...[...cfg.rpc.protocols.keys()].map(rpcProtocolFeature), RPC_MAX_MESSAGE_PREFIX + cfg.rpc.maxMessageBytes);
   }
@@ -70,6 +82,7 @@ const toBytes = (d) => (typeof d === 'string' ? enc.encode(d) : d);
  * @param {object | null} cfg.hostKey - from loadHostKey
  * @param {object | null} cfg.mcp - from createMcpHost
  * @param {{ protocols: Map<string, Function>, maxMessageBytes: number, maxInflight: number } | null} cfg.rpc - from createRpcHost
+ * @param {{ signKey: CryptoKey | null, hybrid?: boolean } | null} cfg.e2e - answer KeyExchange and seal virtual-mode traffic (bridged connections only)
  * @param {object | null} cfg.relay - a RelayHub (see relay.mjs); null = not a relay
  * @param {object | null} cfg.sessions - a SessionRegistry (see sessions.mjs); null = sessions die with their connection
  * @param {{ limiter: object, key: Function, failureDelayMs: number } | null} cfg.rateLimit - password throttle
@@ -86,8 +99,9 @@ export function createConnectionFactory(cfg) {
    *    for a transport that already frames messages itself (a relay bridge, WebTransport). `authenticated`
    *    (`{ username, fingerprint }`) skips the handshake for a peer something else has already authenticated;
    *    `dataStreams: false` means there are no client-opened streams to bind, so exec output rides SessionData.
+   *    `e2e` overrides `cfg.e2e` for this connection.
    */
-  return function attach({ send: sendBytes, sendMessage = null, remote = {}, authenticated = null, dataStreams = true, closeTransport = null }) {
+  return function attach({ send: sendBytes, sendMessage = null, remote = {}, authenticated = null, dataStreams = true, closeTransport = null, e2e: e2eOverride }) {
     const cid = ++counter;
     const log = (m) => cfg.log(`[conn ${cid}] ${m}`);
     const state = {
@@ -111,7 +125,39 @@ export function createConnectionFactory(cfg) {
 
     const qmux = sendMessage ? null : new QMuxConnection({ isClient: false, send: (b) => { if (!state.closed) sendBytes(b); } });
 
+    // ── End-to-end encryption (#90) ──────────────────────────────────
+    // An operator that sent KeyExchange for one of this connection's sessions gets that session's output sealed
+    // into EncryptedFrames and its input opened from them; plaintext input for it is ignored from then on (a relay
+    // must not be able to inject stdin around the encryption).
+    /** channel id -> the session id its OpenOk named (the AAD of its frames). */
+    const sessionOfChannel = new Map();
+    /** session id -> { key, sendCounter, recvCounter } for sessions with E2E on. */
+    const e2eSessions = new Map();
+    const MAX_E2E_SESSIONS = 64;
+    const e2eConfig = e2eOverride === undefined ? cfg.e2e : e2eOverride;
+    const responder = e2eConfig && !dataStreams ? new E2EResponder({ signKey: e2eConfig.signKey ?? null, hybrid: e2eConfig.hybrid !== false }) : null;
+    /** Outgoing messages in order, once E2E is on (sealing is async; order on the wire must not change). */
+    let outChain = Promise.resolve();
+    const e2eFor = (channelId) => e2eSessions.get(sessionOfChannel.get(channelId));
+
     const send = (msg) => {
+      if (state.closed) return Promise.resolve();
+      if (msg.type === MSG.OPEN_OK && typeof msg.session_id === 'string') sessionOfChannel.set(msg.channel_id, msg.session_id);
+      if (e2eSessions.size === 0) return deliver(msg);
+      // E2E is on for something: everything goes through one chain, so a sealed frame cannot be overtaken
+      // by (or overtake) the Exit that follows it.
+      const done = outChain.then(async () => {
+        const st = msg.type === MSG.SESSION_DATA ? e2eFor(msg.channel_id) : undefined;
+        if (!st) return deliver(msg);
+        const sessionId = sessionOfChannel.get(msg.channel_id);
+        const { nonce, ciphertext } = await sealFrame(st.key, sessionId, ROLE_TAGS.responder, st.sendCounter++, msg.data);
+        return deliver(encryptedFrameMsg({ channelId: msg.channel_id, nonce, ciphertext, sessionId }));
+      });
+      outChain = done.catch(() => {});
+      return done.catch(() => {});
+    };
+
+    const deliver = (msg) => {
       if (state.closed) return Promise.resolve();
       if (sendMessage) {
         try { return Promise.resolve(sendMessage(msg)).catch(() => {}); } catch { return Promise.resolve(); }
@@ -129,6 +175,8 @@ export function createConnectionFactory(cfg) {
       // everything else dies with it.
       for (const ch of [...channels.values()]) (ch.detach ?? ch.kill)?.call(ch);
       channels.clear();
+      responder?.close();
+      e2eSessions.clear();
       for (const c of mcpCalls) c.abort(new Error('connection closed'));
       mcpCalls.clear();
       for (const w of fileWrites.values()) clearTimeout(w.timer);
@@ -179,7 +227,12 @@ export function createConnectionFactory(cfg) {
       switch (m.type) {
         case MSG.PING: return send(pong({ id: m.id }));
         case MSG.OPEN: return handleOpen(m);
-        case MSG.SESSION_DATA: channels.get(m.channel_id)?.input?.(m.data); return;
+        case MSG.SESSION_DATA:
+          if (e2eFor(m.channel_id)) { log(`ignored plaintext SessionData on E2E channel ${m.channel_id}`); return; }
+          channels.get(m.channel_id)?.input?.(m.data);
+          return;
+        case MSG.KEY_EXCHANGE: return handleKeyExchange(m);
+        case MSG.ENCRYPTED_FRAME: return handleEncryptedFrame(m);
         case MSG.RESIZE: channels.get(m.channel_id)?.resize?.(m.cols, m.rows); return;
         case MSG.SIGNAL: channels.get(m.channel_id)?.signal?.(String(m.signal ?? '')); return;
         case MSG.CLOSE: { const ch = channels.get(m.channel_id); if (ch) { ch.kill?.(); } return; }
@@ -197,6 +250,37 @@ export function createConnectionFactory(cfg) {
       }
     }
 
+    async function handleKeyExchange(m) {
+      if (!responder) { log('ignored KeyExchange (end-to-end encryption is not enabled here)'); return; }
+      // Only for a session this connection actually has: the id is the AAD of every frame, and it bounds the state kept.
+      const known = typeof m.session_id === 'string' && [...sessionOfChannel.values()].includes(m.session_id);
+      if (!known) { log('ignored KeyExchange for a session this connection does not have'); return; }
+      if (!e2eSessions.has(m.session_id) && e2eSessions.size >= MAX_E2E_SESSIONS) { log('ignored KeyExchange (too many encrypted sessions)'); return; }
+      let result;
+      try { result = await responder.respond(m, send); } catch (e) { log(`KeyExchange refused: ${e.message}`); return; }
+      // Sealing starts after the reply is on its way, so the operator never gets a frame before it has the key.
+      if (result.key) {
+        e2eSessions.set(result.sessionId, { key: result.key, sendCounter: 0, recvCounter: 0 });
+        log(`end-to-end encryption on for session ${result.sessionId.slice(0, 8)}${result.hybrid ? ' (hybrid)' : ''}`);
+      }
+    }
+
+    async function handleEncryptedFrame(m) {
+      const st = e2eSessions.get(m.session_id);
+      // The frame must name the session its channel belongs to, or it could be spliced onto another channel.
+      if (!st || sessionOfChannel.get(m.channel_id) !== m.session_id) { log('dropped an EncryptedFrame for a session without E2E'); return; }
+      let plain;
+      try {
+        plain = await openFrame(st.key, m.session_id, st.recvCounter, { nonce: m.nonce, ciphertext: m.ciphertext, expectedRoleTag: ROLE_TAGS.initiator });
+      } catch (e) {
+        // Garbage (a replay, a tampered or foreign frame) neither reaches the process nor consumes a counter.
+        log(`dropped an EncryptedFrame: ${e.message}`);
+        return;
+      }
+      st.recvCounter += 1;
+      channels.get(m.channel_id)?.input?.(plain);
+    }
+
     async function refuse(reason) {
       log(`AUTH_FAIL: ${reason}`);
       await send(authFail({ reason }));

package/src/server/index.d.ts

@@ -198,6 +198,12 @@ export interface WshServerRelayOptions {
   ) => boolean | Promise<boolean>;
   /** Registered peers at once (default 1024). */
   maxPeers?: number;
+  /**
+   * Operators bridged to one peer at once (default 1). Above 1, a peer is held to one operator unless it stated
+   * `relay-multi-operator` in its `ReverseAccept.features` (it then addresses replies with
+   * `RelayForward.to_fingerprint` and is sent a `ReverseClose` when one operator leaves instead of being closed).
+   */
+  maxOperatorsPerPeer?: number;
   /** How long a `ReverseConnect` waits for its peer to answer before the operator is rejected (default 8000). */
   connectTimeoutMs?: number;
 }
@@ -421,6 +427,23 @@ export interface WshReverseHostOptions {
   /** Extra options for the relay connection (`expectHostKey`, `knownHosts`, `trustOnFirstUse`, ...). */
   connect?: Record<string, unknown>;
   peerType?: string;
+  /**
+   * Operators served at once (default 1). Above 1 the host states `relay-multi-operator`, addresses each reply
+   * to its operator and handles `ReverseClose`; each operator gets its own connection state and is passed through
+   * `accept` on its own identity. Needs `reportFeatures` and a relay with `maxOperatorsPerPeer` above 1; any other
+   * relay serves it one operator, as before.
+   */
+  maxOperators?: number;
+  /** State this host's features in `ReverseAccept.features` (default true). */
+  reportFeatures?: boolean;
+  /**
+   * End-to-end encryption for the bridged operator (default on; `false` turns it off). The host answers an
+   * operator's `KeyExchange` for a session it opened, seals that session's output into `EncryptedFrame`s and opens its
+   * input from them. `sign`: sign the reply with `keyPair` so the operator can authenticate this host
+   * (`initiateE2E(..., { verifyPeer })`); `'auto'` (default) signs when the relay's `ServerHello` carries `e2e-sign`.
+   * `hybrid` (default true) offers X25519+ML-KEM-768 when asked and available.
+   */
+  e2e?: false | { sign?: boolean | 'auto'; hybrid?: boolean };
   onLog?: (line: string) => void;
 }
 
@@ -433,6 +456,8 @@ export interface WshReverseHost {
   readonly fingerprint: string | null;
   /** Is the relay connection up right now? */
   readonly connected: boolean;
+  /** Fingerprints of the operators bridged right now. */
+  readonly operators: string[];
 }
 
 /**

package/src/server/index.mjs

@@ -39,10 +39,12 @@ export { generateSelfSignedCertificate } from './self-signed.mjs';
  *   SHA-256 through `webTransport()` / `certificateHashes()` for the client's `serverCertificateHashes`, and rotates it
  *   before it expires (`rotateCertificate()`). Needs the optional peers
  *   `@fails-components/webtransport` and `@fails-components/webtransport-transport-http3-quiche`. Off unless given.
- * @param {{ canRegister?: Function, canConnect?: Function, maxPeers?: number, connectTimeoutMs?: number }} [options.relay]
+ * @param {{ canRegister?: Function, canConnect?: Function, maxPeers?: number, maxOperatorsPerPeer?: number, connectTimeoutMs?: number }} [options.relay]
  *   Act as a relay: peers register (`ReverseRegister`, a signed record), operators list and connect to them
  *   (`ReverseList` / `ReverseConnect`) and traffic is carried between them as `RelayForward`. Default deny:
  *   `canRegister(who, record)` and `canConnect(from, to)` must both be given and return true. See `createReverseHost`.
+ *   One operator per peer unless `maxOperatorsPerPeer` is raised -- and then only for peers that state
+ *   `relay-multi-operator` (`createReverseHost({ maxOperators })`).
  * @param {true | { detachTtlMs?: number, maxDetached?: number, ringBytes?: number, sessionSecret?: string | Uint8Array }} [options.sessions]
  *   Keep pty/exec sessions alive across disconnects so a client can `resumeSession()` / `attachSession()`
  *   them: a per-server registry, a ring buffer of the newest `ringBytes` (default 1 MiB) of output, and
@@ -76,6 +78,9 @@ export function createWshServer({
     key: typeof rl.key === 'function' ? rl.key : ({ address }) => String(address ?? 'unknown'),
     failureDelayMs: rl.failureDelayMs ?? 250,
   };
+  if (relay && relay.maxOperatorsPerPeer !== undefined && (!Number.isSafeInteger(relay.maxOperatorsPerPeer) || relay.maxOperatorsPerPeer < 1)) {
+    throw new TypeError('createWshServer: relay.maxOperatorsPerPeer must be a positive integer');
+  }
   const rpcHost = rpc ? createRpcHost(rpc, { files, maxMessageBytes: rpcMaxMessageBytes, maxInflight: rpcMaxInflight }) : null;
   let host_ = null;
   let registry = null;

package/src/server/relay.mjs

@@ -12,15 +12,22 @@
  * Nothing is open by default: `canRegister` and `canConnect` both default to
  * "no". `from_fingerprint` is always the sender's own authenticated key, set
  * here and never taken from a payload; only the spec's `forwardable` message
- * types cross a bridge. One bridge per peer at a time, and it lasts as long as
+ * types cross a bridge.
+ *
+ * By default there is one operator per peer, and the bridge lasts as long as
  * both connections: when either ends the hub closes the other, so no
  * half-bridged state (and no processes belonging to a departed operator)
- * outlives it.
+ * outlives it. With `maxOperatorsPerPeer` > 1 a peer that states
+ * `relay-multi-operator` in its ReverseAccept may be bridged to several
+ * operators at once (#89): it addresses each reply with
+ * `RelayForward.to_fingerprint`, and is told with a `ReverseClose` when one
+ * operator leaves (instead of being closed). A peer that never stated it is
+ * still held to one operator, and still closed when its operator leaves.
  */
 
 import { cborEncode, cborDecode } from '../cbor.mjs';
 import { fingerprint as fingerprintOf, importPublicKeyRaw, verifyPeerRecord } from '../auth.mjs';
-import { MSG, isRelayForwardable, relayForward, reversePeers, reverseReject } from '../messages.gen.mjs';
+import { MSG, isRelayForwardable, relayForward, reversePeers, reverseReject, reverseClose } from '../messages.gen.mjs';
 
 export const DEFAULT_MAX_PEERS = 1024;
 export const DEFAULT_CONNECT_TIMEOUT_MS = 8000;
@@ -32,7 +39,16 @@ const MAX_FEATURE_LENGTH = 128;
 const SHORT_PREFIX_MIN = 8;
 
 /** `ReverseReject.reason` for a peer that already has an operator (or a request awaiting its answer). */
+/** `ReverseAccept.features` entry by which a peer says it serves several operators (addresses replies, understands `ReverseClose`). */
+export const MULTI_OPERATOR_FEATURE = 'relay-multi-operator';
+/** `ReverseReject.reason` for a peer that is answering another operator's request right now. */
+export const BUSY_ANSWERING = 'busy: this peer is answering another operator\'s request (retry in a moment)';
+
 export const BUSY_PEER = 'busy: this peer already has an operator (the relay bridges one operator per peer at a time; retry when it leaves)';
+/** `ReverseReject.reason` for a multi-operator peer already at the relay's operator cap. */
+export function busyPeerFull(max) {
+  return `busy: this peer already has its maximum of ${max} operators (retry when one leaves)`;
+}
 /** `ReverseReject.reason` for an operator that already has a bridge, or a request in flight. */
 export const BUSY_OPERATOR = 'busy: you already have a bridge or a pending connect on this relay (one at a time)';
 
@@ -50,17 +66,23 @@ const isStr = (v, max) => typeof v === 'string' && v.length <= max;
 export class RelayHub {
   /** @type {Map<string, { handle: RelayHandle, meta: object }>} */
   #peers = new Map();
-  /** @type {Map<RelayHandle, RelayHandle>} */
-  #pairs = new Map();
+  /** operator -> the peer it is bridged to (an operator has one bridge). @type {Map<RelayHandle, RelayHandle>} */
+  #operatorPeer = new Map();
+  /** peer -> its bridged operators, by the operator's fingerprint. @type {Map<RelayHandle, Map<string, RelayHandle>>} */
+  #peerOperators = new Map();
+  /** peers that stated `relay-multi-operator` in their last ReverseAccept. @type {WeakSet<RelayHandle>} */
+  #multiPeers = new WeakSet();
   /** target handle -> { operator, timer } while a ReverseConnect awaits its answer */
   #pending = new Map();
   /** fingerprint -> the highest record seq ever accepted (survives disconnects, bounded). */
   #lastSeq = new Map();
 
-  constructor({ canRegister, canConnect, maxPeers = DEFAULT_MAX_PEERS, connectTimeoutMs = DEFAULT_CONNECT_TIMEOUT_MS, log = () => {} } = {}) {
+  constructor({ canRegister, canConnect, maxPeers = DEFAULT_MAX_PEERS, maxOperatorsPerPeer = 1, connectTimeoutMs = DEFAULT_CONNECT_TIMEOUT_MS, log = () => {} } = {}) {
     for (const [name, fn] of [['canRegister', canRegister], ['canConnect', canConnect]]) {
       if (fn !== undefined && typeof fn !== 'function') throw new TypeError(`createWshServer: relay.${name} must be a function`);
     }
+    if (!Number.isSafeInteger(maxOperatorsPerPeer) || maxOperatorsPerPeer < 1) throw new TypeError('createWshServer: relay.maxOperatorsPerPeer must be a positive integer');
+    this.maxOperatorsPerPeer = maxOperatorsPerPeer;
     this.canRegister = canRegister ?? (() => false);
     this.canConnect = canConnect ?? (() => false);
     this.maxPeers = maxPeers;
@@ -193,14 +215,19 @@ export class RelayHub {
         fingerprint: target.meta.fingerprint, username: target.meta.username, capabilities: target.meta.capabilities,
       }))) return reject('no such peer');
     if (this.#peers.get(target.meta.fingerprint) !== target) return reject('no such peer');
-    // One bridge per peer, and one per operator: a RelayForward names its sender but not its recipient, so a
-    // peer's reply could not be addressed to one of several operators (see the README's relay section).
-    if (this.#pairs.has(target.handle) || this.#pending.has(target.handle)) {
-      return reject(BUSY_PEER);
-    }
-    if (this.#pairs.has(operator) || [...this.#pending.values()].some((p) => p.operator === operator)) {
+    // One bridge per operator. A peer serves one operator unless the relay allows more AND the peer said it can
+    // address them (RelayForward.to_fingerprint) and be told when one leaves (ReverseClose); it also answers one
+    // request at a time, since a ReverseAccept does not name the operator it answers.
+    if (this.bridged(operator) || [...this.#pending.values()].some((p) => p.operator === operator)) {
       return reject(BUSY_OPERATOR);
     }
+    if (this.#pending.has(target.handle)) return reject(this.maxOperatorsPerPeer > 1 ? BUSY_ANSWERING : BUSY_PEER);
+    const serving = this.#peerOperators.get(target.handle);
+    if (serving?.size) {
+      if (this.maxOperatorsPerPeer < 2 || !this.#multiPeers.has(target.handle)) return reject(BUSY_PEER);
+      if (serving.has(operator.fingerprint)) return reject(BUSY_OPERATOR);
+      if (serving.size >= this.maxOperatorsPerPeer) return reject(busyPeerFull(this.maxOperatorsPerPeer));
+    }
     const timer = setTimeout(() => {
       if (this.#pending.get(target.handle)?.operator !== operator) return;
       this.#pending.delete(target.handle);
@@ -232,26 +259,39 @@ export class RelayHub {
     }
     p.operator.send(answer);
     if (msg.type === MSG.REVERSE_ACCEPT) {
-      this.#pairs.set(p.operator, peer);
-      this.#pairs.set(peer, p.operator);
+      if (Array.isArray(answer.features) && answer.features.includes(MULTI_OPERATOR_FEATURE)) this.#multiPeers.add(peer);
+      else this.#multiPeers.delete(peer);
+      this.#operatorPeer.set(p.operator, peer);
+      let ops = this.#peerOperators.get(peer);
+      if (!ops) this.#peerOperators.set(peer, ops = new Map());
+      ops.set(p.operator.fingerprint, p.operator);
       this.log(`relay: bridged ${p.operator.fingerprint.slice(0, 12)} <-> ${peer.fingerprint.slice(0, 12)}`);
     }
   }
 
   /** Is `handle` one end of a bridge? */
   bridged(handle) {
-    return this.#pairs.has(handle);
+    return this.#operatorPeer.has(handle) || (this.#peerOperators.get(handle)?.size ?? 0) > 0;
+  }
+
+  /** Fingerprints of the operators bridged to the peer with this fingerprint. */
+  operatorsOf(peerFingerprint) {
+    const peer = this.#peers.get(peerFingerprint)?.handle;
+    return peer ? [...(this.#peerOperators.get(peer)?.keys() ?? [])] : [];
   }
 
   /**
    * Carry a message across `handle`'s bridge. A plain forwardable message is
    * wrapped; a `RelayForward` the client wrote itself is unwrapped, its inner
    * checked, and re-wrapped -- either way `from_fingerprint` is the sender's
-   * authenticated key. Returns false when there is no bridge.
+   * authenticated key. A peer serving several operators names the recipient
+   * with `to_fingerprint` (never set on what is delivered); with one operator
+   * it may leave it out. Returns false when there is no bridge.
    */
   forward(handle, msg) {
-    const partner = this.#pairs.get(handle);
-    if (!partner) return false;
+    const partner = this.#recipient(handle, msg);
+    if (partner === undefined) return false;
+    if (partner === null) return true; // bridged, but the message could not be addressed: dropped (and logged)
     let inner;
     if (msg.type === MSG.RELAY_FORWARD) {
       if (!(msg.inner instanceof Uint8Array) || msg.inner.byteLength > MAX_FORWARD_BYTES) { this.log('relay: dropped an oversized or malformed RelayForward'); return true; }
@@ -270,6 +310,27 @@ export class RelayHub {
     return true;
   }
 
+  /**
+   * Who `msg` from `handle` goes to: the partner handle; `undefined` if `handle` has no bridge (the caller
+   * treats the message as an ordinary one); `null` if it has one but the recipient cannot be determined.
+   */
+  #recipient(handle, msg) {
+    const peer = this.#operatorPeer.get(handle);
+    if (peer) return peer; // an operator has exactly one bridge
+    const operators = this.#peerOperators.get(handle);
+    if (!operators || operators.size === 0) return undefined;
+    const to = msg.type === MSG.RELAY_FORWARD ? msg.to_fingerprint : undefined;
+    if (to !== undefined) {
+      const operator = typeof to === 'string' ? operators.get(to) : undefined;
+      if (operator) return operator;
+      this.log(`relay: dropped a forward from ${handle.fingerprint.slice(0, 12)} addressed to ${typeof to === 'string' ? to.slice(0, 12) : '?'}, which is not one of its operators`);
+      return null;
+    }
+    if (operators.size === 1) return operators.values().next().value;
+    this.log(`relay: dropped an unaddressed forward from ${handle.fingerprint.slice(0, 12)}, which serves ${operators.size} operators (to_fingerprint is required)`);
+    return null;
+  }
+
   /** A connection ended: forget it, and end the bridge it was part of. */
   drop(handle) {
     const entry = this.#peers.get(handle.fingerprint);
@@ -279,11 +340,28 @@ export class RelayHub {
     for (const [target, pend] of this.#pending) {
       if (pend.operator === handle) { clearTimeout(pend.timer); this.#pending.delete(target); }
     }
-    const partner = this.#pairs.get(handle);
-    if (partner) {
-      this.#pairs.delete(handle);
-      this.#pairs.delete(partner);
-      partner.close();
+    // An operator leaving.
+    const peer = this.#operatorPeer.get(handle);
+    if (peer) {
+      this.#operatorPeer.delete(handle);
+      const operators = this.#peerOperators.get(peer);
+      operators?.delete(handle.fingerprint);
+      if (this.maxOperatorsPerPeer > 1 && this.#multiPeers.has(peer)) {
+        // The peer serves others too (or may serve more): tell it which operator left instead of closing it.
+        peer.send(reverseClose({ targetFingerprint: handle.fingerprint, reason: 'operator left' }));
+      } else {
+        this.#peerOperators.delete(peer);
+        peer.close();
+      }
+    }
+    // A peer leaving: every operator bridged to it goes with it.
+    const operators = this.#peerOperators.get(handle);
+    if (operators) {
+      this.#peerOperators.delete(handle);
+      for (const operator of operators.values()) {
+        this.#operatorPeer.delete(operator);
+        operator.close();
+      }
     }
   }
 
@@ -291,7 +369,8 @@ export class RelayHub {
   closeAll() {
     for (const p of this.#pending.values()) clearTimeout(p.timer);
     this.#pending.clear();
-    this.#pairs.clear();
+    this.#operatorPeer.clear();
+    this.#peerOperators.clear();
     this.#peers.clear();
   }
 }