Shade/packages/shade-recovery/tests/helpers.ts

/**
 * Test helpers — boot a local prekey server, mint Shade instances, and
 * pair them with an in-process delivery transport that calls
 * `shade.receive` on the recipient. Mirrors the pattern other Shade
 * test suites use, kept private to this package so we don't pull in
 * `@shade/server` at runtime.
 */

import { createPrekeyServer, MemoryPrekeyStore, PrekeyServerEvents } from '@shade/server';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { createShade, type Shade } from '@shade/sdk';
import type { ShadeEnvelope } from '@shade/core';
import type { RecoveryDeliver } from '../src/setup.js';

export interface TestEnv {
  prekeyUrl: string;
  stop: () => void;
}

export async function startTestPrekeyServer(): Promise<TestEnv> {
  const crypto = new SubtleCryptoProvider();
  const prekey = createPrekeyServer({
    crypto,
    store: new MemoryPrekeyStore(),
    disableRateLimit: true,
    events: new PrekeyServerEvents(),
  });
  const server = Bun.serve({ port: 0, fetch: prekey.fetch });
  const port = (server as unknown as { port: number }).port;
  return {
    prekeyUrl: `http://localhost:${port}`,
    stop: () => server.stop(),
  };
}

export async function spawnShade(prekeyUrl: string, address: string): Promise<Shade> {
  return createShade({ prekeyServer: prekeyUrl, address });
}

/**
 * In-process transport that delivers `(to, envelope)` to the named
 * Shade by calling `shade.receive(from, envelope)` on it. Matches the
 * `RecoveryDeliver` shape so it can be plugged into setup/guardian/
 * request flows.
 *
 * Construction order matters: register every party before delivering
 * so the lookup never fails. Use `addr` to keep the from-address
 * symmetric with what `Shade.send` uses.
 */
export class MemoryRecoveryTransport {
  private readonly directory = new Map<string, Shade>();
  /** Per-pair pending-deliveries chain to preserve ordering. */
  private readonly chains = new Map<string, Promise<unknown>>();
  /** Counters of how many envelopes flowed in each direction (for tests). */
  public readonly delivered: Array<{ from: string; to: string }> = [];
  /**
   * Optional drop-after-N policy used to simulate guardians that go
   * unreachable. Keyed on `to` address. Set with `dropAfter(addr, n)`.
   */
  private readonly dropPolicies = new Map<string, number>();

  add(shade: Shade): void {
    this.directory.set(shade.myAddress, shade);
  }

  /** Drop envelopes addressed to `to` after the first `n` have flowed. */
  dropAfter(to: string, n: number): void {
    this.dropPolicies.set(to, n);
  }

  /**
   * Build a `RecoveryDeliver` callback bound to `from`. Call this once
   * per Shade so its outbound sends route through the transport.
   */
  bind(from: Shade): RecoveryDeliver {
    return async (to: string, envelope: ShadeEnvelope) => {
      const recipient = this.directory.get(to);
      if (recipient === undefined) {
        throw new Error(`MemoryRecoveryTransport: unknown recipient "${to}"`);
      }
      // Apply drop policy.
      const policy = this.dropPolicies.get(to);
      if (policy !== undefined) {
        if (policy <= 0) throw new Error(`MemoryRecoveryTransport: dropping for "${to}"`);
        this.dropPolicies.set(to, policy - 1);
      }
      this.delivered.push({ from: from.myAddress, to });
      // Serialize per (from→to) pair to preserve ordering.
      const key = `${from.myAddress}→${to}`;
      const prev = this.chains.get(key) ?? Promise.resolve();
      const next = prev.then(async () => {
        await recipient.receive(from.myAddress, envelope);
      });
      this.chains.set(
        key,
        next.catch(() => {
          // chain even on failures so subsequent sends don't deadlock
          return undefined;
        }),
      );
      await next;
    };
  }
}