(feat): crypto migrations

test/e2ee.mjs

@@ -0,0 +1,503 @@
+import { initSync } from "../dist/raw/index.js";
+import fs from "fs";
+import path from "path";
+import { fileURLToPath } from "url";
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+const wasmPath = path.resolve(__dirname, "../wasm/pkg/mtp_wasm_bg.wasm");
+const wasmBytes = fs.readFileSync(wasmPath);
+const wasmModule = new WebAssembly.Module(wasmBytes);
+initSync(wasmModule);
+
+const sdk = await import("../dist/sdk/index.js");
+const { MTPRatchet } = await import("../dist/sdk/ratchet.js");
+const {
+  serializeEncryptedMessage,
+  parseEncryptedMessage,
+  encryptPayload,
+  decryptPayload,
+  FLAG_INIT,
+  MTP_E2EE_VERSION,
+} = await import("../dist/sdk/encrypted-message.js");
+const {
+  MTPSessionManager,
+  InMemorySessionStorage,
+  deriveSessionKeys,
+  getConversationId,
+} = await import("../dist/sdk/session.js");
+
+const bindings = sdk.raw;
+
+function concat(...arrays) {
+  const totalLen = arrays.reduce((sum, a) => sum + a.length, 0);
+  const result = new Uint8Array(totalLen);
+  let offset = 0;
+  for (const a of arrays) {
+    result.set(a, offset);
+    offset += a.length;
+  }
+  return result;
+}
+
+function setupSessions(sharedSecret, aliceId = 1n, bobId = 2n) {
+  const aliceStorage = new InMemorySessionStorage();
+  const aliceManager = new MTPSessionManager(aliceStorage);
+  const bobStorage = new InMemorySessionStorage();
+  const bobManager = new MTPSessionManager(bobStorage);
+
+  return {
+    aliceManager,
+    aliceStorage,
+    bobManager,
+    bobStorage,
+    async initSessions() {
+      const { initiatorSend, initiatorRecv } = await deriveSessionKeys(
+        sharedSecret,
+        new Uint8Array(0),
+      );
+      const aliceSession = await aliceManager.createSession({
+        ownClientId: aliceId,
+        peerClientId: bobId,
+        peerPublicKey: new Uint8Array(32),
+        sharedSecret,
+        role: "initiator",
+      });
+      const bobSession = await bobManager.createSession({
+        ownClientId: bobId,
+        peerClientId: aliceId,
+        peerPublicKey: new Uint8Array(32),
+        sharedSecret,
+        role: "receiver",
+      });
+      return { aliceSession, bobSession, initiatorSend, initiatorRecv };
+    },
+  };
+}
+
+await describe("E2EE Session Derivation", async () => {
+  await it("Both sides derive same shared secret", async () => {
+    const sharedSecret = sdk.crypto.sha256(new Uint8Array([1, 2, 3, 4, 5]));
+    const transcript = new Uint8Array(0);
+
+    const aliceKeys = await deriveSessionKeys(sharedSecret, transcript);
+    const bobKeys = await deriveSessionKeys(sharedSecret, transcript);
+
+    // Deterministic: same inputs → same outputs
+    assert.deepEqual(aliceKeys.initiatorSend, bobKeys.initiatorSend);
+    assert.deepEqual(aliceKeys.initiatorRecv, bobKeys.initiatorRecv);
+
+    // Init and recv keys are different
+    assert.notDeepEqual(aliceKeys.initiatorSend, aliceKeys.initiatorRecv);
+  });
+
+  await it("Session manager assigns correct chain keys per role", async () => {
+    const ss = sdk.crypto.sha256(new Uint8Array([1]));
+    const { initSessions } = setupSessions(ss);
+    const { aliceSession, bobSession, initiatorSend, initiatorRecv } =
+      await initSessions();
+
+    // Alice (initiator): send = initiatorSend, recv = initiatorRecv
+    assert.deepEqual(aliceSession.sendChainKey, initiatorSend);
+    assert.deepEqual(aliceSession.recvChainKey, initiatorRecv);
+
+    // Bob (receiver): send = initiatorRecv, recv = initiatorSend
+    assert.deepEqual(bobSession.sendChainKey, initiatorRecv);
+    assert.deepEqual(bobSession.recvChainKey, initiatorSend);
+
+    // Alice's send chain = Bob's recv chain
+    assert.deepEqual(aliceSession.sendChainKey, bobSession.recvChainKey);
+    // Alice's recv chain = Bob's send chain
+    assert.deepEqual(aliceSession.recvChainKey, bobSession.sendChainKey);
+  });
+
+  await it("Different transcripts produce different keys", async () => {
+    const ss = sdk.crypto.sha256(new Uint8Array([99]));
+    const aliceKeys1 = await deriveSessionKeys(ss, new Uint8Array(0));
+    const aliceKeys2 = await deriveSessionKeys(
+      ss,
+      sdk.crypto.sha256(new Uint8Array([42])),
+    );
+    assert.notDeepEqual(aliceKeys1.initiatorSend, aliceKeys2.initiatorSend);
+  });
+});
+
+await describe("E2EE Ratchet", async () => {
+  await it("Repeated sends produce different message keys", async () => {
+    const chainKey = sdk.crypto.sha256(new Uint8Array([42]));
+    const step1 = await MTPRatchet.step(chainKey);
+    const step2 = await MTPRatchet.step(step1.chainKey);
+    const step3 = await MTPRatchet.step(step2.chainKey);
+
+    assert.notDeepEqual(step1.key, step2.key);
+    assert.notDeepEqual(step2.key, step3.key);
+    assert.notDeepEqual(step1.key, step3.key);
+    assert.notDeepEqual(chainKey, step1.chainKey);
+  });
+
+  await it("Receiver can decrypt messages sent by sender in order", async () => {
+    const chainKey = sdk.crypto.sha256(new Uint8Array([7]));
+
+    const send1 = await MTPRatchet.step(chainKey);
+    const send2 = await MTPRatchet.step(send1.chainKey);
+    const send3 = await MTPRatchet.step(send2.chainKey);
+
+    const recv1 = await MTPRatchet.step(chainKey);
+    const recv2 = await MTPRatchet.step(recv1.chainKey);
+    const recv3 = await MTPRatchet.step(recv2.chainKey);
+
+    assert.deepEqual(send1.key, recv1.key);
+    assert.deepEqual(send2.key, recv2.key);
+    assert.deepEqual(send3.key, recv3.key);
+  });
+});
+
+await describe("E2EE Serialization", async () => {
+  await it("Roundtrips a basic message", () => {
+    const msg = {
+      header: {
+        version: 1,
+        flags: 0,
+        senderClientId: 0x1234567890abcdefn,
+        recipientClientId: 0xfedcba0987654321n,
+        messageNumber: 42,
+      },
+      aeadPayload: new Uint8Array([1, 2, 3, 4, 5]),
+    };
+    const bytes = serializeEncryptedMessage(msg);
+    const parsed = parseEncryptedMessage(bytes);
+    assert.equal(parsed.header.version, 1);
+    assert.equal(parsed.header.flags, 0);
+    assert.equal(parsed.header.senderClientId, msg.header.senderClientId);
+    assert.equal(parsed.header.recipientClientId, msg.header.recipientClientId);
+    assert.equal(parsed.header.messageNumber, 42);
+    assert.equal(parsed.header.kemCiphertext, undefined);
+    assert.deepEqual(parsed.aeadPayload, msg.aeadPayload);
+  });
+
+  await it("Roundtrips an init message with KEM ciphertext", () => {
+    const msg = {
+      header: {
+        version: 1,
+        flags: FLAG_INIT,
+        senderClientId: 1n,
+        recipientClientId: 2n,
+        messageNumber: 0,
+        kemCiphertext: new Uint8Array([0xde, 0xad, 0xbe, 0xef]),
+      },
+      aeadPayload: new Uint8Array([10, 20, 30]),
+    };
+    const bytes = serializeEncryptedMessage(msg);
+    const parsed = parseEncryptedMessage(bytes);
+    assert.equal(parsed.header.flags & FLAG_INIT, FLAG_INIT);
+    assert.deepEqual(parsed.header.kemCiphertext, msg.header.kemCiphertext);
+  });
+
+  await it("Roundtrip: serialize(parse(x)) === x", () => {
+    const msg = {
+      header: {
+        version: 1,
+        flags: 0,
+        senderClientId: 0xaaaabbbbccccddddn,
+        recipientClientId: 0xffff000011112222n,
+        messageNumber: 65535,
+      },
+      aeadPayload: new Uint8Array(100).fill(0x42),
+    };
+    const bytes = serializeEncryptedMessage(msg);
+    const parsed = parseEncryptedMessage(bytes);
+    const bytes2 = serializeEncryptedMessage(parsed);
+    assert.deepEqual(bytes, bytes2);
+  });
+
+  await it("Rejects malformed payloads", () => {
+    assert.throws(() => parseEncryptedMessage(new Uint8Array(0)));
+    assert.throws(() => parseEncryptedMessage(new Uint8Array([0x01])));
+    assert.throws(() =>
+      parseEncryptedMessage(
+        new Uint8Array([0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]),
+      ),
+    );
+  });
+
+  await it("Rejects unsupported version", () => {
+    const msg = {
+      header: {
+        version: 1,
+        flags: 0,
+        senderClientId: 0n,
+        recipientClientId: 0n,
+        messageNumber: 0,
+      },
+      aeadPayload: new Uint8Array([1]),
+    };
+    const bytes = serializeEncryptedMessage(msg);
+    bytes[0] = 99;
+    assert.throws(() => parseEncryptedMessage(bytes));
+  });
+
+  await it("Rejects trailing data", () => {
+    const msg = {
+      header: {
+        version: 1,
+        flags: 0,
+        senderClientId: 0n,
+        recipientClientId: 0n,
+        messageNumber: 0,
+      },
+      aeadPayload: new Uint8Array([1]),
+    };
+    const bytes = concat(
+      serializeEncryptedMessage(msg),
+      new Uint8Array([0xff]),
+    );
+    assert.throws(() => parseEncryptedMessage(bytes));
+  });
+});
+
+await describe("E2EE Encrypt/Decrypt", async () => {
+  await it("Alice encrypts and Bob decrypts successfully", async () => {
+    const keyring = sdk.crypto.generateKeyring();
+    const bobKeys = sdk.crypto.keyringToKeys(keyring);
+
+    // Alice encapsulates to Bob's KEM public key
+    const enc = sdk.crypto.encapsulate(bobKeys.kemPublicKey);
+    // Bob decapsulates the ciphertext
+    const bobSS = sdk.crypto.decapsulate(bobKeys.kemSecretKey, enc.ciphertext);
+    assert.deepEqual(bobSS, enc.shared_secret);
+
+    const ss = enc.shared_secret;
+    const { initSessions } = setupSessions(ss);
+    const { aliceSession, bobSession } = await initSessions();
+
+    // Alice encrypts a message to Bob
+    const plaintext = sdk.codec.encode(
+      "Ping",
+      { Version: "hello from Alice" },
+      { sender: 1n, receiver: 2n },
+    );
+
+    const { payload, session: aliceNewSession } = await encryptPayload({
+      plaintext,
+      session: aliceSession,
+      kemCiphertext: enc.ciphertext,
+    });
+
+    // Bob decrypts the message
+    const { plaintext: decrypted, session: bobNewSession } =
+      await decryptPayload({
+        payload,
+        session: bobSession,
+      });
+
+    const frame = sdk.codec.decode(decrypted);
+    assert.equal(frame.type, "Ping");
+    assert.equal(frame.data["Version"], "hello from Alice");
+
+    // Chain keys advanced correctly
+    assert.deepEqual(aliceNewSession.sendChainKey, bobNewSession.recvChainKey);
+    assert.equal(aliceNewSession.sendCount, 1);
+    assert.equal(bobNewSession.recvCount, 1);
+    assert.notDeepEqual(
+      aliceNewSession.sendChainKey,
+      aliceSession.sendChainKey,
+    );
+  });
+
+  await it("Encrypt-decrypt multiple messages with chain advance", async () => {
+    const ss = sdk.crypto.sha256(new Uint8Array([1, 2, 3]));
+    const { initSessions } = setupSessions(ss);
+    let { aliceSession, bobSession } = await initSessions();
+
+    // Message 1
+    const { payload: p1, session: aliceAfter1 } = await encryptPayload({
+      plaintext: sdk.codec.encode("Ping", { Version: "msg1" }),
+      session: aliceSession,
+    });
+    const { plaintext: d1, session: bobAfter1 } = await decryptPayload({
+      payload: p1,
+      session: bobSession,
+    });
+    assert.equal(sdk.codec.decode(d1).data["Version"], "msg1");
+    assert.equal(bobAfter1.recvCount, 1);
+
+    // Message 2
+    const { payload: p2, session: aliceAfter2 } = await encryptPayload({
+      plaintext: sdk.codec.encode("Ping", { Version: "msg2" }),
+      session: aliceAfter1,
+    });
+    const { plaintext: d2, session: bobAfter2 } = await decryptPayload({
+      payload: p2,
+      session: bobAfter1,
+    });
+    assert.equal(sdk.codec.decode(d2).data["Version"], "msg2");
+    assert.equal(bobAfter2.recvCount, 2);
+
+    // Chain keys match after two messages
+    assert.deepEqual(aliceAfter2.sendChainKey, bobAfter2.recvChainKey);
+    assert.equal(aliceAfter2.sendCount, 2);
+  });
+});
+
+await describe("E2EE Public Key Bundle", async () => {
+  await it("parses public key bundles from GetUserData.PublicKey", () => {
+    const keyring = sdk.crypto.generateKeyring();
+    const keys = sdk.crypto.keyringToKeys(keyring);
+    const bundle = concat(
+      new Uint8Array([keys.kemPublicKey.length >> 8, keys.kemPublicKey.length & 0xff]),
+      keys.kemPublicKey,
+      new Uint8Array([keys.sigPqPublicKey.length >> 8, keys.sigPqPublicKey.length & 0xff]),
+      keys.sigPqPublicKey,
+      new Uint8Array([keys.sigClPublicKey.length >> 8, keys.sigClPublicKey.length & 0xff]),
+      keys.sigClPublicKey,
+    );
+
+    const parsed = sdk.crypto.publicKeyBundleToKeys(bundle);
+    assert.deepEqual(parsed.kemPublicKey, keys.kemPublicKey);
+    assert.deepEqual(parsed.sigPqPublicKey, keys.sigPqPublicKey);
+    assert.deepEqual(parsed.sigClPublicKey, keys.sigClPublicKey);
+  });
+});
+
+await describe("E2EE Session Manager", async () => {
+  await it("getConversationId is consistent regardless of order", () => {
+    const id1 = getConversationId(5n, 10n);
+    const id2 = getConversationId(10n, 5n);
+    assert.equal(id1, id2);
+  });
+
+  await it("MTPSessionManager creates, retrieves, and deletes sessions", async () => {
+    const ss = sdk.crypto.sha256(new Uint8Array([1, 2, 3]));
+    const storage = new InMemorySessionStorage();
+    const manager = new MTPSessionManager(storage);
+
+    assert.equal(await manager.getSession(1n, 2n), null);
+
+    const session = await manager.createSession({
+      ownClientId: 1n,
+      peerClientId: 2n,
+      peerPublicKey: new Uint8Array(32),
+      sharedSecret: ss,
+      role: "initiator",
+    });
+    assert.equal(session.version, 1);
+    assert.equal(session.sendCount, 0);
+    assert.equal(session.recvCount, 0);
+
+    await manager.saveSession(session);
+    const retrieved = await manager.getSession(1n, 2n);
+    assert.notEqual(retrieved, null);
+    assert.equal(retrieved.conversationId, session.conversationId);
+
+    await manager.deleteSession(1n, 2n);
+    assert.equal(await manager.getSession(1n, 2n), null);
+  });
+});
+
+await describe("E2EE Full Flow: KEM + Session + Ratchet + AEAD", async () => {
+  await it("Alice encapsulates to Bob, both derive matching sessions, encrypt-decrypt works", async () => {
+    // Bob generates keyring
+    const bobKeyring = sdk.crypto.generateKeyring();
+    const bobKeys = sdk.crypto.keyringToKeys(bobKeyring);
+
+    // Alice encapsulates to Bob's KEM public key
+    const enc = sdk.crypto.encapsulate(bobKeys.kemPublicKey);
+
+    // Bob decapsulates
+    const bobSharedSecret = sdk.crypto.decapsulate(
+      bobKeys.kemSecretKey,
+      enc.ciphertext,
+    );
+    assert.deepEqual(bobSharedSecret, enc.shared_secret);
+
+    const ss = enc.shared_secret;
+    const { initSessions } = setupSessions(ss);
+    const { aliceSession, bobSession } = await initSessions();
+
+    // Alice sends encrypted init message with KEM ciphertext
+    const msg1 = new Uint8Array([0x48, 0x65, 0x6c, 0x6c, 0x6f]); // "Hello"
+    const { payload: p1, session: aliceAfter1 } = await encryptPayload({
+      plaintext: msg1,
+      session: aliceSession,
+      kemCiphertext: enc.ciphertext,
+    });
+
+    // Bob receives and decrypts
+    const { plaintext: d1, session: bobAfter1 } = await decryptPayload({
+      payload: p1,
+      session: bobSession,
+    });
+    assert.deepEqual(d1, msg1);
+    assert.deepEqual(aliceAfter1.sendChainKey, bobAfter1.recvChainKey);
+    assert.equal(aliceAfter1.sendCount, 1);
+    assert.equal(bobAfter1.recvCount, 1);
+
+    // Second message (no KEM ciphertext)
+    const msg2 = new Uint8Array([0x57, 0x6f, 0x72, 0x6c, 0x64]); // "World"
+    const { payload: p2, session: aliceAfter2 } = await encryptPayload({
+      plaintext: msg2,
+      session: aliceAfter1,
+    });
+
+    const { plaintext: d2, session: bobAfter2 } = await decryptPayload({
+      payload: p2,
+      session: bobAfter1,
+    });
+    assert.deepEqual(d2, msg2);
+    assert.deepEqual(aliceAfter2.sendChainKey, bobAfter2.recvChainKey);
+    assert.equal(aliceAfter2.sendCount, 2);
+    assert.equal(bobAfter2.recvCount, 2);
+  });
+});
+
+await describe("E2EE Tamper Detection", async () => {
+  await it("Rejects modified ciphertext", async () => {
+    const ss = sdk.crypto.sha256(new Uint8Array([42]));
+    const { initSessions } = setupSessions(ss);
+    const { aliceSession, bobSession } = await initSessions();
+
+    const plaintext = new Uint8Array([0x01, 0x02, 0x03]);
+    const { payload } = await encryptPayload({
+      plaintext,
+      session: aliceSession,
+    });
+
+    // Tamper with AEAD payload
+    const tampered = new Uint8Array(payload);
+    tampered[tampered.length - 1] ^= 0xff;
+
+    await assert.rejects(
+      () => decryptPayload({ payload: tampered, session: bobSession }),
+      /decrypt failed/,
+    );
+  });
+
+  await it("Rejects out-of-order message numbers", async () => {
+    const ss = sdk.crypto.sha256(new Uint8Array([7]));
+    const { initSessions } = setupSessions(ss);
+    const { aliceSession, bobSession } = await initSessions();
+
+    // Send two messages
+    const { payload: p1, session: aliceAfter1 } = await encryptPayload({
+      plaintext: sdk.codec.encode("Ping", { Version: "a" }),
+      session: aliceSession,
+    });
+    await encryptPayload({
+      plaintext: sdk.codec.encode("Ping", { Version: "b" }),
+      session: aliceAfter1,
+    });
+
+    // Bob decrypts p1
+    const { session: bobAfter1 } = await decryptPayload({
+      payload: p1,
+      session: bobSession,
+    });
+
+    // Now bob expects msgNumber 1, but we try to replay msgNumber 0
+    await assert.rejects(
+      () => decryptPayload({ payload: p1, session: bobAfter1 }),
+      /replay|out of order/,
+    );
+  });
+});