Shade/packages/shade-key-transparency/src/index-tree.ts

/**
 * Address-index commitment.
 *
 * The Merkle log itself records mutation events (`address → bundle_hash`
 * at time T), but doesn't natively answer "what's the *current* state for
 * `address`?" or "does `address` exist?".
 *
 * The address index is a **lexicographically sorted snapshot** of the
 * current `(address, latest_leaf_index)` mapping. Its commitment hash —
 * `index_root` — is part of every Signed Tree Head.
 *
 * Inclusion proof:  the entry exists at sorted index `i`, prove it via
 *                   audit path (same Merkle construction as the main log).
 *
 * Absence proof:    the address would sort between two adjacent existing
 *                   entries; prove inclusion of those two adjacent entries
 *                   and that the queried address sorts strictly between them.
 *
 * V1 representation: a flat sorted array. We re-hash the whole index per
 * STH (cheap up to ~1M entries). V2 will move to a sparse Merkle tree if
 * the dataset grows enough that flat re-hash becomes a bottleneck.
 */

import { leafHash, nodeHash, emptyRootHash } from './hashes.js';
import { sha256Sync } from './sha256.js';
import { constantTimeEqual } from './util.js';
import { mth, auditPath, recomputeRootFromAuditPath } from './log.js';

export interface AddressIndexEntry {
  /** The address (UTF-8 string). */
  address: string;
  /** Most recent log leaf index that mutated this address. */
  latestLeafIndex: number;
  /** Most recent bundle hash committed for this address. */
  bundleHash: Uint8Array;
  /** Whether the latest event was a delete (tombstone). */
  deleted: boolean;
}

/** Encode an index entry into the bytes that go into a Merkle leaf. */
export function encodeIndexEntry(entry: AddressIndexEntry): Uint8Array {
  const addrBytes = new TextEncoder().encode(entry.address);
  if (addrBytes.length > 0xffff) throw new Error('address too long');
  if (entry.bundleHash.length > 0xffff) throw new Error('bundleHash too long');
  const len = 2 + addrBytes.length + 4 + 1 + 2 + entry.bundleHash.length;
  const out = new Uint8Array(len);
  const view = new DataView(out.buffer);
  let off = 0;
  view.setUint16(off, addrBytes.length);
  off += 2;
  out.set(addrBytes, off);
  off += addrBytes.length;
  view.setUint32(off, entry.latestLeafIndex >>> 0);
  off += 4;
  out[off++] = entry.deleted ? 1 : 0;
  view.setUint16(off, entry.bundleHash.length);
  off += 2;
  out.set(entry.bundleHash, off);
  return out;
}

/** Compute index_root over a sorted entry list. */
export function computeIndexRoot(sortedEntries: AddressIndexEntry[]): Uint8Array {
  if (sortedEntries.length === 0) return emptyRootHash();
  const leaves = sortedEntries.map((e) => leafHash(encodeIndexEntry(e)));
  return mth(leaves, 0, leaves.length);
}

/** Compare two addresses lexicographically (by UTF-8 byte order). */
export function compareAddresses(a: string, b: string): number {
  const ab = new TextEncoder().encode(a);
  const bb = new TextEncoder().encode(b);
  const len = Math.min(ab.length, bb.length);
  for (let i = 0; i < len; i++) {
    if (ab[i]! !== bb[i]!) return ab[i]! - bb[i]!;
  }
  return ab.length - bb.length;
}

/**
 * In-memory address index. Maintains the canonical sorted ordering; on
 * mutate, the operator re-computes index_root for the next STH.
 */
export class AddressIndex {
  private entries: AddressIndexEntry[] = [];
  private positionByAddress = new Map<string, number>();

  get size(): number {
    return this.entries.length;
  }

  /** Idempotently set an entry; re-sorts only when a new address is added. */
  upsert(entry: AddressIndexEntry): void {
    const existingPos = this.positionByAddress.get(entry.address);
    if (existingPos !== undefined) {
      this.entries[existingPos] = { ...entry };
      return;
    }
    // Insert keeping sort order
    let lo = 0;
    let hi = this.entries.length;
    while (lo < hi) {
      const mid = (lo + hi) >>> 1;
      if (compareAddresses(this.entries[mid]!.address, entry.address) < 0) lo = mid + 1;
      else hi = mid;
    }
    this.entries.splice(lo, 0, { ...entry });
    // Rebuild position map (positions shift after insert)
    this.positionByAddress.clear();
    for (let i = 0; i < this.entries.length; i++) {
      this.positionByAddress.set(this.entries[i]!.address, i);
    }
  }

  /** Mark an address tombstoned. Keeps the entry in sorted order. */
  tombstone(address: string, latestLeafIndex: number): void {
    const pos = this.positionByAddress.get(address);
    if (pos === undefined) return;
    const e = this.entries[pos]!;
    this.entries[pos] = {
      ...e,
      deleted: true,
      latestLeafIndex,
      bundleHash: new Uint8Array(0),
    };
  }

  /** Snapshot ordered list (defensive copy). */
  snapshot(): AddressIndexEntry[] {
    return this.entries.map((e) => ({ ...e, bundleHash: new Uint8Array(e.bundleHash) }));
  }

  /** Compute the index commitment root over the current sorted list. */
  rootHash(): Uint8Array {
    return computeIndexRoot(this.entries);
  }

  /** Look up an entry. */
  get(address: string): AddressIndexEntry | undefined {
    const pos = this.positionByAddress.get(address);
    if (pos === undefined) return undefined;
    return { ...this.entries[pos]!, bundleHash: new Uint8Array(this.entries[pos]!.bundleHash) };
  }

  /** Build inclusion proof: returns sorted-position + audit path. */
  inclusionProof(address: string): IndexInclusionProof | null {
    const pos = this.positionByAddress.get(address);
    if (pos === undefined) return null;
    const leaves = this.entries.map((e) => leafHash(encodeIndexEntry(e)));
    return {
      kind: 'inclusion',
      position: pos,
      treeSize: this.entries.length,
      entry: { ...this.entries[pos]!, bundleHash: new Uint8Array(this.entries[pos]!.bundleHash) },
      auditPath: auditPath(leaves, pos, leaves.length),
    };
  }

  /**
   * Build absence proof: returns the two adjacent entries that bracket the
   * queried address (or boundary case for first/last).
   */
  absenceProof(address: string): IndexAbsenceProof | null {
    if (this.positionByAddress.has(address)) return null;

    if (this.entries.length === 0) {
      return {
        kind: 'absence',
        treeSize: 0,
        queryAddress: address,
        prev: null,
        next: null,
      };
    }

    // Find insertion position
    let lo = 0;
    let hi = this.entries.length;
    while (lo < hi) {
      const mid = (lo + hi) >>> 1;
      if (compareAddresses(this.entries[mid]!.address, address) < 0) lo = mid + 1;
      else hi = mid;
    }

    const leaves = this.entries.map((e) => leafHash(encodeIndexEntry(e)));
    const prevPos = lo - 1;
    const nextPos = lo;

    const prev =
      prevPos >= 0
        ? {
            position: prevPos,
            entry: {
              ...this.entries[prevPos]!,
              bundleHash: new Uint8Array(this.entries[prevPos]!.bundleHash),
            },
            auditPath: auditPath(leaves, prevPos, leaves.length),
          }
        : null;

    const next =
      nextPos < this.entries.length
        ? {
            position: nextPos,
            entry: {
              ...this.entries[nextPos]!,
              bundleHash: new Uint8Array(this.entries[nextPos]!.bundleHash),
            },
            auditPath: auditPath(leaves, nextPos, leaves.length),
          }
        : null;

    return {
      kind: 'absence',
      treeSize: this.entries.length,
      queryAddress: address,
      prev,
      next,
    };
  }

  /** Hot-load from a sorted entry array (used by persistent stores). */
  static fromEntries(sortedEntries: AddressIndexEntry[]): AddressIndex {
    const idx = new AddressIndex();
    for (const e of sortedEntries) {
      idx.entries.push({ ...e, bundleHash: new Uint8Array(e.bundleHash) });
    }
    for (let i = 0; i < idx.entries.length; i++) {
      idx.positionByAddress.set(idx.entries[i]!.address, i);
    }
    return idx;
  }
}

export interface IndexInclusionProof {
  kind: 'inclusion';
  position: number;
  treeSize: number;
  entry: AddressIndexEntry;
  auditPath: Uint8Array[];
}

export interface IndexAbsenceProof {
  kind: 'absence';
  treeSize: number;
  queryAddress: string;
  /**
   * Largest existing entry less than the query (null if the query would
   * be the first entry).
   */
  prev: { position: number; entry: AddressIndexEntry; auditPath: Uint8Array[] } | null;
  /**
   * Smallest existing entry greater than the query (null if the query
   * would be appended after the last entry).
   */
  next: { position: number; entry: AddressIndexEntry; auditPath: Uint8Array[] } | null;
}

export type IndexProof = IndexInclusionProof | IndexAbsenceProof;

/**
 * Verify an inclusion proof against an `index_root` commitment.
 */
export function verifyInclusionProof(
  proof: IndexInclusionProof,
  indexRoot: Uint8Array,
): boolean {
  const lh = leafHash(encodeIndexEntry(proof.entry));
  let recomputed: Uint8Array;
  try {
    recomputed = recomputeRootFromAuditPath(lh, proof.position, proof.treeSize, proof.auditPath);
  } catch {
    return false;
  }
  return constantTimeEqual(recomputed, indexRoot);
}

/**
 * Verify an absence proof:
 *  - prev exists and address(prev) < query
 *  - next exists and query < address(next)
 *  - prev.position + 1 === next.position (they are adjacent)
 *  - both inclusion sub-proofs verify against indexRoot
 *
 * Boundary cases:
 *  - tree empty (treeSize === 0): valid if both prev and next are null
 *  - query smaller than all entries: prev is null, next.position === 0
 *  - query larger than all entries:  next is null, prev.position === treeSize - 1
 */
export function verifyAbsenceProof(
  proof: IndexAbsenceProof,
  indexRoot: Uint8Array,
): boolean {
  if (proof.treeSize === 0) {
    if (proof.prev !== null || proof.next !== null) return false;
    return constantTimeEqual(emptyRootHash(), indexRoot);
  }

  const queryAddr = proof.queryAddress;

  if (proof.prev) {
    if (compareAddresses(proof.prev.entry.address, queryAddr) >= 0) return false;
    const lh = leafHash(encodeIndexEntry(proof.prev.entry));
    let r: Uint8Array;
    try {
      r = recomputeRootFromAuditPath(lh, proof.prev.position, proof.treeSize, proof.prev.auditPath);
    } catch {
      return false;
    }
    if (!constantTimeEqual(r, indexRoot)) return false;
  }

  if (proof.next) {
    if (compareAddresses(queryAddr, proof.next.entry.address) >= 0) return false;
    const lh = leafHash(encodeIndexEntry(proof.next.entry));
    let r: Uint8Array;
    try {
      r = recomputeRootFromAuditPath(lh, proof.next.position, proof.treeSize, proof.next.auditPath);
    } catch {
      return false;
    }
    if (!constantTimeEqual(r, indexRoot)) return false;
  }

  // Boundary checks
  if (proof.prev === null) {
    if (proof.next === null) return false; // already handled treeSize===0
    if (proof.next.position !== 0) return false;
  } else if (proof.next === null) {
    if (proof.prev.position !== proof.treeSize - 1) return false;
  } else {
    if (proof.prev.position + 1 !== proof.next.position) return false;
  }

  return true;
}

/** sha256 helper export for callers that need the same hash function. */
export { sha256Sync };