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

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';

const crypto = new SubtleCryptoProvider();

interface TestRig {
  alice: Shade;
  bob: Shade;
  bobBaseUrl: string;
  prekeyStop: () => 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 = 21000 + 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<TestRig> {
  const prekey = await startPrekeyServer();
  const alice = await createShade({ prekeyServer: prekey.url, address: 'alice' });
  const bob = await createShade({ prekeyServer: prekey.url, address: 'bob' });

  // Bob's transferRoute lazily creates a TransferEngine, which requires
  // configureTransfers first.
  bob.configureTransfers({
    resolveBaseUrl: async () => {
      throw new Error('bob is receive-only in this test');
    },
  });

  // Spin up Bob's HTTP transfer endpoint.
  const bobApp = await bob.transferRoute();
  const port = 21500 + Math.floor(Math.random() * 500);
  const bobServer = Bun.serve({ port, fetch: bobApp.fetch });
  const bobBaseUrl = `http://localhost:${port}`;

  // Wire up Alice's outgoing transfer routing.
  alice.configureTransfers({
    resolveBaseUrl: async (addr) => {
      if (addr === 'bob') return bobBaseUrl;
      throw new Error(`unknown peer ${addr}`);
    },
  });

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

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

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

async function uploadAndAwait(
  rig: TestRig,
  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: 'integration-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('Shade SDK end-to-end E2EE transfer', () => {
  let rig: TestRig;
  beforeAll(async () => {
    rig = await setupRig();
  });
  afterAll(async () => {
    await teardownRig(rig);
  });

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

  test('512 KiB payload — 4 lanes range partition', async () => {
    const input = crypto.randomBytes(512 * 1024);
    const { senderResult, received } = await uploadAndAwait(rig, input, {
      lanes: 4,
      chunkSize: 32 * 1024,
    });
    expect(received).toEqual(input);
    expect(senderResult.sha256).toBe(hex(sha256Once(input)));
  });

  test(
    '4 MiB payload — 4 lanes, simulates Dispatch upload',
    async () => {
      const input = crypto.randomBytes(4 * 1024 * 1024);
      const { senderResult, received } = await uploadAndAwait(rig, input, {
        lanes: 4,
        chunkSize: 128 * 1024,
      });
      expect(received).toEqual(input);
      expect(senderResult.sha256).toBe(hex(sha256Once(input)));
    },
    30_000,
  );

  test('listTransfers returns empty before transfer (memory storage)', async () => {
    // MemoryStorage's listActiveStreamStates is implemented but never
    // populated by the engine in v0.2.0 (resume persistence is M-Stream-6).
    const list = await rig.alice.listTransfers();
    expect(Array.isArray(list)).toBe(true);
  });
});