Shade/packages/shade-sdk/tests/webrtc-integration.test.ts

/**
 * V3.11 — full SDK integration: two Shade instances exchange a transfer
 * over the in-process `MemoryRtcFactory`. The WebRTC transport sits on
 * top of `MultiTransportFallback([webrtc, http])`, so this also verifies
 * the SDK wires the fallback chain correctly.
 */
import { describe, test, expect, beforeAll, afterAll } from 'bun:test';
import {
  createShade,
  type Shade,
  type TransferHandle,
  type TransferResult,
} from '../src/index.js';
import {
  createPrekeyServer,
  MemoryPrekeyStore,
  PrekeyServerEvents,
} from '@shade/server';
import { SubtleCryptoProvider } from '@shade/crypto-web';
import { sha256Once } from '@shade/streams';
import { MemoryRtcFactory } from '@shade/transport-webrtc';

const crypto = new SubtleCryptoProvider();

interface Rig {
  alice: Shade;
  bob: Shade;
  aliceBaseUrl: string;
  bobBaseUrl: string;
  prekeyStop: () => void;
  aliceServerStop: () => void;
  bobServerStop: () => void;
}

async function startPrekeyServer(): Promise<{ url: string; stop: () => void }> {
  const events = new PrekeyServerEvents();
  const server = createPrekeyServer({
    crypto,
    store: new MemoryPrekeyStore(),
    disableRateLimit: true,
    events,
  });
  const port = 22000 + Math.floor(Math.random() * 500);
  const handle = Bun.serve({ port, fetch: server.fetch });
  return { url: `http://localhost:${port}`, stop: () => handle.stop() };
}

async function setupRig(): Promise<Rig> {
  const prekey = await startPrekeyServer();
  const alice = await createShade({ prekeyServer: prekey.url, address: 'alice' });
  const bob = await createShade({ prekeyServer: prekey.url, address: 'bob' });

  // Both peers need bidirectional resolveBaseUrl since signaling envelopes
  // ride the control plane in BOTH directions (offer one way, answer
  // back). Static map for this test rig.
  const baseUrls = new Map<string, string>();
  const resolveBaseUrl = async (addr: string): Promise<string> => {
    const url = baseUrls.get(addr);
    if (url === undefined) throw new Error(`unknown peer ${addr}`);
    return url;
  };
  alice.configureTransfers({ resolveBaseUrl });
  bob.configureTransfers({ resolveBaseUrl });

  // V3.11: opt-in to WebRTC BEFORE the engine is built (transferRoute
  // builds it lazily). Both peers use the same in-process factory so
  // their PCs can pair up via the registry.
  const factory = new MemoryRtcFactory();
  alice.configureWebRTC({ factory, connectTimeoutMs: 10_000 });
  bob.configureWebRTC({ factory, connectTimeoutMs: 10_000 });

  const bobApp = await bob.transferRoute();
  const bobPort = 22500 + Math.floor(Math.random() * 500);
  const bobServer = Bun.serve({ port: bobPort, fetch: bobApp.fetch });
  const bobBaseUrl = `http://localhost:${bobPort}`;

  const aliceApp = await alice.transferRoute();
  const alicePort = 22000 + Math.floor(Math.random() * 500);
  const aliceServer = Bun.serve({ port: alicePort, fetch: aliceApp.fetch });
  const aliceBaseUrl = `http://localhost:${alicePort}`;

  baseUrls.set('alice', aliceBaseUrl);
  baseUrls.set('bob', bobBaseUrl);

  return {
    alice,
    bob,
    aliceBaseUrl,
    bobBaseUrl,
    prekeyStop: prekey.stop,
    aliceServerStop: () => aliceServer.stop(),
    bobServerStop: () => bobServer.stop(),
  };
}

async function teardownRig(rig: Rig): Promise<void> {
  await rig.alice.shutdown();
  await rig.bob.shutdown();
  rig.bobServerStop();
  rig.aliceServerStop();
  rig.prekeyStop();
  MemoryRtcFactory.reset();
}

function hex(b: Uint8Array): string {
  return Array.from(b, (x) => x.toString(16).padStart(2, '0')).join('');
}

async function uploadAndAwait(
  rig: Rig,
  input: Uint8Array,
  opts?: { lanes?: number; chunkSize?: number },
): Promise<{ senderResult: TransferResult; received: Uint8Array }> {
  let resolveRecv!: (h: TransferHandle) => void;
  const recvHandlePromise = new Promise<TransferHandle>((r) => {
    resolveRecv = r;
  });
  const unsubscribe = await rig.bob.onIncomingTransfer(async (incoming) => {
    const h = await incoming.accept({ output: { kind: 'buffer' } });
    resolveRecv(h);
  });

  const handle = await rig.alice.upload({
    to: 'bob',
    input,
    ...(opts?.lanes !== undefined ? { lanes: opts.lanes } : {}),
    ...(opts?.chunkSize !== undefined ? { chunkSize: opts.chunkSize } : {}),
    metadata: { name: 'webrtc-test.bin' },
  });
  const recvHandle = await recvHandlePromise;
  const [senderResult, recvResult] = await Promise.all([
    handle.done(),
    recvHandle.done(),
  ]);
  unsubscribe();
  const received =
    (recvResult as TransferResult & { bytes?: Uint8Array }).bytes ?? new Uint8Array();
  return { senderResult, received };
}

describe('V3.11 WebRTC integration via MemoryRtcFactory', () => {
  let rig: Rig;
  beforeAll(async () => {
    rig = await setupRig();
  });
  afterAll(async () => {
    await teardownRig(rig);
  });

  test('256 KiB payload over WebRTC primary', async () => {
    const input = crypto.randomBytes(256 * 1024);
    const { senderResult, received } = await uploadAndAwait(rig, input, {
      lanes: 1,
      chunkSize: 64 * 1024,
    });
    expect(received).toEqual(input);
    expect(senderResult.sha256).toBe(hex(sha256Once(input)));

    // Verify the WebRTC runtime is alive and the multi-fallback hasn't
    // demoted away from webrtc.
    const runtime = rig.alice.getWebRtcRuntime();
    expect(runtime).not.toBeNull();
    expect(runtime!.fallback.activeName).toBe('webrtc');
    expect(runtime!.fallback.hasFallenBack).toBe(false);
    expect(runtime!.manager.isConnected('bob')).toBe(true);
  });

  test('1 MiB payload — 4 lanes range partition over WebRTC', async () => {
    const input = crypto.randomBytes(1024 * 1024);
    const { received } = await uploadAndAwait(rig, input, {
      lanes: 4,
      chunkSize: 64 * 1024,
    });
    expect(received).toEqual(input);

    const runtime = rig.alice.getWebRtcRuntime();
    expect(runtime!.fallback.activeName).toBe('webrtc');
  });
});