Shade/packages/shade-core/tests/sesame.test.ts

import { describe, test, expect } from 'bun:test';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import { ShadeSessionManager } from '../src/index.js';
import {
  DeviceListManager,
  fanOutEncrypt,
  observeIncoming,
  userOfDevice,
  deviceIdOf,
} from '../src/sesame.js';

const crypto = new SubtleCryptoProvider();

describe('DeviceListManager', () => {
  test('add and get devices', () => {
    const mgr = new DeviceListManager();
    mgr.addDevice('bob', 'bob:phone');
    mgr.addDevice('bob', 'bob:laptop');
    mgr.addDevice('bob', 'bob:tablet');

    const devices = mgr.getDevices('bob');
    expect(devices.length).toBe(3);
    expect(devices).toContain('bob:phone');
    expect(devices).toContain('bob:laptop');
    expect(devices).toContain('bob:tablet');
  });

  test('remove device', () => {
    const mgr = new DeviceListManager();
    mgr.addDevice('bob', 'bob:phone');
    mgr.addDevice('bob', 'bob:laptop');
    mgr.removeDevice('bob', 'bob:phone');
    expect(mgr.getDevices('bob')).toEqual(['bob:laptop']);
  });

  test('setDevices replaces list', () => {
    const mgr = new DeviceListManager();
    mgr.addDevice('bob', 'bob:phone');
    mgr.setDevices('bob', ['bob:laptop', 'bob:tablet']);
    expect(mgr.getDevices('bob').sort()).toEqual(['bob:laptop', 'bob:tablet']);
  });

  test('empty list for unknown user', () => {
    const mgr = new DeviceListManager();
    expect(mgr.getDevices('nobody')).toEqual([]);
  });

  test('toJSON and fromJSON roundtrip', () => {
    const mgr = new DeviceListManager();
    mgr.addDevice('bob', 'bob:phone');
    mgr.addDevice('alice', 'alice:laptop');

    const json = mgr.toJSON();
    const restored = new DeviceListManager();
    restored.fromJSON(json);

    expect(restored.getDevices('bob')).toEqual(['bob:phone']);
    expect(restored.getDevices('alice')).toEqual(['alice:laptop']);
  });
});

describe('Address parsing helpers', () => {
  test('userOfDevice extracts the user part', () => {
    expect(userOfDevice('bob:phone')).toBe('bob');
    expect(userOfDevice('alice@example.com:tablet')).toBe('alice@example.com');
    expect(userOfDevice('noColon')).toBe('noColon');
  });

  test('deviceIdOf extracts the device part', () => {
    expect(deviceIdOf('bob:phone')).toBe('phone');
    expect(deviceIdOf('alice@example.com:tablet')).toBe('tablet');
    expect(deviceIdOf('noColon')).toBe('');
  });
});

describe('observeIncoming', () => {
  test('auto-registers a new device', () => {
    const mgr = new DeviceListManager();
    observeIncoming(mgr, 'bob:newPhone');
    expect(mgr.getDevices('bob')).toEqual(['bob:newPhone']);
  });

  test('second message from same device is idempotent', () => {
    const mgr = new DeviceListManager();
    observeIncoming(mgr, 'bob:phone');
    observeIncoming(mgr, 'bob:phone');
    expect(mgr.getDevices('bob').length).toBe(1);
  });
});

describe('fanOutEncrypt: multi-device fan-out', () => {
  async function setupAliceToBobDevices(devices: string[]) {
    // Alice has one SessionManager
    const alice = new ShadeSessionManager(crypto, new MemoryStorage());
    await alice.initialize();

    // Each Bob device has its own SessionManager (separate storage)
    const bobs = new Map<string, ShadeSessionManager>();
    for (const device of devices) {
      const mgr = new ShadeSessionManager(crypto, new MemoryStorage());
      await mgr.initialize();
      bobs.set(device, mgr);
    }

    // Alice establishes a session with each Bob device
    for (const [device, bob] of bobs) {
      const otpks = await bob.generateOneTimePreKeys(5);
      const bundle = await bob.createPreKeyBundle();
      bundle.oneTimePreKey = { keyId: otpks[0].keyId, publicKey: otpks[0].keyPair.publicKey };
      await alice.initSessionFromBundle(device, bundle);
    }

    // Alice tracks Bob's device list
    const deviceList = new DeviceListManager();
    for (const device of devices) {
      deviceList.addDevice('bob', device);
    }

    return { alice, bobs, deviceList };
  }

  test('fans out to all devices', async () => {
    const devices = ['bob:phone', 'bob:laptop', 'bob:tablet'];
    const { alice, bobs, deviceList } = await setupAliceToBobDevices(devices);

    const fanOut = await fanOutEncrypt(alice, deviceList, 'bob', 'hello all my devices');

    expect(fanOut.length).toBe(3);
    for (const { deviceAddress, envelope } of fanOut) {
      const bob = bobs.get(deviceAddress)!;
      const plain = await bob.decrypt('alice', envelope);
      expect(plain).toBe('hello all my devices');
    }
  });

  test('each device gets an independent session (DH ratchet per device)', async () => {
    const devices = ['bob:phone', 'bob:laptop'];
    const { alice, bobs, deviceList } = await setupAliceToBobDevices(devices);

    // Send two rounds
    for (let i = 0; i < 2; i++) {
      const fanOut = await fanOutEncrypt(alice, deviceList, 'bob', `round ${i}`);
      for (const { deviceAddress, envelope } of fanOut) {
        const bob = bobs.get(deviceAddress)!;
        const plain = await bob.decrypt('alice', envelope);
        expect(plain).toBe(`round ${i}`);
      }
    }
  });

  test('throws when user has no known devices', async () => {
    const alice = new ShadeSessionManager(crypto, new MemoryStorage());
    await alice.initialize();
    const deviceList = new DeviceListManager();

    expect(
      fanOutEncrypt(alice, deviceList, 'unknown-user', 'hello'),
    ).rejects.toThrow(/No known devices/);
  });
});