Shade/packages/shade-key-transparency/src/manager.ts

/**
 * KTLogManager — server-side orchestration of the log + address index.
 *
 * Wraps a `KTLogStore` with the algorithmic primitives so that callers
 * (the prekey-server integration layer) never have to think about Merkle
 * paths, index commitments, or STH signing in isolation.
 *
 * Lifecycle:
 *   const mgr = await KTLogManager.create({ crypto, store, signingKey });
 *   await mgr.recordRegister(address, bundleHash, timestampMs);
 *   const sth = await mgr.publishSTH();
 *   const proof = await mgr.buildBundleInclusionProof(address);
 *
 * Concurrency: the manager is single-writer. Server callers must serialize
 * mutations behind a mutex (the integration layer does this with an in-process
 * lock that is sufficient for the single-instance default deployment;
 * multi-instance HA requires external coordination — documented in
 * docs/key-transparency.md §"Recovery and HA").
 */

import type { CryptoProvider } from '@shade/core';
import { OP_DELETE, OP_REGISTER, OP_REPLENISH, encodeLeafData, leafHash } from './hashes.js';
import { MerkleLog, auditPath } from './log.js';
import {
  AddressIndex,
  type AddressIndexEntry,
  type IndexAbsenceProof,
  type IndexInclusionProof,
} from './index-tree.js';
import {
  type SignedTreeHead,
  computeLogId,
  signSth,
} from './sth.js';
import type { KTLogStore } from './store.js';
import type {
  KTBundleAbsenceProof,
  KTBundleInclusionProof,
  KTBundleTombstoneProof,
  KTProof,
} from './proof.js';
import { consistencyProof, verifyConsistencyProof } from './log.js';

export interface KTLogManagerOptions {
  crypto: CryptoProvider;
  store: KTLogStore;
  /** Operator's Ed25519 signing key (private, 32-byte seed). */
  signingPrivateKey: Uint8Array;
  /** Operator's Ed25519 signing public key (pinned by clients). */
  signingPublicKey: Uint8Array;
  /** Time source — defaults to `Date.now()`. */
  now?: () => number;
}

export class KTLogManager {
  private readonly crypto: CryptoProvider;
  private readonly store: KTLogStore;
  private readonly signingPrivateKey: Uint8Array;
  private readonly signingPublicKey: Uint8Array;
  private readonly logId: Uint8Array;
  private readonly now: () => number;

  // In-memory mirror of the persistent state for fast proof generation.
  private merkleLog: MerkleLog;
  private addressIndex: AddressIndex;

  private constructor(opts: KTLogManagerOptions, log: MerkleLog, idx: AddressIndex) {
    this.crypto = opts.crypto;
    this.store = opts.store;
    this.signingPrivateKey = opts.signingPrivateKey;
    this.signingPublicKey = opts.signingPublicKey;
    this.logId = computeLogId(opts.signingPublicKey);
    this.now = opts.now ?? (() => Date.now());
    this.merkleLog = log;
    this.addressIndex = idx;
  }

  static async create(opts: KTLogManagerOptions): Promise<KTLogManager> {
    const size = await opts.store.size();
    const leaves = await opts.store.getLeaves(0, size);
    const log = MerkleLog.fromLeaves(leaves.map((l) => l.leafHash));
    const idx = AddressIndex.fromEntries(await opts.store.getAllIndexEntries());
    return new KTLogManager(opts, log, idx);
  }

  /** Operator's pinned signing public key, for callers to ship to clients. */
  getSigningPublicKey(): Uint8Array {
    return new Uint8Array(this.signingPublicKey);
  }

  /** log_id (== sha256(signingPublicKey)). */
  getLogId(): Uint8Array {
    return new Uint8Array(this.logId);
  }

  /** Current tree size (number of leaves). */
  getTreeSize(): number {
    return this.merkleLog.size;
  }

  /** Record a register/rotate event. Bundle hash is the canonical bundle commit. */
  async recordRegister(
    address: string,
    bundleHash: Uint8Array,
    timestampMs?: number,
  ): Promise<{ leafIndex: number }> {
    return this.recordOperation(address, OP_REGISTER, bundleHash, timestampMs);
  }

  /**
   * Record a "replenish" event. Per §11 of the design notat, replenish does
   * NOT mutate bundle_hash (one-time prekeys aren't part of the commitment).
   * In practice this method is rarely called — kept for cases where the
   * operator wants liveness-evidence of OTP top-ups in the log. When used,
   * the leaf's bundle_hash equals the current index entry's bundle_hash.
   */
  async recordReplenish(
    address: string,
    timestampMs?: number,
  ): Promise<{ leafIndex: number } | null> {
    const existing = await this.store.getIndexEntry(address);
    if (!existing) return null;
    return this.recordOperation(address, OP_REPLENISH, existing.bundleHash, timestampMs);
  }

  /** Record an unregister/tombstone event. */
  async recordDelete(address: string, timestampMs?: number): Promise<{ leafIndex: number }> {
    const result = await this.recordOperation(
      address,
      OP_DELETE,
      new Uint8Array(0),
      timestampMs,
    );
    await this.store.tombstoneIndexEntry(address, result.leafIndex);
    this.addressIndex.tombstone(address, result.leafIndex);
    return result;
  }

  private async recordOperation(
    address: string,
    operation: number,
    bundleHash: Uint8Array,
    timestampMsOpt?: number,
  ): Promise<{ leafIndex: number }> {
    const timestampMs = timestampMsOpt ?? this.now();
    const data = encodeLeafData(timestampMs, operation, address, bundleHash);
    const lh = leafHash(data);
    const index = await this.store.appendLeaf({
      leafHash: lh,
      timestampMs,
      operation,
      address,
      bundleHash,
    });
    this.merkleLog.appendLeafHash(lh);

    if (operation !== OP_DELETE) {
      const entry: AddressIndexEntry = {
        address,
        latestLeafIndex: index,
        bundleHash,
        deleted: false,
      };
      await this.store.upsertIndexEntry(entry);
      this.addressIndex.upsert(entry);
    }

    return { leafIndex: index };
  }

  /** Re-sign and persist the current STH. Idempotent if no change since last call. */
  async publishSTH(timestampMsOpt?: number): Promise<SignedTreeHead> {
    const treeSize = this.merkleLog.size;
    const rootHash = this.merkleLog.rootHash();
    const indexRoot = this.addressIndex.rootHash();
    const timestampMs = timestampMsOpt ?? this.now();
    const sth = await signSth(this.crypto, this.signingPrivateKey, {
      treeSize,
      timestampMs,
      rootHash,
      indexRoot,
      logId: this.logId,
    });
    await this.store.saveSTH(sth);
    return sth;
  }

  /** Build an inclusion proof for the address's *latest* event. */
  async buildBundleInclusionProof(
    address: string,
    sth: SignedTreeHead,
  ): Promise<KTProof | null> {
    const indexEntry = this.addressIndex.get(address);
    if (!indexEntry) return null;
    const leaf = await this.store.getLeaf(indexEntry.latestLeafIndex);
    if (!leaf) return null;
    if (sth.treeSize <= indexEntry.latestLeafIndex) return null;

    // Audit path is over the snapshot at sth.treeSize. The current in-memory
    // log is exactly that size when the manager produced this STH.
    const path = this.auditPathAt(indexEntry.latestLeafIndex, sth.treeSize);

    const indexProof = this.addressIndex.inclusionProof(address);
    if (!indexProof) return null;

    if (indexEntry.deleted) {
      const body: KTBundleTombstoneProof = {
        kind: 'tombstone',
        leafIndex: indexEntry.latestLeafIndex,
        leafTimestampMs: leaf.timestampMs,
        operation: leaf.operation,
        auditPath: path,
        indexProof,
      };
      return { sth, body };
    }

    const body: KTBundleInclusionProof = {
      kind: 'inclusion',
      leafIndex: indexEntry.latestLeafIndex,
      leafTimestampMs: leaf.timestampMs,
      operation: leaf.operation,
      auditPath: path,
      indexProof,
    };
    return { sth, body };
  }

  /** Build an absence proof for an address that does not exist. */
  buildBundleAbsenceProof(address: string, sth: SignedTreeHead): KTProof | null {
    const indexProof = this.addressIndex.absenceProof(address);
    if (!indexProof) return null; // address actually exists
    const body: KTBundleAbsenceProof = { kind: 'absence', indexProof };
    return { sth, body };
  }

  /**
   * Compute a consistency proof from `oldTreeSize` to current. Works against
   * the in-memory log; for very-large logs this becomes O(N) and a future
   * persistent-only path may be needed.
   */
  async buildConsistencyProof(oldTreeSize: number): Promise<{
    proof: Uint8Array[];
    fromTreeSize: number;
    toTreeSize: number;
  }> {
    const size = this.merkleLog.size;
    const leaves = this.merkleLog.exportLeaves();
    const proof = consistencyProof(leaves, oldTreeSize, size);
    return { proof, fromTreeSize: oldTreeSize, toTreeSize: size };
  }

  /**
   * Compute a consistency proof between two arbitrary historical sizes.
   * Reads leaves from the persistent store (so older snapshots can be proven
   * even if the in-memory log has grown further).
   */
  async buildHistoricalConsistencyProof(oldSize: number, newSize: number): Promise<Uint8Array[]> {
    if (newSize > this.merkleLog.size) {
      throw new Error(`newSize ${newSize} exceeds current tree size ${this.merkleLog.size}`);
    }
    const leaves = this.merkleLog.exportLeaves();
    return consistencyProof(leaves, oldSize, newSize);
  }

  // ─── Helpers ──────────────────────────────────────────

  private auditPathAt(leafIndex: number, treeSize: number): Uint8Array[] {
    const leaves = this.merkleLog.exportLeaves().slice(0, treeSize);
    return auditPath(leaves, leafIndex, leaves.length);
  }
}

export { verifyConsistencyProof };