HealthLog/src/lib/crypto/encryption-corpus.ts

/**
 * v1.23 — read-only scan + in-place re-encrypt over the whole encrypted-column
 * corpus, driven by the canonical registry (`encrypted-columns.ts`).
 *
 * Two consumers:
 *   - the admin encryption-status view (`GET /api/admin/encryption/status`)
 *     calls `scanCorpus()` to bucket every encrypted column's rows by key id;
 *   - the admin-triggered rotation pg-boss job (`encryption-key-rotate`) calls
 *     `rotateCorpus()` to re-encrypt every row that is not already on the
 *     active key.
 *
 * The standalone CLI (`scripts/rotate-encryption-key.ts`) remains the canonical
 * rotation path and stays independent (its own Prisma client). This module is
 * the in-app convenience that reuses the SAME registry, so the guard test keeps
 * both in lock-step.
 *
 * GUARANTEES (the security review must confirm these on the rotation path):
 *  - ACTIVE-KEY-ONLY. Re-encryption is `encrypt(decrypt(value))`; `encrypt()`
 *    always writes the configured active key id. There is no code path here
 *    that selects any other write key.
 *  - NEVER ADDS / DROPS A KEY. This module never reads or mutates
 *    `ENCRYPTION_KEYS` / `ENCRYPTION_ACTIVE_KEY_ID`. The operator's env key map
 *    is the only place keys live; a key drop stays a deliberate env + redeploy
 *    act, never a button.
 *  - IDEMPOTENT. `shouldRotate()` skips rows already on the active key, so a
 *    second pass (or two racing workers) re-encrypts zero rows.
 *  - FAIL-CLOSED. A row written under a key id that is no longer configured
 *    throws on decrypt (counted as an error, the row is left untouched) rather
 *    than being silently dropped or overwritten — exactly the property that
 *    protects against dropping a legacy key too early.
 */
import { Buffer } from "node:buffer";
import {
  decrypt,
  encrypt,
  extractKeyId,
  extractKeyIdFromBytes,
  getActiveKeyId,
  reencryptBytesToActive,
} from "@/lib/crypto";
import {
  ENCRYPTED_COLUMNS,
  type EncryptedColumn,
} from "@/lib/crypto/encrypted-columns";

/** Sentinel bucket for legacy (unversioned) ciphertext under `byKeyId`. */
export const LEGACY_BUCKET = "legacy";

/**
 * Batch size for blob columns (`codecField` or `batched` set). The document
 * vault's rows are up to cap-sized ciphertexts and a backup row is the whole
 * account compressed, so the walk is id-cursor paginated — at most this many
 * blobs are in memory at once.
 */
export const BLOB_ROTATION_BATCH_SIZE = 25;

/** Minimal Prisma delegate shape this module needs. */
interface ColumnDelegate {
  findMany: (args: {
    select: Record<string, true>;
    orderBy?: Record<string, "asc" | "desc">;
    take?: number;
    cursor?: Record<string, unknown>;
    skip?: number;
  }) => Promise<Array<Record<string, unknown>>>;
  update: (args: {
    where: Record<string, string>;
    data: Record<string, unknown>;
  }) => Promise<unknown>;
  /** Only ever called for a `disposable` column's unreadable rows. */
  delete: (args: { where: Record<string, string> }) => Promise<unknown>;
}

/** The subset of the Prisma client we touch: one delegate per model. */
export type CorpusClient = Record<string, ColumnDelegate>;

/** PascalCase model name -> camelCase Prisma delegate key. */
function delegateKey(model: string): string {
  return model.charAt(0).toLowerCase() + model.slice(1);
}

function getDelegate(client: CorpusClient, model: string): ColumnDelegate {
  const delegate = client[delegateKey(model)];
  if (!delegate) {
    throw new Error(`No Prisma delegate for model '${model}'`);
  }
  return delegate;
}

/** Read a registry column's value as a ciphertext string (Bytes -> utf8). */
function toCiphertext(
  value: unknown,
  kind: EncryptedColumn["kind"],
): string | null {
  if (value == null) return null;
  if (kind === "bytes") {
    const buf = value as Uint8Array;
    if (buf.byteLength === 0) return null;
    return Buffer.from(buf).toString("utf8");
  }
  const s = value as string;
  return s.length === 0 ? null : s;
}

/** Encode a re-encrypted ciphertext string back into the column's storage shape. */
function fromCiphertext(
  value: string,
  kind: EncryptedColumn["kind"],
): string | Uint8Array {
  if (kind !== "bytes") return value;
  const encoded = Buffer.from(value, "utf8");
  const next = new Uint8Array(new ArrayBuffer(encoded.byteLength));
  next.set(encoded);
  return next;
}

function shouldRotate(ciphertext: string): boolean {
  return extractKeyId(ciphertext) !== getActiveKeyId();
}

// ─── Codec-dispatched blob columns (document vault) ─────────────────────────

/**
 * The key id a codec-dispatched blob row was written under, or null when
 * legacy/unparsable. "binary2" parses the binary header; every other codec
 * value (notably "base64v1") is the `encrypt()`-string-as-UTF-8 shape.
 */
function blobKeyId(value: Uint8Array, codec: string): string | null {
  const buf = Buffer.from(value);
  if (codec === "binary2") return extractKeyIdFromBytes(buf);
  return extractKeyId(buf.toString("utf8"));
}

/** Re-encrypt one codec-dispatched blob under its OWN codec (never converts). */
function reencryptBlob(value: Uint8Array, codec: string): Uint8Array {
  const buf = Buffer.from(value);
  if (codec === "binary2") {
    const rotated = reencryptBytesToActive(buf);
    const next = new Uint8Array(new ArrayBuffer(rotated.byteLength));
    next.set(rotated);
    return next;
  }
  if (codec === "base64v1") {
    const rotated = encrypt(decrypt(buf.toString("utf8")));
    const encoded = Buffer.from(rotated, "utf8");
    const next = new Uint8Array(new ArrayBuffer(encoded.byteLength));
    next.set(encoded);
    return next;
  }
  // FAIL-CLOSED: an unknown codec is counted as an error and left untouched.
  throw new Error(`Unknown content codec '${codec}'`);
}

/**
 * Walk a blob column in bounded id-cursor batches, invoking `onRow` per
 * non-empty row. At most `BLOB_ROTATION_BATCH_SIZE` rows are in memory per
 * step; an interrupted run resumes safely on re-invocation because processing
 * is idempotent (already-active rows are skipped by the callers).
 *
 * `codec` is the row's own codec label for a codec-dispatched column and null
 * for a plain `batched` one — a backup blob has one layout, just a large one.
 */
async function walkBlobColumn(
  delegate: ColumnDelegate,
  col: EncryptedColumn,
  onRow: (row: {
    id: string;
    value: unknown;
    codec: string | null;
  }) => Promise<void> | void,
): Promise<void> {
  const codecField = col.codecField;
  const pk = pkField(col);
  let cursor: string | null = null;
  for (;;) {
    const rows = await delegate.findMany({
      select: {
        [pk]: true,
        [col.field]: true,
        ...(codecField ? { [codecField]: true } : {}),
      },
      orderBy: { [pk]: "asc" },
      take: BLOB_ROTATION_BATCH_SIZE,
      ...(cursor ? { cursor: { [pk]: cursor }, skip: 1 } : {}),
    });
    if (rows.length === 0) break;
    for (const row of rows) {
      const value = row[col.field];
      if (value == null) continue;
      if (value instanceof Uint8Array && value.byteLength === 0) continue;
      if (typeof value === "string" && value.length === 0) continue;
      await onRow({
        id: row[pk] as string,
        value,
        codec: codecField ? String(row[codecField] ?? "") : null,
      });
    }
    cursor = rows[rows.length - 1]![pk] as string;
    if (rows.length < BLOB_ROTATION_BATCH_SIZE) break;
  }
}

/** The model's primary-key field: `id` unless the registry says otherwise. */
function pkField(col: EncryptedColumn): string {
  return col.pkField ?? "id";
}

/** True when the column's rows must be walked in bounded batches. */
function isBatched(col: EncryptedColumn): boolean {
  return Boolean(col.codecField ?? col.batched);
}

/** The key id a walked row was written under, dispatching on its codec. */
function walkedKeyId(
  value: unknown,
  codec: string | null,
  kind: EncryptedColumn["kind"],
): string | null {
  if (codec !== null) return blobKeyId(value as Uint8Array, codec);
  const ciphertext = toCiphertext(value, kind);
  return ciphertext == null ? null : extractKeyId(ciphertext);
}

export interface ColumnScan {
  model: string;
  field: string;
  kind: EncryptedColumn["kind"];
  /** Non-null ciphertext rows. */
  total: number;
  /** Rows per key id; legacy/unversioned rows land under `LEGACY_BUCKET`. */
  byKeyId: Record<string, number>;
  /** Rows under the legacy/unversioned format (= `byKeyId[LEGACY_BUCKET]`). */
  legacy: number;
}

/** Scan one encrypted column: bucket every non-null ciphertext by key id. */
export async function scanColumn(
  client: CorpusClient,
  col: EncryptedColumn,
): Promise<ColumnScan> {
  const delegate = getDelegate(client, col.model);
  const byKeyId: Record<string, number> = {};
  let total = 0;

  if (isBatched(col)) {
    // Blob column: bounded batches, per-row codec where the column has one.
    await walkBlobColumn(delegate, col, ({ value, codec }) => {
      total += 1;
      const id = walkedKeyId(value, codec, col.kind) ?? LEGACY_BUCKET;
      byKeyId[id] = (byKeyId[id] ?? 0) + 1;
    });
  } else {
    const rows = await delegate.findMany({
      select: { [pkField(col)]: true, [col.field]: true },
    });
    for (const row of rows) {
      const ciphertext = toCiphertext(row[col.field], col.kind);
      if (ciphertext == null) continue;
      total += 1;
      const id = extractKeyId(ciphertext) ?? LEGACY_BUCKET;
      byKeyId[id] = (byKeyId[id] ?? 0) + 1;
    }
  }
  return {
    model: col.model,
    field: col.field,
    kind: col.kind,
    total,
    byKeyId,
    legacy: byKeyId[LEGACY_BUCKET] ?? 0,
  };
}

export interface CorpusScan {
  activeKeyId: string;
  columns: ColumnScan[];
  /** Total non-null ciphertext rows across the corpus. */
  totalRows: number;
  /** Rows already on the active key. */
  activeRows: number;
  /** Rows NOT on the active key (legacy + any non-active versioned). */
  staleRows: number;
  /**
   * True iff every column has zero rows that are not on the active key — the
   * single signal an operator needs before dropping a legacy key.
   */
  rotationComplete: boolean;
}

/** Scan the whole corpus. Read-only; never writes. */
export async function scanCorpus(client: CorpusClient): Promise<CorpusScan> {
  const activeKeyId = getActiveKeyId();
  const columns: ColumnScan[] = [];
  for (const col of ENCRYPTED_COLUMNS) {
    columns.push(await scanColumn(client, col));
  }
  let totalRows = 0;
  let activeRows = 0;
  for (const c of columns) {
    totalRows += c.total;
    activeRows += c.byKeyId[activeKeyId] ?? 0;
  }
  const staleRows = totalRows - activeRows;
  return {
    activeKeyId,
    columns,
    totalRows,
    activeRows,
    staleRows,
    rotationComplete: staleRows === 0,
  };
}

export interface RotationResult {
  model: string;
  field: string;
  scanned: number;
  rotated: number;
  errors: number;
  /** Unreadable rows deleted from a `disposable` column. */
  dropped: number;
}

/** Re-encrypt one column's stale rows to the active key. */
export async function rotateColumn(
  client: CorpusClient,
  col: EncryptedColumn,
): Promise<RotationResult> {
  const delegate = getDelegate(client, col.model);
  const result: RotationResult = {
    model: col.model,
    field: col.field,
    scanned: 0,
    rotated: 0,
    errors: 0,
    dropped: 0,
  };

  /**
   * A row that could not be re-encrypted. FAIL-CLOSED by default: count it and
   * leave it untouched rather than dropping data. A `disposable` column takes
   * the other branch — the value is a reproducible cache entry, and keeping an
   * unreadable one only guarantees the next reader gets a body it cannot
   * parse, so the row goes instead.
   */
  const onUnreadable = async (id: string): Promise<void> => {
    if (!col.disposable) {
      result.errors += 1;
      return;
    }
    try {
      await delegate.delete({ where: { [pkField(col)]: id } });
      result.dropped += 1;
    } catch {
      result.errors += 1;
    }
  };

  if (isBatched(col)) {
    // Blob column: bounded id-cursor batches (never an unbounded findMany),
    // re-encrypted under each row's OWN codec where it has one. Idempotent —
    // rows already on the active key are skipped, so an interrupted run
    // resumes cleanly on the next invocation.
    await walkBlobColumn(delegate, col, async ({ id, value, codec }) => {
      result.scanned += 1;
      if (walkedKeyId(value, codec, col.kind) === getActiveKeyId()) return;
      try {
        // The plaintext is never inspected, only re-sealed, so every envelope
        // a stored value can legitimately carry survives rotation untouched.
        const next =
          codec !== null
            ? reencryptBlob(value as Uint8Array, codec)
            : fromCiphertext(
                encrypt(decrypt(toCiphertext(value, col.kind)!)),
                col.kind,
              );
        await delegate.update({
          where: { [pkField(col)]: id },
          data: { [col.field]: next },
        });
        result.rotated += 1;
      } catch {
        await onUnreadable(id);
      }
    });
    return result;
  }

  const pk = pkField(col);
  const rows = await delegate.findMany({
    select: { [pk]: true, [col.field]: true },
  });
  result.scanned = rows.length;
  for (const row of rows) {
    const ciphertext = toCiphertext(row[col.field], col.kind);
    if (ciphertext == null || !shouldRotate(ciphertext)) continue;
    const id = row[pk] as string;
    try {
      // ACTIVE-KEY-ONLY: encrypt() always writes the active key id.
      const reencrypted = encrypt(decrypt(ciphertext));
      await delegate.update({
        where: { [pk]: id },
        data: { [col.field]: fromCiphertext(reencrypted, col.kind) },
      });
      result.rotated += 1;
    } catch {
      await onUnreadable(id);
    }
  }
  return result;
}

export interface CorpusRotation {
  activeKeyId: string;
  results: RotationResult[];
  totalScanned: number;
  totalRotated: number;
  totalErrors: number;
  /** Unreadable rows deleted from `disposable` columns. */
  totalDropped: number;
}

/** Re-encrypt the whole corpus to the active key. Idempotent + active-key-only. */
export async function rotateCorpus(
  client: CorpusClient,
): Promise<CorpusRotation> {
  const activeKeyId = getActiveKeyId();
  const results: RotationResult[] = [];
  for (const col of ENCRYPTED_COLUMNS) {
    results.push(await rotateColumn(client, col));
  }
  let totalScanned = 0;
  let totalRotated = 0;
  let totalErrors = 0;
  let totalDropped = 0;
  for (const r of results) {
    totalScanned += r.scanned;
    totalRotated += r.rotated;
    totalErrors += r.errors;
    totalDropped += r.dropped;
  }
  return {
    activeKeyId,
    results,
    totalScanned,
    totalRotated,
    totalErrors,
    totalDropped,
  };
}