feat: M4 Session Manager + demo

ShadeSessionManager wraps X3DH + Double Ratchet into a simple API: - initialize(), createPreKeyBundle(), encrypt(), decrypt() - Automatic PreKeyMessage for first message, RatchetMessage after - Signed prekey rotation, multi-peer sessions, one-time prekey mgmt - Interactive demo.ts showing full frontend↔backend E2EE flow 80 tests, 0 failures across M1-M4. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-09 20:12:01 +02:00
parent bd6452044f
commit a60ff9d6e8
4 changed files with 642 additions and 0 deletions
--- a/packages/shade-core/src/index.ts
+++ b/packages/shade-core/src/index.ts
@@ -5,3 +5,4 @@ export * from './keys.js';
 export * from './errors.js';
 export * from './x3dh.js';
 export * from './ratchet.js';
+export { ShadeSessionManager } from './session.js';
--- a/packages/shade-core/src/session.ts
+++ b/packages/shade-core/src/session.ts
@@ -0,0 +1,254 @@
+import type { CryptoProvider } from './crypto.js';
+import type { StorageProvider } from './storage.js';
+import type {
+  IdentityKeyPair,
+  SignedPreKey,
+  OneTimePreKey,
+  PreKeyBundle,
+  PreKeyMessage,
+  RatchetMessage,
+  ShadeEnvelope,
+} from './types.js';
+import {
+  generateIdentityKeyPair,
+  generateSignedPreKey,
+  generateOneTimePreKeys,
+  createPreKeyBundle,
+  processPreKeyBundle,
+  processPreKeyMessage,
+} from './x3dh.js';
+import {
+  initSenderSession,
+  initReceiverSession,
+  ratchetEncrypt,
+  ratchetDecrypt,
+} from './ratchet.js';
+import { NoSessionError } from './errors.js';
+
+const enc = new TextEncoder();
+const dec = new TextDecoder();
+
+/**
+ * ShadeSessionManager — the high-level API for using Shade.
+ *
+ * Wraps X3DH key agreement and Double Ratchet into a simple interface:
+ *   - `initialize()` — generate or load identity keys
+ *   - `createPreKeyBundle()` — publish to prekey server
+ *   - `encrypt(address, plaintext)` — encrypt for a peer
+ *   - `decrypt(address, envelope)` — decrypt from a peer
+ *
+ * Usage:
+ * ```ts
+ * const manager = new ShadeSessionManager(crypto, storage);
+ * await manager.initialize();
+ *
+ * // To initiate: fetch bundle, then encrypt
+ * await manager.initSessionFromBundle('bob', bundle);
+ * const envelope = await manager.encrypt('bob', 'Hello!');
+ *
+ * // To receive: decrypt handles everything
+ * const plaintext = await manager.decrypt('alice', envelope);
+ * ```
+ */
+export class ShadeSessionManager {
+  private identity: IdentityKeyPair | null = null;
+  private registrationId: number = 0;
+  private currentSignedPreKeyId: number = 0;
+
+  constructor(
+    private readonly crypto: CryptoProvider,
+    private readonly storage: StorageProvider,
+  ) {}
+
+  // ─── Initialization ────────────────────────────────────────
+
+  /** Initialize: load or generate identity keys and a signed prekey */
+  async initialize(): Promise<void> {
+    // Load or generate identity
+    this.identity = await this.storage.getIdentityKeyPair();
+    if (!this.identity) {
+      this.identity = await generateIdentityKeyPair(this.crypto);
+      await this.storage.saveIdentityKeyPair(this.identity);
+    }
+
+    // Load or generate registration ID
+    this.registrationId = await this.storage.getLocalRegistrationId();
+    if (this.registrationId === 0) {
+      this.registrationId = Math.floor(Math.random() * 0xffffffff) + 1;
+      await this.storage.saveLocalRegistrationId(this.registrationId);
+    }
+
+    // Generate initial signed prekey if none exists
+    const spk = await this.storage.getSignedPreKey(1);
+    if (!spk) {
+      const signedPreKey = await generateSignedPreKey(this.crypto, this.identity, 1);
+      await this.storage.saveSignedPreKey(signedPreKey);
+      this.currentSignedPreKeyId = 1;
+    } else {
+      this.currentSignedPreKeyId = spk.keyId;
+    }
+  }
+
+  /** Get our identity's DH public key (for addressing) */
+  getPublicIdentity(): { signingKey: Uint8Array; dhKey: Uint8Array } {
+    if (!this.identity) throw new Error('Not initialized');
+    return {
+      signingKey: this.identity.signingPublicKey,
+      dhKey: this.identity.dhPublicKey,
+    };
+  }
+
+  // ─── Prekey Management ─────────────────────────────────────
+
+  /** Create a prekey bundle to publish to the prekey server */
+  async createPreKeyBundle(): Promise<PreKeyBundle> {
+    if (!this.identity) throw new Error('Not initialized');
+    const spk = await this.storage.getSignedPreKey(this.currentSignedPreKeyId);
+    if (!spk) throw new Error('No signed prekey');
+
+    // Try to include a one-time prekey
+    // (In real usage, the prekey server would pick one — here we just check if any exist)
+    return createPreKeyBundle(this.registrationId, this.identity, spk);
+  }
+
+  /** Generate and store a batch of one-time prekeys */
+  async generateOneTimePreKeys(count: number): Promise<OneTimePreKey[]> {
+    const existingCount = await this.storage.getOneTimePreKeyCount();
+    const startId = existingCount + 1;
+    const keys = await generateOneTimePreKeys(this.crypto, startId, count);
+    for (const key of keys) {
+      await this.storage.saveOneTimePreKey(key);
+    }
+    return keys;
+  }
+
+  /** Rotate the signed prekey (recommended: every 1-7 days) */
+  async rotateSignedPreKey(): Promise<SignedPreKey> {
+    if (!this.identity) throw new Error('Not initialized');
+    const newId = this.currentSignedPreKeyId + 1;
+    const spk = await generateSignedPreKey(this.crypto, this.identity, newId);
+    await this.storage.saveSignedPreKey(spk);
+    // Keep old one for a grace period (sessions may still reference it)
+    this.currentSignedPreKeyId = newId;
+    return spk;
+  }
+
+  // ─── Session Establishment ─────────────────────────────────
+
+  /**
+   * Initiate a session with a peer by processing their prekey bundle.
+   * Call this before the first `encrypt()` to a new peer.
+   */
+  async initSessionFromBundle(address: string, bundle: PreKeyBundle): Promise<void> {
+    const x3dhResult = await processPreKeyBundle(this.crypto, this.storage, bundle);
+
+    const session = await initSenderSession(
+      this.crypto,
+      x3dhResult.rootKey,
+      x3dhResult.remoteIdentityKey,
+      x3dhResult.remoteSignedPreKey,
+    );
+
+    await this.storage.saveSession(address, session);
+    await this.storage.saveTrustedIdentity(address, x3dhResult.remoteIdentityKey);
+
+    // Store X3DH metadata for the first message
+    // We stash this on the session object for the first encrypt call
+    (session as any).__x3dh = {
+      ephemeralPublicKey: x3dhResult.ephemeralPublicKey,
+      signedPreKeyId: x3dhResult.signedPreKeyId,
+      preKeyId: x3dhResult.preKeyId,
+      identityDHKey: this.identity!.dhPublicKey,
+      registrationId: this.registrationId,
+    };
+    await this.storage.saveSession(address, session);
+  }
+
+  // ─── Encrypt / Decrypt ─────────────────────────────────────
+
+  /**
+   * Encrypt a message for a peer. Returns a ShadeEnvelope ready to send.
+   *
+   * The first message to a new peer will be a PreKeyMessage (includes X3DH info).
+   * Subsequent messages are standard RatchetMessages.
+   */
+  async encrypt(address: string, plaintext: string): Promise<ShadeEnvelope> {
+    const session = await this.storage.getSession(address);
+    if (!session) throw new NoSessionError(address);
+
+    const ratchetMsg = await ratchetEncrypt(this.crypto, session, enc.encode(plaintext));
+
+    // Check if this is the first message (X3DH metadata attached)
+    const x3dh = (session as any).__x3dh;
+    if (x3dh) {
+      delete (session as any).__x3dh;
+      await this.storage.saveSession(address, session);
+
+      const preKeyMsg: PreKeyMessage = {
+        registrationId: x3dh.registrationId,
+        preKeyId: x3dh.preKeyId,
+        signedPreKeyId: x3dh.signedPreKeyId,
+        ephemeralKey: x3dh.ephemeralPublicKey,
+        identityDHKey: x3dh.identityDHKey,
+        message: ratchetMsg,
+      };
+      return {
+        type: 'prekey',
+        content: preKeyMsg,
+        timestamp: Date.now(),
+        senderAddress: address,
+      };
+    }
+
+    await this.storage.saveSession(address, session);
+    return {
+      type: 'ratchet',
+      content: ratchetMsg,
+      timestamp: Date.now(),
+      senderAddress: address,
+    };
+  }
+
+  /**
+   * Decrypt a message from a peer. Handles both PreKeyMessage and RatchetMessage.
+   */
+  async decrypt(address: string, envelope: ShadeEnvelope): Promise<string> {
+    if (envelope.type === 'prekey') {
+      return this.decryptPreKeyMessage(address, envelope.content as PreKeyMessage);
+    }
+    return this.decryptRatchetMessage(address, envelope.content as RatchetMessage);
+  }
+
+  private async decryptPreKeyMessage(address: string, message: PreKeyMessage): Promise<string> {
+    // Process X3DH to establish session
+    const x3dhResult = await processPreKeyMessage(this.crypto, this.storage, message);
+
+    // Find the signed prekey that was used
+    const spk = await this.storage.getSignedPreKey(message.signedPreKeyId);
+    if (!spk) throw new Error(`Signed prekey ${message.signedPreKeyId} not found`);
+
+    const session = initReceiverSession(
+      x3dhResult.rootKey,
+      x3dhResult.remoteIdentityKey,
+      spk.keyPair,
+    );
+
+    // Decrypt the embedded first ratchet message
+    const plaintext = await ratchetDecrypt(this.crypto, session, x3dhResult.initialMessage);
+
+    await this.storage.saveSession(address, session);
+    await this.storage.saveTrustedIdentity(address, x3dhResult.remoteIdentityKey);
+
+    return dec.decode(plaintext);
+  }
+
+  private async decryptRatchetMessage(address: string, message: RatchetMessage): Promise<string> {
+    const session = await this.storage.getSession(address);
+    if (!session) throw new NoSessionError(address);
+
+    const plaintext = await ratchetDecrypt(this.crypto, session, message);
+    await this.storage.saveSession(address, session);
+
+    return dec.decode(plaintext);
+  }
+}
--- a/packages/shade-core/tests/session.test.ts
+++ b/packages/shade-core/tests/session.test.ts
@@ -0,0 +1,232 @@
+import { describe, test, expect, beforeEach } from 'bun:test';
+import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
+import {
+  ShadeSessionManager,
+  generateOneTimePreKeys,
+  createPreKeyBundle,
+} from '../src/index.js';
+
+const crypto = new SubtleCryptoProvider();
+
+describe('ShadeSessionManager', () => {
+  let alice: ShadeSessionManager;
+  let bob: ShadeSessionManager;
+  let aliceStorage: MemoryStorage;
+  let bobStorage: MemoryStorage;
+
+  beforeEach(async () => {
+    aliceStorage = new MemoryStorage();
+    bobStorage = new MemoryStorage();
+    alice = new ShadeSessionManager(crypto, aliceStorage);
+    bob = new ShadeSessionManager(crypto, bobStorage);
+    await alice.initialize();
+    await bob.initialize();
+  });
+
+  /** Helper: establish a session from Alice to Bob */
+  async function establishSession() {
+    // Bob generates one-time prekeys
+    const otpks = await bob.generateOneTimePreKeys(10);
+
+    // Bob creates a bundle (with one-time prekey)
+    const bobBundle = await bob.createPreKeyBundle();
+    // Add a one-time prekey to the bundle (in real life the prekey server does this)
+    const otpk = otpks[0];
+    bobBundle.oneTimePreKey = { keyId: otpk.keyId, publicKey: otpk.keyPair.publicKey };
+
+    // Alice initiates session
+    await alice.initSessionFromBundle('bob', bobBundle);
+  }
+
+  // ─── Initialization ────────────────────────────────────────
+
+  describe('initialization', () => {
+    test('generates identity keys on first init', async () => {
+      const pub = alice.getPublicIdentity();
+      expect(pub.signingKey.length).toBe(32);
+      expect(pub.dhKey.length).toBe(32);
+    });
+
+    test('reuses identity keys on second init', async () => {
+      const pub1 = alice.getPublicIdentity();
+      const alice2 = new ShadeSessionManager(crypto, aliceStorage);
+      await alice2.initialize();
+      const pub2 = alice2.getPublicIdentity();
+      expect(pub1.signingKey).toEqual(pub2.signingKey);
+      expect(pub1.dhKey).toEqual(pub2.dhKey);
+    });
+
+    test('creates a prekey bundle', async () => {
+      const bundle = await alice.createPreKeyBundle();
+      expect(bundle.identitySigningKey.length).toBe(32);
+      expect(bundle.identityDHKey.length).toBe(32);
+      expect(bundle.signedPreKey.keyId).toBe(1);
+      expect(bundle.signedPreKey.publicKey.length).toBe(32);
+      expect(bundle.signedPreKey.signature.length).toBe(64);
+    });
+
+    test('generates one-time prekeys', async () => {
+      const keys = await alice.generateOneTimePreKeys(5);
+      expect(keys.length).toBe(5);
+      expect(await aliceStorage.getOneTimePreKeyCount()).toBe(5);
+    });
+  });
+
+  // ─── Full Conversation ─────────────────────────────────────
+
+  describe('full conversation via managers', () => {
+    test('Alice sends to Bob, Bob replies', async () => {
+      await establishSession();
+
+      // Alice → Bob (first message = PreKeyMessage)
+      const env1 = await alice.encrypt('bob', 'Hello Bob!');
+      expect(env1.type).toBe('prekey');
+
+      const plain1 = await bob.decrypt('alice', env1);
+      expect(plain1).toBe('Hello Bob!');
+
+      // Alice → Bob (second message = RatchetMessage)
+      const env2 = await alice.encrypt('bob', 'Still me');
+      expect(env2.type).toBe('ratchet');
+
+      const plain2 = await bob.decrypt('alice', env2);
+      expect(plain2).toBe('Still me');
+
+      // Bob → Alice (reply)
+      const env3 = await bob.encrypt('alice', 'Hi Alice!');
+      expect(env3.type).toBe('ratchet');
+
+      const plain3 = await alice.decrypt('bob', env3);
+      expect(plain3).toBe('Hi Alice!');
+    });
+
+    test('extended conversation with many turns', async () => {
+      await establishSession();
+
+      // Alice sends first message to establish Bob's session
+      const first = await alice.encrypt('bob', 'init');
+      await bob.decrypt('alice', first);
+
+      for (let i = 0; i < 20; i++) {
+        const senderMgr = i % 2 === 0 ? alice : bob;
+        const receiverMgr = i % 2 === 0 ? bob : alice;
+        const senderAddr = i % 2 === 0 ? 'bob' : 'alice';
+        const receiverAddr = i % 2 === 0 ? 'alice' : 'bob';
+
+        const text = `Turn ${i}`;
+        const env = await senderMgr.encrypt(senderAddr, text);
+        const plain = await receiverMgr.decrypt(receiverAddr, env);
+        expect(plain).toBe(text);
+      }
+    });
+
+    test('burst messages then reply', async () => {
+      await establishSession();
+
+      // Alice sends 5 messages
+      const envelopes = [];
+      for (let i = 0; i < 5; i++) {
+        envelopes.push(await alice.encrypt('bob', `msg-${i}`));
+      }
+
+      // Bob decrypts all 5
+      for (let i = 0; i < 5; i++) {
+        expect(await bob.decrypt('alice', envelopes[i])).toBe(`msg-${i}`);
+      }
+
+      // Bob replies
+      const reply = await bob.encrypt('alice', 'Got all 5!');
+      expect(await alice.decrypt('bob', reply)).toBe('Got all 5!');
+    });
+  });
+
+  // ─── Prekey Rotation ──────────────────────────────────────
+
+  describe('prekey rotation', () => {
+    test('rotated signed prekey works for new sessions', async () => {
+      // Bob rotates his signed prekey
+      await bob.rotateSignedPreKey();
+
+      // New bundle uses the new signed prekey
+      const bundle = await bob.createPreKeyBundle();
+      expect(bundle.signedPreKey.keyId).toBe(2);
+
+      // Alice can still establish a session with the new bundle
+      await alice.initSessionFromBundle('bob', bundle);
+      const env = await alice.encrypt('bob', 'After rotation');
+      const plain = await bob.decrypt('alice', env);
+      expect(plain).toBe('After rotation');
+    });
+
+    test('old sessions continue working after rotation', async () => {
+      await establishSession();
+
+      // Establish session and exchange messages
+      const env1 = await alice.encrypt('bob', 'Before rotation');
+      await bob.decrypt('alice', env1);
+
+      // Bob rotates
+      await bob.rotateSignedPreKey();
+
+      // Existing session still works (uses ratchet keys, not prekeys)
+      const env2 = await bob.encrypt('alice', 'After rotation, same session');
+      expect(await alice.decrypt('bob', env2)).toBe('After rotation, same session');
+    });
+  });
+
+  // ─── Multi-Peer ───────────────────────────────────────────
+
+  describe('multi-peer sessions', () => {
+    test('Alice talks to Bob and Charlie simultaneously', async () => {
+      const charlieStorage = new MemoryStorage();
+      const charlie = new ShadeSessionManager(crypto, charlieStorage);
+      await charlie.initialize();
+      await charlie.generateOneTimePreKeys(5);
+
+      // Establish Alice → Bob
+      await establishSession();
+      const envB = await alice.encrypt('bob', 'Hi Bob');
+      await bob.decrypt('alice', envB);
+
+      // Establish Alice → Charlie
+      const charlieBundle = await charlie.createPreKeyBundle();
+      const otpks = await charlie.generateOneTimePreKeys(1);
+      charlieBundle.oneTimePreKey = { keyId: otpks[0].keyId, publicKey: otpks[0].keyPair.publicKey };
+      await alice.initSessionFromBundle('charlie', charlieBundle);
+
+      const envC = await alice.encrypt('charlie', 'Hi Charlie');
+      expect(await charlie.decrypt('alice', envC)).toBe('Hi Charlie');
+
+      // Both sessions work independently
+      const envB2 = await alice.encrypt('bob', 'Still talking to you Bob');
+      expect(await bob.decrypt('alice', envB2)).toBe('Still talking to you Bob');
+
+      const envC2 = await alice.encrypt('charlie', 'And you Charlie');
+      expect(await charlie.decrypt('alice', envC2)).toBe('And you Charlie');
+    });
+  });
+
+  // ─── Error Cases ──────────────────────────────────────────
+
+  describe('error cases', () => {
+    test('encrypt to unknown peer throws NoSessionError', async () => {
+      expect(alice.encrypt('nobody', 'test')).rejects.toThrow('No session');
+    });
+
+    test('decrypt from unknown peer throws NoSessionError', async () => {
+      const fakeEnvelope = {
+        type: 'ratchet' as const,
+        content: {
+          dhPublicKey: crypto.randomBytes(32),
+          previousCounter: 0,
+          counter: 0,
+          ciphertext: crypto.randomBytes(48),
+          nonce: crypto.randomBytes(12),
+        },
+        timestamp: Date.now(),
+        senderAddress: 'nobody',
+      };
+      expect(alice.decrypt('nobody', fakeEnvelope)).rejects.toThrow('No session');
+    });
+  });
+});