Shade/packages/shade-core/tests/identity-rotation.test.ts

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

const crypto = new SubtleCryptoProvider();

describe('Identity rotation', () => {
  test('rotateIdentity generates new identity and archives old', async () => {
    const storage = new MemoryStorage();
    const mgr = new ShadeSessionManager(crypto, storage);
    await mgr.initialize();

    const oldFp = await mgr.getIdentityFingerprint();
    const oldPub = mgr.getPublicIdentity();

    await mgr.rotateIdentity();

    const newFp = await mgr.getIdentityFingerprint();
    const newPub = mgr.getPublicIdentity();

    // New identity should differ
    expect(newFp).not.toBe(oldFp);
    expect(newPub.signingKey).not.toEqual(oldPub.signingKey);
    expect(newPub.dhKey).not.toEqual(oldPub.dhKey);

    // Old identity is archived
    const retired = await storage.getRetiredIdentities();
    expect(retired.length).toBe(1);
    expect(retired[0].keyPair.signingPublicKey).toEqual(oldPub.signingKey);
  });

  test('rotation returns a new prekey bundle', async () => {
    const mgr = new ShadeSessionManager(crypto, new MemoryStorage());
    await mgr.initialize();

    const bundle = await mgr.rotateIdentity();
    expect(bundle.identitySigningKey.length).toBe(32);
    expect(bundle.identityDHKey.length).toBe(32);
    expect(bundle.signedPreKey.keyId).toBe(2); // advanced
    expect(bundle.signedPreKey.signature.length).toBe(64);
  });

  test('getActiveRetiredIdentities returns entries within grace period', async () => {
    const mgr = new ShadeSessionManager(crypto, new MemoryStorage());
    await mgr.initialize();

    await mgr.rotateIdentity();
    await mgr.rotateIdentity();

    const active = await mgr.getActiveRetiredIdentities();
    expect(active.length).toBe(2);
  });

  test('getActiveRetiredIdentities filters out expired entries', async () => {
    const storage = new MemoryStorage();
    const mgr = new ShadeSessionManager(crypto, storage);
    await mgr.initialize();

    // Manually add a very old retired identity
    await storage.addRetiredIdentity({
      keyPair: (await mgr.getPublicIdentity()) as any, // placeholder
      retiredAt: Date.now() - (GRACE_PERIOD_MS + 1000),
    });

    // And a fresh one
    await mgr.rotateIdentity();

    const active = await mgr.getActiveRetiredIdentities();
    expect(active.length).toBe(1);
  });

  test('pruneExpiredIdentities removes old entries', async () => {
    const storage = new MemoryStorage();
    const mgr = new ShadeSessionManager(crypto, storage);
    await mgr.initialize();

    // Rotate twice
    await mgr.rotateIdentity();
    await mgr.rotateIdentity();

    expect((await storage.getRetiredIdentities()).length).toBe(2);

    // Default grace period: nothing is expired yet
    await mgr.pruneExpiredIdentities();
    expect((await storage.getRetiredIdentities()).length).toBe(2);

    // Force prune with 0 grace → everything goes
    await mgr.pruneExpiredIdentities(0);
    expect((await storage.getRetiredIdentities()).length).toBe(0);
  });

  test('rotation persists across manager restart', async () => {
    const storage = new MemoryStorage();
    const mgr1 = new ShadeSessionManager(crypto, storage);
    await mgr1.initialize();
    await mgr1.rotateIdentity();
    const fp1 = await mgr1.getIdentityFingerprint();

    const mgr2 = new ShadeSessionManager(crypto, storage);
    await mgr2.initialize();
    const fp2 = await mgr2.getIdentityFingerprint();

    expect(fp1).toBe(fp2);
  });

  test('existing sessions survive identity rotation', async () => {
    // Set up Alice-Bob conversation
    const alice = new ShadeSessionManager(crypto, new MemoryStorage());
    const bob = new ShadeSessionManager(crypto, new MemoryStorage());
    await alice.initialize();
    await bob.initialize();

    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('bob', bundle);

    // Exchange messages to establish bidirectional session
    const env1 = await alice.encrypt('bob', 'hello');
    expect(await bob.decrypt('alice', env1)).toBe('hello');
    const env2 = await bob.encrypt('alice', 'hi');
    expect(await alice.decrypt('bob', env2)).toBe('hi');

    // Alice rotates her identity — but her session with Bob should still work
    // because the Double Ratchet uses ephemeral DH keys, not identity keys
    await alice.rotateIdentity();

    const env3 = await alice.encrypt('bob', 'after rotation');
    expect(await bob.decrypt('alice', env3)).toBe('after rotation');
  });
});