Shade/packages/shade-server/src/kt-integration.ts

/**
 * Key-Transparency integration for the Shade prekey server.
 *
 * The prekey server is the *source of truth* for which prekey bundle is
 * currently published for each address. Without KT a malicious server
 * could swap a bundle without anyone noticing. With KT enabled:
 *
 *   - Every register / delete operation appends a leaf to an append-only
 *     Merkle log via `KTLogManager.recordRegister` / `recordDelete`.
 *   - After each mutation the manager re-signs and publishes a fresh STH.
 *   - GET /v1/keys/bundle/:address attaches a `ktProof` to its response,
 *     so the client can verify inclusion + freshness.
 *   - GET /v1/kt/sth and friends expose the log to witnesses.
 *
 * KT is **opt-in**: pass `keyTransparency` to `createPrekeyServer`. When
 * absent, the server behaves exactly as before — proof fields are simply
 * not added to the bundle response.
 */

import type { CryptoProvider } from '@shade/core';
import {
  KTLogManager,
  type KTLogStore,
  computeBundleHash,
  ktProofToWire,
  sthToWire,
  type KTProof,
  type KTProofWire,
  type STHWire,
  type SignedTreeHead,
} from '@shade/key-transparency';

export interface KeyTransparencyConfig {
  /** Persistent store for the log + index + STH set. */
  store: KTLogStore;
  /** Operator's STH signing key (32-byte Ed25519 seed). */
  signingPrivateKey: Uint8Array;
  /** Operator's STH signing public key (32-byte Ed25519). */
  signingPublicKey: Uint8Array;
  /**
   * Heartbeat interval — minimum gap between fresh STHs even when no
   * mutations occur. Default 10 minutes; set to 0 to disable.
   */
  heartbeatIntervalMs?: number;
  /** Time source override (testing). */
  now?: () => number;
}

/**
 * Wraps a `KTLogManager` with the bookkeeping the server cares about:
 *  - Serializes mutations (single-writer guarantee).
 *  - Caches the latest STH so bundle-fetch is hot-path-fast.
 *  - Schedules / surfaces heartbeats.
 *  - Lazily backfills index entries from the prekey-server's existing
 *    state when KT is first turned on.
 */
export class KeyTransparencyService {
  private readonly mgr: KTLogManager;
  private readonly store: KTLogStore;
  private readonly heartbeatIntervalMs: number;
  private readonly now: () => number;
  private latest: SignedTreeHead | null = null;
  private mutex: Promise<unknown> = Promise.resolve();

  private constructor(
    mgr: KTLogManager,
    store: KTLogStore,
    opts: { heartbeatIntervalMs?: number; now?: () => number },
  ) {
    this.mgr = mgr;
    this.store = store;
    this.heartbeatIntervalMs = opts.heartbeatIntervalMs ?? 10 * 60 * 1000;
    this.now = opts.now ?? (() => Date.now());
  }

  static async create(crypto: CryptoProvider, cfg: KeyTransparencyConfig): Promise<KeyTransparencyService> {
    const mgr = await KTLogManager.create({
      crypto,
      store: cfg.store,
      signingPrivateKey: cfg.signingPrivateKey,
      signingPublicKey: cfg.signingPublicKey,
      ...(cfg.now ? { now: cfg.now } : {}),
    });
    const svc = new KeyTransparencyService(mgr, cfg.store, {
      ...(cfg.heartbeatIntervalMs !== undefined ? { heartbeatIntervalMs: cfg.heartbeatIntervalMs } : {}),
      ...(cfg.now ? { now: cfg.now } : {}),
    });

    // Cache or generate the initial STH so bundle responses always have one.
    const existing = await cfg.store.getLatestSTH();
    if (existing) {
      svc.latest = existing;
    } else {
      svc.latest = await mgr.publishSTH();
    }
    return svc;
  }

  /**
   * Run a mutation under the manager's serial lock and refresh the STH.
   */
  private async withLock<T>(fn: () => Promise<T>): Promise<T> {
    const prev = this.mutex;
    let resolveNext: () => void;
    const next = new Promise<void>((res) => {
      resolveNext = res;
    });
    this.mutex = next;
    try {
      await prev.catch(() => {});
      return await fn();
    } finally {
      resolveNext!();
    }
  }

  async recordRegister(address: string, bundle: {
    identitySigningKey: Uint8Array;
    identityDHKey: Uint8Array;
    signedPreKey: { keyId: number; publicKey: Uint8Array; signature: Uint8Array };
  }): Promise<SignedTreeHead> {
    return this.withLock(async () => {
      await this.mgr.recordRegister(address, computeBundleHash(bundle));
      this.latest = await this.mgr.publishSTH();
      return this.latest!;
    });
  }

  async recordDelete(address: string): Promise<SignedTreeHead> {
    return this.withLock(async () => {
      await this.mgr.recordDelete(address);
      this.latest = await this.mgr.publishSTH();
      return this.latest!;
    });
  }

  /**
   * Build a proof for a freshly-fetched bundle. Returns null if the
   * address has no live entry (caller can request an absence proof
   * via `buildAbsenceProof` instead).
   */
  async buildBundleInclusion(address: string): Promise<KTProof | null> {
    const sth = await this.maybeHeartbeat();
    return this.mgr.buildBundleInclusionProof(address, sth);
  }

  async buildAbsence(address: string): Promise<KTProof | null> {
    const sth = await this.maybeHeartbeat();
    return this.mgr.buildBundleAbsenceProof(address, sth);
  }

  /** Latest STH — issuing a heartbeat first if the cached one is stale. */
  async getLatestSTH(): Promise<SignedTreeHead> {
    return this.maybeHeartbeat();
  }

  /** Historical STH at a specific tree size. */
  async getSTHByTreeSize(treeSize: number): Promise<SignedTreeHead | null> {
    return this.store.getSTHByTreeSize(treeSize);
  }

  /** All persisted STHs in a time window — used by witness backfill. */
  async listSTHs(fromTimestampMs?: number, toTimestampMs?: number): Promise<SignedTreeHead[]> {
    return this.store.listSTHs(fromTimestampMs, toTimestampMs);
  }

  /** Build a consistency proof for `from → to`. */
  async buildConsistencyProof(fromTreeSize: number, toTreeSize?: number): Promise<{
    fromTreeSize: number;
    toTreeSize: number;
    proof: Uint8Array[];
  }> {
    return this.withLock(async () => {
      const targetSize = toTreeSize ?? this.mgr.getTreeSize();
      const proof = await this.mgr.buildHistoricalConsistencyProof(fromTreeSize, targetSize);
      return { fromTreeSize, toTreeSize: targetSize, proof };
    });
  }

  /** STH signing public key — operators expose this to clients OOB. */
  getSigningPublicKey(): Uint8Array {
    return this.mgr.getSigningPublicKey();
  }

  getLogId(): Uint8Array {
    return this.mgr.getLogId();
  }

  private async maybeHeartbeat(): Promise<SignedTreeHead> {
    if (!this.latest) {
      return this.withLock(async () => {
        this.latest = await this.mgr.publishSTH();
        return this.latest!;
      });
    }
    if (this.heartbeatIntervalMs <= 0) return this.latest;
    const age = this.now() - this.latest.timestampMs;
    if (age < this.heartbeatIntervalMs) return this.latest;
    return this.withLock(async () => {
      // Re-check age inside the lock — another caller may have published.
      const ageNow = this.now() - (this.latest?.timestampMs ?? 0);
      if (ageNow < this.heartbeatIntervalMs) return this.latest!;
      this.latest = await this.mgr.publishSTH();
      return this.latest!;
    });
  }
}

/** Helpers to encode an STH for the wire (base64). */
export function encodeSthForWire(sth: SignedTreeHead): STHWire {
  return sthToWire(sth, (b) => Buffer.from(b).toString('base64'));
}

export function encodeProofForWire(proof: KTProof): KTProofWire {
  return ktProofToWire(proof);
}