Shade/packages/shade-core/tests/session-aliases.test.ts

import { describe, test, expect, beforeEach } from 'bun:test';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import { ShadeSessionManager } from '../src/index.js';

const crypto = new SubtleCryptoProvider();

/**
 * Durable session-label aliases (V4.12).
 *
 * THE BUG THIS FILE EXISTS TO KILL — diagnosed live in Prism:
 *
 *   A phone pairs with a host. First contact forces the host to label
 *   the session by the only sender hint the relay surfaces, an 8-byte
 *   signing-key fingerprint (`fp:<hex>`). The pair handshake then
 *   announces the phone's canonical address and the host calls
 *   `aliasSession(fp:<hex> → device:<addr>)`, which moved the session
 *   on disk and dropped the binding.
 *
 *   The phone, however, keeps sending under `fp:<hex>` — its transport
 *   derives the same label from the same relay hint every time. While
 *   the host process lived, an in-memory map papered over the gap.
 *   After a restart that map was empty, every inbound ratchet frame
 *   resolved to `fp:<hex>`, found no session, and failed. The phone
 *   held a perfectly valid session so it never re-ran X3DH — meaning
 *   the failure was permanent and self-inflicted, not transient.
 *
 *   Observed as `No session for address: fp:579c3b335d66e2c0` on every
 *   receive for three days, with the phone's RPCs timing out forever.
 *
 * The fix: `aliasSession` persists the binding, and session lookup
 * follows it. These tests pin the restart behaviour specifically —
 * a same-process test cannot fail the way production did.
 */
describe('session label aliases', () => {
  let alice: ShadeSessionManager;
  let bob: ShadeSessionManager;
  let aliceStorage: MemoryStorage;
  let bobStorage: MemoryStorage;

  /** The first-contact label Alice is forced to use for Bob. */
  const FP = 'fp:579c3b335d66e2c0';

  beforeEach(async () => {
    aliceStorage = new MemoryStorage();
    bobStorage = new MemoryStorage();
    alice = new ShadeSessionManager(crypto, aliceStorage);
    bob = new ShadeSessionManager(crypto, bobStorage);
    await alice.initialize();
    await bob.initialize();
  });

  /**
   * Bob initiates X3DH against Alice, exactly like a phone reaching a
   * host it just scanned. Returns Bob's first (prekey) envelope.
   */
  async function bobInitiates(target: ShadeSessionManager, initiator: ShadeSessionManager) {
    const otpks = await target.generateOneTimePreKeys(10);
    const bundle = await target.createPreKeyBundle();
    const otpk = otpks[0]!;
    bundle.oneTimePreKey = { keyId: otpk.keyId, publicKey: otpk.keyPair.publicKey };
    await initiator.initSessionFromBundle('alice', bundle);
  }

  /** Simulate a host restart: fresh manager, same durable storage. */
  async function restartAlice(): Promise<ShadeSessionManager> {
    const revived = new ShadeSessionManager(crypto, aliceStorage);
    await revived.initialize();
    return revived;
  }

  test('an aliased session still decrypts under the old label after a restart', async () => {
    await bobInitiates(alice, bob);

    // First contact lands under the fingerprint label.
    const env1 = await bob.encrypt('alice', 'hello, my address is bob');
    expect(await alice.decrypt(FP, env1)).toBe('hello, my address is bob');

    // Alice canonicalizes to Bob's announced address.
    await alice.aliasSession(FP, 'bob');

    // The host restarts. Storage survives; every in-memory map does not.
    const alice2 = await restartAlice();

    // Bob has a valid session and keeps sending under the same label he
    // always has. Before the fix this threw NoSessionError forever.
    const env2 = await bob.encrypt('alice', 'still here after restart');
    expect(await alice2.decrypt(FP, env2)).toBe('still here after restart');
  });

  test('the host can reply under the old label after a restart', async () => {
    await bobInitiates(alice, bob);
    const env1 = await bob.encrypt('alice', 'hi');
    await alice.decrypt(FP, env1);
    await alice.aliasSession(FP, 'bob');

    const alice2 = await restartAlice();

    // Decrypting is only half of it — a host that cannot encrypt back
    // leaves every RPC hanging just the same.
    const reply = await alice2.encrypt(FP, 'reply from the host');
    expect(await bob.decrypt('alice', reply)).toBe('reply from the host');
  });

  test('a live session under the label wins over an alias (re-link)', async () => {
    await bobInitiates(alice, bob);
    const env1 = await bob.encrypt('alice', 'first pairing');
    await alice.decrypt(FP, env1);
    await alice.aliasSession(FP, 'bob');
    const alice2 = await restartAlice();

    // Bob reinstalls: brand-new identity, same relay fingerprint label.
    const bob2Storage = new MemoryStorage();
    const bob2 = new ShadeSessionManager(crypto, bob2Storage);
    await bob2.initialize();
    await bobInitiates(alice2, bob2);

    // The prekey envelope must establish a FRESH session under FP rather
    // than being redirected into the stale aliased one.
    const fresh1 = await bob2.encrypt('alice', 'fresh contact');
    expect(await alice2.decrypt(FP, fresh1)).toBe('fresh contact');

    // And subsequent ratchet frames must keep using that new session.
    const fresh2 = await bob2.encrypt('alice', 'second message');
    expect(await alice2.decrypt(FP, fresh2)).toBe('second message');
  });

  test('resolveSessionLabel reports where the state actually lives', async () => {
    await bobInitiates(alice, bob);
    const env1 = await bob.encrypt('alice', 'hi');
    await alice.decrypt(FP, env1);

    expect(await alice.resolveSessionLabel(FP)).toBe(FP);
    await alice.aliasSession(FP, 'bob');

    const alice2 = await restartAlice();
    expect(await alice2.resolveSessionLabel(FP)).toBe('bob');
    // An unaliased label resolves to itself.
    expect(await alice2.resolveSessionLabel('carol')).toBe('carol');
  });

  test('resetSession drops aliases pointing at the cleared session', async () => {
    await bobInitiates(alice, bob);
    const env1 = await bob.encrypt('alice', 'hi');
    await alice.decrypt(FP, env1);
    await alice.aliasSession(FP, 'bob');
    expect(await aliceStorage.getSessionAlias(FP)).toBe('bob');

    await alice.resetSession('bob');

    // A dangling alias would redirect the next first-contact frame into
    // a session that no longer exists, defeating the reset.
    expect(await aliceStorage.getSessionAlias(FP)).toBeNull();
    expect(await alice.resolveSessionLabel(FP)).toBe(FP);
  });

  test('aliasing persists the binding to storage', async () => {
    await bobInitiates(alice, bob);
    const env1 = await bob.encrypt('alice', 'hi');
    await alice.decrypt(FP, env1);

    expect(await aliceStorage.getSessionAlias(FP)).toBeNull();
    await alice.aliasSession(FP, 'bob');
    expect(await aliceStorage.getSessionAlias(FP)).toBe('bob');
  });
});