dchain/client-app/lib/ws.ts

/**
 * DChain WebSocket client — replaces balance / inbox / contacts polling with
 * server-push. Matches `node/ws.go` exactly.
 *
 * Usage:
 *   const ws = getWSClient();
 *   ws.connect();                              // idempotent
 *   const off = ws.subscribe('addr:ab12…', ev => { ... });
 *   // later:
 *   off();                                     // unsubscribe + stop handler
 *   ws.disconnect();
 *
 * Features:
 *   - Auto-reconnect with exponential backoff (1s → 30s cap).
 *   - Re-subscribes all topics after a reconnect.
 *   - `hello` frame exposes chain_id + tip_height for connection state UI.
 *   - Degrades silently if the endpoint returns 501 (old node without WS).
 */

import { getNodeUrl, onNodeUrlChange } from './api';
import { sign } from './crypto';

export type WSEventName =
  | 'hello'
  | 'block'
  | 'tx'
  | 'contract_log'
  | 'inbox'
  | 'typing'
  | 'pong'
  | 'error'
  | 'subscribed'
  | 'submit_ack'
  | 'lag';

export interface WSFrame {
  event:       WSEventName;
  data?:       unknown;
  topic?:      string;
  msg?:        string;
  chain_id?:   string;
  tip_height?: number;
  /** Server-issued nonce in the hello frame; client signs it for auth. */
  auth_nonce?: string;
  // submit_ack fields
  id?:         string;
  status?:     'accepted' | 'rejected';
  reason?:     string;
}

type Handler = (frame: WSFrame) => void;

class WSClient {
  private ws:          WebSocket | null = null;
  private url:         string | null = null;
  private reconnectMs: number = 1000;
  private closing:     boolean = false;

  /** topic → set of handlers interested in frames for this topic */
  private handlers: Map<string, Set<Handler>> = new Map();
  /** topics we want the server to push — replayed on every reconnect */
  private wantedTopics: Set<string> = new Set();

  private connectionListeners: Set<(ok: boolean, err?: string) => void> = new Set();
  private helloInfo: { chainId?: string; tipHeight?: number; authNonce?: string } = {};

  /**
   * Credentials used for auto-auth on every (re)connect. The signer runs on
   * each hello frame so scoped subscriptions (addr:*, inbox:*) are accepted.
   * Without these, subscribe requests to scoped topics get rejected by the
   * server; global topics (blocks, tx, …) still work unauthenticated.
   */
  private authCreds: { pubKey: string; privKey: string } | null = null;

  /** Current connection state (read-only for UI). */
  isConnected(): boolean {
    return this.ws?.readyState === WebSocket.OPEN;
  }
  getHelloInfo(): { chainId?: string; tipHeight?: number } {
    return this.helloInfo;
  }

  /** Subscribe to a connection-state listener — fires on connect/disconnect. */
  onConnectionChange(cb: (ok: boolean, err?: string) => void): () => void {
    this.connectionListeners.add(cb);
    return () => this.connectionListeners.delete(cb) as unknown as void;
  }
  private fireConnectionChange(ok: boolean, err?: string) {
    for (const cb of this.connectionListeners) {
      try { cb(ok, err); } catch { /* noop */ }
    }
  }

  /**
   * Register the Ed25519 keypair used for auto-auth. The signer runs on each
   * (re)connect against the server-issued nonce so the connection is bound
   * to this identity. Pass null to disable auth (only global topics will
   * work — useful for observers).
   */
  setAuthCreds(creds: { pubKey: string; privKey: string } | null): void {
    this.authCreds = creds;
    // If we're already connected, kick off auth immediately.
    if (creds && this.isConnected() && this.helloInfo.authNonce) {
      this.sendAuth(this.helloInfo.authNonce);
    }
  }

  /** Idempotent connect. Call once on app boot. */
  connect(): void {
    const base = getNodeUrl();
    const newURL = base.replace(/^http/, 'ws') + '/api/ws';
    if (this.ws) {
      const state = this.ws.readyState;
      // Already pointing at this URL and connected / connecting — nothing to do.
      if (this.url === newURL && (state === WebSocket.OPEN || state === WebSocket.CONNECTING)) {
        return;
      }
      // URL changed (operator flipped nodes in settings) — tear down and
      // re-dial. Existing subscriptions live in wantedTopics and will be
      // replayed after the new onopen fires.
      if (this.url !== newURL && (state === WebSocket.OPEN || state === WebSocket.CONNECTING)) {
        try { this.ws.close(); } catch { /* noop */ }
        this.ws = null;
      }
    }
    this.closing = false;
    this.url = newURL;
    try {
      this.ws = new WebSocket(this.url);
    } catch (e: any) {
      this.fireConnectionChange(false, e?.message ?? 'ws construct failed');
      this.scheduleReconnect();
      return;
    }

    this.ws.onopen = () => {
      this.reconnectMs = 1000;  // reset backoff
      this.fireConnectionChange(true);
      // Replay all wanted subscriptions.
      for (const topic of this.wantedTopics) {
        this.sendRaw({ op: 'subscribe', topic });
      }
    };

    this.ws.onmessage = (ev) => {
      let frame: WSFrame;
      try {
        frame = JSON.parse(typeof ev.data === 'string' ? ev.data : '');
      } catch {
        return;
      }
      if (frame.event === 'hello') {
        this.helloInfo = {
          chainId:   frame.chain_id,
          tipHeight: frame.tip_height,
          authNonce: frame.auth_nonce,
        };
        // Auto-authenticate if credentials are set. The server binds this
        // connection to the signed pubkey so scoped subscriptions (addr:*,
        // inbox:*) get through. On reconnect a new nonce is issued, so the
        // auth dance repeats transparently.
        if (this.authCreds && frame.auth_nonce) {
          this.sendAuth(frame.auth_nonce);
        }
      }
      // Dispatch to all handlers for any topic that could match this frame.
      // We use a simple predicate: look at the frame to decide which topics it
      // was fanned out to, then fire every matching handler.
      for (const topic of this.topicsForFrame(frame)) {
        const set = this.handlers.get(topic);
        if (!set) continue;
        for (const h of set) {
          try { h(frame); } catch (e) { console.warn('[ws] handler error', e); }
        }
      }
    };

    this.ws.onerror = (e: any) => {
      this.fireConnectionChange(false, 'ws error');
    };

    this.ws.onclose = () => {
      this.ws = null;
      this.fireConnectionChange(false);
      if (!this.closing) this.scheduleReconnect();
    };
  }

  disconnect(): void {
    this.closing = true;
    if (this.ws) {
      try { this.ws.close(); } catch { /* noop */ }
      this.ws = null;
    }
  }

  /**
   * Subscribe to a topic. Returns an `off()` function that unsubscribes AND
   * removes the handler. If multiple callers subscribe to the same topic,
   * the server is only notified on the first and last caller.
   */
  subscribe(topic: string, handler: Handler): () => void {
    let set = this.handlers.get(topic);
    if (!set) {
      set = new Set();
      this.handlers.set(topic, set);
    }
    set.add(handler);

    // Notify server only on the first handler for this topic.
    if (!this.wantedTopics.has(topic)) {
      this.wantedTopics.add(topic);
      if (this.isConnected()) {
        this.sendRaw({ op: 'subscribe', topic });
      } else {
        this.connect();  // lazy-connect on first subscribe
      }
    }

    return () => {
      const s = this.handlers.get(topic);
      if (!s) return;
      s.delete(handler);
      if (s.size === 0) {
        this.handlers.delete(topic);
        this.wantedTopics.delete(topic);
        if (this.isConnected()) {
          this.sendRaw({ op: 'unsubscribe', topic });
        }
      }
    };
  }

  /** Force a keepalive ping. Useful for debugging. */
  ping(): void {
    this.sendRaw({ op: 'ping' });
  }

  /**
   * Send a typing indicator to another user. Recipient is their X25519 pubkey
   * (the one used for inbox encryption). Ephemeral — no ack, no retry; just
   * fire and forget. Call on each keystroke but throttle to once per 2-3s
   * at the caller side so we don't flood the WS with frames.
   */
  sendTyping(recipientX25519: string): void {
    if (!this.isConnected()) return;
    try {
      this.ws!.send(JSON.stringify({ op: 'typing', to: recipientX25519 }));
    } catch { /* best-effort */ }
  }

  /**
   * Submit a signed transaction over the WebSocket and resolve once the
   * server returns a `submit_ack`. Saves the HTTP round-trip on every tx
   * and gives the UI immediate accept/reject feedback.
   *
   * Rejects if:
   *   - WS is not connected (caller should fall back to HTTP)
   *   - Server returns `status: "rejected"` — `reason` is surfaced as error msg
   *   - No ack within `timeoutMs` (default 10 s)
   */
  submitTx(tx: unknown, timeoutMs = 10_000): Promise<{ id: string }> {
    if (!this.isConnected()) {
      return Promise.reject(new Error('WS not connected'));
    }
    const reqId = 's_' + Date.now() + '_' + Math.random().toString(36).slice(2, 8);

    return new Promise((resolve, reject) => {
      const off = this.subscribe('$system', (frame) => {
        if (frame.event !== 'submit_ack' || frame.id !== reqId) return;
        off();
        clearTimeout(timer);
        if (frame.status === 'accepted') {
          // `msg` carries the server-confirmed tx id.
          resolve({ id: typeof frame.msg === 'string' ? frame.msg : '' });
        } else {
          reject(new Error(frame.reason || 'submit_tx rejected'));
        }
      });
      const timer = setTimeout(() => {
        off();
        reject(new Error('submit_tx timeout (' + timeoutMs + 'ms)'));
      }, timeoutMs);

      try {
        this.ws!.send(JSON.stringify({ op: 'submit_tx', tx, id: reqId }));
      } catch (e: any) {
        off();
        clearTimeout(timer);
        reject(new Error('WS send failed: ' + (e?.message ?? 'unknown')));
      }
    });
  }

  // ── internals ───────────────────────────────────────────────────────────

  private scheduleReconnect(): void {
    if (this.closing) return;
    const delay = Math.min(this.reconnectMs, 30_000);
    this.reconnectMs = Math.min(this.reconnectMs * 2, 30_000);
    setTimeout(() => {
      if (!this.closing) this.connect();
    }, delay);
  }

  private sendRaw(cmd: { op: string; topic?: string }): void {
    if (!this.isConnected()) return;
    try { this.ws!.send(JSON.stringify(cmd)); } catch { /* noop */ }
  }

  /**
   * Sign the server nonce with our Ed25519 private key and send the `auth`
   * op. The server binds this connection to `authCreds.pubKey`; subsequent
   * subscribe requests to `addr:<pubKey>` / `inbox:<my_x25519>` are accepted.
   */
  private sendAuth(nonce: string): void {
    if (!this.authCreds || !this.isConnected()) return;
    try {
      const bytes = new TextEncoder().encode(nonce);
      const sig = sign(bytes, this.authCreds.privKey);
      this.ws!.send(JSON.stringify({
        op:     'auth',
        pubkey: this.authCreds.pubKey,
        sig,
      }));
    } catch (e) {
      console.warn('[ws] auth send failed:', e);
    }
  }

  /**
   * Given an incoming frame, enumerate every topic that handlers could have
   * subscribed to and still be interested. This mirrors the fan-out logic in
   * node/ws.go:EmitBlock / EmitTx / EmitContractLog.
   */
  private topicsForFrame(frame: WSFrame): string[] {
    switch (frame.event) {
      case 'block':
        return ['blocks'];
      case 'tx': {
        const d = frame.data as { from?: string; to?: string } | undefined;
        const topics = ['tx'];
        if (d?.from) topics.push('addr:' + d.from);
        if (d?.to && d.to !== d.from) topics.push('addr:' + d.to);
        return topics;
      }
      case 'contract_log': {
        const d = frame.data as { contract_id?: string } | undefined;
        const topics = ['contract_log'];
        if (d?.contract_id) topics.push('contract:' + d.contract_id);
        return topics;
      }
      case 'inbox': {
        // Node fans inbox events to `inbox` + `inbox:<recipient_x25519>`;
        // we mirror that here so both firehose listeners and address-scoped
        // subscribers see the event.
        const d = frame.data as { recipient_pub?: string } | undefined;
        const topics = ['inbox'];
        if (d?.recipient_pub) topics.push('inbox:' + d.recipient_pub);
        return topics;
      }
      case 'typing': {
        // Server fans to `typing:<to>` only (the recipient).
        const d = frame.data as { to?: string } | undefined;
        return d?.to ? ['typing:' + d.to] : [];
      }
      // Control-plane events — no topic fan-out; use a pseudo-topic so UI
      // can listen for them via subscribe('$system', ...).
      case 'hello':
      case 'pong':
      case 'error':
      case 'subscribed':
      case 'submit_ack':
      case 'lag':
        return ['$system'];
      default:
        return [];
    }
  }
}

let _singleton: WSClient | null = null;

/**
 * Return the app-wide WebSocket client. Safe to call from any component;
 * `.connect()` is idempotent.
 *
 * On first creation we register a node-URL listener so flipping the node
 * in Settings tears down the existing socket and dials the new one — the
 * user's active subscriptions (addr:*, inbox:*) replay automatically.
 */
export function getWSClient(): WSClient {
  if (!_singleton) {
    _singleton = new WSClient();
    onNodeUrlChange(() => {
      // Fire and forget — connect() is idempotent and handles stale URLs.
      _singleton!.connect();
    });
  }
  return _singleton;
}