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

import { describe, test, expect, beforeEach, afterEach } from 'bun:test';
import { createShade, type Shade } from '../src/index.js';
import {
  createPrekeyServer,
  MemoryPrekeyStore,
  PrekeyServerEvents,
} from '@shade/server';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';

const crypto = new SubtleCryptoProvider();

/**
 * Spin up a real prekey server on a random port and return its URL
 * + a teardown function.
 */
async function startPrekeyServer(): Promise<{
  url: string;
  stop: () => void;
  events: PrekeyServerEvents;
}> {
  const events = new PrekeyServerEvents();
  const server = createPrekeyServer({
    crypto,
    store: new MemoryPrekeyStore(),
    disableRateLimit: true,
    events,
  });
  const port = 19500 + Math.floor(Math.random() * 500);
  const handle = Bun.serve({ port, fetch: server.fetch });
  return {
    url: `http://localhost:${port}`,
    stop: () => handle.stop(),
    events,
  };
}

describe('createShade — happy path', () => {
  let server: Awaited<ReturnType<typeof startPrekeyServer>>;
  let alice: Shade;
  let bob: Shade;

  beforeEach(async () => {
    server = await startPrekeyServer();
  });

  afterEach(async () => {
    await alice?.shutdown();
    await bob?.shutdown();
    server.stop();
  });

  test('one-liner creation and initialization', async () => {
    alice = await createShade({
      prekeyServer: server.url,
      address: 'alice',
    });

    expect(alice.myAddress).toBe('alice');
    const fp = await alice.fingerprint;
    expect(fp.split(' ').length).toBe(12);
  });

  test('auto-publishes bundle on init', async () => {
    const registered: string[] = [];
    server.events.on((e) => {
      if (e.name === 'server.identity_registered') {
        registered.push(e.data.address);
      }
    });

    alice = await createShade({
      prekeyServer: server.url,
      address: 'alice',
    });

    expect(registered).toContain('alice');
  });

  test('two-process conversation: Alice ↔ Bob via SDK only', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    bob = await createShade({ prekeyServer: server.url, address: 'bob' });

    // Alice sends to Bob — SDK auto-establishes session
    const env1 = await alice.send('bob', 'hello Bob');
    const plain1 = await bob.receive('alice', env1);
    expect(plain1).toBe('hello Bob');

    // Bob replies (DH ratchet triggers)
    const env2 = await bob.send('alice', 'hi Alice');
    const plain2 = await alice.receive('bob', env2);
    expect(plain2).toBe('hi Alice');
  });

  test('onMessage handler fires on receive', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    bob = await createShade({ prekeyServer: server.url, address: 'bob' });

    const received: Array<{ from: string; msg: string }> = [];
    bob.onMessage((from, msg) => received.push({ from, msg }));

    const env = await alice.send('bob', 'callback test');
    await bob.receive('alice', env);

    expect(received.length).toBe(1);
    expect(received[0]!.from).toBe('alice');
    expect(received[0]!.msg).toBe('callback test');
  });

  test('concurrent sends to same new peer establish exactly one session', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    bob = await createShade({ prekeyServer: server.url, address: 'bob' });

    // Fire 3 parallel sends to Bob
    const results = await Promise.all([
      alice.send('bob', 'msg1'),
      alice.send('bob', 'msg2'),
      alice.send('bob', 'msg3'),
    ]);

    // All 3 should succeed and be decryptable in order
    expect(results.length).toBe(3);
    const decrypted: string[] = [];
    for (const env of results) {
      decrypted.push(await bob.receive('alice', env));
    }
    expect(decrypted.sort()).toEqual(['msg1', 'msg2', 'msg3']);
  });

  test('send to unknown address throws clear error', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    await expect(alice.send('nobody', 'ghost')).rejects.toThrow();
  });

  test('aliasSession migrates a session from fp:<hex> to a canonical address label (V4.8.3)', async () => {
    // Reproduces the Prism FR `session-label-asymmetry-v4.8.2`. Bob
    // initiates X3DH against Alice using Alice's prekey-server
    // address. Alice receives the prekey envelope under the relay's
    // sender-fingerprint hint (`fp:<bobfp>`), because that's the only
    // sender label the bridge surfaces at first contact. The
    // post-decrypt plaintext announces Bob's real address; Alice then
    // canonicalizes the session by aliasing `fp:<bobfp>` → `bob` and
    // every subsequent send/receive operates symmetrically.
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    bob = await createShade({ prekeyServer: server.url, address: 'bob' });

    // First contact — Bob sends, Alice receives under the fp-label.
    const env1 = await bob.send('alice', 'hello, my address is bob');
    const fpLabel = 'fp:bobfingerprint16';
    expect(await alice.receive(fpLabel, env1)).toBe('hello, my address is bob');

    // Alice canonicalizes: move the session from the fp-label to bob's
    // real address.
    await alice.aliasSession(fpLabel, 'bob');

    // Subsequent ratchet messages flow under the canonical label both
    // directions. Bob's session for Alice is keyed under `alice`
    // (Bob's send target); Alice's session for Bob is now keyed under
    // `bob` (post-alias). Symmetry restored.
    const env2 = await bob.send('alice', 'reply 1');
    expect(await alice.receive('bob', env2)).toBe('reply 1');

    const env3 = await alice.send('bob', 'reply 2');
    expect(await bob.receive('alice', env3)).toBe('reply 2');

    // The old fp-label has no session — receive under it would now
    // fail. (We don't assert the error shape, only that the label is
    // gone.)
    await expect(alice.receive(fpLabel, env3)).rejects.toThrow();
  });

  test('aliasSession refuses to overwrite an existing session', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    bob = await createShade({ prekeyServer: server.url, address: 'bob' });
    const carol = await createShade({ prekeyServer: server.url, address: 'carol' });
    try {
      // Two distinct first-contact prekey envelopes — one from Bob,
      // one from Carol — let Alice end up with two real sessions in
      // storage at two different labels.
      const env1 = await bob.send('alice', 'one');
      await alice.receive('fp:bobfp', env1);
      const env2 = await carol.send('alice', 'two');
      await alice.receive('fp:carolfp', env2);

      await expect(alice.aliasSession('fp:carolfp', 'fp:bobfp')).rejects.toThrow(
        /refusing to overwrite/i,
      );
    } finally {
      await carol.shutdown();
    }
  });

  test('aliasSession is a no-op when oldLabel === newLabel', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    bob = await createShade({ prekeyServer: server.url, address: 'bob' });
    const env = await bob.send('alice', 'hi');
    await alice.receive('fp:bobfp', env);
    // Same-label alias is a no-op; session must still decrypt the next message.
    await alice.aliasSession('fp:bobfp', 'fp:bobfp');
    const env2 = await bob.send('alice', 'hi again');
    expect(await alice.receive('fp:bobfp', env2)).toBe('hi again');
  });

  test('concurrent receive(from, env) for same `from` does not race the ratchet (V4.8.2)', async () => {
    // Reproduces the Prism FR scenario: a single PUT is fanned out
    // multiple times by the relay (or any duplicating transport), the
    // receiver dispatches several `shade.receive(from, env)` in
    // parallel, and the underlying SessionManager + StorageProvider
    // would race on the ratchet (and on storage writes — sqlite throws
    // "database is locked", IDB throws transaction conflicts) without
    // per-`from` serialization. We pre-establish a session, then fire
    // the same envelope at `bob.receive` from many concurrent callers
    // and verify all of them either decrypt to the same plaintext or
    // surface a benign "already-consumed" error. Crucially: no
    // unhandled storage races, no ratchet corruption, and the next
    // legitimate message still decrypts.
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    bob = await createShade({ prekeyServer: server.url, address: 'bob' });

    const env1 = await alice.send('bob', 'first');
    expect(await bob.receive('alice', env1)).toBe('first');

    const env2 = await alice.send('bob', 'second');
    // Fan the same envelope out to 8 concurrent receives — exactly the
    // shape of the relay duplicate fan-out described in the FR.
    const dispatches = await Promise.allSettled(
      Array.from({ length: 8 }, () => bob.receive('alice', env2)),
    );
    // At least one must have succeeded with the right plaintext; the
    // others may legitimately reject (replay protection / OTPK
    // already-consumed) but MUST NOT corrupt the ratchet or throw
    // "database is locked".
    const fulfilled = dispatches.filter((d) => d.status === 'fulfilled') as Array<
      PromiseFulfilledResult<string>
    >;
    expect(fulfilled.length).toBeGreaterThan(0);
    expect(fulfilled[0]!.value).toBe('second');

    for (const d of dispatches) {
      if (d.status === 'rejected') {
        const msg = String((d.reason as Error)?.message ?? d.reason);
        expect(msg).not.toMatch(/database is locked/i);
      }
    }

    // Ratchet must still advance — the next legitimate message decrypts.
    const env3 = await alice.send('bob', 'third');
    expect(await bob.receive('alice', env3)).toBe('third');
  });

  test('verify fingerprint matches pinned identity', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    bob = await createShade({ prekeyServer: server.url, address: 'bob' });

    // Establish session
    const env = await alice.send('bob', 'init');
    await bob.receive('alice', env);

    const bobFp = await bob.fingerprint;
    // Alice knows Bob via session; but remote fingerprint is derived from
    // stored DH key only (not full identity), so we just check it's returned
    const remoteFp = await alice.getFingerprintFor('bob');
    expect(remoteFp.split(' ').length).toBe(12);
  });

  test('shutdown clears background timers and closes storage', async () => {
    alice = await createShade({
      prekeyServer: server.url,
      address: 'alice',
      autoReplenish: { min: 5, target: 20, intervalMs: 100 },
    });
    await alice.shutdown();
    // If background timer wasn't cleared, this test would hang after the
    // final afterEach. We rely on bun test's cleanup to catch that.
  });

  test('manual replenish is callable', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    // Initially has 20 prekeys, so replenish is a no-op
    const n = await alice.replenish();
    expect(n).toBe(0);
  });

  test('auto-replenish is disabled when set to false', async () => {
    alice = await createShade({
      prekeyServer: server.url,
      address: 'alice',
      autoReplenish: false,
    });
    expect(alice.myAddress).toBe('alice');
  });

  test('rotate regenerates identity', async () => {
    alice = await createShade({ prekeyServer: server.url, address: 'alice' });
    const oldFp = await alice.fingerprint;

    await alice.rotate();

    const newFp = await alice.fingerprint;
    expect(newFp).not.toBe(oldFp);
  });

  test('identityPublicKey exposes the device Ed25519 key and tracks rotation', async () => {
    const storage = new MemoryStorage();
    alice = await createShade({ prekeyServer: server.url, address: 'alice', storage });

    const pk = await alice.identityPublicKey;
    expect(pk).toBeInstanceOf(Uint8Array);
    expect(pk.length).toBe(32);

    // Matches what the underlying storage holds
    const stored = await storage.getIdentityKeyPair();
    expect(stored).not.toBeNull();
    expect(pk).toEqual(stored!.signingPublicKey);

    // Reflects the new key after rotate (acceptance criteria #3)
    await alice.rotate();
    const pkAfter = await alice.identityPublicKey;
    expect(pkAfter.length).toBe(32);
    expect(pkAfter).not.toEqual(pk);
  });
});

describe('createShade — validation', () => {
  test('throws when prekeyServer is missing', async () => {
    await expect(createShade({} as any)).rejects.toThrow(/prekeyServer is required/);
  });

  test('throws when observer token is too short', async () => {
    await expect(
      createShade({
        prekeyServer: 'http://localhost:9999',
        observer: { token: 'short' },
      }),
    ).rejects.toThrow(/at least 16/);
  });
});