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

import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import { ShadeSessionManager } from '@shade/core';
import { createPrekeyServer, MemoryPrekeyStore } from '@shade/server';
import { ShadeFetchTransport } from '../src/fetch-transport.js';

const crypto = new SubtleCryptoProvider();

describe('ShadeFetchTransport', () => {
  test('full flow: register → fetch bundle → establish session → talk', async () => {
    const store = new MemoryPrekeyStore();
    const server = createPrekeyServer({ crypto, store, disableRateLimit: true });

    const port = 19000 + Math.floor(Math.random() * 1000);
    const handle = Bun.serve({ port, fetch: server.fetch });

    try {
      const baseUrl = `http://localhost:${port}`;

      // ─── Bob: set up session manager, register with transport ──
      const bobStorage = new MemoryStorage();
      const bobManager = new ShadeSessionManager(crypto, bobStorage);
      await bobManager.initialize();

      const bobIdentity = await bobStorage.getIdentityKeyPair();
      const bobTransport = new ShadeFetchTransport({
        baseUrl,
        crypto,
        signingPrivateKey: bobIdentity!.signingPrivateKey,
      });

      const bobOTPKs = await bobManager.generateOneTimePreKeys(5);
      const bobBundle = await bobManager.createPreKeyBundle();

      await bobTransport.register(
        'bob',
        bobManager.getPublicIdentity(),
        bobBundle.signedPreKey,
        bobOTPKs,
      );

      expect(await bobTransport.getKeyCount('bob')).toBe(5);

      // ─── Alice: anonymous fetch + establish session ───────
      const aliceStorage = new MemoryStorage();
      const aliceManager = new ShadeSessionManager(crypto, aliceStorage);
      await aliceManager.initialize();

      // Alice doesn't need a signing key to fetch bundles
      const aliceTransport = new ShadeFetchTransport({ baseUrl, crypto });

      const fetchedBundle = await aliceTransport.fetchBundle('bob');
      expect(fetchedBundle.identityDHKey).toEqual(bobManager.getPublicIdentity().dhKey);

      expect(await aliceTransport.getKeyCount('bob')).toBe(4);

      await aliceManager.initSessionFromBundle('bob', fetchedBundle);

      // ─── Alice → Bob encrypted ───────────────────────────
      const env1 = await aliceManager.encrypt('bob', 'Hello via transport!');
      const plain1 = await bobManager.decrypt('alice', env1);
      expect(plain1).toBe('Hello via transport!');

      // ─── Bob → Alice reply ───────────────────────────────
      const env2 = await bobManager.encrypt('alice', 'Got it!');
      const plain2 = await aliceManager.decrypt('bob', env2);
      expect(plain2).toBe('Got it!');

      // ─── Replenish (signed) ──────────────────────────────
      const remaining = await bobTransport.replenish('bob', [
        { keyId: 200, publicKey: crypto.randomBytes(32) },
        { keyId: 201, publicKey: crypto.randomBytes(32) },
      ]);
      expect(remaining).toBe(6);

      // ─── Unregister (signed) ─────────────────────────────
      await bobTransport.unregister('bob');
      await expect(aliceTransport.fetchBundle('bob')).rejects.toThrow();

    } finally {
      handle.stop();
    }
  });
});