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

/**
 * KeyManager — owns the masterKey lifecycle for at-rest encryption.
 *
 * Three sources are supported:
 *   1. passphrase — scrypt-derived from a developer-supplied secret
 *   2. keychain   — fetched from OS keychain via @shade/keychain (optional dep)
 *   3. injected   — caller supplies the 32-byte raw key directly
 *
 * 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, deriveStorageKey, type ScryptParams, DEFAULT_SCRYPT } from './kdf.js';

export type KeySource =
  | {
      kind: 'passphrase';
      passphrase: string;
      /** Stable 16+ byte salt persisted alongside the DB. */
      salt: Uint8Array;
      params?: ScryptParams;
    }
  | {
      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;
    };

/** 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 '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;
    }
  }
}