Shade/packages/shade-recovery/src/request.ts

/**
 * New-device flow: rebuild the original identity from threshold-many
 * guardian shares.
 *
 * Sequence:
 *
 *   1. The new device boots a Shade with a temporary identity (caller
 *      must do this BEFORE invoking `requestRecovery`; the new device
 *      needs a published prekey bundle so guardians can reply).
 *   2. The new device's safety number is read off `shade.fingerprint`
 *      and embedded in every `recovery-request` envelope so the
 *      guardian's user can OOB-confirm before approving.
 *   3. For each guardian in the supplied list, we send one
 *      `recovery-request` envelope. We register a transient
 *      `Shade.onMessage` handler that collects the matching
 *      `share-grant` and `share-decline` replies.
 *   4. When `threshold` distinct grants have arrived, we Shamir-combine
 *      them, re-derive the backup passphrase, and call
 *      `Shade.importBackup` — which atomically swaps the temporary
 *      identity for the recovered one.
 *   5. If too many guardians decline (so the threshold can no longer
 *      be reached) or the timeout elapses, we abort with a typed
 *      error.
 *
 * The reconstruction is authenticated end-to-end: a forged share is
 * detected when the AES-GCM tag inside the backup blob fails to
 * verify. We retry combination with subsets of size `threshold` from
 * the received-grants pool until one succeeds OR every subset fails;
 * the latter is a {@link RecoveryReconstructionError} the caller
 * MUST treat as adversarial (do not retry blindly — at least one
 * guardian is malicious).
 */

import type { ShadeEnvelope } from '@shade/core';
import type { Shade } from '@shade/sdk';
import {
  base64UrlToBytes,
  bytesToBase64Url,
  recoveryKeyToBackupPassphrase,
} from './encoding.js';
import {
  RecoveryDeclinedError,
  RecoveryReconstructionError,
  RecoveryTimeoutError,
} from './errors.js';
import {
  encodeRecoveryEnvelope,
  tryParseRecoveryEnvelope,
  type RecoveryRequestEnvelope,
} from './protocol.js';
import type { RecoveryDeliver } from './setup.js';
import { combineShares, decodeShare, type ShamirShare } from './shamir.js';

export interface RequestRecoveryOptions {
  /** Initialized Shade with a temporary identity. */
  shade: Shade;
  /** Address of the original (lost) identity to recover. */
  originalAddress: string;
  /**
   * Guardians to query, in any order. Caller can supply a superset of
   * the threshold — we'll continue collecting until we have enough
   * grants, then stop early. Length must be ≥ threshold.
   */
  guardians: ReadonlyArray<string>;
  /** Reconstruction threshold `k`. Must match the value used at setup time. */
  threshold: number;
  /** Setup id from the original setupRecovery() call. */
  setupId: string;
  /** Outbound transport callback, same shape as setupRecovery. */
  deliver: RecoveryDeliver;
  /**
   * Timeout for the entire flow in milliseconds. Defaults to 5 min.
   * Pass `Infinity` to disable.
   */
  timeoutMs?: number;
  /**
   * Optional progress callback. Fired once per guardian response. The
   * caller wires this to UI so the user sees "3/5 guardians responded".
   */
  onProgress?: (progress: RecoveryProgress) => void;
  /** Wall-clock source. Defaults to Date.now. */
  now?: () => number;
}

export interface RecoveryProgress {
  granted: number;
  declined: number;
  pending: number;
  threshold: number;
  /** Latest guardian to respond. */
  fromAddress?: string;
  /** Outcome of the latest response. */
  latest?: 'grant' | 'decline';
}

export interface RecoveryResult {
  /** True iff `Shade.importBackup` ran and returned. */
  applied: boolean;
  /** Guardians that returned grants. */
  granted: ReadonlyArray<string>;
  /** Guardians that explicitly declined; absent guardians are not listed here. */
  declined: ReadonlyArray<{ address: string; reason: string }>;
  /** Setup fingerprint of the original device at setup time. Hand to UI. */
  setupFingerprint: string | null;
  /**
   * Recovered safety number. Will equal the original device's pre-loss
   * fingerprint (since we restored the same identity). Compare to the
   * pre-loss value the user has on a recovery card / second device for
   * sanity-checking.
   */
  restoredFingerprint: string;
}

interface PendingShare {
  share: ShamirShare;
  backupBlob: string;
  setupFingerprint?: string;
  fromAddress: string;
}

/**
 * Drive the recovery flow to completion. On success the new Shade
 * instance now hosts the original identity (its prior temporary
 * identity is overwritten). On failure throws a typed
 * {@link RecoveryError} subclass — never `applied: false`.
 */
export async function requestRecovery(opts: RequestRecoveryOptions): Promise<RecoveryResult> {
  validateOptions(opts);
  const now = opts.now ?? Date.now;
  const timeoutMs = opts.timeoutMs ?? 5 * 60 * 1000;
  const requesterFingerprint = await opts.shade.fingerprint;
  const flowId = `recover:${opts.setupId}:${bytesToBase64Url(crypto.getRandomValues(new Uint8Array(8)))}`;

  // Buckets for collecting responses.
  const grants = new Map<string, PendingShare>();
  const declines: Array<{ address: string; reason: string }> = [];
  const guardianSet = new Set(opts.guardians);

  let resolveCollection!: (value: 'enough' | 'declined-out' | 'timeout') => void;
  const collected = new Promise<'enough' | 'declined-out' | 'timeout'>((resolve) => {
    resolveCollection = resolve;
  });

  // Fingerprint we're waiting to see at setup-time. Filled as soon as
  // any grant lands so subsequent grants can be sanity-checked.
  let firstSetupFingerprint: string | null = null;

  const detach = opts.shade.onMessage(async (from, plaintext) => {
    if (!guardianSet.has(from)) return;
    const env = safeParse(plaintext);
    if (env === null) return;
    if (env.flowId !== flowId) return;
    if (env.type === 'share-grant') {
      if (grants.has(from)) return; // ignore duplicates
      try {
        const shareBytes = base64UrlToBytes(env.shareBytes);
        const share = decodeShare(shareBytes);
        const pending: PendingShare = {
          share,
          backupBlob: env.backupBlob,
          fromAddress: from,
        };
        if (firstSetupFingerprint === null) {
          // Shape from the share-grant doesn't include setupFingerprint;
          // we'll fall back to whatever importBackup reports.
          firstSetupFingerprint = null;
        }
        grants.set(from, pending);
      } catch {
        // Malformed share — count this as an implicit decline.
        declines.push({ address: from, reason: 'malformed share-bytes' });
      }
    } else if (env.type === 'share-decline') {
      if (grants.has(from)) return;
      declines.push({ address: from, reason: env.reason });
    } else {
      return;
    }
    opts.onProgress?.({
      granted: grants.size,
      declined: declines.length,
      pending: opts.guardians.length - grants.size - declines.length,
      threshold: opts.threshold,
      fromAddress: from,
      latest: env.type === 'share-grant' ? 'grant' : 'decline',
    });
    if (grants.size >= opts.threshold) {
      resolveCollection('enough');
      return;
    }
    if (opts.guardians.length - declines.length < opts.threshold) {
      resolveCollection('declined-out');
    }
  });

  // Send recovery-request to every guardian. Failures count as
  // implicit declines (the network couldn't reach them); we keep
  // going so partial reachability still clears the threshold.
  const requestedAt = now();
  for (const guardian of opts.guardians) {
    const env: RecoveryRequestEnvelope = {
      shadeRecovery: 1,
      type: 'recovery-request',
      flowId,
      originalAddress: opts.originalAddress,
      setupId: opts.setupId,
      requesterFingerprint,
      requestedAt,
    };
    const plaintext = encodeRecoveryEnvelope(env);
    try {
      const envelope: ShadeEnvelope = await opts.shade.send(guardian, plaintext);
      await opts.deliver(guardian, envelope);
    } catch (err) {
      declines.push({
        address: guardian,
        reason: `delivery failed: ${(err as Error).message}`,
      });
      opts.onProgress?.({
        granted: grants.size,
        declined: declines.length,
        pending: opts.guardians.length - grants.size - declines.length,
        threshold: opts.threshold,
        fromAddress: guardian,
        latest: 'decline',
      });
    }
  }

  // Already over-quota on declines from delivery failures alone.
  if (opts.guardians.length - declines.length < opts.threshold) {
    resolveCollection('declined-out');
  }

  // Race the collection against the timeout.
  let timer: ReturnType<typeof setTimeout> | null = null;
  if (Number.isFinite(timeoutMs)) {
    timer = setTimeout(() => resolveCollection('timeout'), timeoutMs);
  }
  let outcome: 'enough' | 'declined-out' | 'timeout';
  try {
    outcome = await collected;
  } finally {
    if (timer !== null) clearTimeout(timer);
    detach();
  }

  if (outcome === 'declined-out') {
    throw new RecoveryDeclinedError(
      declines.map((d) => d.address),
      opts.threshold,
      opts.guardians.length - declines.length,
    );
  }
  if (outcome === 'timeout') {
    throw new RecoveryTimeoutError(grants.size, opts.threshold);
  }

  // Reconstruct. Try the natural subset first; if a guardian forged
  // their share, retry with `threshold`-sized subsets until one
  // authenticates against the AES-GCM tag inside the backup blob.
  const granted = Array.from(grants.values());
  const result = await tryReconstruct(opts.shade, granted, opts.threshold);
  if (result === null) {
    throw new RecoveryReconstructionError(
      'none of the threshold-sized share subsets authenticated the backup blob — ' +
        'at least one guardian likely supplied a forged share',
    );
  }

  const restoredFingerprint = await opts.shade.fingerprint;
  return {
    applied: true,
    granted: granted.map((g) => g.fromAddress),
    declined: declines.slice(),
    setupFingerprint: firstSetupFingerprint,
    restoredFingerprint,
  };
}

function safeParse(plaintext: string): ReturnType<typeof tryParseRecoveryEnvelope> {
  try {
    return tryParseRecoveryEnvelope(plaintext);
  } catch {
    return null;
  }
}

/**
 * Attempt reconstruction with successively-chosen `threshold`-sized
 * subsets of `granted`. The first subset whose combined recoveryKey
 * authenticates the backup blob wins. Returns `null` if every subset
 * is rejected by the AEAD.
 *
 * In the honest case the very first subset (the natural-order
 * threshold-many) succeeds and we exit immediately. In the adversarial
 * case where one or more guardians supplied forged shares, we
 * exhaustively try all C(granted, threshold) subsets — bounded above
 * by C(255, 128) which is finite but large, so callers SHOULD cap the
 * size of the granted pool to a sane maximum (e.g. threshold + 2)
 * before invoking this. Inside `requestRecovery` we never collect more
 * than `threshold` grants because we stop the listener as soon as the
 * threshold lands; the iterative branch only fires when the caller
 * passes a pre-collected set.
 */
async function tryReconstruct(
  shade: Shade,
  granted: ReadonlyArray<PendingShare>,
  threshold: number,
): Promise<{ backupBlob: string } | null> {
  if (granted.length < threshold) return null;
  // All grants must agree on the backupBlob — if a guardian shipped a
  // different blob, we treat it as adversarial and skip its share.
  const groups = new Map<string, PendingShare[]>();
  for (const g of granted) {
    const bucket = groups.get(g.backupBlob);
    if (bucket === undefined) groups.set(g.backupBlob, [g]);
    else bucket.push(g);
  }

  for (const [backupBlob, group] of groups) {
    if (group.length < threshold) continue;
    const subsets = subsetsOfSize(group, threshold);
    for (const subset of subsets) {
      const recoveryKey = combineShares(subset.map((s) => s.share));
      try {
        const passphrase = recoveryKeyToBackupPassphrase(recoveryKey);
        await shade.importBackup(backupBlob, passphrase);
        return { backupBlob };
      } catch {
        // Wrong combination — keep trying other subsets.
        continue;
      } finally {
        recoveryKey.fill(0);
      }
    }
  }
  return null;
}

function* subsetsOfSize<T>(items: ReadonlyArray<T>, size: number): IterableIterator<T[]> {
  if (size <= 0) {
    yield [];
    return;
  }
  if (items.length < size) return;
  const indices = new Array<number>(size);
  for (let i = 0; i < size; i++) indices[i] = i;
  while (true) {
    yield indices.map((i) => items[i]!);
    // advance
    let i = size - 1;
    while (i >= 0 && indices[i]! === items.length - size + i) i--;
    if (i < 0) return;
    indices[i] = indices[i]! + 1;
    for (let j = i + 1; j < size; j++) indices[j] = indices[j - 1]! + 1;
  }
}

function validateOptions(opts: RequestRecoveryOptions): void {
  if (typeof opts.threshold !== 'number' || !Number.isInteger(opts.threshold) || opts.threshold < 1) {
    throw new RangeError('requestRecovery: threshold must be an integer ≥ 1');
  }
  if (!Array.isArray(opts.guardians) || opts.guardians.length === 0) {
    throw new RangeError('requestRecovery: guardians must be a non-empty array');
  }
  if (opts.guardians.length < opts.threshold) {
    throw new RangeError(
      `requestRecovery: guardians.length (${opts.guardians.length}) must be ≥ threshold (${opts.threshold})`,
    );
  }
  if (new Set(opts.guardians).size !== opts.guardians.length) {
    throw new RangeError('requestRecovery: guardians must be unique');
  }
  if (opts.guardians.includes(opts.shade.myAddress)) {
    throw new RangeError('requestRecovery: cannot include your own address as a guardian');
  }
  if (typeof opts.deliver !== 'function') {
    throw new TypeError('requestRecovery: deliver must be a function');
  }
  if (typeof opts.originalAddress !== 'string' || opts.originalAddress.length === 0) {
    throw new TypeError('requestRecovery: originalAddress must be a non-empty string');
  }
  if (typeof opts.setupId !== 'string' || opts.setupId.length === 0) {
    throw new TypeError('requestRecovery: setupId must be a non-empty string');
  }
}