rspace-online/shared/local-first/crypto.ts

/**
 * Layer 1: Crypto — Document-level encryption for local-first data.
 *
 * Key hierarchy:
 *   Master Key (PRF output / passphrase) → HKDF
 *     → Space Key (info: "rspace:{spaceId}") → HKDF
 *       → Doc Key (info: "rspace:{spaceId}:{docId}") → AES-256-GCM
 *
 * Server never sees plaintext — only encrypted Automerge binary blobs.
 * Extends the existing EncryptIDKeyManager key derivation pattern.
 */

// ============================================================================
// TYPES
// ============================================================================

export interface EncryptedBlob {
  /** AES-256-GCM ciphertext (includes auth tag) */
  ciphertext: Uint8Array;
  /** 12-byte random nonce */
  nonce: Uint8Array;
}

// ============================================================================
// DocCrypto
// ============================================================================

const encoder = new TextEncoder();

/**
 * DocCrypto — derives per-space and per-document encryption keys from a
 * master key (typically the EncryptIDKeyManager's AES-256-GCM key material).
 *
 * Usage:
 *   const crypto = new DocCrypto();
 *   await crypto.init(masterKey);           // from EncryptIDKeyManager PRF
 *   const spaceKey = await crypto.deriveSpaceKey('my-space');
 *   const docKey = await crypto.deriveDocKey(spaceKey, 'notes:items');
 *   const blob = await crypto.encrypt(docKey, automergeBytes);
 *   const plain = await crypto.decrypt(docKey, blob);
 */
export class DocCrypto {
  #masterKeyMaterial: CryptoKey | null = null;

  /**
   * Initialize from a master key. Accepts either:
   *  - A CryptoKey with HKDF usage (from EncryptIDKeyManager.initFromPRF)
   *  - Raw key bytes (Uint8Array / ArrayBuffer) that will be imported as HKDF material
   */
  async init(masterKey: CryptoKey | Uint8Array | ArrayBuffer): Promise<void> {
    if (masterKey instanceof CryptoKey) {
      // If the key already supports deriveBits/deriveKey, use it directly
      if (masterKey.algorithm.name === 'HKDF') {
        this.#masterKeyMaterial = masterKey;
      } else {
        // It's an AES-GCM key — export raw bits and re-import as HKDF
        const raw = await crypto.subtle.exportKey('raw', masterKey);
        this.#masterKeyMaterial = await crypto.subtle.importKey(
          'raw',
          raw as ArrayBuffer,
          { name: 'HKDF' },
          false,
          ['deriveKey', 'deriveBits']
        );
      }
    } else {
      // Raw bytes → import as HKDF
      const buf = masterKey instanceof Uint8Array ? masterKey.buffer as ArrayBuffer : masterKey;
      this.#masterKeyMaterial = await crypto.subtle.importKey(
        'raw',
        buf,
        { name: 'HKDF' },
        false,
        ['deriveKey', 'deriveBits']
      );
    }
  }

  /**
   * Initialize directly from WebAuthn PRF output (convenience shortcut).
   */
  async initFromPRF(prfOutput: ArrayBuffer): Promise<void> {
    await this.init(new Uint8Array(prfOutput));
  }

  get isInitialized(): boolean {
    return this.#masterKeyMaterial !== null;
  }

  /**
   * Derive a space-level HKDF key.
   * info = "rspace:{spaceId}"
   */
  async deriveSpaceKey(spaceId: string): Promise<CryptoKey> {
    this.#assertInit();

    // Derive 256 bits of key material for the space
    const bits = await crypto.subtle.deriveBits(
      {
        name: 'HKDF',
        hash: 'SHA-256',
        salt: encoder.encode('rspace-space-key-v1'),
        info: encoder.encode(`rspace:${spaceId}`),
      },
      this.#masterKeyMaterial!,
      256
    );

    // Re-import as HKDF for further derivation (space → doc)
    return crypto.subtle.importKey(
      'raw',
      bits,
      { name: 'HKDF' },
      false,
      ['deriveKey', 'deriveBits']
    );
  }

  /**
   * Derive a document-level AES-256-GCM key from a space key.
   * info = "rspace:{spaceId}:{docId}"
   */
  async deriveDocKey(spaceKey: CryptoKey, docId: string): Promise<CryptoKey> {
    return crypto.subtle.deriveKey(
      {
        name: 'HKDF',
        hash: 'SHA-256',
        salt: encoder.encode('rspace-doc-key-v1'),
        info: encoder.encode(`doc:${docId}`),
      },
      spaceKey,
      { name: 'AES-GCM', length: 256 },
      false, // non-extractable
      ['encrypt', 'decrypt']
    );
  }

  /**
   * Encrypt data with a document key.
   * Returns ciphertext + 12-byte random nonce.
   */
  async encrypt(docKey: CryptoKey, data: Uint8Array): Promise<EncryptedBlob> {
    const nonce = crypto.getRandomValues(new Uint8Array(12));
    const ciphertext = await crypto.subtle.encrypt(
      { name: 'AES-GCM', iv: nonce },
      docKey,
      data.buffer as ArrayBuffer
    );
    return {
      ciphertext: new Uint8Array(ciphertext),
      nonce,
    };
  }

  /**
   * Decrypt an encrypted blob with a document key.
   */
  async decrypt(docKey: CryptoKey, blob: EncryptedBlob): Promise<Uint8Array> {
    const iv = new Uint8Array(blob.nonce) as unknown as ArrayBuffer;
    const ct = new Uint8Array(blob.ciphertext) as unknown as ArrayBuffer;
    const plaintext = await crypto.subtle.decrypt(
      { name: 'AES-GCM', iv },
      docKey,
      ct
    );
    return new Uint8Array(plaintext);
  }

  /**
   * Convenience: derive doc key directly from master key (space + doc in one call).
   */
  async deriveDocKeyDirect(spaceId: string, docId: string): Promise<CryptoKey> {
    const spaceKey = await this.deriveSpaceKey(spaceId);
    return this.deriveDocKey(spaceKey, docId);
  }

  /**
   * Serialize an EncryptedBlob for storage (nonce prepended to ciphertext).
   * Format: [12-byte nonce][ciphertext...]
   */
  static pack(blob: EncryptedBlob): Uint8Array {
    const packed = new Uint8Array(12 + blob.ciphertext.length);
    packed.set(blob.nonce, 0);
    packed.set(blob.ciphertext, 12);
    return packed;
  }

  /**
   * Deserialize a packed blob.
   */
  static unpack(packed: Uint8Array): EncryptedBlob {
    return {
      nonce: packed.slice(0, 12),
      ciphertext: packed.slice(12),
    };
  }

  /**
   * Clear master key material from memory.
   */
  clear(): void {
    this.#masterKeyMaterial = null;
  }

  #assertInit(): void {
    if (!this.#masterKeyMaterial) {
      throw new Error('DocCrypto not initialized — call init() first');
    }
  }
}