Shade/packages/shade-storage-encrypted/src/crypto/key-manager.ts

/**
 * KeyManager — owns the masterKey lifecycle for at-rest encryption.
 *
 * Five sources are supported:
 *   1. passphrase — scrypt-derived from a developer-supplied secret
 *   2. argon2id   — memory-hard KDF for low-entropy secrets (PINs)
 *   3. keychain   — fetched from OS keychain via @shade/keychain (optional dep)
 *   4. injected   — caller supplies the 32-byte raw key directly
 *   5. composite  — HKDF-combine N sub-sources into one master key
 *                   (multi-factor unlock — every source is required)
 *
 * The KeyManager pre-derives storageKey at construction time and caches the
 * per-(table, column) field keys. masterKey is zeroized after storageKey
 * derivation to limit residency.
 */

import {
  deriveFieldKey,
  deriveMasterKey,
  deriveMasterKeyArgon2id,
  deriveStorageKey,
  hkdfDerive,
  DEFAULT_SCRYPT,
  type Argon2idParams,
  type ScryptParams,
} from './kdf.js';

export type KeySource =
  | {
      kind: 'passphrase';
      passphrase: string;
      /** Stable 16+ byte salt persisted alongside the DB. */
      salt: Uint8Array;
      params?: ScryptParams;
    }
  | {
      kind: 'argon2id';
      /** Low-entropy secret (PIN, short password). */
      secret: string | Uint8Array;
      /** Stable 16+ byte salt persisted alongside the DB. */
      salt: Uint8Array;
      params?: Argon2idParams;
    }
  | {
      kind: 'keychain';
      /** Service identifier (e.g. "shade.storage"). */
      service: string;
      /** Account / key name within the keychain. */
      account: string;
      /** If true, generate + store a new 32-byte key when one is absent. */
      createIfMissing?: boolean;
    }
  | {
      kind: 'injected';
      /** Raw 32-byte master key. */
      key: Uint8Array;
    }
  | {
      kind: 'composite';
      /**
       * Sub-sources HKDF-combined into the master key, in order. Every
       * source is mandatory: omitting or substituting any source yields
       * a different master key and the storage open() will fail.
       *
       * Order is significant by design — `[pwd, pin]` and `[pin, pwd]`
       * derive different master keys.
       */
      sources: KeySource[];
      /**
       * HKDF info string for domain separation. Defaults to
       * `shade-composite-master-v1`. Apps that want their composite key
       * to be cryptographically distinct from any other Shade composite
       * should override this with an app-specific string.
       */
      info?: string;
    };

/** Pluggable keychain backend. Implementations live in @shade/keychain. */
export interface KeychainBackend {
  get(service: string, account: string): Promise<Uint8Array | null>;
  set(service: string, account: string, value: Uint8Array): Promise<void>;
  delete(service: string, account: string): Promise<void>;
}

export interface KeyManagerOptions {
  /** Required when KeySource.kind === 'keychain'. */
  keychain?: KeychainBackend;
}

export class KeyManager {
  private readonly storageKey: Uint8Array;
  private readonly fieldKeyCache = new Map<string, Uint8Array>();

  private constructor(storageKey: Uint8Array) {
    this.storageKey = storageKey;
  }

  /**
   * Open a KeyManager from any supported source. May call the OS keychain
   * (async) or run scrypt (slow on cold start).
   */
  static async open(source: KeySource, opts: KeyManagerOptions = {}): Promise<KeyManager> {
    const masterKey = await resolveMasterKey(source, opts);
    try {
      const storageKey = deriveStorageKey(masterKey);
      return new KeyManager(storageKey);
    } finally {
      // Zeroize whichever buffer we can — never the caller's own buffer
      // when source.kind === 'injected', since the caller may still hold it.
      if (source.kind !== 'injected') masterKey.fill(0);
    }
  }

  /** Look up (and cache) the AEAD key for a given table+column. */
  fieldKey(table: string, column: string): Uint8Array {
    const cacheKey = `${table}\x1f${column}`;
    let k = this.fieldKeyCache.get(cacheKey);
    if (!k) {
      k = deriveFieldKey(this.storageKey, table, column);
      this.fieldKeyCache.set(cacheKey, k);
    }
    return k;
  }

  /** Sanity helper for verifying that a passphrase decrypts the active DB. */
  storageKeyFingerprint(): Uint8Array {
    // Deterministic 8-byte tag for "is this the same masterKey?" without
    // exposing the storageKey itself. Tag-only — not a MAC.
    const tag = new Uint8Array(8);
    for (let i = 0; i < this.storageKey.length; i++) {
      tag[i % 8]! ^= this.storageKey[i]!;
    }
    return tag;
  }

  /** Forget all derived keys — call on shutdown. */
  destroy(): void {
    this.storageKey.fill(0);
    for (const v of this.fieldKeyCache.values()) v.fill(0);
    this.fieldKeyCache.clear();
  }
}

async function resolveMasterKey(source: KeySource, opts: KeyManagerOptions): Promise<Uint8Array> {
  switch (source.kind) {
    case 'passphrase':
      return deriveMasterKey(source.passphrase, source.salt, source.params ?? DEFAULT_SCRYPT);

    case 'argon2id':
      return deriveMasterKeyArgon2id(source.secret, source.salt, source.params);

    case 'injected':
      if (source.key.length !== 32) throw new Error('injected key must be exactly 32 bytes');
      return new Uint8Array(source.key); // copy, in case caller mutates

    case 'keychain': {
      if (!opts.keychain) {
        throw new Error('keychain source requires opts.keychain (install @shade/keychain)');
      }
      const existing = await opts.keychain.get(source.service, source.account);
      if (existing) {
        if (existing.length !== 32) throw new Error('keychain returned a non-32-byte key');
        return existing;
      }
      if (!source.createIfMissing) {
        throw new Error(`no key in keychain for ${source.service}/${source.account}`);
      }
      const fresh = new Uint8Array(32);
      globalThis.crypto.getRandomValues(fresh);
      await opts.keychain.set(source.service, source.account, fresh);
      return fresh;
    }

    case 'composite': {
      if (source.sources.length === 0) {
        throw new Error('composite source requires at least one sub-source');
      }
      const subKeys: Uint8Array[] = [];
      try {
        for (const sub of source.sources) {
          if (sub.kind === 'composite') {
            // Composite-of-composite would silently flatten the unlock
            // semantics ("any of N" vs "all of N") if the inner is misread.
            // Forbidding nesting keeps the contract clear.
            throw new Error('composite sources cannot be nested');
          }
          subKeys.push(await resolveMasterKey(sub, opts));
        }
        let total = 0;
        for (const k of subKeys) total += k.length;
        const ikm = new Uint8Array(total);
        let off = 0;
        for (const k of subKeys) {
          ikm.set(k, off);
          off += k.length;
        }
        const info = source.info ?? 'shade-composite-master-v1';
        try {
          return hkdfDerive(ikm, info, 32);
        } finally {
          ikm.fill(0);
        }
      } finally {
        for (const k of subKeys) k.fill(0);
      }
    }
  }
}