trails/apps/planner/app/lib/yjs-server.ts

import { WebSocketServer, type WebSocket } from "ws";
import * as Y from "yjs";
import * as syncProtocol from "y-protocols/sync";
import * as awarenessProtocol from "y-protocols/awareness";
import * as encoding from "lib0/encoding";
import * as decoding from "lib0/decoding";
import type { IncomingMessage, Server } from "node:http";
import { saveSessionState, loadSessionState, touchSession, getSession } from "./sessions.ts";
import { plannerActiveSessions, plannerConnectedClients } from "./metrics.server.ts";

const messageSync = 0;
const messageAwareness = 1;

// Bound the per-session doc size — the real abuse guard, stopping a client
// from filling memory or the persisted `yjs_state` blob without limit. A
// route with hundreds of waypoints is well under this once its BRouter
// geometry is stored in the doc, which is what actually drives the size.
export const MAX_DOC_BYTES = 5 * 1024 * 1024; // 5 MB per session doc
// Per-WebSocket-frame cap. It MUST exceed the doc cap: a full-state sync
// (initial load, or any reconnect) puts the entire doc into a single frame,
// so a frame limit below MAX_DOC_BYTES makes any legal-but-larger-than-the-
// -frame doc permanently unsyncable — the client's sync frame is rejected,
// it reconnects, resends the same frame, and loops forever (observed at the
// old 256 KB cap: a 305 KB / 4-waypoint route froze in a reconnect storm).
// The doc cap remains the size limit; this just leaves room to transmit it.
export const MAX_MESSAGE_BYTES = MAX_DOC_BYTES + 256 * 1024; // doc + protocol overhead
// WebSocket close code for policy-violation responses (1008 is the
// standard RFC 6455 code for "received a message that violates policy").
const POLICY_VIOLATION = 1008;

const docs = new Map<string, Y.Doc>();
const awarenessMap = new Map<string, awarenessProtocol.Awareness>();
const conns = new Map<WebSocket, { sessionId: string; clientIds: Set<number> }>();
const oversizedSessions = new Set<string>();

/**
 * Current persisted size of a session's Yjs doc — same bytes that
 * would land in `sessions.yjs_state` if we saved right now. Exported
 * so tests can assert the size-cap accounting directly.
 */
export function docByteSize(sessionId: string): number {
  const doc = docs.get(sessionId);
  if (!doc) return 0;
  return Y.encodeStateAsUpdate(doc).byteLength;
}

/**
 * Refuse further activity on a session whose doc exceeded the cap.
 * Close every connected client and mark the session so the next
 * persisted save is skipped (avoid writing an oversized blob to DB).
 */
function quarantineSession(sessionId: string, reason: string): void {
  oversizedSessions.add(sessionId);
  for (const [ws, meta] of conns) {
    if (meta.sessionId === sessionId && ws.readyState === ws.OPEN) {
      try { ws.close(POLICY_VIOLATION, reason); } catch { /* ignore */ }
    }
  }
}

/**
 * Count the number of distinct session ids represented in an iterable
 * of connection metadata. Pure, for testability — the server calls it
 * with the `conns` map's values.
 */
export function countActiveSessions(
  connsIter: Iterable<{ sessionId: string }>,
): number {
  const sessionIds = new Set<string>();
  for (const { sessionId } of connsIter) sessionIds.add(sessionId);
  return sessionIds.size;
}

/**
 * Set the active-sessions gauge from the actual `conns` state. Derived
 * each time, so it can't drift negative — the previous inc/dec pattern
 * leaked decrements whenever a reconnect found the doc cached in
 * `docs` and skipped the matching increment.
 */
function refreshActiveSessions(): void {
  plannerActiveSessions.set(countActiveSessions(conns.values()));
}

export function getOrCreateDoc(sessionId: string): Y.Doc {
  let doc = docs.get(sessionId);
  if (!doc) {
    doc = new Y.Doc();
    docs.set(sessionId, doc);
  }
  return doc;
}

export async function getOrLoadDoc(sessionId: string): Promise<Y.Doc> {
  let doc = docs.get(sessionId);
  if (!doc) {
    doc = new Y.Doc();
    const savedState = await loadSessionState(sessionId);
    if (savedState) {
      Y.applyUpdate(doc, savedState);
    }
    docs.set(sessionId, doc);
  }
  return doc;
}

function getAwareness(sessionId: string, doc: Y.Doc): awarenessProtocol.Awareness {
  let awareness = awarenessMap.get(sessionId);
  if (!awareness) {
    awareness = new awarenessProtocol.Awareness(doc);
    awarenessMap.set(sessionId, awareness);

    awareness.on("update", ({ added, updated, removed }: { added: number[]; updated: number[]; removed: number[] }) => {
      const changedClients = added.concat(updated, removed);
      const encoder = encoding.createEncoder();
      encoding.writeVarUint(encoder, messageAwareness);
      encoding.writeVarUint8Array(encoder, awarenessProtocol.encodeAwarenessUpdate(awareness!, changedClients));
      const message = encoding.toUint8Array(encoder);

      for (const [ws, meta] of conns) {
        if (meta.sessionId === sessionId && ws.readyState === ws.OPEN) {
          ws.send(message);
        }
      }
    });
  }
  return awareness;
}

export function deleteDoc(sessionId: string): boolean {
  const doc = docs.get(sessionId);
  if (doc) {
    doc.destroy();
    docs.delete(sessionId);
    const awareness = awarenessMap.get(sessionId);
    if (awareness) {
      awareness.destroy();
      awarenessMap.delete(sessionId);
    }
    return true;
  }
  return false;
}

export function getDocCount(): number {
  return docs.size;
}

export function getClientCount(sessionId: string): number {
  let count = 0;
  for (const meta of conns.values()) {
    if (meta.sessionId === sessionId) count++;
  }
  return count;
}

// Debounced persistence — save at most every 5 seconds per session
const saveTimers = new Map<string, ReturnType<typeof setTimeout>>();

function debouncedSave(sessionId: string): void {
  if (saveTimers.has(sessionId)) return;
  saveTimers.set(
    sessionId,
    setTimeout(async () => {
      saveTimers.delete(sessionId);
      try {
        await saveSessionState(sessionId);
      } catch (_e) {
        // Log but don't crash on save failure
      }
    }, 5000),
  );
}

function handleMessage(ws: WebSocket, message: Uint8Array, sessionId: string) {
  const doc = docs.get(sessionId);
  if (!doc) return;
  if (oversizedSessions.has(sessionId)) return;

  // Per-frame cap. Anything larger than 256 KB on a Yjs session is
  // pathological — a single waypoint add is hundreds of bytes — so the
  // sender is either buggy or hostile.
  if (message.byteLength > MAX_MESSAGE_BYTES) {
    try { ws.close(POLICY_VIOLATION, "message too large"); } catch { /* ignore */ }
    return;
  }

  const awareness = awarenessMap.get(sessionId);
  const decoder = decoding.createDecoder(message);
  const messageType = decoding.readVarUint(decoder);

  switch (messageType) {
    case messageSync: {
      const encoder = encoding.createEncoder();
      encoding.writeVarUint(encoder, messageSync);
      syncProtocol.readSyncMessage(decoder, encoder, doc, ws);
      const resp = encoding.toUint8Array(encoder);
      if (encoding.length(encoder) > 1) {
        ws.send(resp);
      }
      // Per-doc cap. Apply succeeded above; if the doc just grew past
      // the limit, quarantine before debouncing the next save so we
      // never write an oversized blob to Postgres.
      if (docByteSize(sessionId) > MAX_DOC_BYTES) {
        quarantineSession(sessionId, "doc size cap exceeded");
        break;
      }
      debouncedSave(sessionId);
      break;
    }
    case messageAwareness: {
      if (awareness) {
        const update = decoding.readVarUint8Array(decoder);
        awarenessProtocol.applyAwarenessUpdate(awareness, update, ws);

        const meta = conns.get(ws);
        if (meta) {
          awareness.getStates().forEach((_state, clientId) => {
            meta.clientIds.add(clientId);
          });
        }
      }
      break;
    }
  }
}

export function setupYjsWebSocket(server: Server): WebSocketServer {
  const wss = new WebSocketServer({ noServer: true });

  server.on("upgrade", async (request, socket, head) => {
    const url = new URL(request.url ?? "", `http://${request.headers.host}`);
    const match = url.pathname.match(/^\/sync\/(.+)$/);

    if (!match) {
      socket.destroy();
      return;
    }

    const sessionId = match[1]!;
    // Reject upgrades for sessions that don't exist or are already
    // closed/expired. Before this guard, `getOrLoadDoc` happily
    // *created* a fresh Y.Doc for any sessionId, so an attacker
    // hitting /sync/<random-uuid> repeatedly could exhaust memory
    // and resurrect closed sessions. `getSession()` already filters
    // out `closed = true` rows.
    try {
      const session = await getSession(sessionId);
      if (!session) {
        socket.destroy();
        return;
      }
    } catch {
      // DB unreachable — fail closed (refuse) rather than open. The
      // client will reconnect after backoff once the DB recovers.
      socket.destroy();
      return;
    }

    wss.handleUpgrade(request, socket, head, (ws) => {
      wss.emit("connection", ws, request, sessionId);
    });
  });

  wss.on("connection", async (ws: WebSocket, _request: IncomingMessage, sessionId: string) => {
    const doc = await getOrLoadDoc(sessionId);
    const awareness = getAwareness(sessionId, doc);

    conns.set(ws, { sessionId, clientIds: new Set() });
    plannerConnectedClients.inc();
    refreshActiveSessions();

    // Broadcast doc updates to all connections in this session
    const onUpdate = (update: Uint8Array, origin: unknown) => {
      if (origin === ws) return;
      const encoder = encoding.createEncoder();
      encoding.writeVarUint(encoder, messageSync);
      syncProtocol.writeUpdate(encoder, update);
      const message = encoding.toUint8Array(encoder);

      for (const [client, meta] of conns) {
        if (meta.sessionId === sessionId && client !== ws && client.readyState === ws.OPEN) {
          client.send(message);
        }
      }
    };
    doc.on("update", onUpdate);

    // Send sync step 1
    const encoder = encoding.createEncoder();
    encoding.writeVarUint(encoder, messageSync);
    syncProtocol.writeSyncStep1(encoder, doc);
    ws.send(encoding.toUint8Array(encoder));

    // Send current awareness states
    const awarenessStates = awareness.getStates();
    if (awarenessStates.size > 0) {
      const awarenessEncoder = encoding.createEncoder();
      encoding.writeVarUint(awarenessEncoder, messageAwareness);
      encoding.writeVarUint8Array(
        awarenessEncoder,
        awarenessProtocol.encodeAwarenessUpdate(awareness, Array.from(awarenessStates.keys())),
      );
      ws.send(encoding.toUint8Array(awarenessEncoder));
    }

    await touchSession(sessionId);

    ws.on("message", (data: ArrayBuffer) => {
      try {
        handleMessage(ws, new Uint8Array(data), sessionId);
      } catch (_e) {
        // Ignore malformed messages
      }
    });

    ws.on("close", () => {
      const meta = conns.get(ws);
      if (meta && meta.clientIds.size > 0) {
        awarenessProtocol.removeAwarenessStates(
          awareness,
          Array.from(meta.clientIds),
          null,
        );
      }
      conns.delete(ws);
      plannerConnectedClients.dec();
      refreshActiveSessions();
      doc.off("update", onUpdate);

      // Save when last client leaves
      const hasClients = Array.from(conns.values()).some((m) => m.sessionId === sessionId);
      if (!hasClients) {
        saveSessionState(sessionId).catch(() => {});
      }
    });
  });

  return wss;
}