(feat): crypto migrations

src/sdk/session.ts

@@ -0,0 +1,249 @@
+import * as bindings from "mtp/raw";
+
+function utf8Encode(text: string): Uint8Array {
+  if (typeof TextEncoder !== "undefined") {
+    return new TextEncoder().encode(text);
+  }
+  if (typeof Buffer !== "undefined") {
+    return new Uint8Array(Buffer.from(text, "utf-8"));
+  }
+  const bytes = new Uint8Array(text.length * 4);
+  let len = 0;
+  for (let i = 0; i < text.length; i += 1) {
+    const code = text.codePointAt(i) as number;
+    if (code < 0x80) {
+      bytes[len++] = code;
+    } else if (code < 0x800) {
+      bytes[len++] = 0xc0 | (code >> 6);
+      bytes[len++] = 0x80 | (code & 0x3f);
+    } else if (code < 0x10000) {
+      bytes[len++] = 0xe0 | (code >> 12);
+      bytes[len++] = 0x80 | ((code >> 6) & 0x3f);
+      bytes[len++] = 0x80 | (code & 0x3f);
+    } else {
+      bytes[len++] = 0xf0 | (code >> 18);
+      bytes[len++] = 0x80 | ((code >> 12) & 0x3f);
+      bytes[len++] = 0x80 | ((code >> 6) & 0x3f);
+      bytes[len++] = 0x80 | (code & 0x3f);
+      i += 1;
+    }
+  }
+  return bytes.subarray(0, len);
+}
+
+export const HKDF_SALT_ROOT = "mtp-e2ee-v1-root";
+const HKDF_INITIATOR_SEND = "mtp-e2ee-v1-initiator-send";
+const HKDF_INITIATOR_RECV = "mtp-e2ee-v1-initiator-recv";
+
+export interface MTPSessionTranscriptContext {
+  senderUserId?: string;
+  senderClientId: bigint;
+  recipientUserId?: string;
+  recipientClientId: bigint;
+  recipientPublicKey: Uint8Array;
+  kemCiphertext: Uint8Array;
+  conversationId: string;
+}
+
+export interface MTPSessionState {
+  version: 1;
+  conversationId: string;
+  ownClientId: bigint;
+  peerClientId: bigint;
+  peerPublicKey: Uint8Array;
+  sendChainKey: Uint8Array;
+  recvChainKey: Uint8Array;
+  sendCount: number;
+  recvCount: number;
+  createdAt: number;
+  updatedAt: number;
+}
+
+export interface MTPSessionStorage {
+  getSession(conversationId: string): Promise<MTPSessionState | null>;
+  setSession(state: MTPSessionState): Promise<void>;
+  deleteSession(conversationId: string): Promise<void>;
+}
+
+export class InMemorySessionStorage implements MTPSessionStorage {
+  private store = new Map<string, MTPSessionState>();
+
+  async getSession(conversationId: string): Promise<MTPSessionState | null> {
+    return this.store.get(conversationId) ?? null;
+  }
+
+  async setSession(state: MTPSessionState): Promise<void> {
+    this.store.set(state.conversationId, { ...state });
+  }
+
+  async deleteSession(conversationId: string): Promise<void> {
+    this.store.delete(conversationId);
+  }
+}
+
+function writeU32BE(value: number): Uint8Array {
+  return new Uint8Array([
+    (value >>> 24) & 0xff,
+    (value >>> 16) & 0xff,
+    (value >>> 8) & 0xff,
+    value & 0xff,
+  ]);
+}
+
+function writeU64BE(value: bigint): Uint8Array {
+  if (value < 0n || value > 0xffff_ffff_ffff_ffffn)
+    throw new Error("u64 out of range");
+  const buf = new Uint8Array(8);
+  for (let i = 7; i >= 0; i -= 1) {
+    buf[i] = Number(value & 0xffn);
+    value >>= 8n;
+  }
+  return buf;
+}
+
+function concatBytes(parts: Uint8Array[]): Uint8Array {
+  const out = new Uint8Array(parts.reduce((sum, part) => sum + part.length, 0));
+  let offset = 0;
+  for (const part of parts) {
+    out.set(part, offset);
+    offset += part.length;
+  }
+  return out;
+}
+
+function transcriptField(label: string, value: Uint8Array): Uint8Array {
+  const labelBytes = utf8Encode(label);
+  return concatBytes([
+    writeU32BE(labelBytes.length),
+    labelBytes,
+    writeU32BE(value.length),
+    value,
+  ]);
+}
+
+export function buildSessionTranscript(
+  args: MTPSessionTranscriptContext,
+): Uint8Array {
+  const recipientPublicKeyHash = bindings.wasm_sha256(args.recipientPublicKey);
+  const kemHash = bindings.wasm_sha256(args.kemCiphertext);
+  return concatBytes([
+    transcriptField("domain", utf8Encode("mtp-e2ee-session-transcript-v1")),
+    transcriptField("version", utf8Encode("1")),
+    transcriptField("senderUserId", utf8Encode(args.senderUserId ?? "")),
+    transcriptField("senderClientId", writeU64BE(args.senderClientId)),
+    transcriptField("recipientUserId", utf8Encode(args.recipientUserId ?? "")),
+    transcriptField("recipientClientId", writeU64BE(args.recipientClientId)),
+    transcriptField("recipientPublicKeyHash", recipientPublicKeyHash),
+    transcriptField("kemCiphertextHash", kemHash),
+    transcriptField("conversationId", utf8Encode(args.conversationId)),
+  ]);
+}
+
+export async function deriveSessionKeys(
+  sharedSecret: Uint8Array,
+  transcript: Uint8Array = new Uint8Array(0),
+): Promise<{
+  root: Uint8Array;
+  initiatorSend: Uint8Array;
+  initiatorRecv: Uint8Array;
+}> {
+  const rootInfo = concatBytes([utf8Encode(HKDF_SALT_ROOT), transcript]);
+  const root = bindings.wasm_hkdf_expand(
+    sharedSecret,
+    new Uint8Array(0),
+    rootInfo,
+    32,
+  );
+  const initiatorSend = bindings.wasm_hkdf_expand(
+    root,
+    new Uint8Array(0),
+    utf8Encode(HKDF_INITIATOR_SEND),
+    32,
+  );
+  const initiatorRecv = bindings.wasm_hkdf_expand(
+    root,
+    new Uint8Array(0),
+    utf8Encode(HKDF_INITIATOR_RECV),
+    32,
+  );
+  return { root, initiatorSend, initiatorRecv };
+}
+
+export function getConversationId(
+  ownClientId: bigint,
+  peerClientId: bigint,
+): string {
+  const ids = [ownClientId, peerClientId].sort((a, b) =>
+    a < b ? -1 : a > b ? 1 : 0,
+  );
+  return `${ids[0].toString(16)}:${ids[1].toString(16)}`;
+}
+
+export class MTPSessionManager {
+  constructor(private storage: MTPSessionStorage) {}
+
+  getConversationId(
+    ownClientId: bigint,
+    peerClientId: bigint,
+  ): Promise<string> {
+    return Promise.resolve(getConversationId(ownClientId, peerClientId));
+  }
+
+  async getSession(
+    ownClientId: bigint,
+    peerClientId: bigint,
+  ): Promise<MTPSessionState | null> {
+    return this.storage.getSession(
+      getConversationId(ownClientId, peerClientId),
+    );
+  }
+
+  async saveSession(state: MTPSessionState): Promise<void> {
+    await this.storage.setSession({ ...state, updatedAt: Date.now() });
+  }
+
+  async deleteSession(
+    ownClientId: bigint,
+    peerClientId: bigint,
+  ): Promise<void> {
+    await this.storage.deleteSession(
+      getConversationId(ownClientId, peerClientId),
+    );
+  }
+
+  async createSession(args: {
+    ownClientId: bigint;
+    peerClientId: bigint;
+    peerPublicKey: Uint8Array;
+    sharedSecret: Uint8Array;
+    role: "initiator" | "receiver";
+    transcript?: Uint8Array;
+    transcriptContext?: MTPSessionTranscriptContext;
+  }): Promise<MTPSessionState> {
+    const transcript =
+      args.transcript ??
+      (args.transcriptContext
+        ? buildSessionTranscript(args.transcriptContext)
+        : undefined);
+    const { root, initiatorSend, initiatorRecv } = await deriveSessionKeys(
+      args.sharedSecret,
+      transcript,
+    );
+    const now = Date.now();
+    const state: MTPSessionState = {
+      version: 1,
+      conversationId: getConversationId(args.ownClientId, args.peerClientId),
+      ownClientId: args.ownClientId,
+      peerClientId: args.peerClientId,
+      peerPublicKey: args.peerPublicKey,
+      sendChainKey: args.role === "initiator" ? initiatorSend : initiatorRecv,
+      recvChainKey: args.role === "initiator" ? initiatorRecv : initiatorSend,
+      sendCount: 0,
+      recvCount: 0,
+      createdAt: now,
+      updatedAt: now,
+    };
+    root.fill(0);
+    return state;
+  }
+}