@yourrealm/sdk 0.0.0-alpha.1378.ce03c8dc

package/src/worker-protocol/uplink/session.ts

@@ -0,0 +1,540 @@
+// One uplink session: a WebSocket, a completed Noise XX handshake, the
+// transport cipher pair, and a `Mux` + `Rpc` riding the plaintext.
+//
+// Framing is one Noise message per binary WebSocket message. A text frame is
+// a protocol error; `permessage-deflate` is never negotiated, because Deno's
+// `WebSocket` does not offer it.
+//
+// The session owns no schedule beyond the handshake deadline and the
+// backpressure poll. `todo/ideas/uplink.md` § Session lifecycle puts the ping
+// cadence and the reconnect on the worker's loop and the idle deadline on
+// Home's session registry; this layer only exposes `ping()` and
+// `lastReceivedAt` for them to read.
+
+import {
+  type CipherState,
+  HandshakeState,
+  type Keypair,
+  NoiseError,
+  type TransportPair,
+} from "@yourrealm/noise";
+import {
+  HANDSHAKE_TIMEOUT_MS,
+  NOISE_MAX_PAYLOAD,
+  SOCKET_HIGH_WATER,
+  SOCKET_INBOUND_CAP,
+  UPLINK_PROLOGUE,
+} from "./constants.ts";
+import { HomeKeyMismatch, SessionClosed, UplinkProtocolError } from "./errors.ts";
+import { EMPTY_PAYLOAD } from "./frame.ts";
+import { type ByteDuplex, type Limits, Mux, type MuxSide } from "./mux.ts";
+import { Rpc, type RpcHandler } from "./rpc.ts";
+
+/** Noise transport messages carry no associated data; the key pair is the AD. */
+const EMPTY_AD = EMPTY_PAYLOAD;
+
+const CLOSE_NORMAL = 1000;
+const CLOSE_PROTOCOL = 1002;
+/** 1008 (policy violation) is what a Noise failure closes with. */
+const CLOSE_POLICY = 1008;
+
+/** A close reason is capped at 123 UTF-8 bytes; stay well inside it. */
+const MAX_CLOSE_REASON_CHARS = 100;
+
+/** There is no drain event in the WHATWG WebSocket API, so the gate polls. */
+const DRAIN_POLL_MS = 5;
+
+/** `setTimeout` is a `number` under Deno's lib and a `Timeout` under node's. */
+type TimerId = ReturnType<typeof setTimeout>;
+
+export interface SessionOptions {
+  /** Defaults to `UPLINK_PROLOGUE`. Only a test has reason to change it. */
+  prologue?: Uint8Array;
+  handshakeTimeoutMs?: number;
+  limits?: Partial<Limits>;
+}
+
+export interface UplinkSession {
+  /** The peer's Noise static key. Home looks it up; the worker pinned it. */
+  readonly peerPublicKey: Uint8Array;
+  readonly rpc: Rpc;
+  /** The close reason. Never rejects. */
+  readonly closed: Promise<string>;
+  /** ms epoch of the last inbound frame, for whoever owns the idle deadline. */
+  readonly lastReceivedAt: number;
+  /** Resolves on the matching PONG. The caller owns the deadline. */
+  ping(): Promise<void>;
+  close(reason: string): void;
+}
+
+/** Worker side. Dials, handshakes as initiator, pins Home's static key. */
+export async function connectUplink(
+  opts: {
+    url: string;
+    identity: Keypair;
+    homePublicKey: Uint8Array;
+    handlers: Record<string, RpcHandler>;
+  } & SessionOptions,
+): Promise<UplinkSession> {
+  const socket = new WebSocket(opts.url);
+  const feed = new SocketFeed(socket);
+  try {
+    const { peer, transport } = await withDeadline(
+      opts.handshakeTimeoutMs ?? HANDSHAKE_TIMEOUT_MS,
+      async () => {
+        await feed.opened;
+        const handshake = new HandshakeState({
+          initiator: true,
+          prologue: opts.prologue ?? UPLINK_PROLOGUE,
+          staticKeypair: opts.identity,
+        });
+        sendBinary(socket, handshake.writeMessage(EMPTY_AD));
+        handshake.readMessage(await feed.next());
+        // Message 2 carries Home's static. Compare before message 3, which is
+        // what keeps the worker's own static off an impostor's wire.
+        const peer = requirePeerKey(handshake);
+        if (!sameKey(peer, opts.homePublicKey)) {
+          throw new HomeKeyMismatch(opts.homePublicKey, peer);
+        }
+        sendBinary(socket, handshake.writeMessage(EMPTY_AD));
+        return { peer, transport: handshake.split() };
+      },
+    );
+    return new Session(socket, feed, transport, peer, "initiator", opts.handlers, opts.limits);
+  } catch (error) {
+    closeSocket(socket, failureCode(error), asError(error).message);
+    throw error;
+  }
+}
+
+/**
+ * Home side. Takes an already-upgraded socket and handshakes as responder.
+ *
+ * Called from inside the upgrade handler, which returns its response without
+ * waiting, so the caller owns this promise's rejection — a failed handshake
+ * has already closed the socket by the time it rejects.
+ */
+export async function acceptUplink(
+  opts: {
+    socket: WebSocket;
+    identity: Keypair;
+    /** Chosen once the worker's static is known: an unknown key gets the
+     * enroll-only map (`todo/ideas/uplink.md` § 2 Noise). */
+    handlers: (peer: Uint8Array) => Record<string, RpcHandler>;
+  } & SessionOptions,
+): Promise<UplinkSession> {
+  const socket = opts.socket;
+  const feed = new SocketFeed(socket);
+  try {
+    const { peer, transport } = await withDeadline(
+      opts.handshakeTimeoutMs ?? HANDSHAKE_TIMEOUT_MS,
+      async () => {
+        await feed.opened;
+        const handshake = new HandshakeState({
+          initiator: false,
+          prologue: opts.prologue ?? UPLINK_PROLOGUE,
+          staticKeypair: opts.identity,
+        });
+        handshake.readMessage(await feed.next());
+        sendBinary(socket, handshake.writeMessage(EMPTY_AD));
+        handshake.readMessage(await feed.next());
+        return { peer: requirePeerKey(handshake), transport: handshake.split() };
+      },
+    );
+    return new Session(
+      socket,
+      feed,
+      transport,
+      peer,
+      "responder",
+      opts.handlers(peer),
+      opts.limits,
+    );
+  } catch (error) {
+    closeSocket(socket, failureCode(error), asError(error).message);
+    throw error;
+  }
+}
+
+class Session implements UplinkSession {
+  readonly peerPublicKey: Uint8Array;
+  readonly rpc: Rpc;
+  readonly closed: Promise<string>;
+
+  readonly #socket: WebSocket;
+  readonly #feed: SocketFeed;
+  readonly #mux: Mux;
+  readonly #closedResolve: (reason: string) => void;
+  /** Every armed drain poll, so `closed` can never resolve over a live timer. */
+  readonly #polls = new Map<TimerId, () => void>();
+  #lastReceivedAt = Date.now();
+  #closeReason: string | undefined;
+
+  constructor(
+    socket: WebSocket,
+    feed: SocketFeed,
+    transport: TransportPair,
+    peer: Uint8Array,
+    side: MuxSide,
+    handlers: Record<string, RpcHandler>,
+    limits: Partial<Limits> | undefined,
+  ) {
+    this.peerPublicKey = peer;
+    this.#socket = socket;
+    this.#feed = feed;
+
+    let resolveClosed: (reason: string) => void = () => {};
+    this.closed = new Promise<string>((resolve) => {
+      resolveClosed = resolve;
+    });
+    this.#closedResolve = resolveClosed;
+
+    const duplex: ByteDuplex = {
+      readable: this.#makeReadable(transport.recv),
+      writable: this.#makeWritable(transport.send),
+    };
+    this.#mux = new Mux(duplex, { side, ...limits });
+    this.rpc = new Rpc(this.#mux, handlers);
+
+    // The socket dying, and the mux giving up on the transport, are the two
+    // ends the session cannot see from inside a stream.
+    feed.onEnd = (reason, code) => {
+      this.#shutdown(reason, code);
+    };
+    void this.#mux.closed.then((reason) => {
+      this.#shutdown(reason, CLOSE_PROTOCOL);
+    });
+  }
+
+  get lastReceivedAt(): number {
+    return this.#lastReceivedAt;
+  }
+
+  ping(): Promise<void> {
+    return this.#mux.ping();
+  }
+
+  close(reason: string): void {
+    this.#shutdown(reason, CLOSE_NORMAL);
+  }
+
+  // --- transport -------------------------------------------------------
+
+  // One binary WebSocket message per Noise message. The mux writes one frame
+  // per `write()`, at most 7 + 16384 bytes, so the split below is a guard
+  // rather than a hot path.
+  #makeWritable(send: CipherState): WritableStream<Uint8Array> {
+    return new WritableStream<Uint8Array>({
+      write: async (chunk) => {
+        for (let offset = 0; offset < chunk.length; offset += NOISE_MAX_PAYLOAD) {
+          const end = Math.min(offset + NOISE_MAX_PAYLOAD, chunk.length);
+          await this.#awaitDrain();
+          const reason = this.#closeReason;
+          if (reason !== undefined) throw new SessionClosed(reason);
+          try {
+            sendBinary(this.#socket, send.encryptWithAd(EMPTY_AD, chunk.subarray(offset, end)));
+          } catch (e) {
+            const error = asError(e);
+            this.#shutdown(error.message, failureCode(error));
+            throw error;
+          }
+        }
+      },
+      // The mux aborts the writable and cancels the readable on its own
+      // shutdown; both mean the same thing here.
+      abort: (reason) => {
+        this.#shutdown(reasonText(reason), CLOSE_PROTOCOL);
+      },
+      close: () => {
+        this.#shutdown("mux closed the transport", CLOSE_NORMAL);
+      },
+    });
+  }
+
+  // A zero high-water mark is the inbound backpressure: one message is
+  // decrypted per outstanding read, so nothing is decrypted ahead of the mux.
+  // What bounds the peer beyond that is the mux's own in-flight cap, which a
+  // flooding peer trips as a protocol error.
+  #makeReadable(recv: CipherState): ReadableStream<Uint8Array> {
+    return new ReadableStream<Uint8Array>(
+      {
+        pull: async (controller) => {
+          const message = await this.#feed.messages.next();
+          if (message === undefined) {
+            controller.close();
+            return;
+          }
+          this.#lastReceivedAt = Date.now();
+          let plaintext: Uint8Array;
+          try {
+            plaintext = recv.decryptWithAd(EMPTY_AD, message);
+          } catch (e) {
+            const error = asError(e);
+            this.#shutdown(error.message, failureCode(error));
+            throw error;
+          }
+          if (plaintext.length > 0) controller.enqueue(plaintext);
+        },
+        cancel: (reason) => {
+          this.#shutdown(reasonText(reason), CLOSE_PROTOCOL);
+        },
+      },
+      { highWaterMark: 0 },
+    );
+  }
+
+  // `WebSocket.send` buffers without bound, so the mux's `write()` waits here
+  // until the socket has drained. That is what makes a resolved `write()` mean
+  // "handed off", which is what keeps control frames ahead of a body.
+  async #awaitDrain(): Promise<void> {
+    while (this.#closeReason === undefined && this.#socket.bufferedAmount > SOCKET_HIGH_WATER) {
+      await this.#poll();
+    }
+  }
+
+  #poll(): Promise<void> {
+    return new Promise<void>((resolve) => {
+      const id = setTimeout(() => {
+        this.#polls.delete(id);
+        resolve();
+      }, DRAIN_POLL_MS);
+      this.#polls.set(id, resolve);
+    });
+  }
+
+  // --- teardown --------------------------------------------------------
+
+  #shutdown(reason: string, code: number): void {
+    if (this.#closeReason !== undefined) return;
+    this.#closeReason = reason;
+    for (const [id, resolve] of this.#polls) {
+      clearTimeout(id);
+      resolve();
+    }
+    this.#polls.clear();
+    this.#mux.close(reason);
+    closeSocket(this.#socket, code, reason);
+    this.#closedResolve(reason);
+  }
+}
+
+/**
+ * The socket's events, turned into one ordered queue of whole messages plus
+ * an `opened` promise. Listeners are attached before the handshake so nothing
+ * that arrives between the upgrade and the first read is lost.
+ */
+class SocketFeed {
+  readonly messages = new MessageQueue();
+  readonly opened: Promise<void>;
+  /** Set once the session exists; before that, `messages` carries the failure. */
+  onEnd: ((reason: string, code: number) => void) | undefined;
+
+  #open = false;
+
+  constructor(socket: WebSocket) {
+    socket.binaryType = "arraybuffer";
+    let resolveOpen: () => void = () => {};
+    let rejectOpen: (error: Error) => void = () => {};
+    this.opened = new Promise<void>((resolve, reject) => {
+      resolveOpen = resolve;
+      rejectOpen = reject;
+    });
+    if (socket.readyState === WebSocket.OPEN) {
+      this.#open = true;
+      resolveOpen();
+    }
+
+    socket.addEventListener("open", () => {
+      this.#open = true;
+      resolveOpen();
+    });
+    socket.addEventListener("message", (event: MessageEvent) => {
+      const data: unknown = event.data;
+      if (data instanceof ArrayBuffer) {
+        this.messages.push(new Uint8Array(data));
+        if (this.messages.queuedBytes > SOCKET_INBOUND_CAP) {
+          const error = new UplinkProtocolError("the peer outran the inbound queue");
+          this.#end(error.message, error, CLOSE_PROTOCOL);
+        }
+        return;
+      }
+      // A text frame carries no Noise message, so there is nothing to do with
+      // it but close: `todo/ideas/uplink.md` § 1 Byte stream.
+      const error = new UplinkProtocolError("the peer sent a text frame");
+      this.#end(error.message, error, CLOSE_PROTOCOL);
+    });
+    socket.addEventListener("error", () => {
+      // The event carries nothing useful; `close` follows with the detail.
+      this.#end("socket error", new SessionClosed("socket error"), CLOSE_NORMAL);
+    });
+    socket.addEventListener("close", (event: CloseEvent) => {
+      const detail = event.reason === "" ? `code ${event.code}` : event.reason;
+      const reason = `socket closed: ${detail}`;
+      if (!this.#open) rejectOpen(new SessionClosed(reason));
+      this.messages.end();
+      this.onEnd?.(reason, CLOSE_NORMAL);
+    });
+  }
+
+  /** The next whole message, or a rejection if the socket ended first. */
+  async next(): Promise<Uint8Array> {
+    const message = await this.messages.next();
+    if (message === undefined) {
+      throw new UplinkProtocolError("the socket closed mid-handshake");
+    }
+    return message;
+  }
+
+  #end(reason: string, error: Error, code: number): void {
+    this.messages.fail(error);
+    this.onEnd?.(reason, code);
+  }
+}
+
+/** Whole messages, in order, handed out one at a time. `undefined` is EOF. */
+class MessageQueue {
+  readonly #queue: Uint8Array[] = [];
+  #queuedBytes = 0;
+  #pending:
+    | { resolve: (value: Uint8Array | undefined) => void; reject: (error: Error) => void }
+    | undefined;
+  #error: Error | undefined;
+  #ended = false;
+
+  /** Bytes waiting for a `next()`, the number the inbound cap bounds. */
+  get queuedBytes(): number {
+    return this.#queuedBytes;
+  }
+
+  push(bytes: Uint8Array): void {
+    const pending = this.#pending;
+    if (pending === undefined) {
+      this.#queue.push(bytes);
+      this.#queuedBytes += bytes.length;
+      return;
+    }
+    this.#pending = undefined;
+    pending.resolve(bytes);
+  }
+
+  end(): void {
+    if (this.#ended) return;
+    this.#ended = true;
+    // A pending waiter means the queue is empty, so nothing is dropped.
+    const pending = this.#pending;
+    if (pending !== undefined) {
+      this.#pending = undefined;
+      pending.resolve(undefined);
+    }
+  }
+
+  fail(error: Error): void {
+    this.#error ??= error;
+    // Whatever was queued belongs to a session that is closing.
+    this.#queue.length = 0;
+    this.#queuedBytes = 0;
+    const pending = this.#pending;
+    if (pending !== undefined) {
+      this.#pending = undefined;
+      pending.reject(error);
+    }
+  }
+
+  next(): Promise<Uint8Array | undefined> {
+    const head = this.#queue.shift();
+    if (head !== undefined) {
+      this.#queuedBytes -= head.length;
+      return Promise.resolve(head);
+    }
+    if (this.#error !== undefined) return Promise.reject(this.#error);
+    if (this.#ended) return Promise.resolve(undefined);
+    return new Promise<Uint8Array | undefined>((resolve, reject) => {
+      this.#pending = { resolve, reject };
+    });
+  }
+}
+
+// --- helpers ------------------------------------------------------------
+
+// The whole handshake, socket open included, under one deadline. `run` is
+// caught separately so its own later failure is not an unhandled rejection
+// once the race has already reported the expiry.
+async function withDeadline<T>(ms: number, run: () => Promise<T>): Promise<T> {
+  const running = run();
+  running.catch(() => {});
+  let timer: TimerId | undefined;
+  const expiry = new Promise<never>((_, reject) => {
+    timer = setTimeout(() => {
+      reject(new UplinkProtocolError(`the handshake did not finish inside ${ms}ms`));
+    }, ms);
+  });
+  try {
+    return await Promise.race([running, expiry]);
+  } finally {
+    if (timer !== undefined) clearTimeout(timer);
+  }
+}
+
+function requirePeerKey(handshake: HandshakeState): Uint8Array {
+  const peer = handshake.remoteStaticPublicKey;
+  // XX transmits both statics, so this is unreachable on a message the
+  // handshake accepted; it is here so the type is not an assertion.
+  if (peer === undefined) throw new UplinkProtocolError("the handshake carried no static key");
+  return peer;
+}
+
+// Public keys, so there is nothing to leak by timing.
+function sameKey(a: Uint8Array, b: Uint8Array): boolean {
+  if (a.length !== b.length) return false;
+  for (let i = 0; i < a.length; i++) {
+    if (a[i] !== b[i]) return false;
+  }
+  return true;
+}
+
+function failureCode(error: unknown): number {
+  // A Noise failure — a bad tag, a rejected key, the pin — is 1008; anything
+  // else the session refuses is framing, which is 1002.
+  return error instanceof NoiseError || error instanceof HomeKeyMismatch
+    ? CLOSE_POLICY
+    : CLOSE_PROTOCOL;
+}
+
+// `send` is typed for an `ArrayBuffer`-backed view; every Noise message is a
+// fresh allocation, so this is a type-level fact rather than a claim.
+function sendBinary(socket: WebSocket, message: Uint8Array): void {
+  socket.send(message as Uint8Array<ArrayBuffer>);
+}
+
+// Both fallbacks earn their place. The WHATWG API lets application code send
+// only 1000 or 3000-4999, which Deno enforces on a *client* socket while
+// allowing the reserved codes on an upgraded server one — so the worker's end
+// cannot spell 1002 or 1008 and falls back to 1000, where the reason string
+// still carries the detail. The last resort covers a reason that measured
+// longer in UTF-8 than its character count suggested.
+function closeSocket(socket: WebSocket, code: number, reason: string): void {
+  if (socket.readyState === WebSocket.CLOSING || socket.readyState === WebSocket.CLOSED) return;
+  const text = reason.slice(0, MAX_CLOSE_REASON_CHARS);
+  const attempts = [
+    () => socket.close(code, text),
+    () => socket.close(CLOSE_NORMAL, text),
+    () => socket.close(),
+  ];
+  for (const attempt of attempts) {
+    try {
+      attempt();
+      return;
+    } catch {
+      /* try the next spelling */
+    }
+  }
+}
+
+function asError(value: unknown): Error {
+  return value instanceof Error ? value : new Error(String(value));
+}
+
+function reasonText(reason: unknown): string {
+  if (reason === undefined || reason === null) return "transport closed";
+  return reason instanceof Error ? reason.message : String(reason);
+}