Kez/kez-chat/web/src/lib/identity-store.ts

// Local persistence for the current user's identity.
//
// v0.1: stores the ed25519 seed in IndexedDB encrypted under a
// passphrase the user enters at unlock time. Standard pattern: derive
// a key from the passphrase via PBKDF2-SHA-256 → AES-GCM-256 wrap the
// seed → store the wrapped blob + salt + nonce. To unlock, the user
// types the passphrase again; if PBKDF2 produces the right key,
// AES-GCM unwraps the seed; if not, decryption fails.
//
// Caveats spelled out in the spec doc: browsers have no Keychain
// equivalent, so this is the best we can do client-side. The CLI and
// future native GUI use OS keychain and don't have this limitation.

import { get, set, del } from "idb-keyval";
import { bytesToHex, hexToBytes } from "@noble/hashes/utils";
import type { Identity } from "./kez.js";

const IDB_KEY = "kez-chat:identity";

// TS 5.6+ defaults Uint8Array's generic param to ArrayBufferLike, which
// isn't assignable to BufferSource (= ArrayBufferView<ArrayBuffer>) that
// WebCrypto expects. Copy to a plain ArrayBuffer so the types line up.
function asBuffer(u: Uint8Array): ArrayBuffer {
  return u.buffer.slice(u.byteOffset, u.byteOffset + u.byteLength) as ArrayBuffer;
}

interface StoredIdentity {
  version: 1;
  handle: string;
  server: string;        // e.g. "kez.lat"
  primary: Identity;     // ed25519:<hex>
  // Encrypted seed:
  salt: string;          // hex, 16 bytes
  nonce: string;         // hex, 12 bytes
  ciphertext: string;    // hex; AES-GCM(seed) under PBKDF2(passphrase)
  // Optional: encrypted 12-word recovery phrase (added in the mnemonic
  // rollout). New accounts have both; pre-mnemonic accounts have only
  // the seed. Encrypted under the SAME PBKDF2 key as the seed; uses its
  // own nonce because AES-GCM reuse is unsafe.
  phrase_nonce?: string;       // hex, 12 bytes
  phrase_ciphertext?: string;  // hex; AES-GCM(utf8(phrase))
  // Metadata:
  created_at: string;    // RFC3339
}

export interface UnlockedIdentity {
  handle: string;
  server: string;
  primary: Identity;
  seed: Uint8Array;
  /** 12-word recovery phrase. Absent for pre-mnemonic accounts. */
  phrase?: string;
}

const PBKDF2_ITERATIONS = 600_000; // OWASP 2024 SHA-256 guidance

async function deriveKey(
  passphrase: string,
  salt: Uint8Array,
): Promise<CryptoKey> {
  const baseKey = await crypto.subtle.importKey(
    "raw",
    asBuffer(new TextEncoder().encode(passphrase)),
    "PBKDF2",
    false,
    ["deriveKey"],
  );
  return crypto.subtle.deriveKey(
    {
      name: "PBKDF2",
      salt: asBuffer(salt),
      iterations: PBKDF2_ITERATIONS,
      hash: "SHA-256",
    },
    baseKey,
    { name: "AES-GCM", length: 256 },
    false,
    ["encrypt", "decrypt"],
  );
}

export async function hasStoredIdentity(): Promise<boolean> {
  const stored = await get<StoredIdentity>(IDB_KEY);
  return !!stored;
}

/** True iff the stored identity carries an encrypted recovery phrase
 *  (created via the 12-word-mnemonic flow). Used to distinguish a
 *  truly-legacy hex-only account from a phrase account that just isn't
 *  available in the current session (e.g. biometric unlock didn't decrypt it). */
export async function hasStoredPhrase(): Promise<boolean> {
  const stored = await get<StoredIdentity>(IDB_KEY);
  return !!(stored?.phrase_ciphertext && stored?.phrase_nonce);
}

export async function loadStoredIdentityMeta(): Promise<
  Pick<StoredIdentity, "handle" | "server" | "primary" | "created_at"> | null
> {
  const stored = await get<StoredIdentity>(IDB_KEY);
  if (!stored) return null;
  const { handle, server, primary, created_at } = stored;
  return { handle, server, primary, created_at };
}

export async function saveIdentity(opts: {
  handle: string;
  server: string;
  primary: Identity;
  seed: Uint8Array;
  passphrase: string;
  /** Optional 12-word phrase — stored encrypted so the user can re-display
   *  it later. The seed is derived from the phrase one-way (SHA-256 with
   *  domain tag); we can't recover the phrase from the seed, hence
   *  storing it explicitly. */
  phrase?: string;
}): Promise<void> {
  const salt = crypto.getRandomValues(new Uint8Array(16));
  const nonce = crypto.getRandomValues(new Uint8Array(12));
  const key = await deriveKey(opts.passphrase, salt);
  const ciphertext = new Uint8Array(
    await crypto.subtle.encrypt(
      { name: "AES-GCM", iv: asBuffer(nonce) },
      key,
      asBuffer(opts.seed),
    ),
  );

  const record: StoredIdentity = {
    version: 1,
    handle: opts.handle,
    server: opts.server,
    primary: opts.primary,
    salt: bytesToHex(salt),
    nonce: bytesToHex(nonce),
    ciphertext: bytesToHex(ciphertext),
    created_at: new Date().toISOString(),
  };

  if (opts.phrase) {
    // Fresh nonce — AES-GCM nonce reuse under the same key is fatal.
    const phraseNonce = crypto.getRandomValues(new Uint8Array(12));
    const phraseCt = new Uint8Array(
      await crypto.subtle.encrypt(
        { name: "AES-GCM", iv: asBuffer(phraseNonce) },
        key,
        asBuffer(new TextEncoder().encode(opts.phrase)),
      ),
    );
    record.phrase_nonce = bytesToHex(phraseNonce);
    record.phrase_ciphertext = bytesToHex(phraseCt);
  }

  await set(IDB_KEY, record);
}

export async function unlockIdentity(
  passphrase: string,
): Promise<UnlockedIdentity> {
  const stored = await get<StoredIdentity>(IDB_KEY);
  if (!stored) throw new Error("no stored identity");

  const salt = hexToBytes(stored.salt);
  const nonce = hexToBytes(stored.nonce);
  const ciphertext = hexToBytes(stored.ciphertext);
  const key = await deriveKey(passphrase, salt);

  let plaintext: Uint8Array;
  try {
    plaintext = new Uint8Array(
      await crypto.subtle.decrypt(
        { name: "AES-GCM", iv: asBuffer(nonce) },
        key,
        asBuffer(ciphertext),
      ),
    );
  } catch {
    throw new Error("wrong passphrase");
  }
  if (plaintext.length !== 32) {
    throw new Error(
      `unlocked seed is ${plaintext.length} bytes, expected 32`,
    );
  }

  let phrase: string | undefined;
  if (stored.phrase_nonce && stored.phrase_ciphertext) {
    try {
      const pn = hexToBytes(stored.phrase_nonce);
      const pc = hexToBytes(stored.phrase_ciphertext);
      const pBytes = new Uint8Array(
        await crypto.subtle.decrypt(
          { name: "AES-GCM", iv: asBuffer(pn) },
          key,
          asBuffer(pc),
        ),
      );
      phrase = new TextDecoder().decode(pBytes);
    } catch {
      // The seed unlocked fine, so the passphrase is right; a phrase
      // decrypt failure here would point at IDB corruption. Surface it
      // by logging, but don't block unlock — the user can still chat.
      console.error("identity-store: phrase decrypt failed (seed OK)");
    }
  }

  return {
    handle: stored.handle,
    server: stored.server,
    primary: stored.primary,
    seed: plaintext,
    phrase,
  };
}

export async function deleteStoredIdentity(): Promise<void> {
  await del(IDB_KEY);
}

/**
 * Reconstruct an `UnlockedIdentity` from a seed + the stored metadata.
 * Used by the biometric flow: WebAuthn hands us back the seed, but
 * handle/server/primary live in the same record the passphrase flow
 * writes — we just merge them. Throws if no stored identity exists or
 * if the seed isn't 32 bytes.
 */
export async function unlockWithSeed(seed: Uint8Array): Promise<UnlockedIdentity> {
  if (seed.length !== 32) {
    throw new Error(`seed is ${seed.length} bytes, expected 32`);
  }
  const stored = await get<StoredIdentity>(IDB_KEY);
  if (!stored) throw new Error("no stored identity");
  return {
    handle: stored.handle,
    server: stored.server,
    primary: stored.primary,
    seed,
  };
}