HiveGrid-Wireframes/hivegrid-cloud/apps/web/app/utils/auth/mfa/recovery-codes.server.ts

import { createHmac, randomBytes, timingSafeEqual } from "node:crypto";
import { loadRequiredWebConfig } from "@toolctrl/config";

export const DEFAULT_RECOVERY_CODE_COUNT = 10;
export const DEFAULT_RECOVERY_CODE_SEGMENTS = 3;
export const DEFAULT_RECOVERY_CODE_SEGMENT_LENGTH = 4;
export const RECOVERY_CODE_ALPHABET = "23456789ABCDEFGHJKLMNPQRSTUVWXYZ";

export type RecoveryCodeHashEntry = {
  hash: string;
  usedAt: Date | null;
  createdAt: Date;
};

export type RecoveryCodeGeneration = {
  codes: string[];
  entries: RecoveryCodeHashEntry[];
  generatedAt: Date;
};

export type RecoveryCodeStatus = {
  configured: boolean;
  totalCount: number;
  remainingCount: number;
  usedCount: number;
  generatedAt: Date | null;
  lastUsedAt: Date | null;
};

type GenerateRecoveryCodesOptions = {
  count?: number;
  segments?: number;
  segmentLength?: number;
  now?: Date;
};

type NormalizeOptions = {
  segments?: number;
  segmentLength?: number;
};

function assertPositiveInteger(value: number, name: string) {
  if (!Number.isInteger(value) || value <= 0) {
    throw new Error(`${name} must be a positive integer`);
  }
}

function getRecoveryCodeSecret() {
  const { cryptoSecret, sessionSecret } = loadRequiredWebConfig();
  const secret = cryptoSecret || sessionSecret;
  if (!secret) {
    throw new Error("CRYPTO_SECRET or SESSION_SECRET env variable not set");
  }
  return secret;
}

function randomAlphabetCharacter() {
  const byteLimit = Math.floor(256 / RECOVERY_CODE_ALPHABET.length) * RECOVERY_CODE_ALPHABET.length;

  while (true) {
    const value = randomBytes(1)[0];
    if (value < byteLimit) {
      return RECOVERY_CODE_ALPHABET[value % RECOVERY_CODE_ALPHABET.length];
    }
  }
}

export function normalizeRecoveryCode(input: string, options: NormalizeOptions = {}) {
  const segments = options.segments ?? DEFAULT_RECOVERY_CODE_SEGMENTS;
  const segmentLength = options.segmentLength ?? DEFAULT_RECOVERY_CODE_SEGMENT_LENGTH;
  assertPositiveInteger(segments, "segments");
  assertPositiveInteger(segmentLength, "segmentLength");

  const compact = input.trim().toUpperCase().replace(/[\s-]/g, "");
  const expectedLength = segments * segmentLength;
  if (compact.length !== expectedLength || [...compact].some((character) => !RECOVERY_CODE_ALPHABET.includes(character))) {
    return null;
  }

  return Array.from({ length: segments }, (_, index) =>
    compact.slice(index * segmentLength, (index + 1) * segmentLength)
  ).join("-");
}

export function generateRecoveryCode(options: NormalizeOptions = {}) {
  const segments = options.segments ?? DEFAULT_RECOVERY_CODE_SEGMENTS;
  const segmentLength = options.segmentLength ?? DEFAULT_RECOVERY_CODE_SEGMENT_LENGTH;
  assertPositiveInteger(segments, "segments");
  assertPositiveInteger(segmentLength, "segmentLength");

  return Array.from({ length: segments }, () =>
    Array.from({ length: segmentLength }, () => randomAlphabetCharacter()).join("")
  ).join("-");
}

export function hashRecoveryCode(code: string) {
  const normalized = normalizeRecoveryCode(code);
  if (!normalized) {
    return null;
  }

  return createHmac("sha256", getRecoveryCodeSecret()).update(normalized).digest("hex");
}

function hashMatches(candidateHash: string, storedHash: string) {
  const candidateBuffer = Buffer.from(candidateHash, "hex");
  const storedBuffer = Buffer.from(storedHash, "hex");
  if (candidateBuffer.length !== storedBuffer.length) {
    return false;
  }
  return timingSafeEqual(candidateBuffer, storedBuffer);
}

export function verifyRecoveryCodeHash(code: string, storedHash: string) {
  const candidateHash = hashRecoveryCode(code);
  if (!candidateHash) {
    return false;
  }
  return hashMatches(candidateHash, storedHash);
}

export function generateRecoveryCodes(options: GenerateRecoveryCodesOptions = {}): RecoveryCodeGeneration {
  const count = options.count ?? DEFAULT_RECOVERY_CODE_COUNT;
  const segments = options.segments ?? DEFAULT_RECOVERY_CODE_SEGMENTS;
  const segmentLength = options.segmentLength ?? DEFAULT_RECOVERY_CODE_SEGMENT_LENGTH;
  assertPositiveInteger(count, "count");
  assertPositiveInteger(segments, "segments");
  assertPositiveInteger(segmentLength, "segmentLength");

  const generatedAt = options.now ?? new Date();
  // 32^12 (~10^18) code space means in-batch collisions are astronomically
  // unlikely. Generate `count` codes upfront and only retry once if we ever
  // see a duplicate (instead of looping forever on every call).
  const MAX_RETRIES = 2;
  let codeList: string[] = [];
  for (let attempt = 0; attempt < MAX_RETRIES; attempt += 1) {
    const candidates = Array.from({ length: count }, () => generateRecoveryCode({ segments, segmentLength }));
    if (new Set(candidates).size === count) {
      codeList = candidates;
      break;
    }
  }
  if (codeList.length !== count) {
    throw new Error("Failed to generate unique recovery codes after retry; suspected entropy source failure.");
  }
  return {
    codes: codeList,
    entries: codeList.map((code) => ({
      hash: hashRecoveryCode(code)!,
      usedAt: null,
      createdAt: generatedAt,
    })),
    generatedAt,
  };
}

export function consumeRecoveryCode(
  entries: RecoveryCodeHashEntry[],
  code: string,
  options: { now?: Date } = {}
): { matched: boolean; entries: RecoveryCodeHashEntry[]; usedAt: Date | null } {
  const candidateHash = hashRecoveryCode(code);
  if (!candidateHash) {
    return { matched: false, entries, usedAt: null };
  }

  const usedAt = options.now ?? new Date();
  let matched = false;
  const nextEntries = entries.map((entry) => {
    if (entry.usedAt || matched || !hashMatches(candidateHash, entry.hash)) {
      return entry;
    }
    matched = true;
    return { ...entry, usedAt };
  });

  return { matched, entries: matched ? nextEntries : entries, usedAt: matched ? usedAt : null };
}

export function getRecoveryCodeStatus(
  entries: RecoveryCodeHashEntry[] | null | undefined,
  options: { generatedAt?: Date | null; lastUsedAt?: Date | null } = {}
): RecoveryCodeStatus {
  const configured = Boolean(entries?.length);
  const totalCount = entries?.length ?? 0;
  const usedEntries = entries?.filter((entry) => entry.usedAt) ?? [];
  const remainingCount = entries?.filter((entry) => !entry.usedAt).length ?? 0;
  const lastEntryUsedAt = usedEntries.reduce<Date | null>((latest, entry) => {
    if (!entry.usedAt) {
      return latest;
    }
    if (!latest || entry.usedAt.getTime() > latest.getTime()) {
      return entry.usedAt;
    }
    return latest;
  }, null);

  return {
    configured,
    totalCount,
    remainingCount,
    usedCount: usedEntries.length,
    generatedAt: options.generatedAt ?? (configured ? entries?.[0]?.createdAt ?? null : null),
    lastUsedAt: options.lastUsedAt ?? lastEntryUsedAt,
  };
}